From 840c83b4034a266cdd9e0858b9a568e6c348cf22 Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Fri, 30 Oct 2026 11:19:08 -0700 Subject: [PATCH 01/46] test: bedrock claude prompt caching on /v1/messages reports and charges cache tokens (#49265) * test: bedrock claude prompt caching on /v1/messages reports and charges cache tokens * ci: re-run after ci-coverage and assert-shard-coverage were cancelled by the GitHub Actions incident --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...st_bedrock_messages_prompt_caching_wire.py | 191 ++++++++++++++++++ 1 file changed, 191 insertions(+) create mode 100644 tests/integration/messages_endpoint/providers/bedrock/test_bedrock_messages_prompt_caching_wire.py diff --git a/tests/integration/messages_endpoint/providers/bedrock/test_bedrock_messages_prompt_caching_wire.py b/tests/integration/messages_endpoint/providers/bedrock/test_bedrock_messages_prompt_caching_wire.py new file mode 100644 index 00000000000..f2622da1bf6 --- /dev/null +++ b/tests/integration/messages_endpoint/providers/bedrock/test_bedrock_messages_prompt_caching_wire.py @@ -0,0 +1,191 @@ +import base64 +import json +import uuid +from typing import Final + +import anthropic +import httpx +import pytest +from integration._support import claude_code as cc +from integration._support.client import Gateway, eventually +from integration._support.database import read_rows +from integration._support.upstream import _aws_event_frame +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +_MODEL_ID: Final = "us.anthropic.claude-sonnet-4-5-20250929-v1:0" +_TOKEN: Final = "synthetic-bedrock-bearer" +_PROMPT: Final = "Reply with exactly the word PONG" +_INPUT_TOKENS: Final = 10 +_OUTPUT_TOKENS: Final = 4 +_CACHE_READ_TOKENS: Final = 900 +_CACHE_WRITE_TOKENS: Final = 300 +_INPUT_RATE: Final = 0.001 +_OUTPUT_RATE: Final = 0.002 +_CACHE_READ_RATE: Final = 0.0001 +_CACHE_WRITE_RATE: Final = 0.0015 +_EXPECTED_SPEND: Final = ( + _INPUT_TOKENS * _INPUT_RATE + + _CACHE_READ_TOKENS * _CACHE_READ_RATE + + _CACHE_WRITE_TOKENS * _CACHE_WRITE_RATE + + _OUTPUT_TOKENS * _OUTPUT_RATE +) +_PROMPT_TOKENS: Final = _INPUT_TOKENS + _CACHE_READ_TOKENS + _CACHE_WRITE_TOKENS +_USAGE: Final = { + "input_tokens": _INPUT_TOKENS, + "output_tokens": _OUTPUT_TOKENS, + "cache_read_input_tokens": _CACHE_READ_TOKENS, + "cache_creation_input_tokens": _CACHE_WRITE_TOKENS, +} +_USER_CONTENT: Final = [{"type": "text", "text": _PROMPT, "cache_control": cc.CACHE}] +_CACHED_SYSTEM_BLOCKS: Final = [block for block in cc.system_blocks() if "cache_control" in block] + + +def _assert_claude_code_breakpoints_reach_bedrock(request: Request, target: str) -> dict[str, JsonValue]: + assert request.method == "POST" and request.target == target, request.target + assert request.headers["authorization"] == f"Bearer {_TOKEN}" + body: Final = cc.JSON_OBJECT.validate_json(request.body) + assert body["anthropic_version"] == "bedrock-2023-05-31", body + assert body["system"] == _CACHED_SYSTEM_BLOCKS, f"billing-header block dropped, breakpoints kept: {body}" + assert body["messages"] == [{"role": "user", "content": _USER_CONTENT}], body + return body + + +def _invoke_chunk(event: dict[str, JsonValue]) -> bytes: + return _aws_event_frame("chunk", {"bytes": base64.b64encode(json.dumps(event).encode()).decode()}, "sc", "u") + + +def _stream_frames(message_id: str) -> tuple[bytes, ...]: + return ( + _invoke_chunk( + { + "type": "message_start", + "message": { + "id": message_id, + "type": "message", + "role": "assistant", + "content": [], + "model": "claude-sonnet-4-5-20250929", + "stop_reason": None, + "stop_sequence": None, + "usage": {**_USAGE, "output_tokens": 1}, + }, + } + ), + _invoke_chunk({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + _invoke_chunk({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": "PONG"}}), + _invoke_chunk({"type": "content_block_stop", "index": 0}), + _invoke_chunk( + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": _OUTPUT_TOKENS}, + } + ), + _invoke_chunk({"type": "message_stop"}), + ) + + +def _spend_row(message_id: str) -> dict[str, object]: + rows: Final = eventually( + lambda: read_rows( + 'SELECT spend, prompt_tokens, completion_tokens FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (message_id,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + return rows[0] + + +def _client(gateway: Gateway, transport: httpx.Client) -> anthropic.Anthropic: + return anthropic.Anthropic( + api_key=gateway.key, base_url=str(gateway.client.base_url), max_retries=0, http_client=transport + ) + + +def test_bedrock_claude_cache_control_reaches_invoke_and_cache_tokens_are_reported_and_charged( + gateway: Gateway, +) -> None: + message_id: Final = f"msg_{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + _assert_claude_code_breakpoints_reach_bedrock(request, f"/model/{_MODEL_ID}/invoke") + return Reply( + body=cc.message_reply( + message_id, "claude-sonnet-4-5-20250929", ({"type": "text", "text": "PONG"},), dict(_USAGE) + ) + ) + + with ( + wire_server(respond) as wire, + gateway.scenario() as scenario, + httpx.Client(timeout=15, trust_env=False) as transport, + ): + model: Final = scenario.model( + model=f"bedrock/invoke/{_MODEL_ID}", + api_key=_TOKEN, + aws_region_name="us-east-1", + api_base=wire.url, + input_cost_per_token=_INPUT_RATE, + output_cost_per_token=_OUTPUT_RATE, + cache_read_input_token_cost=_CACHE_READ_RATE, + cache_creation_input_token_cost=_CACHE_WRITE_RATE, + ) + final: Final = _client(gateway, transport).messages.create( + model=model, + max_tokens=32, + system=list(cc.system_blocks()), + messages=[{"role": "user", "content": _USER_CONTENT}], + ) + assert len(wire.drain()) == 1 + assert final.id == message_id, final + assert [(block.type, block.text) for block in final.content] == [("text", "PONG")], final + assert final.usage.model_dump(exclude_none=True, include=set(_USAGE)) == _USAGE, final.usage + row: Final = _spend_row(message_id) + + assert float(str(row["spend"])) == pytest.approx(_EXPECTED_SPEND), row + assert (row["prompt_tokens"], row["completion_tokens"]) == (_PROMPT_TOKENS, _OUTPUT_TOKENS), row + + +def test_bedrock_claude_stream_cache_tokens_from_message_start_are_reported_and_charged(gateway: Gateway) -> None: + message_id: Final = f"msg_{uuid.uuid4().hex}" + + def respond(request: Request) -> Reply: + body: Final = _assert_claude_code_breakpoints_reach_bedrock( + request, f"/model/{_MODEL_ID}/invoke-with-response-stream" + ) + assert "stream" not in body, body + return Reply(chunks=_stream_frames(message_id), content_type="application/vnd.amazon.eventstream") + + with ( + wire_server(respond) as wire, + gateway.scenario() as scenario, + httpx.Client(timeout=15, trust_env=False) as transport, + ): + model: Final = scenario.model( + model=f"bedrock/invoke/{_MODEL_ID}", + api_key=_TOKEN, + aws_region_name="us-east-1", + api_base=wire.url, + input_cost_per_token=_INPUT_RATE, + output_cost_per_token=_OUTPUT_RATE, + cache_read_input_token_cost=_CACHE_READ_RATE, + cache_creation_input_token_cost=_CACHE_WRITE_RATE, + ) + with _client(gateway, transport).messages.stream( + model=model, + max_tokens=32, + system=list(cc.system_blocks()), + messages=[{"role": "user", "content": _USER_CONTENT}], + ) as stream: + deltas: Final = tuple(stream.text_stream) + final: Final = stream.get_final_message() + assert len(wire.drain()) == 1 + assert deltas == ("PONG",) + assert final.id == message_id, final + assert final.usage.model_dump(exclude_none=True, include=set(_USAGE)) == _USAGE, final.usage + row: Final = _spend_row(message_id) + + assert float(str(row["spend"])) == pytest.approx(_EXPECTED_SPEND), row + assert (row["prompt_tokens"], row["completion_tokens"]) == (_PROMPT_TOKENS, _OUTPUT_TOKENS), row From d23af8bdd08943fe1a0e1d9e528fcb12000b287a Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Fri, 30 Oct 2026 11:24:49 -0700 Subject: [PATCH 02/46] fix(cost): bill the margin once on rebuilt streams and streamed Responses through the chat bridge (#49254) * fix(cost): bill the margin once on rebuilt streams and streamed Responses through the chat bridge The last usage chunk and the Responses terminal event carry litellm's bill in usage.cost, and two readers took that bill for the upstream's own charge and margined it again: completion_cost on a stream_chunk_builder result, and stamp_responses_usage_cost on the chat bridge. The stamp now keeps the upstream's usage.cost under the hidden upstream_usage_cost key, the chunk builder carries it onto the rebuilt response, and the Responses stamp prices that saved figure * fix(cost): hand the Responses stamp the upstream charge through upstream_usage_cost only, never the response-cost header response_cost_calculator returns the response-cost header as the final bill without the margin, so a priced copy carrying the upstream's raw usage.cost there billed the client the unmargined figure. The hidden upstream_usage_cost key already feeds billed_usage_cost, which margins it once. The Responses test now prices through a real Logging object instead of a mocked calculator * fix(cost): leave stamp_responses_usage_cost as on main The Responses chat bridge copies the inner stream's final usage chunk, so the wrapper stamp keeping upstream_usage_cost already prices it once. Changing the Responses stamp was not needed for that and re-priced xai MCP rounds whose usage.cost is already litellm's bill (0.00512 for 0.00312 in test_xai_stream_auto_execute_bills_each_round_at_its_own_price_with_one_fixed_margin) --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../streaming_chunk_builder_utils.py | 18 ++++++++++- .../litellm_core_utils/streaming_handler.py | 18 ++++++++--- .../test_streaming_chunk_builder_utils.py | 24 ++++++++++++++ .../test_streaming_handler.py | 31 +++++++++++++++++++ 4 files changed, 86 insertions(+), 5 deletions(-) diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 574f1103c1a..0ddfa1210c2 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -14,12 +14,18 @@ import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.core_helpers import ( PROVIDERS_BILLED_FROM_THEIR_OWN_USAGE_FIELDS, + UPSTREAM_USAGE_COST, chunk_usage, finite_cost, provider_reported_usage_cost, stamp_litellm_bill_on_usage, ) -from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params +from litellm.litellm_core_utils.hidden_params import ( + HIDDEN_PARAMS_ATTR, + get_hidden_params, + set_hidden_param, + set_hidden_params, +) from litellm.litellm_core_utils.llm_cost_calc.guardrail_cost import BEDROCK_TRACE_FIELDS, object_items, object_mapping from litellm.litellm_core_utils.llm_response_utils.convert_dict_to_response import partition_annotations from litellm.litellm_core_utils.prompt_templates.common_utils import ( @@ -1444,6 +1450,14 @@ def _set_stream_builder_response_cost( _set_hidden_response_cost(response, response_cost) +def _carry_upstream_usage_cost(response: ModelResponse, chunks: Sequence[object]) -> None: + chunk_hidden_params: Final = (get_hidden_params(chunk) for chunk in reversed(chunks)) + stamped: Final = (hidden for hidden in chunk_hidden_params if hidden is not None and UPSTREAM_USAGE_COST in hidden) + for hidden in stamped: + set_hidden_param(response, UPSTREAM_USAGE_COST, hidden[UPSTREAM_USAGE_COST]) + return + + def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_obj: Optional["Logging"]) -> None: if logging_obj is None: return @@ -1794,6 +1808,7 @@ def stream_chunk_builder( response, logging_obj, ) + _carry_upstream_usage_cost(response, chunks) # pyright: ignore[reportUnknownArgumentType] # chunks is an untyped list param _stamp_streaming_usage_cost(usage, response, logging_obj) _set_stream_builder_response_cost( response, @@ -1836,6 +1851,7 @@ def stream_chunk_builder( ) break + _carry_upstream_usage_cost(response, chunks) # pyright: ignore[reportUnknownArgumentType] # chunks is an untyped list param _stamp_streaming_usage_cost(usage, response, logging_obj) _set_stream_builder_response_cost( response, diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index b1b3e20f218..9710956c97e 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -24,7 +24,12 @@ import litellm from litellm import verbose_logger from litellm._uuid import uuid from litellm.litellm_core_utils.asyncify import asyncify -from litellm.litellm_core_utils.core_helpers import chunk_usage, provider_reported_usage_cost +from litellm.litellm_core_utils.core_helpers import ( + UPSTREAM_USAGE_COST, + chunk_usage, + provider_reported_usage_cost, + stamp_litellm_bill_on_usage, +) from litellm.litellm_core_utils.hidden_params import ( HIDDEN_PARAMS_ATTR, get_hidden_params, @@ -98,12 +103,17 @@ def _stamp_usage_cost( usage_chunk: ModelResponse | TextCompletionResponse | ModelResponseStream, complete_response: ModelResponse | TextCompletionResponse | ModelResponseStream, ) -> None: - bill: Final = (get_hidden_params(complete_response) or {}).get("response_cost") + complete_hidden: Final = get_hidden_params(complete_response) or {} + bill: Final = complete_hidden.get("response_cost") usage: Final[object] = getattr(usage_chunk, "usage", None) if usage is None or isinstance(bill, bool) or not isinstance(bill, (int, float)): return - if bill > 0 or isinstance(getattr(usage, "cost", None), (int, float)): - setattr(usage, "cost", float(bill)) # noqa: B010 # cost is an optional extra on the usage model + if bill <= 0 and not isinstance(getattr(usage, "cost", None), (int, float)): + return + chunk_hidden: Final = get_hidden_params(usage_chunk) + if chunk_hidden is not None and UPSTREAM_USAGE_COST in complete_hidden: + chunk_hidden[UPSTREAM_USAGE_COST] = complete_hidden[UPSTREAM_USAGE_COST] + stamp_litellm_bill_on_usage(usage_chunk, usage, float(bill)) def _stamp_served_service_tier(response: ModelResponseStream, complete_streaming_response: ModelResponse) -> None: diff --git a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index c8fa9c94c17..9df45b92686 100644 --- a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -1077,6 +1077,30 @@ def test_cost_field_in_usage_chunks(): assert usage.completion_tokens == 5 +def test_the_rebuilt_response_keeps_the_upstream_cost_a_stamped_usage_chunk_saved(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setattr(litellm, "cost_margin_config", {"global": 0.5}) + text_chunk = ModelResponseStream( + id="chatcmpl-1", + model="gpt-4.1", + choices=[StreamingChoices(finish_reason="stop", index=0, delta=Delta(content="Hi"))], + ) + usage_chunk = ModelResponseStream( + id="chatcmpl-1", + model="gpt-4.1", + choices=[], + usage=Usage(completion_tokens=5, prompt_tokens=10, total_tokens=15, cost=0.0015), + ) + usage_chunk._hidden_params = {"response_cost": 0.0015, "upstream_usage_cost": 0.001} + + rebuilt = stream_chunk_builder([text_chunk, usage_chunk], messages=[{"role": "user", "content": "hi"}]) + + assert rebuilt is not None + assert rebuilt._hidden_params["upstream_usage_cost"] == 0.001 + assert litellm.completion_cost(rebuilt, custom_llm_provider="openai") == pytest.approx(0.0015), ( + "the upstream's 0.001 gets the 50% margin once, litellm's 0.0015 bill is not margined again" + ) + + def _perplexity_usage_chunk(cost: Mapping[str, float], finish_reason: str | None) -> ModelResponseStream: return ModelResponseStream( diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 148b7193ef6..6df4731e335 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -6847,6 +6847,37 @@ def test_the_usage_chunk_carries_the_bill_the_builder_priced_for_the_spend_row() assert usage_chunk.usage is not None and usage_chunk.usage.cost == 0.00032, "the client sees what the row bills" +@pytest.mark.parametrize("upstream_cost", [None, 0.001]) +def test_rebuilding_a_priced_stream_bills_the_margin_once(monkeypatch: pytest.MonkeyPatch, upstream_cost: float | None) -> None: + from litellm.litellm_core_utils.core_helpers import stamp_litellm_bill_on_usage + from litellm.litellm_core_utils.streaming_handler import ( + _stamp_usage_cost, # pyright: ignore[reportPrivateUsage] # the helper under test + ) + + monkeypatch.setattr(litellm, "cost_margin_config", {"global": 0.5}) + model: Final = "anthropic/claude-sonnet-4-5" + shared_usage: Final = Usage(prompt_tokens=9, completion_tokens=16, total_tokens=25) + if upstream_cost is not None: + setattr(shared_usage, "cost", upstream_cost) + reference: Final = ModelResponse(model=model, usage=shared_usage.model_copy()) + billed_once: Final = litellm.completion_cost(reference, custom_llm_provider="anthropic") + complete: Final = ModelResponse(model=model, usage=shared_usage) + stamp_litellm_bill_on_usage(complete, shared_usage, billed_once) + text_chunk: Final = ModelResponseStream( + id="chunk-1", model=model, choices=[StreamingChoices(index=0, delta=Delta(content="hi"), finish_reason="stop")] + ) + usage_chunk: Final = ModelResponseStream(id="chunk-1", model=model, choices=[], usage=shared_usage) + + _stamp_usage_cost(usage_chunk, complete) + rebuilt: Final = litellm.stream_chunk_builder([text_chunk, usage_chunk], messages=[{"role": "user", "content": "hi"}]) + + assert rebuilt is not None and rebuilt.usage is not None # pyright: ignore[reportAttributeAccessIssue] # ModelResponse carries usage at runtime + assert usage_chunk.usage is not None and usage_chunk.usage.cost == pytest.approx(billed_once) + assert litellm.completion_cost(rebuilt, custom_llm_provider="anthropic") == pytest.approx(billed_once), ( + "pricing the rebuilt stream must not margin litellm's own bill again" + ) + + def test_chunk_creator_keeps_provider_model_private_across_stream(): from litellm.router_utils.add_retry_fallback_headers import ( get_hidden_params_dict, From a1c238ff4a1e1e8fc9cf2e9520c3132418e400f6 Mon Sep 17 00:00:00 2001 From: yujonglee Date: Fri, 30 Oct 2026 11:31:04 -0700 Subject: [PATCH 03/46] refactor(rust): extract SecretValue and use it for Redis credentials (#49186) * refactor(rust): move SecretValue into litellm-secret-value * docs(rust): add rust-redaction skill * refactor(redis-client): carry credentials as SecretValue --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../.agents/skills/rust-redaction/SKILL.md | 43 +++++++++ litellm-rust/AGENTS.md | 2 + litellm-rust/Cargo.lock | 50 ++++++++--- litellm-rust/Cargo.toml | 1 + litellm-rust/crates/auth-azure/Cargo.toml | 1 + litellm-rust/crates/auth-azure/src/native.rs | 6 +- litellm-rust/crates/auth-azure/src/resolve.rs | 9 +- litellm-rust/crates/auth-azure/src/types.rs | 4 +- litellm-rust/crates/auth-gcp/Cargo.toml | 1 + litellm-rust/crates/auth-gcp/src/auth.rs | 2 +- litellm-rust/crates/auth-gcp/src/config.rs | 4 +- litellm-rust/crates/auth-types/Cargo.toml | 2 +- .../crates/auth-types/src/credential.rs | 5 +- litellm-rust/crates/auth-types/src/lib.rs | 2 - litellm-rust/crates/auth-types/src/policy.rs | 3 +- litellm-rust/crates/auth-types/src/token.rs | 3 +- litellm-rust/crates/auth-types/tests/token.rs | 6 +- litellm-rust/crates/auth/Cargo.toml | 3 + litellm-rust/crates/auth/tests/facade.rs | 3 +- litellm-rust/crates/cache-gcs/Cargo.toml | 1 + litellm-rust/crates/cache-gcs/src/token.rs | 4 +- .../crates/callbacks-legacy-python/Cargo.toml | 2 +- .../callbacks-legacy-python/src/adapter.rs | 2 +- litellm-rust/crates/config/Cargo.toml | 2 +- litellm-rust/crates/config/src/mcp.rs | 2 +- litellm-rust/crates/config/src/redis.rs | 2 +- litellm-rust/crates/config/src/settings.rs | 2 +- litellm-rust/crates/gateway-auth/Cargo.toml | 2 +- .../crates/gateway-auth/src/authentication.rs | 2 +- litellm-rust/crates/gateway-auth/src/http.rs | 2 +- litellm-rust/crates/gateway-auth/src/keys.rs | 2 +- litellm-rust/crates/gateway-auth/src/ui.rs | 2 +- .../crates/gateway-auth/tests/auth.rs | 2 +- .../crates/gateway-auth/tests/keys.rs | 2 +- litellm-rust/crates/gateway-auth/tests/ui.rs | 2 +- .../crates/gateway-inference/Cargo.toml | 1 + .../crates/gateway-inference/src/ocr.rs | 2 +- .../gateway-inference/tests/support/mod.rs | 4 +- .../crates/gateway-management/Cargo.toml | 2 +- .../crates/gateway-management/src/keys.rs | 2 +- .../crates/gateway-management/tests/keys.rs | 2 +- litellm-rust/crates/gateway-mcp/Cargo.toml | 2 +- .../crates/gateway-mcp/src/configured.rs | 2 +- litellm-rust/crates/gateway-ui/Cargo.toml | 2 +- .../crates/gateway-ui/tests/routes.rs | 2 +- litellm-rust/crates/gateway/Cargo.toml | 2 +- litellm-rust/crates/gateway/src/redis.rs | 10 +-- litellm-rust/crates/gateway/src/secrets.rs | 2 +- litellm-rust/crates/gateway/src/settings.rs | 2 +- litellm-rust/crates/gateway/tests/mcp.rs | 2 +- litellm-rust/crates/gateway/tests/server.rs | 2 +- litellm-rust/crates/host/Cargo.toml | 2 +- litellm-rust/crates/host/src/interceptors.rs | 2 +- litellm-rust/crates/inference-chat/Cargo.toml | 1 + .../crates/inference-chat/src/handler.rs | 2 +- .../crates/inference-chat/tests/caching.rs | 4 +- .../crates/inference-messages/Cargo.toml | 1 + .../crates/inference-messages/src/handler.rs | 4 +- .../inference-messages/tests/caching.rs | 4 +- litellm-rust/crates/inference-ocr/Cargo.toml | 1 + .../crates/inference-ocr/src/handler.rs | 2 +- .../inference-ocr/src/provider_config.rs | 16 ++-- .../crates/inference-ocr/src/types.rs | 4 +- litellm-rust/crates/inference-ocr/src/wire.rs | 4 +- .../inference-ocr/tests/ocr/azure_ai.rs | 6 +- .../inference-ocr/tests/ocr/lifecycle.rs | 2 +- .../crates/inference-ocr/tests/ocr/main.rs | 2 +- .../crates/inference-responses/Cargo.toml | 1 + .../inference-responses/tests/caching.rs | 4 +- .../crates/inference-testing/Cargo.toml | 1 + .../crates/inference-testing/src/lib.rs | 3 +- litellm-rust/crates/llms/Cargo.toml | 1 + .../llms/src/anthropic/chat/transformation.rs | 2 +- .../crates/llms/src/anthropic/common_utils.rs | 4 +- .../src/azure_ai/messages/transformation.rs | 2 +- .../document_intelligence/transformation.rs | 2 +- .../llms/src/azure_ai/ocr/transformation.rs | 4 +- litellm-rust/crates/llms/src/base_llm/auth.rs | 4 +- .../llms/src/base_llm/ocr/transformation.rs | 4 +- .../bedrock/chat/converse_transformation.rs | 4 +- .../messages/transformation.rs | 4 +- .../anthropic_claude_transformation.rs | 4 +- .../bedrock_mantle/messages/transformation.rs | 4 +- .../llms/src/cohere/ocr/transformation.rs | 2 +- .../src/deepseek/messages/transformation.rs | 4 +- .../src/edenai/messages/transformation.rs | 4 +- .../src/minimax/messages/transformation.rs | 4 +- .../llms/src/mistral/ocr/transformation.rs | 2 +- .../src/openai/responses/transformation.rs | 3 +- .../src/openai_like/chat/transformation.rs | 4 +- .../src/openrouter/messages/transformation.rs | 4 +- .../openrouter/responses/transformation.rs | 4 +- .../llms/src/reducto/ocr/transformation.rs | 4 +- .../src/tencent/messages/transformation.rs | 4 +- .../llms/src/vertex_ai/ocr/transformation.rs | 2 +- litellm-rust/crates/python-bridge/Cargo.toml | 1 + .../crates/python-bridge/src/credentials.rs | 4 +- .../python-bridge/src/routes/ocr/project.rs | 2 +- .../python-bridge/src/secrets/callback.rs | 4 +- .../python-bridge/src/secrets/config.rs | 4 +- .../python-bridge/src/secrets/operations.rs | 8 +- .../python-bridge/src/secrets/provider.rs | 4 +- .../python-bridge/src/secrets/python.rs | 4 +- .../python-bridge/src/secrets/resolved.rs | 5 +- .../python-bridge/src/secrets/runtime.rs | 4 +- .../src/secrets/vault/operation.rs | 7 +- litellm-rust/crates/redis-client/Cargo.toml | 1 + .../crates/redis-client/src/client.rs | 11 ++- .../crates/redis-client/src/config.rs | 90 ++++++------------- .../crates/redis-client/src/connection.rs | 4 +- .../crates/redis-client/src/test_support.rs | 3 +- litellm-rust/crates/router-types/Cargo.toml | 1 + .../crates/router-types/src/litellm_params.rs | 4 +- .../router-types/tests/litellm_params.rs | 4 +- litellm-rust/crates/secret-value/Cargo.toml | 14 +++ litellm-rust/crates/secret-value/src/lib.rs | 5 ++ .../src/secret_value.rs} | 6 +- litellm-rust/crates/secrets-aws/Cargo.toml | 1 + .../crates/secrets-aws/src/secret_manager.rs | 5 +- .../secrets-aws/tests/secret_manager.rs | 4 +- litellm-rust/crates/secrets-azure/Cargo.toml | 1 + .../crates/secrets-azure/src/key_vault.rs | 4 +- .../crates/secrets-azure/tests/key_vault.rs | 4 +- .../crates/secrets-cyberark/Cargo.toml | 1 + .../secrets-cyberark/src/secret_manager.rs | 5 +- .../secrets-cyberark/tests/secret_manager.rs | 4 +- litellm-rust/crates/secrets-google/Cargo.toml | 1 + .../crates/secrets-google/src/auth.rs | 2 +- litellm-rust/crates/secrets-google/src/kms.rs | 2 +- .../secrets-google/src/secret_manager.rs | 6 +- .../crates/secrets-hashicorp/Cargo.toml | 1 + .../crates/secrets-hashicorp/src/config.rs | 2 +- .../secrets-hashicorp/src/secret_manager.rs | 5 +- .../secrets-hashicorp/tests/secret_manager.rs | 6 +- litellm-rust/crates/secrets-types/Cargo.toml | 2 +- .../secrets-types/src/base_secret_manager.rs | 3 +- .../crates/secrets-types/src/config.rs | 2 +- .../crates/secrets-types/src/context.rs | 3 +- litellm-rust/crates/secrets-types/src/lib.rs | 1 - .../crates/secrets-types/src/value.rs | 2 +- .../crates/secrets-types/tests/config.rs | 6 +- .../crates/secrets-types/tests/context.rs | 5 +- .../crates/secrets-types/tests/rotation.rs | 6 +- litellm-rust/crates/secrets/Cargo.toml | 1 + .../crates/secrets/src/compatibility.rs | 3 +- litellm-rust/crates/secrets/src/handler.rs | 2 +- litellm-rust/crates/secrets/src/lib.rs | 4 +- litellm-rust/crates/secrets/src/oidc.rs | 3 +- litellm-rust/crates/secrets/src/resolver.rs | 3 +- litellm-rust/crates/secrets/src/source.rs | 3 +- litellm-rust/crates/secrets/tests/aws.rs | 4 +- litellm-rust/crates/secrets/tests/azure.rs | 5 +- .../secrets/tests/common_read_contract.rs | 4 +- litellm-rust/crates/secrets/tests/cyberark.rs | 4 +- litellm-rust/crates/secrets/tests/google.rs | 4 +- litellm-rust/crates/secrets/tests/oidc.rs | 6 +- .../crates/secrets/tests/resolution.rs | 3 +- 157 files changed, 439 insertions(+), 254 deletions(-) create mode 100644 litellm-rust/.agents/skills/rust-redaction/SKILL.md create mode 100644 litellm-rust/crates/secret-value/Cargo.toml create mode 100644 litellm-rust/crates/secret-value/src/lib.rs rename litellm-rust/crates/{auth-types/src/secret.rs => secret-value/src/secret_value.rs} (95%) diff --git a/litellm-rust/.agents/skills/rust-redaction/SKILL.md b/litellm-rust/.agents/skills/rust-redaction/SKILL.md new file mode 100644 index 00000000000..2e906e10607 --- /dev/null +++ b/litellm-rust/.agents/skills/rust-redaction/SKILL.md @@ -0,0 +1,43 @@ +--- +name: rust-redaction +description: Keep secrets out of Debug output in litellm-rust with SecretValue and veil. Use when a struct or enum holds a password, token, key, credential URL or credential provider, when writing or reviewing a Debug impl, or when a type with secrets gets logged or traced +metadata: + internal: true +--- + +# Rust redaction + +Every secret in a Debug output is hidden the same way, so a reviewer can tell at a glance that a field is covered and a new field cannot slip through a hand-written impl + +## Where SecretValue lives + +`SecretValue` lives in `crates/secret-value` and is imported as `litellm_secret_value::SecretValue` everywhere. No crate re-exports it, so there is one import path. Depend on `litellm-secret-value` directly. Do not pull in `litellm-auth-types` or `litellm-auth` just to get it + +`crates/secret-value` holds only the type. Credential acquisition (`TokenProvider`, `CredentialPlan`, provider param tables) stays in `litellm-auth-types`, and generic clients such as `litellm-redis-client` must not depend on it + +## Rules + +- Never hand-write `impl Debug` to hide a value, and never invent markers like `""`, `""` or `.map(|_| "...")`. A hand-written impl silently prints any field added later +- A secret that crosses a crate boundary or lives in config is typed as `SecretValue`. It redacts itself, compares in constant time and only gives up the plaintext through `expose()`. A type whose only secrets are `SecretValue` fields derives plain `Debug` + ```rust + #[derive(Clone, Debug, PartialEq)] + pub struct HostEndpoint { + pub host: String, + pub port: u16, + pub username: Option, + pub password: Option, + pub tls: bool, + pub db: i64, + } + ``` +- Keep a secret as `SecretValue` until the last conversion into the third-party type (redis-rs, reqwest, an SDK) and call `expose()` there. Calling `expose()` early spreads plaintext `String`s through our own types +- A secret that has to stay a raw type, because a serde wire shape or a third-party struct owns the field, is covered by deriving `veil::Redact` on the holder and marking the field `#[redact(with = "[REDACTED]")]`. `Redact` replaces `Debug`, so do not derive both +- A URL that can carry userinfo is a secret as a whole. Do not parse it to show the host. Keep it as a `SecretValue` +- A bag of unparsed settings that may hold credentials (`additional_fields`, raw kwargs) is redacted as a whole with `#[redact(with = "[REDACTED]")]` +- An opaque handle over `Arc` whose trait has no `Debug` gets a newtype whose `Debug` prints only the type name with `finish_non_exhaustive()`. That impl shows nothing, so it is not a redaction impl and the first rule does not apply. A type holding such a handle can derive `Debug` +- Never put a secret in a tracing field or an error message. Log the type with `?value`, or log only non-secret fields. See the rust-tracing and rust-error skills + +## Tests + +- One test per type that holds a secret, built the way production builds it, asserting the plaintext is absent from `format!("{value:?}")` and `[REDACTED]` is present +- Use a distinctive plaintext (`"credential-value"`) so the absence check cannot pass by accident diff --git a/litellm-rust/AGENTS.md b/litellm-rust/AGENTS.md index 9e7a55eb232..5eaf876e3b2 100644 --- a/litellm-rust/AGENTS.md +++ b/litellm-rust/AGENTS.md @@ -1,5 +1,7 @@ # Rust workspace rules +For keeping secrets out of Debug output, follow [.agents/skills/rust-redaction/SKILL.md](.agents/skills/rust-redaction/SKILL.md) + For diagnostic tracing changes, follow [.agents/skills/rust-tracing/SKILL.md](.agents/skills/rust-tracing/SKILL.md) For string-valued enums and their Serde conversions, follow [.agents/skills/rust-string-enums/SKILL.md](.agents/skills/rust-string-enums/SKILL.md) diff --git a/litellm-rust/Cargo.lock b/litellm-rust/Cargo.lock index e959f77917b..f91deaaa401 100644 --- a/litellm-rust/Cargo.lock +++ b/litellm-rust/Cargo.lock @@ -3380,6 +3380,7 @@ dependencies = [ "litellm-auth-azure", "litellm-auth-gcp", "litellm-auth-types", + "litellm-secret-value", ] [[package]] @@ -3416,6 +3417,7 @@ dependencies = [ "azure_identity", "litellm-auth-types", "litellm-error", + "litellm-secret-value", "moka", "rstest", "serde_json", @@ -3435,6 +3437,7 @@ dependencies = [ "http 1.4.2", "litellm-auth-types", "litellm-error", + "litellm-secret-value", "moka", "rstest", "serde", @@ -3450,10 +3453,10 @@ name = "litellm-auth-types" version = "0.1.0" dependencies = [ "litellm-error", + "litellm-secret-value", "rstest", "serde", "serde_json", - "subtle", "thiserror 2.0.19", "tokio", "veil", @@ -3523,6 +3526,7 @@ dependencies = [ "litellm-cache-testing", "litellm-error", "litellm-http", + "litellm-secret-value", "percent-encoding", "reqwest", "rstest", @@ -3668,10 +3672,10 @@ dependencies = [ name = "litellm-callbacks-legacy-python" version = "0.1.0" dependencies = [ - "litellm-auth", "litellm-error", "litellm-host", "litellm-host-python", + "litellm-secret-value", "proptest", "pyo3", "rstest", @@ -3692,9 +3696,9 @@ dependencies = [ name = "litellm-config" version = "0.1.0" dependencies = [ - "litellm-auth-types", "litellm-error", "litellm-router-types", + "litellm-secret-value", "rstest", "serde", "serde_json", @@ -3804,7 +3808,6 @@ dependencies = [ "envy", "futures-util", "http-body-util", - "litellm-auth-types", "litellm-circuit-breaker", "litellm-config", "litellm-error", @@ -3817,6 +3820,7 @@ dependencies = [ "litellm-llms", "litellm-redis-client", "litellm-redis-testing", + "litellm-secret-value", "litellm-secrets", "litellm-tracing", "rstest", @@ -3839,9 +3843,9 @@ dependencies = [ "axum", "axum-login", "futures-util", - "litellm-auth-types", "litellm-config", "litellm-error", + "litellm-secret-value", "litellm-secrets", "rstest", "serde", @@ -3879,6 +3883,7 @@ dependencies = [ "litellm-llms", "litellm-llms-types", "litellm-router", + "litellm-secret-value", "litellm-secrets", "rstest", "serde", @@ -3895,9 +3900,9 @@ name = "litellm-gateway-management" version = "0.1.0" dependencies = [ "base64 0.22.1", - "litellm-auth-types", "litellm-error", "litellm-gateway-auth", + "litellm-secret-value", "rand 0.8.7", "rstest", "thiserror 2.0.19", @@ -3912,9 +3917,9 @@ dependencies = [ "base64 0.22.1", "futures-util", "http 1.4.2", - "litellm-auth-types", "litellm-config", "litellm-error", + "litellm-secret-value", "litellm-secrets", "moka", "rmcp", @@ -3942,9 +3947,9 @@ dependencies = [ "futures-util", "governor", "jsonwebtoken", - "litellm-auth-types", "litellm-error", "litellm-gateway-auth", + "litellm-secret-value", "rand 0.8.7", "rstest", "serde", @@ -3964,8 +3969,8 @@ name = "litellm-host" version = "0.1.0" dependencies = [ "futures-util", - "litellm-auth", "litellm-coroutine", + "litellm-secret-value", "rstest", "serde_json", "thiserror 2.0.19", @@ -4094,6 +4099,7 @@ dependencies = [ "litellm-inference-testing", "litellm-llms", "litellm-llms-types", + "litellm-secret-value", "litellm-secrets", "litellm-tracing", "rstest", @@ -4124,6 +4130,7 @@ dependencies = [ "litellm-llms", "litellm-llms-types", "litellm-router-types", + "litellm-secret-value", "litellm-secrets", "litellm-tracing", "reqwest", @@ -4154,6 +4161,7 @@ dependencies = [ "litellm-inference-testing", "litellm-llms", "litellm-llms-types", + "litellm-secret-value", "litellm-secrets", "litellm-tracing", "mime_guess", @@ -4188,6 +4196,7 @@ dependencies = [ "litellm-inference-testing", "litellm-llms", "litellm-llms-types", + "litellm-secret-value", "litellm-secrets", "litellm-tracing", "reqwest", @@ -4207,6 +4216,7 @@ dependencies = [ "litellm-auth", "litellm-http", "litellm-inference", + "litellm-secret-value", "litellm-secrets", "tokio", ] @@ -4254,6 +4264,7 @@ dependencies = [ "litellm-llms-types", "litellm-python-compat", "litellm-router-types", + "litellm-secret-value", "litellm-secrets", "reqwest", "rstest", @@ -4347,6 +4358,7 @@ dependencies = [ "litellm-llms-types", "litellm-model-catalog", "litellm-router-types", + "litellm-secret-value", "litellm-secrets", "litellm-secrets-aws", "litellm-secrets-types", @@ -4396,6 +4408,7 @@ dependencies = [ "futures-util", "litellm-circuit-breaker", "litellm-redis-testing", + "litellm-secret-value", "redis", "rstest", "thiserror 2.0.19", @@ -4432,12 +4445,23 @@ name = "litellm-router-types" version = "0.1.0" dependencies = [ "litellm-auth-types", + "litellm-secret-value", "rstest", "serde", "serde_json", "serde_with", ] +[[package]] +name = "litellm-secret-value" +version = "0.1.0" +dependencies = [ + "rstest", + "serde", + "subtle", + "veil", +] + [[package]] name = "litellm-secrets" version = "0.1.0" @@ -4451,6 +4475,7 @@ dependencies = [ "litellm-error", "litellm-http", "litellm-python-compat", + "litellm-secret-value", "litellm-secrets-aws", "litellm-secrets-azure", "litellm-secrets-cyberark", @@ -4479,6 +4504,7 @@ dependencies = [ "base64 0.22.1", "litellm-auth-aws", "litellm-core-utils", + "litellm-secret-value", "litellm-secrets-types", "litellm-tracing", "rstest", @@ -4498,6 +4524,7 @@ dependencies = [ "litellm-auth-types", "litellm-core-utils", "litellm-http", + "litellm-secret-value", "litellm-secrets-types", "percent-encoding", "reqwest", @@ -4519,6 +4546,7 @@ dependencies = [ "base64 0.22.1", "litellm-core-utils", "litellm-http", + "litellm-secret-value", "litellm-secrets-types", "litellm-tracing", "moka", @@ -4549,6 +4577,7 @@ dependencies = [ "litellm-auth-types", "litellm-core-utils", "litellm-http", + "litellm-secret-value", "litellm-secrets-types", "moka", "percent-encoding", @@ -4570,6 +4599,7 @@ dependencies = [ "litellm-core-utils", "litellm-error", "litellm-http", + "litellm-secret-value", "litellm-secrets-types", "reqwest", "rstest", @@ -4589,7 +4619,7 @@ dependencies = [ name = "litellm-secrets-types" version = "0.1.0" dependencies = [ - "litellm-auth-types", + "litellm-secret-value", "moka", "rstest", "serde", diff --git a/litellm-rust/Cargo.toml b/litellm-rust/Cargo.toml index 49a33e08456..75cac2a5d63 100644 --- a/litellm-rust/Cargo.toml +++ b/litellm-rust/Cargo.toml @@ -37,6 +37,7 @@ litellm-callbacks-legacy-python = { path = "crates/callbacks-legacy-python" } litellm-framer = { path = "crates/framer" } litellm-auth = { path = "crates/auth" } litellm-auth-types = { path = "crates/auth-types" } +litellm-secret-value = { path = "crates/secret-value" } litellm-auth-aws = { path = "crates/auth-aws" } litellm-auth-azure = { path = "crates/auth-azure" } litellm-auth-gcp = { path = "crates/auth-gcp" } diff --git a/litellm-rust/crates/auth-azure/Cargo.toml b/litellm-rust/crates/auth-azure/Cargo.toml index f99972f813b..7dcb0e513a7 100644 --- a/litellm-rust/crates/auth-azure/Cargo.toml +++ b/litellm-rust/crates/auth-azure/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-auth-types.workspace = true moka.workspace = true diff --git a/litellm-rust/crates/auth-azure/src/native.rs b/litellm-rust/crates/auth-azure/src/native.rs index 9fda67746d1..f40745ab7a2 100644 --- a/litellm-rust/crates/auth-azure/src/native.rs +++ b/litellm-rust/crates/auth-azure/src/native.rs @@ -13,7 +13,8 @@ use azure_identity::{ use sha2::{Digest, Sha256}; use litellm_auth_types::{Error, ErrorDetail}; -use litellm_auth_types::{InputSource, ResolvedCredential, SecretValue, Sourced}; +use litellm_auth_types::{InputSource, ResolvedCredential, Sourced}; +use litellm_secret_value::SecretValue; use super::credential_provider_cache::{ AzureCredentialProviderCache, AzureCredentialProviderCacheKey, @@ -498,7 +499,8 @@ mod tests { use azure_core::{Bytes, Result}; use super::{NativeAzureRequest, NativeAzureTokenAcquirer, ValidatedAzureRequest}; - use litellm_auth_types::{InputSource, SecretValue, Sourced}; + use litellm_auth_types::{InputSource, Sourced}; + use litellm_secret_value::SecretValue; fn deployment(value: T) -> Sourced { Sourced::new(value, InputSource::Deployment) diff --git a/litellm-rust/crates/auth-azure/src/resolve.rs b/litellm-rust/crates/auth-azure/src/resolve.rs index 25b71741342..a92630af26a 100644 --- a/litellm-rust/crates/auth-azure/src/resolve.rs +++ b/litellm-rust/crates/auth-azure/src/resolve.rs @@ -1,7 +1,9 @@ use litellm_auth_types::{ - CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential, - SecretValue, Sourced, TokenProviderHandle, + CredentialFileRef, CredentialLookup, CredentialRef, InputSource, ResolvedCredential, Sourced, + TokenProviderHandle, }; +use litellm_secret_value::SecretValue; + use litellm_auth_types::{Error, ErrorDetail}; use std::future::Future; use std::pin::Pin; @@ -500,8 +502,9 @@ mod tests { use litellm_auth_types::ResolvedCredential; use litellm_auth_types::{ CredentialFileRef, CredentialLookup, CredentialLookupFuture, CredentialRef, - CredentialResolver, CredentialResolverHandle, InputSource, SecretValue, Sourced, + CredentialResolver, CredentialResolverHandle, InputSource, Sourced, }; + use litellm_secret_value::SecretValue; #[derive(Debug)] struct FileResolver; diff --git a/litellm-rust/crates/auth-azure/src/types.rs b/litellm-rust/crates/auth-azure/src/types.rs index a042937a047..6f8290645cf 100644 --- a/litellm-rust/crates/auth-azure/src/types.rs +++ b/litellm-rust/crates/auth-azure/src/types.rs @@ -1,8 +1,10 @@ use std::collections::BTreeMap; use litellm_auth_types::{ - CredentialResolverHandle, Error, InputSource, SecretValue, Sourced, TokenProviderHandle, + CredentialResolverHandle, Error, InputSource, Sourced, TokenProviderHandle, }; +use litellm_secret_value::SecretValue; + use serde_json::{Map, Value}; use strum::EnumString; diff --git a/litellm-rust/crates/auth-gcp/Cargo.toml b/litellm-rust/crates/auth-gcp/Cargo.toml index 7844b35ecc2..ad60f0db817 100644 --- a/litellm-rust/crates/auth-gcp/Cargo.toml +++ b/litellm-rust/crates/auth-gcp/Cargo.toml @@ -9,6 +9,7 @@ repository.workspace = true google-sdk = ["dep:google-cloud-auth", "dep:http"] [dependencies] +litellm-secret-value.workspace = true litellm-auth-types.workspace = true base64.workspace = true diff --git a/litellm-rust/crates/auth-gcp/src/auth.rs b/litellm-rust/crates/auth-gcp/src/auth.rs index 8ac1e0f54ba..650d5831796 100644 --- a/litellm-rust/crates/auth-gcp/src/auth.rs +++ b/litellm-rust/crates/auth-gcp/src/auth.rs @@ -258,7 +258,7 @@ fn auth_acquisition_error(error: gcp_auth::Error) -> Error { #[cfg(test)] mod tests { - use litellm_auth_types::SecretValue; + use litellm_secret_value::SecretValue; use super::*; diff --git a/litellm-rust/crates/auth-gcp/src/config.rs b/litellm-rust/crates/auth-gcp/src/config.rs index 5444897aaaf..133e1c26903 100644 --- a/litellm-rust/crates/auth-gcp/src/config.rs +++ b/litellm-rust/crates/auth-gcp/src/config.rs @@ -4,7 +4,9 @@ use base64::{ Engine, alphabet, engine::{DecodePaddingMode, GeneralPurpose, GeneralPurposeConfig}, }; -use litellm_auth_types::{Error, ErrorDetail, InputSource, SecretValue, Sourced}; +use litellm_auth_types::{Error, ErrorDetail, InputSource, Sourced}; +use litellm_secret_value::SecretValue; + use serde_json::{Map, Value}; use litellm_auth_types::VertexParams; diff --git a/litellm-rust/crates/auth-types/Cargo.toml b/litellm-rust/crates/auth-types/Cargo.toml index 400f2640221..f607f0d1da0 100644 --- a/litellm-rust/crates/auth-types/Cargo.toml +++ b/litellm-rust/crates/auth-types/Cargo.toml @@ -7,9 +7,9 @@ repository.workspace = true [dependencies] litellm-error.workspace = true +litellm-secret-value.workspace = true serde.workspace = true serde_json.workspace = true -subtle.workspace = true thiserror.workspace = true veil.workspace = true diff --git a/litellm-rust/crates/auth-types/src/credential.rs b/litellm-rust/crates/auth-types/src/credential.rs index 4f69c6fffc8..36e21f7794a 100644 --- a/litellm-rust/crates/auth-types/src/credential.rs +++ b/litellm-rust/crates/auth-types/src/credential.rs @@ -5,7 +5,8 @@ use std::sync::Arc; use veil::Redact; -use crate::{Error, ResolvedCredential, SecretValue, TokenProviderHandle}; +use crate::{Error, ResolvedCredential, TokenProviderHandle}; +use litellm_secret_value::SecretValue; #[derive(Clone, Debug, PartialEq, Eq)] pub enum CredentialFileRef { @@ -104,8 +105,8 @@ mod tests { CredentialLookup, CredentialLookupFuture, CredentialPlan, CredentialPlanResolution, CredentialRef, CredentialResolver, CredentialResolverHandle, }; - use crate::SecretValue; use crate::{Error, ErrorDetail}; + use litellm_secret_value::SecretValue; #[derive(Debug)] struct HostResolver; diff --git a/litellm-rust/crates/auth-types/src/lib.rs b/litellm-rust/crates/auth-types/src/lib.rs index 99651977867..ab3bd3febbf 100644 --- a/litellm-rust/crates/auth-types/src/lib.rs +++ b/litellm-rust/crates/auth-types/src/lib.rs @@ -5,7 +5,6 @@ mod error; pub mod http; mod params; mod policy; -mod secret; mod token; use serde::{Deserialize, Serialize}; @@ -55,5 +54,4 @@ pub use error::{Error, ErrorDetail}; pub use http::CredentialPlacement; pub use params::{AwsParams, ParamKind, ParamSpec, VertexParams}; pub use policy::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy}; -pub use secret::SecretValue; pub use token::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; diff --git a/litellm-rust/crates/auth-types/src/policy.rs b/litellm-rust/crates/auth-types/src/policy.rs index 8fff59e284a..2232350965a 100644 --- a/litellm-rust/crates/auth-types/src/policy.rs +++ b/litellm-rust/crates/auth-types/src/policy.rs @@ -71,7 +71,8 @@ impl ProviderAuthPolicy { #[cfg(test)] mod tests { use super::{CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, ProviderAuthPolicy}; - use crate::{CredentialPlacement, ResolvedCredential, SecretValue}; + use crate::{CredentialPlacement, ResolvedCredential}; + use litellm_secret_value::SecretValue; const RULES: &[CredentialRule] = &[CredentialRule { kind: CredentialPlanKind::Static, diff --git a/litellm-rust/crates/auth-types/src/token.rs b/litellm-rust/crates/auth-types/src/token.rs index a3126052065..b948dda44fc 100644 --- a/litellm-rust/crates/auth-types/src/token.rs +++ b/litellm-rust/crates/auth-types/src/token.rs @@ -5,7 +5,8 @@ use std::time::SystemTime; use veil::Redact; -use crate::{Error, SecretValue}; +use crate::Error; +use litellm_secret_value::SecretValue; #[derive(Clone, Debug, PartialEq, Eq)] pub enum ResolvedCredential { diff --git a/litellm-rust/crates/auth-types/tests/token.rs b/litellm-rust/crates/auth-types/tests/token.rs index 7f7027942b0..39908ae7b8a 100644 --- a/litellm-rust/crates/auth-types/tests/token.rs +++ b/litellm-rust/crates/auth-types/tests/token.rs @@ -8,9 +8,9 @@ use std::{ time::{Duration, SystemTime}, }; -use litellm_auth_types::{ - Error, ErrorDetail, ResolvedCredential, SecretValue, TokenProviderHandle, -}; +use litellm_auth_types::{Error, ErrorDetail, ResolvedCredential, TokenProviderHandle}; +use litellm_secret_value::SecretValue; + use rstest::rstest; fn credential(index: usize, access_token: bool) -> ResolvedCredential { diff --git a/litellm-rust/crates/auth/Cargo.toml b/litellm-rust/crates/auth/Cargo.toml index ee4900ebcc2..674ae65a59f 100644 --- a/litellm-rust/crates/auth/Cargo.toml +++ b/litellm-rust/crates/auth/Cargo.toml @@ -16,3 +16,6 @@ litellm-auth-types.workspace = true litellm-auth-aws = { workspace = true, optional = true } litellm-auth-azure = { workspace = true, optional = true } litellm-auth-gcp = { workspace = true, optional = true } + +[dev-dependencies] +litellm-secret-value.workspace = true diff --git a/litellm-rust/crates/auth/tests/facade.rs b/litellm-rust/crates/auth/tests/facade.rs index f1092b15def..5d29bac1072 100644 --- a/litellm-rust/crates/auth/tests/facade.rs +++ b/litellm-rust/crates/auth/tests/facade.rs @@ -1,7 +1,8 @@ use litellm_auth::{ CredentialPlacement, CredentialPlanKind, CredentialRule, ExistingHeaderBehavior, - ProviderAuthPolicy, ResolvedCredential, SecretValue, + ProviderAuthPolicy, ResolvedCredential, }; +use litellm_secret_value::SecretValue; const RULES: &[CredentialRule] = &[CredentialRule { kind: CredentialPlanKind::Static, diff --git a/litellm-rust/crates/cache-gcs/Cargo.toml b/litellm-rust/crates/cache-gcs/Cargo.toml index dbd8c9ac908..878db4b8010 100644 --- a/litellm-rust/crates/cache-gcs/Cargo.toml +++ b/litellm-rust/crates/cache-gcs/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-http.workspace = true futures-util.workspace = true litellm-auth-gcp.workspace = true diff --git a/litellm-rust/crates/cache-gcs/src/token.rs b/litellm-rust/crates/cache-gcs/src/token.rs index 079ec44dbce..852d8bbf035 100644 --- a/litellm-rust/crates/cache-gcs/src/token.rs +++ b/litellm-rust/crates/cache-gcs/src/token.rs @@ -1,7 +1,9 @@ use std::{future::Future, pin::Pin}; use litellm_auth_gcp::{VertexAuth, VertexConfig}; -use litellm_auth_types::{InputSource, SecretValue, Sourced}; +use litellm_auth_types::{InputSource, Sourced}; +use litellm_secret_value::SecretValue; + use litellm_cache::Error; pub trait TokenSource: Send + Sync + 'static { diff --git a/litellm-rust/crates/callbacks-legacy-python/Cargo.toml b/litellm-rust/crates/callbacks-legacy-python/Cargo.toml index 1cd05710366..0133af16b80 100644 --- a/litellm-rust/crates/callbacks-legacy-python/Cargo.toml +++ b/litellm-rust/crates/callbacks-legacy-python/Cargo.toml @@ -16,7 +16,7 @@ serde_json.workspace = true thiserror.workspace = true [dev-dependencies] -litellm-auth.workspace = true +litellm-secret-value.workspace = true proptest.workspace = true rstest.workspace = true serde_json.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 493e91bb5ea..9431b9d1326 100644 --- a/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs +++ b/litellm-rust/crates/callbacks-legacy-python/src/adapter.rs @@ -925,11 +925,11 @@ assert all(value is failure for name, value in logger.calls if name.endswith('_h mod payload_tests { use std::ffi::CStr; - use litellm_auth::SecretValue; use litellm_host::hooks::CallHooks; use litellm_host::interceptors::{RawResponse, RequestContext, WireRequest}; use litellm_host::lifecycle::ExecutionEvent; use litellm_host_python::{HookStep, PythonCallEvent, PythonOwned, to_py}; + use litellm_secret_value::SecretValue; use proptest::prelude::*; use pyo3::gc::{PyTraverseError, PyVisit}; use pyo3::prelude::*; diff --git a/litellm-rust/crates/config/Cargo.toml b/litellm-rust/crates/config/Cargo.toml index 5b0cd22332f..278aea6200e 100644 --- a/litellm-rust/crates/config/Cargo.toml +++ b/litellm-rust/crates/config/Cargo.toml @@ -6,7 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] -litellm-auth-types.workspace = true +litellm-secret-value.workspace = true litellm-router-types.workspace = true serde.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/config/src/mcp.rs b/litellm-rust/crates/config/src/mcp.rs index cbe992f9956..00480fa5be1 100644 --- a/litellm-rust/crates/config/src/mcp.rs +++ b/litellm-rust/crates/config/src/mcp.rs @@ -1,6 +1,6 @@ use std::{collections::BTreeMap, fmt}; -use litellm_auth_types::SecretValue; +use litellm_secret_value::SecretValue; use serde::Deserialize; use crate::Object; diff --git a/litellm-rust/crates/config/src/redis.rs b/litellm-rust/crates/config/src/redis.rs index 5674f993025..7a91691c54b 100644 --- a/litellm-rust/crates/config/src/redis.rs +++ b/litellm-rust/crates/config/src/redis.rs @@ -1,7 +1,7 @@ use std::{collections::BTreeMap, fmt, time::Duration}; -use litellm_auth_types::SecretValue; use litellm_router_types::Spelled; +use litellm_secret_value::SecretValue; use serde::{Deserialize, Deserializer, de::Error as _}; use crate::{AdditionalFields, Error, Object, Value}; diff --git a/litellm-rust/crates/config/src/settings.rs b/litellm-rust/crates/config/src/settings.rs index 9571c4fa5de..1bfc19983f3 100644 --- a/litellm-rust/crates/config/src/settings.rs +++ b/litellm-rust/crates/config/src/settings.rs @@ -1,6 +1,6 @@ use std::fmt; -use litellm_auth_types::SecretValue; +use litellm_secret_value::SecretValue; use serde::Deserialize; use crate::{ diff --git a/litellm-rust/crates/gateway-auth/Cargo.toml b/litellm-rust/crates/gateway-auth/Cargo.toml index 27f20203843..ba9097e85d8 100644 --- a/litellm-rust/crates/gateway-auth/Cargo.toml +++ b/litellm-rust/crates/gateway-auth/Cargo.toml @@ -6,9 +6,9 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true axum = { workspace = true, features = ["matched-path"] } axum-login.workspace = true -litellm-auth-types.workspace = true litellm-config.workspace = true litellm-error.workspace = true litellm-secrets.workspace = true diff --git a/litellm-rust/crates/gateway-auth/src/authentication.rs b/litellm-rust/crates/gateway-auth/src/authentication.rs index 4ecfaa375b4..349913619f5 100644 --- a/litellm-rust/crates/gateway-auth/src/authentication.rs +++ b/litellm-rust/crates/gateway-auth/src/authentication.rs @@ -1,8 +1,8 @@ use std::{future::Future, pin::Pin, sync::Arc, time::SystemTime}; use axum::http::{Request, request::Parts}; -use litellm_auth_types::SecretValue; use litellm_config::Config; +use litellm_secret_value::SecretValue; use litellm_secrets::source::SecretSource; use sha2::{Digest, Sha256}; use subtle::ConstantTimeEq; diff --git a/litellm-rust/crates/gateway-auth/src/http.rs b/litellm-rust/crates/gateway-auth/src/http.rs index ddbf240509a..ad1718161c0 100644 --- a/litellm-rust/crates/gateway-auth/src/http.rs +++ b/litellm-rust/crates/gateway-auth/src/http.rs @@ -4,7 +4,7 @@ use axum::{ middleware::Next, response::Response, }; -use litellm_auth_types::SecretValue; +use litellm_secret_value::SecretValue; use crate::{AccessRequest, Auth, AuthenticatedRequest, Credential, Error}; diff --git a/litellm-rust/crates/gateway-auth/src/keys.rs b/litellm-rust/crates/gateway-auth/src/keys.rs index 1fe8fa60e0e..0deaf3f7a82 100644 --- a/litellm-rust/crates/gateway-auth/src/keys.rs +++ b/litellm-rust/crates/gateway-auth/src/keys.rs @@ -1,6 +1,6 @@ use std::{future::Future, pin::Pin, time::SystemTime}; -use litellm_auth_types::SecretValue; +use litellm_secret_value::SecretValue; use crate::{Error, hash_token}; diff --git a/litellm-rust/crates/gateway-auth/src/ui.rs b/litellm-rust/crates/gateway-auth/src/ui.rs index 3d655bde187..ee4dd670aa7 100644 --- a/litellm-rust/crates/gateway-auth/src/ui.rs +++ b/litellm-rust/crates/gateway-auth/src/ui.rs @@ -3,7 +3,7 @@ use std::{convert::Infallible, sync::Arc}; use axum::extract::FromRequestParts; use axum::http::{header::AUTHORIZATION, request::Parts}; use axum_login::{AuthUser, AuthnBackend, UserId}; -use litellm_auth_types::SecretValue; +use litellm_secret_value::SecretValue; use serde::Deserialize; use sha2::{Digest, Sha256}; use subtle::ConstantTimeEq; diff --git a/litellm-rust/crates/gateway-auth/tests/auth.rs b/litellm-rust/crates/gateway-auth/tests/auth.rs index 781ba385381..54180e30504 100644 --- a/litellm-rust/crates/gateway-auth/tests/auth.rs +++ b/litellm-rust/crates/gateway-auth/tests/auth.rs @@ -8,9 +8,9 @@ use axum::{ routing::get, }; use futures_util::future::BoxFuture; -use litellm_auth_types::SecretValue; use litellm_config::Config; use litellm_gateway_auth::{Auth, RequireMasterKey, hash_token}; +use litellm_secret_value::SecretValue; use litellm_secrets::source::SecretSource; use rstest::{fixture, rstest}; use tower::ServiceExt; diff --git a/litellm-rust/crates/gateway-auth/tests/keys.rs b/litellm-rust/crates/gateway-auth/tests/keys.rs index 4d16d3573bc..1f6fc461a5b 100644 --- a/litellm-rust/crates/gateway-auth/tests/keys.rs +++ b/litellm-rust/crates/gateway-auth/tests/keys.rs @@ -1,10 +1,10 @@ use std::time::{Duration, SystemTime}; -use litellm_auth_types::SecretValue; use litellm_gateway_auth::{ Error, keys::{KeyHash, KeyLookup, KeyStatus, LookupFuture, verify_key}, }; +use litellm_secret_value::SecretValue; use rstest::{fixture, rstest}; struct Lookup { diff --git a/litellm-rust/crates/gateway-auth/tests/ui.rs b/litellm-rust/crates/gateway-auth/tests/ui.rs index 2e5f72d605d..8d465e1198d 100644 --- a/litellm-rust/crates/gateway-auth/tests/ui.rs +++ b/litellm-rust/crates/gateway-auth/tests/ui.rs @@ -1,6 +1,6 @@ use axum_login::{AuthUser, AuthnBackend}; -use litellm_auth_types::SecretValue; use litellm_gateway_auth::{UiBackend, UiCredentials}; +use litellm_secret_value::SecretValue; use rstest::{fixture, rstest}; #[fixture] diff --git a/litellm-rust/crates/gateway-inference/Cargo.toml b/litellm-rust/crates/gateway-inference/Cargo.toml index 034c3a9826e..69ca561e4b5 100644 --- a/litellm-rust/crates/gateway-inference/Cargo.toml +++ b/litellm-rust/crates/gateway-inference/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-cache-response.workspace = true axum = { workspace = true, features = ["json", "multipart", "original-uri"] } base64.workspace = true diff --git a/litellm-rust/crates/gateway-inference/src/ocr.rs b/litellm-rust/crates/gateway-inference/src/ocr.rs index f200fd7b141..b8e6ef4ff4f 100644 --- a/litellm-rust/crates/gateway-inference/src/ocr.rs +++ b/litellm-rust/crates/gateway-inference/src/ocr.rs @@ -2,10 +2,10 @@ use litellm_gateway_auth::AuthenticatedRequest; use std::sync::Arc; use axum::{Json, extract::State, http::HeaderMap, response::IntoResponse}; -use litellm_auth::SecretValue; use litellm_inference_ocr::types::{LiteLLMOcrRequest, OcrConnectionInputs, OcrDocumentInput}; use litellm_llms::base_llm::ocr::transformation::decode_request_value; use litellm_llms_types::formats::ocr::OcrDocument; +use litellm_secret_value::SecretValue; use serde_json::Value; use crate::{ diff --git a/litellm-rust/crates/gateway-inference/tests/support/mod.rs b/litellm-rust/crates/gateway-inference/tests/support/mod.rs index 737dc0c830c..185cbc8594d 100644 --- a/litellm-rust/crates/gateway-inference/tests/support/mod.rs +++ b/litellm-rust/crates/gateway-inference/tests/support/mod.rs @@ -13,7 +13,9 @@ use futures_util::future::BoxFuture; use litellm_gateway_inference::{Deployment, Gateway, router}; use litellm_http::{HttpClientPool, HttpSettings, Resolution, media::PublicDnsResolver}; use litellm_inference::resources::CoreResources; -use litellm_secrets::{SecretValue, source::SecretSource}; +use litellm_secret_value::SecretValue; +use litellm_secrets::source::SecretSource; + use serde_json::Value; use tokio::io::{AsyncReadExt, AsyncWriteExt}; use tower::ServiceExt; diff --git a/litellm-rust/crates/gateway-management/Cargo.toml b/litellm-rust/crates/gateway-management/Cargo.toml index 502760f8fbe..1ff5dd5319d 100644 --- a/litellm-rust/crates/gateway-management/Cargo.toml +++ b/litellm-rust/crates/gateway-management/Cargo.toml @@ -6,8 +6,8 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true base64.workspace = true -litellm-auth-types.workspace = true litellm-error.workspace = true litellm-gateway-auth.workspace = true rand.workspace = true diff --git a/litellm-rust/crates/gateway-management/src/keys.rs b/litellm-rust/crates/gateway-management/src/keys.rs index 45d915b7b3d..feeb40b2403 100644 --- a/litellm-rust/crates/gateway-management/src/keys.rs +++ b/litellm-rust/crates/gateway-management/src/keys.rs @@ -1,8 +1,8 @@ use std::{future::Future, pin::Pin, sync::Arc, time::SystemTime}; use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; -use litellm_auth_types::SecretValue; use litellm_gateway_auth::keys::{KeyHash, KeyStatus}; +use litellm_secret_value::SecretValue; use rand::{RngCore, rngs::OsRng}; use crate::Error; diff --git a/litellm-rust/crates/gateway-management/tests/keys.rs b/litellm-rust/crates/gateway-management/tests/keys.rs index bf0221cf1be..8a970bb407c 100644 --- a/litellm-rust/crates/gateway-management/tests/keys.rs +++ b/litellm-rust/crates/gateway-management/tests/keys.rs @@ -7,7 +7,6 @@ use std::{ }; use base64::{Engine, engine::general_purpose::URL_SAFE_NO_PAD}; -use litellm_auth_types::SecretValue; use litellm_gateway_auth::{ hash_token, keys::{KeyHash, KeyLookup, KeyStatus, verify_key}, @@ -16,6 +15,7 @@ use litellm_gateway_management::{ Error, keys::{Keys, Revocation}, }; +use litellm_secret_value::SecretValue; use rstest::{fixture, rstest}; use support::{MemoryStore, Unavailable}; diff --git a/litellm-rust/crates/gateway-mcp/Cargo.toml b/litellm-rust/crates/gateway-mcp/Cargo.toml index bed48bad58d..65077eec0ea 100644 --- a/litellm-rust/crates/gateway-mcp/Cargo.toml +++ b/litellm-rust/crates/gateway-mcp/Cargo.toml @@ -7,11 +7,11 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true axum = { workspace = true, features = ["json", "query"] } futures-util.workspace = true litellm-config.workspace = true litellm-secrets.workspace = true -litellm-auth-types.workspace = true litellm-error.workspace = true base64.workspace = true http.workspace = true diff --git a/litellm-rust/crates/gateway-mcp/src/configured.rs b/litellm-rust/crates/gateway-mcp/src/configured.rs index cc7a52fc029..33cbd243530 100644 --- a/litellm-rust/crates/gateway-mcp/src/configured.rs +++ b/litellm-rust/crates/gateway-mcp/src/configured.rs @@ -8,8 +8,8 @@ use std::{ use base64::{Engine, engine::general_purpose::STANDARD}; use futures_util::future::try_join_all; use http::{HeaderName, HeaderValue}; -use litellm_auth_types::SecretValue; use litellm_config::{McpAuth, McpServer, McpTransport}; +use litellm_secret_value::SecretValue; use litellm_secrets::source::SecretSource; use rmcp::{ RoleClient, ServiceExt, diff --git a/litellm-rust/crates/gateway-ui/Cargo.toml b/litellm-rust/crates/gateway-ui/Cargo.toml index 1c65665be72..a6104455e96 100644 --- a/litellm-rust/crates/gateway-ui/Cargo.toml +++ b/litellm-rust/crates/gateway-ui/Cargo.toml @@ -23,10 +23,10 @@ tower-http = { version = "0.6.11", features = ["fs", "set-header"] } tower-sessions.workspace = true [dev-dependencies] +litellm-secret-value.workspace = true futures-util.workspace = true serde_json.workspace = true tower = { version = "0.5", features = ["util"] } -litellm-auth-types.workspace = true rstest.workspace = true tempfile.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/gateway-ui/tests/routes.rs b/litellm-rust/crates/gateway-ui/tests/routes.rs index 95246f33a35..86b7db63ed9 100644 --- a/litellm-rust/crates/gateway-ui/tests/routes.rs +++ b/litellm-rust/crates/gateway-ui/tests/routes.rs @@ -16,8 +16,8 @@ use tower_sessions::{ session::{Id, Record}, }; -use litellm_auth_types::SecretValue; use litellm_gateway_auth::UiBackend; +use litellm_secret_value::SecretValue; struct App { router: Router, diff --git a/litellm-rust/crates/gateway/Cargo.toml b/litellm-rust/crates/gateway/Cargo.toml index 4d831aa2314..6eeaa43e013 100644 --- a/litellm-rust/crates/gateway/Cargo.toml +++ b/litellm-rust/crates/gateway/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true axum.workspace = true envy = "0.4.2" http-body-util = "0.1" @@ -17,7 +18,6 @@ thiserror.workspace = true futures-util.workspace = true litellm-gateway-auth.workspace = true litellm-gateway-ui.workspace = true -litellm-auth-types.workspace = true litellm-config.workspace = true litellm-error.workspace = true litellm-circuit-breaker.workspace = true diff --git a/litellm-rust/crates/gateway/src/redis.rs b/litellm-rust/crates/gateway/src/redis.rs index 197f40e7988..c40b289029b 100644 --- a/litellm-rust/crates/gateway/src/redis.rs +++ b/litellm-rust/crates/gateway/src/redis.rs @@ -152,12 +152,12 @@ async fn reachable(client: RedisClient) -> Option { fn client(params: Params, breaker: Option) -> Result { tracing::info!(endpoint = ?params.endpoint, namespace = ?params.namespace, "redis client configured"); let endpoint = match params.endpoint { - RedisEndpoint::Url(url) => Endpoint::Url(url.expose().to_owned()), + RedisEndpoint::Url(url) => Endpoint::Url(url), RedisEndpoint::Host(host) => Endpoint::Host(HostEndpoint { host: host.host, port: host.port, - username: host.username.map(|value| value.expose().to_owned()), - password: host.password.map(|value| value.expose().to_owned()), + username: host.username, + password: host.password, tls: host.tls, db: host.db, }), @@ -170,8 +170,8 @@ fn client(params: Params, breaker: Option) -> Result futures_util::future::BoxFuture< 'a, - Result, litellm_secrets::Error>, + Result, litellm_secrets::Error>, > { Box::pin(async { Ok(None) }) } diff --git a/litellm-rust/crates/gateway/tests/server.rs b/litellm-rust/crates/gateway/tests/server.rs index c876c1cbcbb..7de7a97c870 100644 --- a/litellm-rust/crates/gateway/tests/server.rs +++ b/litellm-rust/crates/gateway/tests/server.rs @@ -159,7 +159,7 @@ async fn mounts_ui_without_exposing_credentials_or_authorizing_inference(inferen std::fs::write(assets.path().join("index.html"), "dashboard").unwrap(); let backend = litellm_gateway_auth::UiBackend::new( "admin".into(), - litellm_auth_types::SecretValue::new("ui-password"), + litellm_secret_value::SecretValue::new("ui-password"), ) .unwrap(); let ui = litellm_gateway_ui::router( diff --git a/litellm-rust/crates/host/Cargo.toml b/litellm-rust/crates/host/Cargo.toml index d3354fa8622..e0ccab4e0ad 100644 --- a/litellm-rust/crates/host/Cargo.toml +++ b/litellm-rust/crates/host/Cargo.toml @@ -6,8 +6,8 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true futures-util.workspace = true -litellm-auth.workspace = true litellm-coroutine.workspace = true serde_json.workspace = true thiserror.workspace = true diff --git a/litellm-rust/crates/host/src/interceptors.rs b/litellm-rust/crates/host/src/interceptors.rs index 3cff781792f..39af7c410e0 100644 --- a/litellm-rust/crates/host/src/interceptors.rs +++ b/litellm-rust/crates/host/src/interceptors.rs @@ -22,7 +22,7 @@ pub struct RequestContext { /// Optional-param names that carry credentials and must be redacted when logged. pub secret_fields: Vec, /// The credential the route resolved for the provider call. - pub api_key: Option, + pub api_key: Option, } #[derive(Clone, Debug, PartialEq, Eq)] diff --git a/litellm-rust/crates/inference-chat/Cargo.toml b/litellm-rust/crates/inference-chat/Cargo.toml index 8a779d44976..bcde3c8d30f 100644 --- a/litellm-rust/crates/inference-chat/Cargo.toml +++ b/litellm-rust/crates/inference-chat/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] } litellm-cache-response.workspace = true litellm-core-utils.workspace = true diff --git a/litellm-rust/crates/inference-chat/src/handler.rs b/litellm-rust/crates/inference-chat/src/handler.rs index 1c724a47956..f948e8fffaf 100644 --- a/litellm-rust/crates/inference-chat/src/handler.rs +++ b/litellm-rust/crates/inference-chat/src/handler.rs @@ -58,7 +58,7 @@ pub(super) async fn execute( custom_llm_provider: provider_name.to_owned(), optional_params: Value::Object(call.optional_params), secret_fields: Vec::new(), - api_key: call.api_key.map(litellm_auth::SecretValue::new), + api_key: call.api_key.map(litellm_secret_value::SecretValue::new), }; let authenticated = resolve_auth(auth, environment, &|key| secrets.get(key)).await?; let identity = litellm_host::interceptors::ProviderIdentity { diff --git a/litellm-rust/crates/inference-chat/tests/caching.rs b/litellm-rust/crates/inference-chat/tests/caching.rs index 4b3e32f04bc..dd229d3d657 100644 --- a/litellm-rust/crates/inference-chat/tests/caching.rs +++ b/litellm-rust/crates/inference-chat/tests/caching.rs @@ -140,7 +140,7 @@ impl litellm_secrets::source::SecretSource for ChangingSecrets { name: &'a str, ) -> futures_util::future::BoxFuture< 'a, - Result, litellm_secrets::Error>, + Result, litellm_secrets::Error>, > { Box::pin(async move { let revision = self.revision.load(Ordering::SeqCst); @@ -154,7 +154,7 @@ impl litellm_secrets::source::SecretSource for ChangingSecrets { } else { None }; - Ok(value.map(litellm_secrets::SecretValue::new)) + Ok(value.map(litellm_secret_value::SecretValue::new)) }) } } diff --git a/litellm-rust/crates/inference-messages/Cargo.toml b/litellm-rust/crates/inference-messages/Cargo.toml index 694d9b4de31..ac52cb65a0d 100644 --- a/litellm-rust/crates/inference-messages/Cargo.toml +++ b/litellm-rust/crates/inference-messages/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true bytes.workspace = true futures-util.workspace = true litellm-auth = { workspace = true, features = ["aws", "azure", "gcp"] } diff --git a/litellm-rust/crates/inference-messages/src/handler.rs b/litellm-rust/crates/inference-messages/src/handler.rs index 9d0bf09197e..2e647688a1a 100644 --- a/litellm-rust/crates/inference-messages/src/handler.rs +++ b/litellm-rust/crates/inference-messages/src/handler.rs @@ -2,7 +2,9 @@ use std::time::Duration; use bytes::Bytes; use futures_util::{FutureExt, StreamExt, TryStreamExt, stream::BoxStream}; -use litellm_auth::{AuthServices, SecretValue}; +use litellm_auth::AuthServices; +use litellm_secret_value::SecretValue; + use litellm_core_utils::{ dot_notation_indexing::delete_nested_value, get_provider_specific_headers::get_provider_specific_headers, settings::Lookup, diff --git a/litellm-rust/crates/inference-messages/tests/caching.rs b/litellm-rust/crates/inference-messages/tests/caching.rs index 480b3c4f3b7..663c87453a9 100644 --- a/litellm-rust/crates/inference-messages/tests/caching.rs +++ b/litellm-rust/crates/inference-messages/tests/caching.rs @@ -107,7 +107,7 @@ impl litellm_secrets::source::SecretSource for ChangingSecrets { name: &'a str, ) -> futures_util::future::BoxFuture< 'a, - Result, litellm_secrets::Error>, + Result, litellm_secrets::Error>, > { Box::pin(async move { let revision = self.revision.load(Ordering::SeqCst); @@ -121,7 +121,7 @@ impl litellm_secrets::source::SecretSource for ChangingSecrets { } else { None }; - Ok(value.map(litellm_secrets::SecretValue::new)) + Ok(value.map(litellm_secret_value::SecretValue::new)) }) } } diff --git a/litellm-rust/crates/inference-ocr/Cargo.toml b/litellm-rust/crates/inference-ocr/Cargo.toml index bbbdaddb050..ad53126b9fe 100644 --- a/litellm-rust/crates/inference-ocr/Cargo.toml +++ b/litellm-rust/crates/inference-ocr/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true base64.workspace = true bytes.workspace = true futures-util.workspace = true diff --git a/litellm-rust/crates/inference-ocr/src/handler.rs b/litellm-rust/crates/inference-ocr/src/handler.rs index 426d4297f28..c126a192153 100644 --- a/litellm-rust/crates/inference-ocr/src/handler.rs +++ b/litellm-rust/crates/inference-ocr/src/handler.rs @@ -128,7 +128,7 @@ mod tests { decode_request(OcrWireRequest { model: model.into(), document, - api_key: Some(litellm_auth::SecretValue::new("test-key")), + api_key: Some(litellm_secret_value::SecretValue::new("test-key")), api_base: Some(base.into()), custom_llm_provider: None, extra_headers: None, diff --git a/litellm-rust/crates/inference-ocr/src/provider_config.rs b/litellm-rust/crates/inference-ocr/src/provider_config.rs index 19e50af0dd1..e76d291254b 100644 --- a/litellm-rust/crates/inference-ocr/src/provider_config.rs +++ b/litellm-rust/crates/inference-ocr/src/provider_config.rs @@ -1,4 +1,6 @@ -use litellm_auth::{InputSource, SecretValue, Sourced}; +use litellm_auth::{InputSource, Sourced}; +use litellm_secret_value::SecretValue; + use litellm_core_utils::get_llm_provider_logic::{ CustomLlmProvider, LlmProviders, get_custom_llm_provider, }; @@ -407,7 +409,7 @@ mod tests { fn connection_resolution_preserves_dynamic_precedence_and_input_sources() { let connection = OcrConfigKind::Mistral.resolve_connection_params(OcrCredentialInputs { api_key: Some(Sourced::new( - litellm_auth::SecretValue::new("explicit-key"), + litellm_secret_value::SecretValue::new("explicit-key"), InputSource::Deployment, )), api_base: Some(Sourced::new( @@ -415,7 +417,7 @@ mod tests { InputSource::Deployment, )), dynamic_api_key: Some(Sourced::new( - litellm_auth::SecretValue::new("dynamic-key"), + litellm_secret_value::SecretValue::new("dynamic-key"), InputSource::Environment, )), dynamic_api_base: Some(Sourced::new( @@ -455,7 +457,7 @@ mod tests { ) { let dynamic_key = dynamic_value.map(|value| { Sourced::new( - litellm_auth::SecretValue::new(value), + litellm_secret_value::SecretValue::new(value), InputSource::Environment, ) }); @@ -463,7 +465,7 @@ mod tests { dynamic_value.map(|value| Sourced::new(value.into(), InputSource::Environment)); let connection = OcrConfigKind::Mistral.resolve_connection_params(OcrCredentialInputs { api_key: Some(Sourced::new( - litellm_auth::SecretValue::new("explicit-key"), + litellm_secret_value::SecretValue::new("explicit-key"), InputSource::Deployment, )), api_base: Some(Sourced::new( @@ -502,14 +504,14 @@ mod tests { OcrCredentialInputs { api_key: explicit_key.map(|value| { Sourced::new( - litellm_auth::SecretValue::new(value), + litellm_secret_value::SecretValue::new(value), InputSource::Deployment, ) }), api_base: explicit_base .map(|value| Sourced::new(value.into(), InputSource::Deployment)), dynamic_api_key: Some(Sourced::new( - litellm_auth::SecretValue::new("dynamic-key"), + litellm_secret_value::SecretValue::new("dynamic-key"), InputSource::Environment, )), dynamic_api_base: Some(Sourced::new( diff --git a/litellm-rust/crates/inference-ocr/src/types.rs b/litellm-rust/crates/inference-ocr/src/types.rs index 98c609a158f..1b2dbb52d77 100644 --- a/litellm-rust/crates/inference-ocr/src/types.rs +++ b/litellm-rust/crates/inference-ocr/src/types.rs @@ -1,7 +1,9 @@ use std::{collections::BTreeMap, path::PathBuf, time::Duration}; use bytes::Bytes; -use litellm_auth::{InputSource, SecretValue, TokenProviderHandle}; +use litellm_auth::{InputSource, TokenProviderHandle}; +use litellm_secret_value::SecretValue; + use litellm_core_utils::call_arguments::CallArguments; use litellm_llms::base_llm::ocr::{ error::Error, diff --git a/litellm-rust/crates/inference-ocr/src/wire.rs b/litellm-rust/crates/inference-ocr/src/wire.rs index 56d22be6523..45a2865823f 100644 --- a/litellm-rust/crates/inference-ocr/src/wire.rs +++ b/litellm-rust/crates/inference-ocr/src/wire.rs @@ -1,6 +1,8 @@ use std::{collections::BTreeMap, time::Duration}; -use litellm_auth::{InputSource, SecretValue}; +use litellm_auth::InputSource; +use litellm_secret_value::SecretValue; + use litellm_llms::base_llm::ocr::{error::Error, transformation::decode_request_value}; use litellm_llms_types::formats::ocr::OcrDocument; use serde::Deserialize; diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/azure_ai.rs b/litellm-rust/crates/inference-ocr/tests/ocr/azure_ai.rs index 6f591af73e0..65c7f309b82 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/azure_ai.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/azure_ai.rs @@ -3,9 +3,9 @@ use std::sync::{ atomic::{AtomicUsize, Ordering}, }; -use litellm_auth::{ - ResolvedCredential, SecretValue, TokenFuture, TokenProvider, TokenProviderHandle, -}; +use litellm_auth::{ResolvedCredential, TokenFuture, TokenProvider, TokenProviderHandle}; +use litellm_secret_value::SecretValue; + use rstest::rstest; use super::*; diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs b/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs index 75a91092daf..9f59bb9a2a2 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/lifecycle.rs @@ -239,7 +239,7 @@ impl litellm_host_native::services::HostCallHandler for CallerTokenHost { OcrOp::AcquireAzureAdToken(reply) => { self.trace.lock().unwrap().push("token".into()); reply.send(litellm_auth::ResolvedCredential::Static( - litellm_auth::SecretValue::new("caller-token"), + litellm_secret_value::SecretValue::new("caller-token"), )); Ok(()) } diff --git a/litellm-rust/crates/inference-ocr/tests/ocr/main.rs b/litellm-rust/crates/inference-ocr/tests/ocr/main.rs index 5a48bba9e1a..6871d1ca1f6 100644 --- a/litellm-rust/crates/inference-ocr/tests/ocr/main.rs +++ b/litellm-rust/crates/inference-ocr/tests/ocr/main.rs @@ -80,7 +80,7 @@ fn wire(model: &str, base: &str, document: Value, options: Value) -> OcrWireRequ OcrWireRequest { model: model.into(), document, - api_key: Some(litellm_auth::SecretValue::new("test-key")), + api_key: Some(litellm_secret_value::SecretValue::new("test-key")), api_base: Some(base.into()), custom_llm_provider: None, extra_headers: None, diff --git a/litellm-rust/crates/inference-responses/Cargo.toml b/litellm-rust/crates/inference-responses/Cargo.toml index 083179c50c8..f90651634c5 100644 --- a/litellm-rust/crates/inference-responses/Cargo.toml +++ b/litellm-rust/crates/inference-responses/Cargo.toml @@ -25,6 +25,7 @@ tokio-tungstenite.workspace = true tracing.workspace = true [dev-dependencies] +litellm-secret-value.workspace = true litellm-http = { workspace = true, features = ["test-support"] } litellm-cache.workspace = true litellm-cache-memory.workspace = true diff --git a/litellm-rust/crates/inference-responses/tests/caching.rs b/litellm-rust/crates/inference-responses/tests/caching.rs index 462461695c3..56e25fbcb98 100644 --- a/litellm-rust/crates/inference-responses/tests/caching.rs +++ b/litellm-rust/crates/inference-responses/tests/caching.rs @@ -98,7 +98,7 @@ impl litellm_secrets::source::SecretSource for ChangingSecrets { name: &'a str, ) -> futures_util::future::BoxFuture< 'a, - Result, litellm_secrets::Error>, + Result, litellm_secrets::Error>, > { Box::pin(async move { let revision = self.revision.load(Ordering::SeqCst); @@ -112,7 +112,7 @@ impl litellm_secrets::source::SecretSource for ChangingSecrets { } else { None }; - Ok(value.map(litellm_secrets::SecretValue::new)) + Ok(value.map(litellm_secret_value::SecretValue::new)) }) } } diff --git a/litellm-rust/crates/inference-testing/Cargo.toml b/litellm-rust/crates/inference-testing/Cargo.toml index 748b1df0d3b..5f690532d9c 100644 --- a/litellm-rust/crates/inference-testing/Cargo.toml +++ b/litellm-rust/crates/inference-testing/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true publish = false [dependencies] +litellm-secret-value.workspace = true litellm-auth.workspace = true futures-util.workspace = true litellm-http.workspace = true diff --git a/litellm-rust/crates/inference-testing/src/lib.rs b/litellm-rust/crates/inference-testing/src/lib.rs index 2e20d5a9e91..f330e02e99a 100644 --- a/litellm-rust/crates/inference-testing/src/lib.rs +++ b/litellm-rust/crates/inference-testing/src/lib.rs @@ -5,7 +5,8 @@ use litellm_http::{ ClientVariant, HttpClientConfig, HttpClientPool, HttpSettings, Resolution, media::PublicDnsResolver, }; -use litellm_secrets::{SecretValue, source::SecretSource}; +use litellm_secret_value::SecretValue; +use litellm_secrets::source::SecretSource; use litellm_inference::resources::CoreResources; diff --git a/litellm-rust/crates/llms/Cargo.toml b/litellm-rust/crates/llms/Cargo.toml index 962f61ef4ec..79ee17e8b97 100644 --- a/litellm-rust/crates/llms/Cargo.toml +++ b/litellm-rust/crates/llms/Cargo.toml @@ -9,6 +9,7 @@ repository.workspace = true test-support = ["litellm-http/test-support"] [dependencies] +litellm-secret-value.workspace = true litellm-llms-types.workspace = true litellm-core-utils.workspace = true litellm-error.workspace = true diff --git a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs index eb796a900f5..164dbce1264 100644 --- a/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/anthropic/chat/transformation.rs @@ -1,4 +1,3 @@ -use litellm_auth::SecretValue; use litellm_core_utils::decline::Decline; use litellm_core_utils::{ core_helpers::{finish_reason_for, unix_now, usage_from_parts}, @@ -10,6 +9,7 @@ use litellm_llms_types::{ }, providers::anthropic::DEFAULT_HEADERS, }; +use litellm_secret_value::SecretValue; use serde::Deserialize; use serde_json::{Map, Value, json}; diff --git a/litellm-rust/crates/llms/src/anthropic/common_utils.rs b/litellm-rust/crates/llms/src/anthropic/common_utils.rs index 7991a57353e..aec5d9ebc27 100644 --- a/litellm-rust/crates/llms/src/anthropic/common_utils.rs +++ b/litellm-rust/crates/llms/src/anthropic/common_utils.rs @@ -1,5 +1,7 @@ use crate::base_llm::messages::context::MessagesModelCapabilities; -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_http::request::{ has_header, header_value, header_values, with_header, without_headers, diff --git a/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs index d578acea68c..a2ba2cee0d3 100644 --- a/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/messages/transformation.rs @@ -1,10 +1,10 @@ -use litellm_auth::SecretValue; use litellm_http::request::{has_bearer_auth, has_header}; use litellm_llms_types::{ formats::messages::{MessagesOptionalParams, MessagesRequest}, providers::anthropic::{BetaProvider, DEFAULT_HEADERS, MESSAGES_PATH}, }; use litellm_router_types::LitellmParams; +use litellm_secret_value::SecretValue; use crate::{ Error, diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs index 36dfe2c890e..870306b1a83 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/document_intelligence/transformation.rs @@ -1250,7 +1250,7 @@ mod tests { #[tokio::test] async fn request_endpoint_accepts_request_owned_key() { let connection = OcrConnection { - api_key: Some(litellm_auth::SecretValue::new("request-key")), + api_key: Some(litellm_secret_value::SecretValue::new("request-key")), api_key_source: InputSource::Request, api_base: Some("https://request.example".into()), api_base_source: InputSource::Request, diff --git a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs index 2f7957fe56e..25f93846967 100644 --- a/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/azure_ai/ocr/transformation.rs @@ -211,7 +211,7 @@ mod tests { #[fixture] fn connection() -> OcrConnection { OcrConnection { - api_key: Some(litellm_auth::SecretValue::new("request-key")), + api_key: Some(litellm_secret_value::SecretValue::new("request-key")), api_base: Some("https://example.com".into()), ..Default::default() } @@ -312,7 +312,7 @@ mod tests { #[tokio::test] async fn request_endpoint_accepts_request_owned_key() { let connection = OcrConnection { - api_key: Some(litellm_auth::SecretValue::new("request-key")), + api_key: Some(litellm_secret_value::SecretValue::new("request-key")), api_key_source: InputSource::Request, api_base: Some("https://request.example".into()), api_base_source: InputSource::Request, diff --git a/litellm-rust/crates/llms/src/base_llm/auth.rs b/litellm-rust/crates/llms/src/base_llm/auth.rs index de271a269d3..ab2cf00614e 100644 --- a/litellm-rust/crates/llms/src/base_llm/auth.rs +++ b/litellm-rust/crates/llms/src/base_llm/auth.rs @@ -5,7 +5,9 @@ //! that split: `validate_environment` shapes the forwarded headers and names the credential //! as an [`AuthScheme`], and [`resolve_auth`] turns the scheme into headers and a signer. -use litellm_auth::{AuthServices, CredentialPlacement, SecretValue, TokenProviderHandle}; +use litellm_auth::{AuthServices, CredentialPlacement, TokenProviderHandle}; +use litellm_secret_value::SecretValue; + use litellm_auth_aws::{AwsCredentialSource, SigV4Signer}; use litellm_auth_gcp::VertexConfig; use litellm_http::request::with_header; diff --git a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs index c90fd99809e..9574e99d357 100644 --- a/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/base_llm/ocr/transformation.rs @@ -5,7 +5,9 @@ use serde::{ use serde_json::{Map, Value}; use std::{collections::BTreeMap, future::Future, sync::Arc, time::Duration}; -use litellm_auth::{InputSource, SecretValue, Sourced, TokenProviderHandle}; +use litellm_auth::{InputSource, Sourced, TokenProviderHandle}; +use litellm_secret_value::SecretValue; + use litellm_core_utils::{call_arguments::CallArguments, settings::ProcessEnvironment}; use litellm_http::outbound::RequestSigner; use litellm_llms_types::formats::ocr::{LiteLLMOcrResponse, OcrDocument, OcrResponseFormat}; 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 976c9161dd2..6cf0167a30f 100644 --- a/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/chat/converse_transformation.rs @@ -1,5 +1,7 @@ use litellm_auth::AwsParams; -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_auth_aws::{ AwsCredentialSource, bedrock_model_id_and_region, constants::{AWS_BEARER_TOKEN_BEDROCK, BEDROCK_RUNTIME_ENDPOINT_TEMPLATE, BEDROCK_SERVICE}, diff --git a/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs index 79c30df9b19..5b89345a1c5 100644 --- a/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/claude_platform/messages/transformation.rs @@ -1,6 +1,8 @@ use std::sync::LazyLock; -use litellm_auth::{AwsParams, CredentialPlacement, ParamKind, ParamSpec, SecretValue}; +use litellm_auth::{AwsParams, CredentialPlacement, ParamKind, ParamSpec}; +use litellm_secret_value::SecretValue; + use litellm_auth_aws::{AwsCredentialSource, resolve_aws_region}; use litellm_core_utils::settings::resolve_non_empty; use litellm_http::request::{header_value, with_header, without_headers}; diff --git a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs index 3ec60f12021..1d380c917f6 100644 --- a/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock/messages/invoke_transformations/anthropic_claude_transformation.rs @@ -7,7 +7,9 @@ use crate::{ }; use futures_util::StreamExt; use litellm_auth::AwsParams; -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_auth_aws::{ AwsCredentialSource, bedrock_model_id_and_region, constants::{ diff --git a/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs b/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs index ae13704f4a0..f139b93e259 100644 --- a/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/bedrock_mantle/messages/transformation.rs @@ -1,6 +1,8 @@ use std::sync::LazyLock; -use litellm_auth::{AwsParams, CredentialPlacement, SecretValue}; +use litellm_auth::{AwsParams, CredentialPlacement}; +use litellm_secret_value::SecretValue; + use litellm_auth_aws::AwsCredentialSource; use litellm_auth_aws::bedrock_model_id_and_region; use litellm_core_utils::settings::resolve_non_empty; diff --git a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs index 76ff27b56fc..5bad06adda2 100644 --- a/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/cohere/ocr/transformation.rs @@ -735,7 +735,7 @@ mod tests { assert!(matches!( CohereParseConfig.resolve_headers( &OcrConnection { - api_key: Some(litellm_auth::SecretValue::new(" ")), + api_key: Some(litellm_secret_value::SecretValue::new(" ")), ..Default::default() }, &|_| None, diff --git a/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs b/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs index 168c3e63920..0dbcea01439 100644 --- a/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/deepseek/messages/transformation.rs @@ -1,4 +1,6 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_llms_types::{ formats::messages::{MessagesOptionalParams, MessagesRequest, MessagesTool}, diff --git a/litellm-rust/crates/llms/src/edenai/messages/transformation.rs b/litellm-rust/crates/llms/src/edenai/messages/transformation.rs index 2c3a7eaef1d..d5944484522 100644 --- a/litellm-rust/crates/llms/src/edenai/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/edenai/messages/transformation.rs @@ -1,4 +1,6 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_llms_types::formats::messages::{ MessagesOptionalParams, MessagesRequest, MessagesResponse, diff --git a/litellm-rust/crates/llms/src/minimax/messages/transformation.rs b/litellm-rust/crates/llms/src/minimax/messages/transformation.rs index 012a657fceb..dafe5fbd16b 100644 --- a/litellm-rust/crates/llms/src/minimax/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/minimax/messages/transformation.rs @@ -1,4 +1,6 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_llms_types::{ formats::messages::{MessagesOptionalParams, MessagesRequest}, diff --git a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs index fa4e2254e86..679aa65f7fd 100644 --- a/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/mistral/ocr/transformation.rs @@ -219,7 +219,7 @@ mod tests { #[default(vec![])] extra_headers: Vec<(String, String)>, ) -> OcrConnection { OcrConnection { - api_key: api_key.map(litellm_auth::SecretValue::new), + api_key: api_key.map(litellm_secret_value::SecretValue::new), extra_headers, ..OcrConnection::default() } diff --git a/litellm-rust/crates/llms/src/openai/responses/transformation.rs b/litellm-rust/crates/llms/src/openai/responses/transformation.rs index cf74871894b..deb163daa07 100644 --- a/litellm-rust/crates/llms/src/openai/responses/transformation.rs +++ b/litellm-rust/crates/llms/src/openai/responses/transformation.rs @@ -3,7 +3,8 @@ use litellm_llms_types::formats::responses::{ }; use serde_json::{Map, Value}; -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; use crate::{ Error, ErrorDetail, diff --git a/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs b/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs index e998f8785cb..83fb8ac3fa2 100644 --- a/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs +++ b/litellm-rust/crates/llms/src/openai_like/chat/transformation.rs @@ -4,7 +4,9 @@ //! the `max_completion_tokens` -> `max_tokens` rename and the usage //! `*_tokens` null-to-zero sanitize. -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::core_helpers::unix_now; use litellm_llms_types::formats::chat_completions::{ ChatCompletionsChoice, ChatCompletionsChoiceMessage, ChatCompletionsResponse, diff --git a/litellm-rust/crates/llms/src/openrouter/messages/transformation.rs b/litellm-rust/crates/llms/src/openrouter/messages/transformation.rs index 5b4e4f77f2a..eeb3d1cb7f8 100644 --- a/litellm-rust/crates/llms/src/openrouter/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/openrouter/messages/transformation.rs @@ -1,4 +1,6 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_http::request::{has_bearer_auth, without_headers}; use litellm_llms_types::{ diff --git a/litellm-rust/crates/llms/src/openrouter/responses/transformation.rs b/litellm-rust/crates/llms/src/openrouter/responses/transformation.rs index 900fba9e916..76155ba55e5 100644 --- a/litellm-rust/crates/llms/src/openrouter/responses/transformation.rs +++ b/litellm-rust/crates/llms/src/openrouter/responses/transformation.rs @@ -1,4 +1,6 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_llms_types::formats::responses::ResponsesApiResponse; use serde_json::{Map, Value}; diff --git a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs index bf0ddc1c7a8..beda6899d04 100644 --- a/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/reducto/ocr/transformation.rs @@ -690,7 +690,7 @@ mod tests { #[test] fn explicit_key_precedes_environment_key() { let connection = OcrConnection { - api_key: Some(litellm_auth::SecretValue::new("passed-key")), + api_key: Some(litellm_secret_value::SecretValue::new("passed-key")), ..Default::default() }; let headers = resolve_headers(&connection, &|_| Some("env-key".into())).unwrap(); @@ -700,7 +700,7 @@ mod tests { #[test] fn blank_explicit_key_uses_environment_key() { let connection = OcrConnection { - api_key: Some(litellm_auth::SecretValue::new(" ")), + api_key: Some(litellm_secret_value::SecretValue::new(" ")), ..Default::default() }; let headers = resolve_headers(&connection, &|_| Some(" env-key ".into())).unwrap(); diff --git a/litellm-rust/crates/llms/src/tencent/messages/transformation.rs b/litellm-rust/crates/llms/src/tencent/messages/transformation.rs index a70336748ad..621a09f4fb0 100644 --- a/litellm-rust/crates/llms/src/tencent/messages/transformation.rs +++ b/litellm-rust/crates/llms/src/tencent/messages/transformation.rs @@ -1,4 +1,6 @@ -use litellm_auth::{CredentialPlacement, SecretValue}; +use litellm_auth::CredentialPlacement; +use litellm_secret_value::SecretValue; + use litellm_core_utils::settings::resolve_non_empty; use litellm_llms_types::{ formats::messages::{MessagesOptionalParams, MessagesRequest}, diff --git a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs index 77881cd2e78..b902967d9dc 100644 --- a/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs +++ b/litellm-rust/crates/llms/src/vertex_ai/ocr/transformation.rs @@ -137,7 +137,7 @@ impl VertexAiOcrConfig { connection .api_key .as_ref() - .map(litellm_auth::SecretValue::expose), + .map(litellm_secret_value::SecretValue::expose), config, &|name: &str| connection.secret(name), ) diff --git a/litellm-rust/crates/python-bridge/Cargo.toml b/litellm-rust/crates/python-bridge/Cargo.toml index a4e1c57ab1e..e3ef45f7876 100644 --- a/litellm-rust/crates/python-bridge/Cargo.toml +++ b/litellm-rust/crates/python-bridge/Cargo.toml @@ -19,6 +19,7 @@ huggingface = ["litellm-token-counter/huggingface"] tiktoken = ["litellm-token-counter/tiktoken"] [dependencies] +litellm-secret-value.workspace = true fancy-regex.workspace = true litellm-tracing.workspace = true litellm-spend-clickhouse.workspace = true diff --git a/litellm-rust/crates/python-bridge/src/credentials.rs b/litellm-rust/crates/python-bridge/src/credentials.rs index a4bb064aae0..adc52a7e77a 100644 --- a/litellm-rust/crates/python-bridge/src/credentials.rs +++ b/litellm-rust/crates/python-bridge/src/credentials.rs @@ -2,7 +2,9 @@ //! keyword arguments and acquired on the host's own thread when the call asks for one. use crate::callable::wrap_failure; -use litellm_auth::{ResolvedCredential, SecretValue}; +use litellm_auth::ResolvedCredential; +use litellm_secret_value::SecretValue; + use pyo3::{ exceptions::PyTypeError, gc::{PyTraverseError, PyVisit}, diff --git a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs index c2c3fe51600..73a2922804a 100644 --- a/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs +++ b/litellm-rust/crates/python-bridge/src/routes/ocr/project.rs @@ -1,4 +1,3 @@ -use litellm_auth::SecretValue; use litellm_auth::VertexParams; use litellm_host_python::{from_py, present}; use litellm_inference_ocr::{ @@ -6,6 +5,7 @@ use litellm_inference_ocr::{ wire::{OcrWireRequest, consumed_optional_params, decode_document, decode_request_input}, }; use litellm_llms::base_llm::ocr::{error::Error, transformation::DROP_PARAMS_PARAM}; +use litellm_secret_value::SecretValue; use pyo3::{prelude::*, types::PyDict}; use serde_json::{Map, Value}; diff --git a/litellm-rust/crates/python-bridge/src/secrets/callback.rs b/litellm-rust/crates/python-bridge/src/secrets/callback.rs index 669c10e602a..af6485a56e1 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/callback.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/callback.rs @@ -2,9 +2,11 @@ use std::{future::Future, pin::Pin, sync::Arc}; use litellm_core_utils::settings::Lookup; use litellm_host_python::{PythonContext, attach_blocking}; +use litellm_secret_value::SecretValue; use litellm_secrets::{ - Error, ExternalSecretManager, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue, + Error, ExternalSecretManager, KeyManagementSettings, KeyManagementSystem, Secret, }; + use pyo3::{ exceptions::PyException, prelude::*, diff --git a/litellm-rust/crates/python-bridge/src/secrets/config.rs b/litellm-rust/crates/python-bridge/src/secrets/config.rs index 7afd06e2373..c1203dcbd89 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/config.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/config.rs @@ -1,7 +1,9 @@ use std::sync::Arc; +use litellm_secret_value::SecretValue; use litellm_secrets::{SecretManager, SecretManagerState}; -use litellm_secrets_types::{AccessMode, KeyManagementSettings, KeyManagementSystem, SecretValue}; +use litellm_secrets_types::{AccessMode, KeyManagementSettings, KeyManagementSystem}; + use pyo3::prelude::*; use serde_json::Value; diff --git a/litellm-rust/crates/python-bridge/src/secrets/operations.rs b/litellm-rust/crates/python-bridge/src/secrets/operations.rs index b458e3f34e8..027da1e4718 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/operations.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/operations.rs @@ -81,7 +81,7 @@ impl From for PythonMutationError { pub(super) async fn write_python_provider( manager: &SecretManager, name: &str, - value: &litellm_secrets::SecretValue, + value: &litellm_secret_value::SecretValue, ) -> Result { match manager { SecretManager::Cyberark(client) => { @@ -128,7 +128,7 @@ pub(super) async fn rotate_python_provider( manager: &SecretManager, current_name: &str, new_name: &str, - value: &litellm_secrets::SecretValue, + value: &litellm_secret_value::SecretValue, ) -> Result { match manager { SecretManager::Cyberark(client) => { @@ -186,7 +186,7 @@ pub(super) enum PythonMutationResponse { pub(super) async fn write_python_provider_with_context( manager: &SecretManager, name: &str, - value: &litellm_secrets::SecretValue, + value: &litellm_secret_value::SecretValue, context: &litellm_secrets_types::SecretWriteContext, ) -> Result { if let (SecretManager::HashicorpVault(client), SecretOperationContext::Hashicorp(operation)) = @@ -235,7 +235,7 @@ pub(super) async fn rotate_python_provider_with_context( manager: &SecretManager, current_name: &str, new_name: &str, - value: &litellm_secrets::SecretValue, + value: &litellm_secret_value::SecretValue, context: &SecretOperationContext, ) -> Result { if let (SecretManager::HashicorpVault(client), SecretOperationContext::Hashicorp(context)) = diff --git a/litellm-rust/crates/python-bridge/src/secrets/provider.rs b/litellm-rust/crates/python-bridge/src/secrets/provider.rs index ff10f596311..55343fb63bb 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/provider.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/provider.rs @@ -1,7 +1,9 @@ use std::time::Duration; use super::operations::PythonReadRequest; -use litellm_secrets::{KeyManagementSystem, SecretValue}; +use litellm_secret_value::SecretValue; +use litellm_secrets::KeyManagementSystem; + use litellm_secrets_types::{ AwsOperationContext, CyberarkOperationContext, GoogleOperationContext, HashicorpOperationContext, SecretOperationContext, diff --git a/litellm-rust/crates/python-bridge/src/secrets/python.rs b/litellm-rust/crates/python-bridge/src/secrets/python.rs index 2f040bff773..c5d9e3565b3 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/python.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/python.rs @@ -2,7 +2,9 @@ use std::sync::Arc; use futures_util::future::BoxFuture; use litellm_host_python::{PythonContext, attach_blocking}; -use litellm_secrets::{Error, SecretValue, source::SecretSource}; +use litellm_secret_value::SecretValue; +use litellm_secrets::{Error, source::SecretSource}; + use pyo3::prelude::*; use super::error::external_error; diff --git a/litellm-rust/crates/python-bridge/src/secrets/resolved.rs b/litellm-rust/crates/python-bridge/src/secrets/resolved.rs index deb07175e7a..752b533886e 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/resolved.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/resolved.rs @@ -4,10 +4,9 @@ use futures_util::future::BoxFuture; use litellm_core_utils::settings::ProcessEnvironment; use litellm_host_python::PythonContext; use litellm_http::Client; +use litellm_secret_value::SecretValue; use litellm_secrets::source::SecretSource; -use litellm_secrets::{ - Error, FailurePolicy, OidcResolver, SecretManagerState, SecretResolver, SecretValue, -}; +use litellm_secrets::{Error, FailurePolicy, OidcResolver, SecretManagerState, SecretResolver}; use super::config::SecretManagerSnapshot; diff --git a/litellm-rust/crates/python-bridge/src/secrets/runtime.rs b/litellm-rust/crates/python-bridge/src/secrets/runtime.rs index b956ec297b6..0d52b0aa629 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/runtime.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/runtime.rs @@ -307,7 +307,7 @@ impl NativeSecretManager { super::operations::write_python_provider_with_context( &backend, &secret_name, - &litellm_secrets::SecretValue::new(secret_value), + &litellm_secret_value::SecretValue::new(secret_value), &context, ) .await, @@ -365,7 +365,7 @@ impl NativeSecretManager { &backend, ¤t_secret_name, &new_secret_name, - &litellm_secrets::SecretValue::new(new_secret_value), + &litellm_secret_value::SecretValue::new(new_secret_value), &context, ) .await, diff --git a/litellm-rust/crates/python-bridge/src/secrets/vault/operation.rs b/litellm-rust/crates/python-bridge/src/secrets/vault/operation.rs index 41658f0baf8..b99dff7aeac 100644 --- a/litellm-rust/crates/python-bridge/src/secrets/vault/operation.rs +++ b/litellm-rust/crates/python-bridge/src/secrets/vault/operation.rs @@ -1,9 +1,8 @@ use std::collections::HashMap; -use litellm_secrets::{ - SecretValue, - hashicorp::{HashicorpVault, RawOperationError}, -}; +use litellm_secret_value::SecretValue; +use litellm_secrets::hashicorp::{HashicorpVault, RawOperationError}; + use litellm_secrets_types::{HashicorpOperationContext, SecretWriteContext}; use serde_json::value::RawValue; diff --git a/litellm-rust/crates/redis-client/Cargo.toml b/litellm-rust/crates/redis-client/Cargo.toml index eab57b36423..3e9447a6bfa 100644 --- a/litellm-rust/crates/redis-client/Cargo.toml +++ b/litellm-rust/crates/redis-client/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true [dependencies] litellm-circuit-breaker.workspace = true +litellm-secret-value.workspace = true futures-util.workspace = true redis = { workspace = true, features = ["cluster-async", "connection-manager", "tokio-rustls-comp"] } thiserror.workspace = true diff --git a/litellm-rust/crates/redis-client/src/client.rs b/litellm-rust/crates/redis-client/src/client.rs index a02b0847663..48dbc2b2dfb 100644 --- a/litellm-rust/crates/redis-client/src/client.rs +++ b/litellm-rust/crates/redis-client/src/client.rs @@ -29,7 +29,7 @@ impl RedisClient { pub fn new(config: RedisConfig) -> Result { let connection = match &config.endpoint { Endpoint::Url(url) => { - Connection::standalone(url.as_str().into_connection_info()?, config.timeout)? + Connection::standalone(url.expose().into_connection_info()?, config.timeout)? } Endpoint::Host(host) => { Connection::standalone(host.connection_info()?, config.timeout)? @@ -204,6 +204,7 @@ mod tests { use std::time::{Duration, Instant}; use litellm_redis_testing::RedisServer; + use litellm_secret_value::SecretValue; use rstest::rstest; use litellm_circuit_breaker::{Breaker, Config, State}; @@ -630,8 +631,12 @@ mod tests { db: 3, })) .unwrap(); - let in_db_0 = connect(Endpoint::Url(server.url().to_owned())).unwrap(); - let via_url = connect(Endpoint::Url(format!("{}/3", server.url()))).unwrap(); + let in_db_0 = connect(Endpoint::Url(SecretValue::new(server.url().to_owned()))).unwrap(); + let via_url = connect(Endpoint::Url(SecretValue::new(format!( + "{}/3", + server.url() + )))) + .unwrap(); in_db_3 .set("k", "3", Duration::from_secs(60)) diff --git a/litellm-rust/crates/redis-client/src/config.rs b/litellm-rust/crates/redis-client/src/config.rs index 64934c68891..c40b9eb6b36 100644 --- a/litellm-rust/crates/redis-client/src/config.rs +++ b/litellm-rust/crates/redis-client/src/config.rs @@ -1,8 +1,9 @@ -use std::{fmt, time::Duration}; +use std::time::Duration; use redis::{ConnectionAddr, ConnectionInfo, IntoConnectionInfo, RedisConnectionInfo, RedisResult}; use litellm_circuit_breaker::Breaker; +use litellm_secret_value::SecretValue; use crate::Namespace; @@ -15,28 +16,28 @@ pub struct RedisConfig { pub breaker: Option, } -#[derive(Clone, PartialEq)] +#[derive(Clone, Debug, PartialEq)] pub enum Endpoint { - Url(String), + Url(SecretValue), Host(HostEndpoint), Cluster(ClusterEndpoint), } -#[derive(Clone, PartialEq)] +#[derive(Clone, Debug, PartialEq)] pub struct HostEndpoint { pub host: String, pub port: u16, - pub username: Option, - pub password: Option, + pub username: Option, + pub password: Option, pub tls: bool, pub db: i64, } -#[derive(Clone, PartialEq)] +#[derive(Clone, Debug, PartialEq)] pub struct ClusterEndpoint { pub nodes: Vec, - pub username: Option, - pub password: Option, + pub username: Option, + pub password: Option, pub tls: bool, } @@ -58,11 +59,11 @@ impl HostEndpoint { pub(crate) fn connection_info(&self) -> RedisResult { let settings = RedisConnectionInfo::default().set_db(self.db); let settings = match &self.username { - Some(username) => settings.set_username(username), + Some(username) => settings.set_username(username.expose()), None => settings, }; let settings = match &self.password { - Some(password) => settings.set_password(password), + Some(password) => settings.set_password(password.expose()), None => settings, }; Ok(addr(&self.host, self.port, self.tls) @@ -83,70 +84,35 @@ fn addr(host: &str, port: u16, tls: bool) -> ConnectionAddr { } } -impl fmt::Debug for Endpoint { - fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result { - match self { - Self::Url(url) => formatter - .debug_tuple("Url") - .field(&redacted_url(url)) - .finish(), - Self::Host(host) => formatter - .debug_struct("Host") - .field("host", &host.host) - .field("port", &host.port) - .field("username", &host.username) - .field("password", &host.password.as_ref().map(|_| "")) - .field("tls", &host.tls) - .field("db", &host.db) - .finish(), - Self::Cluster(cluster) => formatter - .debug_struct("Cluster") - .field("nodes", &cluster.nodes) - .field("username", &cluster.username) - .field("password", &cluster.password.as_ref().map(|_| "")) - .field("tls", &cluster.tls) - .finish(), - } - } -} - -fn redacted_url(url: &str) -> String { - match url - .split_once("://") - .and_then(|(scheme, rest)| Some((scheme, rest.rsplit_once('@')?.1))) - { - Some((scheme, host)) => format!("{scheme}://@{host}"), - None => url.to_owned(), - } -} - #[cfg(test)] mod tests { + use litellm_secret_value::SecretValue; use rstest::rstest; - use super::{Endpoint, HostEndpoint}; + use super::{ClusterEndpoint, Endpoint, HostEndpoint}; #[rstest] - #[case::url_with_password(Endpoint::Url("redis://user:secret@cache:6379/0".to_owned()), "redis://@cache:6379/0")] - #[case::url_with_at_in_password(Endpoint::Url("rediss://:p@ss@cache:6380".to_owned()), "rediss://@cache:6380")] - #[case::host_password( + #[case::url(Endpoint::Url(SecretValue::new("redis://credential-value@cache:6379/0")))] + #[case::host( Endpoint::Host(HostEndpoint { host: "cache".to_owned(), port: 6379, - username: None, - password: Some("secret".to_owned()), + username: Some(SecretValue::new("credential-value")), + password: Some(SecretValue::new("credential-value")), tls: false, db: 0, - }), - "" + }) )] - fn debug_output_never_shows_the_password(#[case] endpoint: Endpoint, #[case] shown: &str) { + #[case::cluster(Endpoint::Cluster(ClusterEndpoint { + nodes: Vec::new(), + username: Some(SecretValue::new("credential-value")), + password: Some(SecretValue::new("credential-value")), + tls: false, + }))] + fn debug_output_never_shows_the_secret(#[case] endpoint: Endpoint) { let debug = format!("{endpoint:?}"); - assert!( - !debug.contains("secret") && !debug.contains("p@ss"), - "{debug}" - ); - assert!(debug.contains(shown), "{debug}"); + assert!(!debug.contains("credential-value"), "{debug}"); + assert!(debug.contains("[REDACTED]"), "{debug}"); } } diff --git a/litellm-rust/crates/redis-client/src/connection.rs b/litellm-rust/crates/redis-client/src/connection.rs index aaf26f105da..41729bfd107 100644 --- a/litellm-rust/crates/redis-client/src/connection.rs +++ b/litellm-rust/crates/redis-client/src/connection.rs @@ -33,11 +33,11 @@ impl Connection { .response_timeout(timeout) .retries(RETRIES); let builder = match &endpoint.username { - Some(username) => builder.username(username), + Some(username) => builder.username(username.expose()), None => builder, }; let builder = match &endpoint.password { - Some(password) => builder.password(password), + Some(password) => builder.password(password.expose()), None => builder, }; let config = ClusterConfig::new() diff --git a/litellm-rust/crates/redis-client/src/test_support.rs b/litellm-rust/crates/redis-client/src/test_support.rs index 64a174b1543..e503739a942 100644 --- a/litellm-rust/crates/redis-client/src/test_support.rs +++ b/litellm-rust/crates/redis-client/src/test_support.rs @@ -2,6 +2,7 @@ use std::time::Duration; use litellm_circuit_breaker::Breaker; use litellm_redis_testing::{RedisCluster, RedisServer}; +use litellm_secret_value::SecretValue; use redis::cluster_routing::{Route, RoutingInfo, SingleNodeRoutingInfo, SlotAddr}; use crate::{ @@ -30,7 +31,7 @@ impl Server { pub(crate) fn endpoint(&self) -> Endpoint { match self { - Self::Standalone(server) => Endpoint::Url(server.url().to_owned()), + Self::Standalone(server) => Endpoint::Url(SecretValue::new(server.url().to_owned())), Self::Cluster(cluster) => Endpoint::Cluster(ClusterEndpoint { nodes: cluster .nodes() diff --git a/litellm-rust/crates/router-types/Cargo.toml b/litellm-rust/crates/router-types/Cargo.toml index 2e9fec28685..f0470ce8a8e 100644 --- a/litellm-rust/crates/router-types/Cargo.toml +++ b/litellm-rust/crates/router-types/Cargo.toml @@ -7,6 +7,7 @@ repository.workspace = true description = "The types of litellm/types/router.py: a deployment's litellm_params as a config or a caller spells them" [dependencies] +litellm-secret-value.workspace = true litellm-auth-types.workspace = true serde.workspace = true serde_json.workspace = true diff --git a/litellm-rust/crates/router-types/src/litellm_params.rs b/litellm-rust/crates/router-types/src/litellm_params.rs index 186a1d98036..aa2c47f6e3c 100644 --- a/litellm-rust/crates/router-types/src/litellm_params.rs +++ b/litellm-rust/crates/router-types/src/litellm_params.rs @@ -1,6 +1,8 @@ use std::{collections::BTreeMap, fmt, ops::Deref}; -use litellm_auth_types::{AwsParams, ParamSpec, SecretValue, VertexParams}; +use litellm_auth_types::{AwsParams, ParamSpec, VertexParams}; +use litellm_secret_value::SecretValue; + use serde::Deserialize; use serde_json::Value; use serde_with::{OneOrMany, formats::PreferMany, serde_as}; diff --git a/litellm-rust/crates/router-types/tests/litellm_params.rs b/litellm-rust/crates/router-types/tests/litellm_params.rs index 7d021fd857a..c5c0d16ee8f 100644 --- a/litellm-rust/crates/router-types/tests/litellm_params.rs +++ b/litellm-rust/crates/router-types/tests/litellm_params.rs @@ -1,4 +1,6 @@ -use litellm_auth_types::{AwsParams, SecretValue, VertexParams}; +use litellm_auth_types::{AwsParams, VertexParams}; +use litellm_secret_value::SecretValue; + use litellm_router_types::{LitellmParams, Spelled}; use rstest::rstest; use serde_json::{Map, Value, json}; diff --git a/litellm-rust/crates/secret-value/Cargo.toml b/litellm-rust/crates/secret-value/Cargo.toml new file mode 100644 index 00000000000..467f852a278 --- /dev/null +++ b/litellm-rust/crates/secret-value/Cargo.toml @@ -0,0 +1,14 @@ +[package] +name = "litellm-secret-value" +version = "0.1.0" +edition.workspace = true +license.workspace = true +repository.workspace = true + +[dependencies] +serde.workspace = true +subtle.workspace = true +veil.workspace = true + +[dev-dependencies] +rstest.workspace = true diff --git a/litellm-rust/crates/secret-value/src/lib.rs b/litellm-rust/crates/secret-value/src/lib.rs new file mode 100644 index 00000000000..f11d22d3701 --- /dev/null +++ b/litellm-rust/crates/secret-value/src/lib.rs @@ -0,0 +1,5 @@ +#![forbid(unsafe_code)] + +mod secret_value; + +pub use secret_value::SecretValue; diff --git a/litellm-rust/crates/auth-types/src/secret.rs b/litellm-rust/crates/secret-value/src/secret_value.rs similarity index 95% rename from litellm-rust/crates/auth-types/src/secret.rs rename to litellm-rust/crates/secret-value/src/secret_value.rs index 7e6789deef7..27e4227080d 100644 --- a/litellm-rust/crates/auth-types/src/secret.rs +++ b/litellm-rust/crates/secret-value/src/secret_value.rs @@ -32,9 +32,11 @@ impl Hash for SecretValue { #[cfg(test)] mod tests { + use rstest::rstest; + use super::SecretValue; - #[test] + #[rstest] fn debug_redacts_plaintext() { let debug = format!("{:?}", SecretValue::new("credential-value")); @@ -42,7 +44,7 @@ mod tests { assert!(debug.contains("REDACTED")); } - #[test] + #[rstest] fn equality_compares_plaintext_values() { assert_eq!(SecretValue::new("same"), SecretValue::new("same")); assert_ne!(SecretValue::new("same"), SecretValue::new("different")); diff --git a/litellm-rust/crates/secrets-aws/Cargo.toml b/litellm-rust/crates/secrets-aws/Cargo.toml index 5d3bd413484..d17dfcfeec2 100644 --- a/litellm-rust/crates/secrets-aws/Cargo.toml +++ b/litellm-rust/crates/secrets-aws/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-auth-aws.workspace = true litellm-secrets-types.workspace = true litellm-core-utils.workspace = true diff --git a/litellm-rust/crates/secrets-aws/src/secret_manager.rs b/litellm-rust/crates/secrets-aws/src/secret_manager.rs index 2220e387838..b25872ab3f6 100644 --- a/litellm-rust/crates/secrets-aws/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-aws/src/secret_manager.rs @@ -21,11 +21,12 @@ use litellm_auth_aws::constants::{ AWS_ACCESS_KEY_ID, AWS_REGION, AWS_REGION_NAME, AWS_SECRET_ACCESS_KEY, }; use litellm_core_utils::settings::Lookup; +use litellm_secret_value::SecretValue; use litellm_secrets_types::{ AwsOperationContext, BaseSecretManager, KeyManagementSettings, RotationError, Secret, - SecretDeleter, SecretRotator, SecretValue, SecretWriteContext, SecretWriter, - async_rotate_secret, + SecretDeleter, SecretRotator, SecretWriteContext, SecretWriter, async_rotate_secret, }; + use serde_json::Value; use crate::{Error, auth}; diff --git a/litellm-rust/crates/secrets-aws/tests/secret_manager.rs b/litellm-rust/crates/secrets-aws/tests/secret_manager.rs index c482a168090..56999633a78 100644 --- a/litellm-rust/crates/secrets-aws/tests/secret_manager.rs +++ b/litellm-rust/crates/secrets-aws/tests/secret_manager.rs @@ -10,11 +10,13 @@ use aws_sdk_secretsmanager::{ Client, config::{BehaviorVersion, Credentials, Region, retry::RetryConfig}, }; +use litellm_secret_value::SecretValue; use litellm_secrets_aws::{AwsSecretsManagerV2, Error, RotationResponse}; use litellm_secrets_types::{ AwsOperationContext, BaseSecretManager, KeyManagementSettings, Secret, SecretDeleter, - SecretValue, SecretWriteContext, SecretWriter, + SecretWriteContext, SecretWriter, }; + use rstest::{fixture, rstest}; use serde_json::json; use wiremock::{ diff --git a/litellm-rust/crates/secrets-azure/Cargo.toml b/litellm-rust/crates/secrets-azure/Cargo.toml index 2aff6129c37..c4bfb24898c 100644 --- a/litellm-rust/crates/secrets-azure/Cargo.toml +++ b/litellm-rust/crates/secrets-azure/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-http.workspace = true tokio.workspace = true litellm-auth-azure.workspace = true diff --git a/litellm-rust/crates/secrets-azure/src/key_vault.rs b/litellm-rust/crates/secrets-azure/src/key_vault.rs index d97b7a22cfc..40028828298 100644 --- a/litellm-rust/crates/secrets-azure/src/key_vault.rs +++ b/litellm-rust/crates/secrets-azure/src/key_vault.rs @@ -3,7 +3,9 @@ use std::sync::Arc; use litellm_auth_azure::{AzureAuthInputs, AzureAuthService, ConfigValue}; use litellm_auth_types::{InputSource, Sourced}; use litellm_core_utils::settings::Lookup; -use litellm_secrets_types::{AzureOperationContext, BaseSecretManager, Secret, SecretValue}; +use litellm_secret_value::SecretValue; +use litellm_secrets_types::{AzureOperationContext, BaseSecretManager, Secret}; + use percent_encoding::{AsciiSet, NON_ALPHANUMERIC}; use serde::Deserialize; diff --git a/litellm-rust/crates/secrets-azure/tests/key_vault.rs b/litellm-rust/crates/secrets-azure/tests/key_vault.rs index a0ba9e6ddf8..c3ba01672a8 100644 --- a/litellm-rust/crates/secrets-azure/tests/key_vault.rs +++ b/litellm-rust/crates/secrets-azure/tests/key_vault.rs @@ -1,7 +1,9 @@ use std::sync::Arc; +use litellm_secret_value::SecretValue; use litellm_secrets_azure::{AzureKeyVault, Error}; -use litellm_secrets_types::{Secret, SecretValue}; +use litellm_secrets_types::Secret; + use rstest::{fixture, rstest}; use serde::Deserialize; use wiremock::{ diff --git a/litellm-rust/crates/secrets-cyberark/Cargo.toml b/litellm-rust/crates/secrets-cyberark/Cargo.toml index 2ae2c1ea5c3..643f322f7fa 100644 --- a/litellm-rust/crates/secrets-cyberark/Cargo.toml +++ b/litellm-rust/crates/secrets-cyberark/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-secrets-types.workspace = true litellm-core-utils.workspace = true litellm-http.workspace = true diff --git a/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs b/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs index c5cb0308b13..cf077a2bbc8 100644 --- a/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-cyberark/src/secret_manager.rs @@ -9,11 +9,12 @@ use litellm_core_utils::settings::Lookup; use litellm_http::{ Client, ClientIdentity, ClientVariant, HttpClientConfig, HttpClientPool, TlsSource, Verify, }; +use litellm_secret_value::SecretValue; use litellm_secrets_types::{ BaseSecretManager, CyberarkOperationContext, RotationError, SecretCache, SecretDeleter, - SecretRotator, SecretValue, SecretWriteContext, SecretWriter, async_rotate_secret, - validate_secret_name, + SecretRotator, SecretWriteContext, SecretWriter, async_rotate_secret, validate_secret_name, }; + use moka::future::Cache; use percent_encoding::{AsciiSet, NON_ALPHANUMERIC, utf8_percent_encode}; diff --git a/litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs b/litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs index 783ba6ff67a..e9b480e026b 100644 --- a/litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs +++ b/litellm-rust/crates/secrets-cyberark/tests/secret_manager.rs @@ -9,8 +9,10 @@ use std::{ use base64::{Engine, engine::general_purpose::STANDARD}; use litellm_core_utils::settings::Lookup; use litellm_http::{HttpClientPool, HttpSettings, Resolution, media::PublicDnsResolver}; +use litellm_secret_value::SecretValue; use litellm_secrets_cyberark::{CyberArkSecretManager, DeleteOutcome, Error}; -use litellm_secrets_types::{BaseSecretManager, CyberarkOperationContext, SecretValue}; +use litellm_secrets_types::{BaseSecretManager, CyberarkOperationContext}; + use rstest::{fixture, rstest}; use serde::Deserialize; use wiremock::{ diff --git a/litellm-rust/crates/secrets-google/Cargo.toml b/litellm-rust/crates/secrets-google/Cargo.toml index 06b3747def5..3c5c2130788 100644 --- a/litellm-rust/crates/secrets-google/Cargo.toml +++ b/litellm-rust/crates/secrets-google/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-http.workspace = true moka.workspace = true tokio.workspace = true diff --git a/litellm-rust/crates/secrets-google/src/auth.rs b/litellm-rust/crates/secrets-google/src/auth.rs index 45fc99d8d5f..7022ecc5ae4 100644 --- a/litellm-rust/crates/secrets-google/src/auth.rs +++ b/litellm-rust/crates/secrets-google/src/auth.rs @@ -3,7 +3,7 @@ use std::sync::Arc; use litellm_auth_gcp::{GoogleCredentials, VertexConfig}; use litellm_auth_types::{InputSource, Sourced}; use litellm_core_utils::settings::Lookup; -use litellm_secrets_types::SecretValue; +use litellm_secret_value::SecretValue; pub(crate) fn credentials( project: Option, diff --git a/litellm-rust/crates/secrets-google/src/kms.rs b/litellm-rust/crates/secrets-google/src/kms.rs index 78ebae008f7..5622a3a4d56 100644 --- a/litellm-rust/crates/secrets-google/src/kms.rs +++ b/litellm-rust/crates/secrets-google/src/kms.rs @@ -2,7 +2,7 @@ use std::sync::Arc; use google_cloud_kms_v1::client::KeyManagementService; use litellm_core_utils::settings::Lookup; -use litellm_secrets_types::SecretValue; +use litellm_secret_value::SecretValue; use crate::{Error, auth}; diff --git a/litellm-rust/crates/secrets-google/src/secret_manager.rs b/litellm-rust/crates/secrets-google/src/secret_manager.rs index 463d1f6fd73..577a2ccbd4d 100644 --- a/litellm-rust/crates/secrets-google/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-google/src/secret_manager.rs @@ -2,9 +2,9 @@ use std::{sync::Arc, time::Duration}; use base64::{Engine, engine::general_purpose::STANDARD}; use litellm_core_utils::settings::Lookup; -use litellm_secrets_types::{ - BaseSecretManager, GoogleOperationContext, Secret, SecretCache, SecretValue, -}; +use litellm_secret_value::SecretValue; +use litellm_secrets_types::{BaseSecretManager, GoogleOperationContext, Secret, SecretCache}; + use serde::Deserialize; use litellm_auth_gcp::GoogleCredentials; diff --git a/litellm-rust/crates/secrets-hashicorp/Cargo.toml b/litellm-rust/crates/secrets-hashicorp/Cargo.toml index 8731a8deb7a..4fcd496931d 100644 --- a/litellm-rust/crates/secrets-hashicorp/Cargo.toml +++ b/litellm-rust/crates/secrets-hashicorp/Cargo.toml @@ -6,6 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] +litellm-secret-value.workspace = true litellm-error.workspace = true litellm-core-utils.workspace = true litellm-http.workspace = true diff --git a/litellm-rust/crates/secrets-hashicorp/src/config.rs b/litellm-rust/crates/secrets-hashicorp/src/config.rs index 4232667ae16..a63161b17b2 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/config.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/config.rs @@ -1,7 +1,7 @@ use std::{path::PathBuf, time::Duration}; use litellm_core_utils::settings::Lookup; -use litellm_secrets_types::SecretValue; +use litellm_secret_value::SecretValue; use crate::Error; diff --git a/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs index 447beacda85..dae14439630 100644 --- a/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs +++ b/litellm-rust/crates/secrets-hashicorp/src/secret_manager.rs @@ -12,11 +12,12 @@ use std::{ }; use litellm_core_utils::settings::Lookup; +use litellm_secret_value::SecretValue; use litellm_secrets_types::{ BaseSecretManager, HashicorpOperationContext, RotationError, SecretCache, SecretDeleter, - SecretRotator, SecretValue, SecretWriteContext, SecretWriter, async_rotate_secret, - validate_secret_name, + SecretRotator, SecretWriteContext, SecretWriter, async_rotate_secret, validate_secret_name, }; + use rustify::errors::ClientError as RustifyClientError; use serde_json::Value; use tokio::sync::Mutex; diff --git a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs index 668871d3fc7..d431e472400 100644 --- a/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs +++ b/litellm-rust/crates/secrets-hashicorp/tests/secret_manager.rs @@ -1,11 +1,13 @@ use std::{collections::HashMap, sync::Arc, time::Duration}; use litellm_core_utils::settings::Lookup; +use litellm_secret_value::SecretValue; use litellm_secrets_hashicorp::{Error, HashicorpVault, HashicorpVaultConfig}; use litellm_secrets_types::{ - BaseSecretManager, HashicorpOperationContext, RotationError, SecretDeleter, SecretValue, - SecretWriteContext, SecretWriter, + BaseSecretManager, HashicorpOperationContext, RotationError, SecretDeleter, SecretWriteContext, + SecretWriter, }; + use rstest::{fixture, rstest}; use serde::Deserialize; use serde_json::json; diff --git a/litellm-rust/crates/secrets-types/Cargo.toml b/litellm-rust/crates/secrets-types/Cargo.toml index 6dd847ec989..b45c41bcf1c 100644 --- a/litellm-rust/crates/secrets-types/Cargo.toml +++ b/litellm-rust/crates/secrets-types/Cargo.toml @@ -6,7 +6,7 @@ license.workspace = true repository.workspace = true [dependencies] -litellm-auth-types.workspace = true +litellm-secret-value.workspace = true moka.workspace = true tokio = { workspace = true, features = ["sync"] } serde.workspace = true diff --git a/litellm-rust/crates/secrets-types/src/base_secret_manager.rs b/litellm-rust/crates/secrets-types/src/base_secret_manager.rs index 737fad28662..41dadd1baa6 100644 --- a/litellm-rust/crates/secrets-types/src/base_secret_manager.rs +++ b/litellm-rust/crates/secrets-types/src/base_secret_manager.rs @@ -1,4 +1,5 @@ -use crate::{Error, SecretValue, SecretWriteContext}; +use crate::{Error, SecretWriteContext}; +use litellm_secret_value::SecretValue; pub fn validate_secret_name(name: &str) -> Result<(), Error> { if name.split('/').any(|segment| segment == "..") diff --git a/litellm-rust/crates/secrets-types/src/config.rs b/litellm-rust/crates/secrets-types/src/config.rs index 9b6e82c3a6b..a45acb4df11 100644 --- a/litellm-rust/crates/secrets-types/src/config.rs +++ b/litellm-rust/crates/secrets-types/src/config.rs @@ -3,7 +3,7 @@ use std::collections::BTreeMap; use serde::{Deserialize, Serialize}; use strum::IntoStaticStr; -use crate::SecretValue; +use litellm_secret_value::SecretValue; #[derive( Clone, diff --git a/litellm-rust/crates/secrets-types/src/context.rs b/litellm-rust/crates/secrets-types/src/context.rs index 1ec493edf42..b71195e2d21 100644 --- a/litellm-rust/crates/secrets-types/src/context.rs +++ b/litellm-rust/crates/secrets-types/src/context.rs @@ -1,6 +1,7 @@ use std::{collections::BTreeMap, time::Duration}; -use crate::{Error, KeyManagementSystem, SecretValue}; +use crate::{Error, KeyManagementSystem}; +use litellm_secret_value::SecretValue; #[derive(Clone, Debug, Default, Eq, PartialEq)] pub enum SecretOperationContext { diff --git a/litellm-rust/crates/secrets-types/src/lib.rs b/litellm-rust/crates/secrets-types/src/lib.rs index 4d7eaf8313e..2563ccbe39d 100644 --- a/litellm-rust/crates/secrets-types/src/lib.rs +++ b/litellm-rust/crates/secrets-types/src/lib.rs @@ -18,5 +18,4 @@ pub use context::{ HashicorpOperationContext, SecretOperationContext, SecretWriteContext, }; pub use error::Error; -pub use litellm_auth_types::SecretValue; pub use value::{PythonSecretRead, Secret}; diff --git a/litellm-rust/crates/secrets-types/src/value.rs b/litellm-rust/crates/secrets-types/src/value.rs index 11f57ef4448..8c92902edad 100644 --- a/litellm-rust/crates/secrets-types/src/value.rs +++ b/litellm-rust/crates/secrets-types/src/value.rs @@ -1,4 +1,4 @@ -use crate::SecretValue; +use litellm_secret_value::SecretValue; #[derive(Clone, PartialEq, Eq, veil::Redact)] pub enum Secret { diff --git a/litellm-rust/crates/secrets-types/tests/config.rs b/litellm-rust/crates/secrets-types/tests/config.rs index 45fd9d0ee03..e7ba63343da 100644 --- a/litellm-rust/crates/secrets-types/tests/config.rs +++ b/litellm-rust/crates/secrets-types/tests/config.rs @@ -1,6 +1,6 @@ -use litellm_secrets_types::{ - AccessMode, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue, -}; +use litellm_secret_value::SecretValue; +use litellm_secrets_types::{AccessMode, KeyManagementSettings, KeyManagementSystem, Secret}; + use rstest::{fixture, rstest}; use serde_json::json; diff --git a/litellm-rust/crates/secrets-types/tests/context.rs b/litellm-rust/crates/secrets-types/tests/context.rs index 8f8eeabdb62..f8d754564dc 100644 --- a/litellm-rust/crates/secrets-types/tests/context.rs +++ b/litellm-rust/crates/secrets-types/tests/context.rs @@ -1,10 +1,11 @@ use std::{collections::BTreeMap, time::Duration}; +use litellm_secret_value::SecretValue; use litellm_secrets_types::{ AwsOperationContext, AzureOperationContext, CyberarkOperationContext, GoogleOperationContext, - HashicorpOperationContext, KeyManagementSystem, SecretOperationContext, SecretValue, - SecretWriteContext, + HashicorpOperationContext, KeyManagementSystem, SecretOperationContext, SecretWriteContext, }; + use rstest::{fixture, rstest}; #[fixture] diff --git a/litellm-rust/crates/secrets-types/tests/rotation.rs b/litellm-rust/crates/secrets-types/tests/rotation.rs index 36c1a065ffa..089890c60bd 100644 --- a/litellm-rust/crates/secrets-types/tests/rotation.rs +++ b/litellm-rust/crates/secrets-types/tests/rotation.rs @@ -1,10 +1,12 @@ use std::sync::atomic::{AtomicUsize, Ordering}; +use litellm_secret_value::SecretValue; use litellm_secrets_types::{ BaseSecretManager, Error, HashicorpOperationContext, RotationError, SecretDeleter, - SecretOperationContext, SecretRotator, SecretValue, SecretWriteContext, SecretWriter, - async_rotate_secret, validate_secret_name, + SecretOperationContext, SecretRotator, SecretWriteContext, SecretWriter, async_rotate_secret, + validate_secret_name, }; + use rstest::{fixture, rstest}; struct Manager { diff --git a/litellm-rust/crates/secrets/Cargo.toml b/litellm-rust/crates/secrets/Cargo.toml index 023d61136e4..21ee1bdc760 100644 --- a/litellm-rust/crates/secrets/Cargo.toml +++ b/litellm-rust/crates/secrets/Cargo.toml @@ -14,6 +14,7 @@ azure = ["dep:litellm-secrets-azure"] cyberark = ["dep:litellm-secrets-cyberark"] [dependencies] +litellm-secret-value.workspace = true futures-util.workspace = true litellm-python-compat.workspace = true litellm-secrets-types.workspace = true diff --git a/litellm-rust/crates/secrets/src/compatibility.rs b/litellm-rust/crates/secrets/src/compatibility.rs index 368c9931b51..3d6ba854325 100644 --- a/litellm-rust/crates/secrets/src/compatibility.rs +++ b/litellm-rust/crates/secrets/src/compatibility.rs @@ -3,9 +3,10 @@ use litellm_python_compat::{Value, literal::literal_eval}; use litellm_secrets_types::PythonSecretRead; use crate::{ - Error, KeyManagementSettings, KeyManagementSystem, Secret, SecretManager, SecretValue, + Error, KeyManagementSettings, KeyManagementSystem, Secret, SecretManager, get_secret_from_manager, }; +use litellm_secret_value::SecretValue; pub async fn get_secret_from_python_manager( manager: &SecretManager, diff --git a/litellm-rust/crates/secrets/src/handler.rs b/litellm-rust/crates/secrets/src/handler.rs index f42e09c6bdb..2d2b81d197f 100644 --- a/litellm-rust/crates/secrets/src/handler.rs +++ b/litellm-rust/crates/secrets/src/handler.rs @@ -5,7 +5,7 @@ use litellm_core_utils::settings::Lookup; use crate::{Error, KeyManagementSettings, KeyManagementSystem, Secret}; #[cfg(any(feature = "aws", feature = "google"))] -use crate::SecretValue; +use litellm_secret_value::SecretValue; #[cfg(any( feature = "google", diff --git a/litellm-rust/crates/secrets/src/lib.rs b/litellm-rust/crates/secrets/src/lib.rs index 8d0f0561608..e4459576e94 100644 --- a/litellm-rust/crates/secrets/src/lib.rs +++ b/litellm-rust/crates/secrets/src/lib.rs @@ -12,9 +12,7 @@ mod state; pub use compatibility::{get_secret_from_python_manager, read_secret_from_python_manager}; pub use error::Error; pub use handler::{ExternalSecretManager, SecretManager, get_secret_from_manager}; -pub use litellm_secrets_types::{ - AccessMode, KeyManagementSettings, KeyManagementSystem, Secret, SecretValue, -}; +pub use litellm_secrets_types::{AccessMode, KeyManagementSettings, KeyManagementSystem, Secret}; pub use native::load_native_manager; pub use oidc::{OidcProvider, OidcReference, OidcResolver}; pub use resolver::{FailurePolicy, SecretResolver, normalize_nonempty_secret_str}; diff --git a/litellm-rust/crates/secrets/src/oidc.rs b/litellm-rust/crates/secrets/src/oidc.rs index 24e6912c797..18fc9479ef1 100644 --- a/litellm-rust/crates/secrets/src/oidc.rs +++ b/litellm-rust/crates/secrets/src/oidc.rs @@ -9,7 +9,8 @@ use litellm_http::Client; use moka::future::Cache; use serde::Deserialize; -use crate::{Error, SecretValue}; +use crate::Error; +use litellm_secret_value::SecretValue; const GOOGLE_TOKEN_MAX_TTL: Duration = Duration::from_secs(3540); const GITHUB_TOKEN_TTL: Duration = Duration::from_secs(295); diff --git a/litellm-rust/crates/secrets/src/resolver.rs b/litellm-rust/crates/secrets/src/resolver.rs index 830e7c22cd5..631b60cba10 100644 --- a/litellm-rust/crates/secrets/src/resolver.rs +++ b/litellm-rust/crates/secrets/src/resolver.rs @@ -4,7 +4,8 @@ use crate::compatibility::python_manager_string; use litellm_core_utils::{serde_compat::parse_str_bool, settings::Lookup}; use crate::state::{LookupTarget, normalize_secret_name}; -use crate::{Error, OidcResolver, Secret, SecretManagerState, SecretValue}; +use crate::{Error, OidcResolver, Secret, SecretManagerState}; +use litellm_secret_value::SecretValue; #[derive(Clone, Copy, Debug, Default, PartialEq, Eq)] pub enum FailurePolicy { diff --git a/litellm-rust/crates/secrets/src/source.rs b/litellm-rust/crates/secrets/src/source.rs index 3c86fd25e9a..1ee224b763a 100644 --- a/litellm-rust/crates/secrets/src/source.rs +++ b/litellm-rust/crates/secrets/src/source.rs @@ -3,7 +3,8 @@ use std::{collections::HashMap, sync::Arc}; use futures_util::future::{BoxFuture, try_join_all}; use litellm_core_utils::settings::Lookup; -use crate::{Error, SecretResolver, SecretValue}; +use crate::{Error, SecretResolver}; +use litellm_secret_value::SecretValue; pub type Secrets = Arc; diff --git a/litellm-rust/crates/secrets/tests/aws.rs b/litellm-rust/crates/secrets/tests/aws.rs index abad04f31c5..6f30d1b5161 100644 --- a/litellm-rust/crates/secrets/tests/aws.rs +++ b/litellm-rust/crates/secrets/tests/aws.rs @@ -2,11 +2,13 @@ use std::sync::Arc; +use litellm_secret_value::SecretValue; use litellm_secrets::{ AccessMode, Error, FailurePolicy, KeyManagementSettings, OidcResolver, Secret, SecretManager, - SecretManagerState, SecretResolver, SecretValue, aws::AwsSecretsManagerV2, + SecretManagerState, SecretResolver, aws::AwsSecretsManagerV2, secret_manager_would_be_consulted, }; + use wiremock::{Mock, MockServer, ResponseTemplate, matchers::method}; fn state(server: &MockServer, settings: KeyManagementSettings) -> SecretManagerState { diff --git a/litellm-rust/crates/secrets/tests/azure.rs b/litellm-rust/crates/secrets/tests/azure.rs index 6e81017b0f6..201e31f44b8 100644 --- a/litellm-rust/crates/secrets/tests/azure.rs +++ b/litellm-rust/crates/secrets/tests/azure.rs @@ -68,10 +68,11 @@ async fn azure_handler_reads_missing_and_failed_secrets() { async fn successful_azure_responses_do_not_fall_back_when_the_value_is_empty_or_null( #[case] body: serde_json::Value, ) { + use litellm_secret_value::SecretValue; use litellm_secrets::{ - OidcResolver, SecretManager, SecretManagerState, SecretResolver, SecretValue, - azure::AzureKeyVault, + OidcResolver, SecretManager, SecretManagerState, SecretResolver, azure::AzureKeyVault, }; + use std::sync::Arc; use wiremock::{Mock, MockServer, ResponseTemplate, matchers::any}; let server = MockServer::start().await; diff --git a/litellm-rust/crates/secrets/tests/common_read_contract.rs b/litellm-rust/crates/secrets/tests/common_read_contract.rs index e6928038799..4cf3513fadf 100644 --- a/litellm-rust/crates/secrets/tests/common_read_contract.rs +++ b/litellm-rust/crates/secrets/tests/common_read_contract.rs @@ -9,8 +9,9 @@ use std::sync::Arc; use base64::{Engine, engine::general_purpose::STANDARD}; use litellm_core_utils::settings::Lookup; +use litellm_secret_value::SecretValue; use litellm_secrets::{ - KeyManagementSettings, SecretManager, SecretValue, + KeyManagementSettings, SecretManager, aws::AwsSecretsManagerV2, azure::AzureKeyVault, cyberark::CyberArkSecretManager, @@ -18,6 +19,7 @@ use litellm_secrets::{ google::GoogleSecretManager, hashicorp::{HashicorpVault, HashicorpVaultConfig}, }; + use rstest::rstest; use serde_json::json; use wiremock::{ diff --git a/litellm-rust/crates/secrets/tests/cyberark.rs b/litellm-rust/crates/secrets/tests/cyberark.rs index 0a6c394937e..09342234118 100644 --- a/litellm-rust/crates/secrets/tests/cyberark.rs +++ b/litellm-rust/crates/secrets/tests/cyberark.rs @@ -4,10 +4,12 @@ async fn cyberark_handler_reads_values_and_surfaces_errors() { use std::time::Duration; + use litellm_secret_value::SecretValue; use litellm_secrets::{ - Error, KeyManagementSettings, SecretManager, SecretValue, cyberark::CyberArkSecretManager, + Error, KeyManagementSettings, SecretManager, cyberark::CyberArkSecretManager, get_secret_from_manager, }; + use wiremock::{ Mock, MockServer, ResponseTemplate, matchers::{body_string, path}, diff --git a/litellm-rust/crates/secrets/tests/google.rs b/litellm-rust/crates/secrets/tests/google.rs index 5542c864dfd..a75e34335de 100644 --- a/litellm-rust/crates/secrets/tests/google.rs +++ b/litellm-rust/crates/secrets/tests/google.rs @@ -7,10 +7,12 @@ use std::sync::Arc; #[case::failure(503)] #[tokio::test] async fn google_resolver_distinguishes_absence_from_failure(#[case] status: u16) { + use litellm_secret_value::SecretValue; use litellm_secrets::{ Error, FailurePolicy, KeyManagementSettings, OidcResolver, SecretManager, - SecretManagerState, SecretResolver, SecretValue, google::GoogleSecretManager, + SecretManagerState, SecretResolver, google::GoogleSecretManager, }; + use wiremock::{Mock, MockServer, ResponseTemplate, matchers::method}; let server = MockServer::start().await; Mock::given(method("GET")) diff --git a/litellm-rust/crates/secrets/tests/oidc.rs b/litellm-rust/crates/secrets/tests/oidc.rs index 867a662cbbb..2d250e1de9d 100644 --- a/litellm-rust/crates/secrets/tests/oidc.rs +++ b/litellm-rust/crates/secrets/tests/oidc.rs @@ -315,7 +315,7 @@ async fn azure_oidc_acquires_the_requested_scope_and_preserves_failures(#[case] Box< dyn std::future::Future< Output = Result< - litellm_secrets::SecretValue, + litellm_secret_value::SecretValue, litellm_secrets::azure::Error, >, > + Send @@ -331,7 +331,7 @@ async fn azure_oidc_acquires_the_requested_scope_and_preserves_failures(#[case] if self.0 { Err(litellm_secrets::azure::Error::MissingCredentials) } else { - Ok(litellm_secrets::SecretValue::new("azure-token")) + Ok(litellm_secret_value::SecretValue::new("azure-token")) } }) } @@ -346,7 +346,7 @@ async fn azure_oidc_acquires_the_requested_scope_and_preserves_failures(#[case] let result = resolver .get_secret_str( "oidc/azure/api://audience/path", - Some(litellm_secrets::SecretValue::new("fallback")), + Some(litellm_secret_value::SecretValue::new("fallback")), ) .await; if failed { diff --git a/litellm-rust/crates/secrets/tests/resolution.rs b/litellm-rust/crates/secrets/tests/resolution.rs index 8435e051315..185f2256780 100644 --- a/litellm-rust/crates/secrets/tests/resolution.rs +++ b/litellm-rust/crates/secrets/tests/resolution.rs @@ -1,9 +1,10 @@ use std::{future::Future, pin::Pin, sync::Arc}; use litellm_core_utils::settings::Lookup; +use litellm_secret_value::SecretValue; use litellm_secrets::{ Error, ExternalSecretManager, FailurePolicy, KeyManagementSettings, KeyManagementSystem, - OidcResolver, Secret, SecretManager, SecretManagerState, SecretResolver, SecretValue, + OidcResolver, Secret, SecretManager, SecretManagerState, SecretResolver, secret_manager_would_be_consulted, }; From 24b8889965a6f86bee6f743d61c77da0896eaa15 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 11:35:18 -0700 Subject: [PATCH 04/46] chore(cost-map): add openai gpt-6-rosalind from the pricing page and changelog (#49268) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 56 +++++++++++++++++++ model_prices_and_context_window.json | 56 +++++++++++++++++++ 2 files changed, 112 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index ea9cacbdabe..dac7af6ab7a 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -81156,5 +81156,61 @@ "output_cost_per_token": 0.0, "source": "https://inference.baseten.co/v1/models", "supports_vision": true + }, + "gpt-6-rosalind": { + "cache_creation_input_token_cost": 1.875e-05, + "cache_creation_input_token_cost_above_12hr": 3.75e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 7.5e-05, + "cache_creation_input_token_cost_above_272k_tokens": 3.75e-05, + "cache_read_input_token_cost": 1.5e-06, + "cache_read_input_token_cost_above_272k_tokens": 3e-06, + "input_cost_per_token": 1.5e-05, + "input_cost_per_token_above_272k_tokens": 3e-05, + "litellm_provider": "openai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.5e-05, + "output_cost_per_token_above_272k_tokens": 0.0001125, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, + "source": "https://developers.openai.com/api/docs/pricing", + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_tool_search": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, + "chatgpt/gpt-6-rosalind": { + "litellm_provider": "chatgpt", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "responses", + "source": "https://developers.openai.com/api/docs/pricing", + "supported_endpoints": [ + "/v1/responses" + ], + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true } } diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index ea9cacbdabe..dac7af6ab7a 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -81156,5 +81156,61 @@ "output_cost_per_token": 0.0, "source": "https://inference.baseten.co/v1/models", "supports_vision": true + }, + "gpt-6-rosalind": { + "cache_creation_input_token_cost": 1.875e-05, + "cache_creation_input_token_cost_above_12hr": 3.75e-05, + "cache_creation_input_token_cost_above_12hr_above_272k_tokens": 7.5e-05, + "cache_creation_input_token_cost_above_272k_tokens": 3.75e-05, + "cache_read_input_token_cost": 1.5e-06, + "cache_read_input_token_cost_above_272k_tokens": 3e-06, + "input_cost_per_token": 1.5e-05, + "input_cost_per_token_above_272k_tokens": 3e-05, + "litellm_provider": "openai", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "chat", + "output_cost_per_token": 7.5e-05, + "output_cost_per_token_above_272k_tokens": 0.0001125, + "regional_processing_uplift_multiplier_eu": 1.1, + "regional_processing_uplift_multiplier_us": 1.1, + "source": "https://developers.openai.com/api/docs/pricing", + "supports_computer_use": true, + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_native_streaming": true, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_pdf_input": true, + "supports_prompt_cache_breakpoint": true, + "supports_prompt_caching": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "supports_tool_search": true, + "supports_vision": true, + "supports_web_search": true, + "supports_xhigh_reasoning_effort": true + }, + "chatgpt/gpt-6-rosalind": { + "litellm_provider": "chatgpt", + "max_input_tokens": 922000, + "max_output_tokens": 128000, + "mode": "responses", + "source": "https://developers.openai.com/api/docs/pricing", + "supported_endpoints": [ + "/v1/responses" + ], + "supports_function_calling": true, + "supports_max_reasoning_effort": true, + "supports_minimal_reasoning_effort": false, + "supports_none_reasoning_effort": false, + "supports_parallel_function_calling": true, + "supports_reasoning": true, + "supports_response_schema": true, + "supports_vision": true, + "supports_xhigh_reasoning_effort": true } } From 1ccb970d520aa24a078d5b2a40ccc7f5249a770c Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Fri, 30 Oct 2026 11:43:13 -0700 Subject: [PATCH 05/46] refactor(ui): remove the deprecated Old Usage page (#49284) Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- tests/e2e/ui/fixtures/menuMappings.ts | 1 - tests/e2e/ui/fixtures/migratedPages.ts | 6 - tests/e2e/ui/fixtures/pages.ts | 1 - .../e2e/ui/tests/logs/logsPagination.spec.ts | 2 +- tests/e2e/ui/tests/usage/usagePage.spec.ts | 2 +- ui/litellm-dashboard/eslint-suppressions.json | 28 - .../app/(dashboard)/legacyPageRoutes.test.ts | 2 +- .../src/app/(dashboard)/legacyPageRoutes.ts | 2 +- .../old-usage/_components/usage.test.tsx | 257 ----- .../old-usage/_components/usage.tsx | 950 ------------------ .../src/app/(dashboard)/old-usage/page.tsx | 22 - .../components/UsagePageView.test.tsx | 4 - .../src/components/leftnav.test.tsx | 26 - .../src/components/leftnav.tsx | 8 - .../src/components/page_metadata.ts | 1 - .../src/components/view_user_spend.tsx | 144 --- .../tests/CreateKeyPage.expiredToken.test.tsx | 1 - 17 files changed, 4 insertions(+), 1453 deletions(-) delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.test.tsx delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx delete mode 100644 ui/litellm-dashboard/src/app/(dashboard)/old-usage/page.tsx delete mode 100644 ui/litellm-dashboard/src/components/view_user_spend.tsx diff --git a/tests/e2e/ui/fixtures/menuMappings.ts b/tests/e2e/ui/fixtures/menuMappings.ts index d6e7ea86982..4cf90c00cd0 100644 --- a/tests/e2e/ui/fixtures/menuMappings.ts +++ b/tests/e2e/ui/fixtures/menuMappings.ts @@ -32,7 +32,6 @@ export const menuLabelToPage: Record = { Budgets: Page.Budgets, "API Playground": Page.TransformRequest, "Tag Management": Page.TagManagement, - "Old Usage": Page.Usage, // Tools submenu items "MCP Servers": Page.McpServers, "Vector Stores": Page.VectorStores, diff --git a/tests/e2e/ui/fixtures/migratedPages.ts b/tests/e2e/ui/fixtures/migratedPages.ts index 5fa1f9d2444..6fa8d0c2587 100644 --- a/tests/e2e/ui/fixtures/migratedPages.ts +++ b/tests/e2e/ui/fixtures/migratedPages.ts @@ -118,12 +118,6 @@ export const MIGRATED_E2E_PAGES: Readonly> = { content: { role: "heading", name: "AI Hub" }, }, new_usage: { segment: "usage", linkName: "Usage", content: { role: "heading", name: "Usage" } }, - usage: { - segment: "old-usage", - linkName: "Old Usage", - group: "Experimental", - content: { role: "tab", name: "All Up" }, - }, agents: { segment: "agents", linkName: "Agents", group: "Agentic", content: { role: "heading", name: "Agents" } }, "router-settings": { segment: "router-settings", diff --git a/tests/e2e/ui/fixtures/pages.ts b/tests/e2e/ui/fixtures/pages.ts index 8210334c166..a850dcb4755 100644 --- a/tests/e2e/ui/fixtures/pages.ts +++ b/tests/e2e/ui/fixtures/pages.ts @@ -30,5 +30,4 @@ export enum Page { TagManagement = "tag-management", VectorStores = "vector-stores", NewUsage = "new_usage", - Usage = "usage", } diff --git a/tests/e2e/ui/tests/logs/logsPagination.spec.ts b/tests/e2e/ui/tests/logs/logsPagination.spec.ts index 585eb31c37c..a611f760e8b 100644 --- a/tests/e2e/ui/tests/logs/logsPagination.spec.ts +++ b/tests/e2e/ui/tests/logs/logsPagination.spec.ts @@ -84,7 +84,7 @@ test.describe("Logs pagination with Aria in the header", () => { await expect(aria).toBeFocused(); await visibleTestId(page, "pagination-first").click(); await expect(visibleTestId(page, "pagination-range")).toHaveText("Showing 1-25 of 26"); - for (const destination of [Page.Usage, Page.LlmPlayground, Page.ApiKeys]) { + for (const destination of [Page.NewUsage, Page.LlmPlayground, Page.ApiKeys]) { await navigateToPage(page, destination); await aria.click(); await expect(page.getByRole("dialog", { name: "Aria Agent", exact: true })).toBeVisible(); diff --git a/tests/e2e/ui/tests/usage/usagePage.spec.ts b/tests/e2e/ui/tests/usage/usagePage.spec.ts index d532042c7af..2a58a1ef59d 100644 --- a/tests/e2e/ui/tests/usage/usagePage.spec.ts +++ b/tests/e2e/ui/tests/usage/usagePage.spec.ts @@ -11,7 +11,7 @@ import { waitForSpendLog, } from "../../helpers/traffic"; -/** Covers /ui/usage. The legacy /ui/old-usage view is deprecated and deliberately not covered. */ +/** Covers /ui/usage. */ const topKeysCard = (page: PlaywrightPage): Locator => page.locator("section").filter({ has: page.getByRole("heading", { name: "Top Virtual Keys", exact: true }) }); diff --git a/ui/litellm-dashboard/eslint-suppressions.json b/ui/litellm-dashboard/eslint-suppressions.json index 9fce93782b2..c14f3e6f036 100644 --- a/ui/litellm-dashboard/eslint-suppressions.json +++ b/ui/litellm-dashboard/eslint-suppressions.json @@ -625,23 +625,6 @@ "count": 1 } }, - "src/app/(dashboard)/old-usage/_components/usage.tsx": { - "local/filename-pascal-case": { - "count": 1 - }, - "max-lines": { - "count": 1 - }, - "prefer-const": { - "count": 6 - }, - "react-hooks/immutability": { - "count": 1 - }, - "react-hooks/purity": { - "count": 1 - } - }, "src/app/(dashboard)/playground/components/chat_ui/AdditionalModelSettings.tsx": { "react-hooks/set-state-in-effect": { "count": 2 @@ -2090,17 +2073,6 @@ "count": 1 } }, - "src/components/view_user_spend.tsx": { - "local/filename-pascal-case": { - "count": 1 - }, - "prefer-const": { - "count": 3 - }, - "react-hooks/set-state-in-effect": { - "count": 2 - } - }, "src/contexts/AuthContext.tsx": { "react-hooks/set-state-in-effect": { "count": 1 diff --git a/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.test.ts b/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.test.ts index 6b02b79b04c..315d988db54 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.test.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.test.ts @@ -10,7 +10,7 @@ describe("legacyPageRedirectHref", () => { expect(redirect("page=models")).toBe("/ui/models-and-endpoints"); expect(redirect("page=llm-playground")).toBe("/ui/playground"); expect(redirect("page=new_usage")).toBe("/ui/usage"); - expect(redirect("page=usage")).toBe("/ui/old-usage"); + expect(redirect("page=usage")).toBe("/ui/usage"); }); it("keeps the older aliases for renamed pages", () => { diff --git a/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts b/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts index e3025266ddb..fe054c7b39a 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts +++ b/ui/litellm-dashboard/src/app/(dashboard)/legacyPageRoutes.ts @@ -36,7 +36,7 @@ const LEGACY_PAGE_ROUTES: ReadonlyMap = new Map( "logging-and-alerts": "logging-and-alerts", "model-hub-table": "model-hub-table", new_usage: "usage", - usage: "old-usage", + usage: "usage", "cost-optimization": "cost-optimization", "model-insights": "model-insights", agents: "agents", diff --git a/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.test.tsx deleted file mode 100644 index 0e4455ee912..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.test.tsx +++ /dev/null @@ -1,257 +0,0 @@ -import React from "react"; -import { describe, it, expect, vi, beforeEach } from "vitest"; -import { act, screen, waitFor, within } from "@testing-library/react"; -import userEvent from "@testing-library/user-event"; -import { renderWithProviders } from "../../../../../tests/test-utils"; -import UsagePage from "./usage"; - -const networking = vi.hoisted(() => ({ - adminSpendLogsCall: vi.fn(), - adminTopKeysCall: vi.fn(), - adminTopModelsCall: vi.fn(), - adminTopEndUsersCall: vi.fn(), - teamSpendLogsCall: vi.fn(), - tagsSpendLogsCall: vi.fn(), - allTagNamesCall: vi.fn(), - adminspendByProvider: vi.fn(), - adminGlobalActivity: vi.fn(), - adminGlobalActivityPerModel: vi.fn(), - getProxyUISettings: vi.fn(), - modelAvailableCall: vi.fn(), - keyInfoV1Call: vi.fn(), -})); - -vi.mock("@/components/networking", () => networking); -vi.mock("../../../../components/networking", () => networking); - -vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ - default: () => ({ - accessToken: "sk-test", - token: "tok", - userRole: "Admin", - userId: "u1", - premiumUser: true, - }), -})); - -const UNLIMITED_SETTINGS = { DISABLE_EXPENSIVE_DB_QUERIES: false, NUM_SPEND_LOGS_ROWS: 10 }; - -const renderUsage = (overrides: Partial> = {}) => - renderWithProviders( - , - ); - -// Width of this window is guarded by "proves the flush window is wide enough". -const flushPendingRequests = async () => { - await act(async () => { - await new Promise((resolve) => setTimeout(resolve, 0)); - await new Promise((resolve) => setTimeout(resolve, 0)); - }); -}; - -beforeEach(() => { - vi.clearAllMocks(); - networking.getProxyUISettings.mockResolvedValue(UNLIMITED_SETTINGS); - networking.adminSpendLogsCall.mockResolvedValue([{ date: "2026-07-01", spend: 12.5 }]); - networking.adminTopKeysCall.mockResolvedValue([ - { api_key: "sk-abcdefghijk", key_alias: "prod-key", total_spend: 9.5 }, - ]); - networking.adminTopModelsCall.mockResolvedValue([{ model: "gpt-5.1", total_spend: 7.25 }]); - networking.adminTopEndUsersCall.mockResolvedValue([ - { end_user: "customer-alpha", total_spend: 3.5, total_count: 42 }, - ]); - networking.teamSpendLogsCall.mockResolvedValue({ - daily_spend: [{ date: "2026-07-01", "team-a": 5 }], - teams: ["team-a"], - total_spend_per_team: [{ team_id: "team-a", total_spend: 5 }], - }); - networking.tagsSpendLogsCall.mockResolvedValue({ spend_per_tag: [{ name: "prod", spend: 4 }] }); - networking.allTagNamesCall.mockResolvedValue({ tag_names: ["prod", "staging"] }); - networking.adminspendByProvider.mockResolvedValue([{ provider: "openai", spend: 6.75 }]); - networking.adminGlobalActivity.mockResolvedValue({ - sum_api_requests: 120, - sum_total_tokens: 4500, - daily_data: [{ date: "2026-07-01", api_requests: 120, total_tokens: 4500 }], - }); - networking.adminGlobalActivityPerModel.mockResolvedValue([]); - networking.modelAvailableCall.mockResolvedValue({ data: [] }); - networking.keyInfoV1Call.mockResolvedValue({ info: {} }); -}); - -describe("old usage page", () => { - describe("when the proxy has disabled expensive DB queries", () => { - beforeEach(() => { - networking.getProxyUISettings.mockResolvedValue({ - DISABLE_EXPENSIVE_DB_QUERIES: true, - NUM_SPEND_LOGS_ROWS: 2500000, - }); - }); - - it("shows the database query limit warning instead of the usage dashboard", async () => { - renderUsage(); - - expect(await screen.findByText("Database Query Limit Reached")).toBeInTheDocument(); - expect(screen.getByText(/SpendLogs in DB has/)).toHaveTextContent("2500000"); - expect(screen.getByText(/Please follow our guide to view usage when SpendLogs has more than 1M rows/i)); - expect(screen.queryByRole("tab", { name: "All Up" })).not.toBeInTheDocument(); - }); - - it("links to the cost tracking guide in a new tab", async () => { - renderUsage(); - - const link = await screen.findByRole("link", { name: "View Usage Guide" }); - expect(link).toHaveAttribute("href", "https://docs.litellm.ai/docs/proxy/cost_tracking"); - expect(link).toHaveAttribute("target", "_blank"); - }); - - it("skips every expensive usage query", async () => { - renderUsage(); - - await screen.findByText("Database Query Limit Reached"); - await waitFor(() => expect(networking.getProxyUISettings).toHaveBeenCalled()); - - expect(networking.adminSpendLogsCall).not.toHaveBeenCalled(); - expect(networking.adminspendByProvider).not.toHaveBeenCalled(); - expect(networking.adminTopKeysCall).not.toHaveBeenCalled(); - expect(networking.adminTopModelsCall).not.toHaveBeenCalled(); - expect(networking.adminGlobalActivity).not.toHaveBeenCalled(); - expect(networking.adminGlobalActivityPerModel).not.toHaveBeenCalled(); - expect(networking.teamSpendLogsCall).not.toHaveBeenCalled(); - expect(networking.adminTopEndUsersCall).not.toHaveBeenCalled(); - expect(networking.tagsSpendLogsCall).not.toHaveBeenCalled(); - }); - }); - - describe("as an admin", () => { - it("renders the admin tabs", async () => { - renderUsage(); - - expect(await screen.findByRole("tab", { name: "All Up" })).toBeInTheDocument(); - expect(screen.getByRole("tab", { name: "Team Based Usage" })).toBeInTheDocument(); - expect(screen.getByRole("tab", { name: "Customer Usage" })).toBeInTheDocument(); - expect(screen.getByRole("tab", { name: "Tag Based Usage" })).toBeInTheDocument(); - }); - - it("renders the cost panel cards", async () => { - renderUsage(); - - expect(await screen.findByText("Monthly Spend")).toBeInTheDocument(); - expect(screen.getByText("Top Virtual Keys")).toBeInTheDocument(); - expect(screen.getByText("Top Models")).toBeInTheDocument(); - expect(screen.getByText("Spend by Provider")).toBeInTheDocument(); - }); - - it("lists spend by provider in a table", async () => { - renderUsage(); - - const providerCell = await screen.findByText("openai"); - const row = providerCell.closest("tr"); - expect(row).not.toBeNull(); - expect(within(row as HTMLElement).getByText("$6.75")).toBeInTheDocument(); - expect(screen.getByRole("columnheader", { name: "Provider" })).toBeInTheDocument(); - }); - - it("shows the customer usage table when its tab is selected", async () => { - const user = userEvent.setup(); - renderUsage(); - - await user.click(await screen.findByRole("tab", { name: "Customer Usage" })); - - const customerCell = await screen.findByText("customer-alpha"); - const row = customerCell.closest("tr"); - expect(row).not.toBeNull(); - expect(within(row as HTMLElement).getByText("$3.50")).toBeInTheDocument(); - expect(within(row as HTMLElement).getByText("42")).toBeInTheDocument(); - expect(screen.getByRole("columnheader", { name: "Total Events" })).toBeInTheDocument(); - }); - - it("shows the tag spend panel when its tab is selected", async () => { - const user = userEvent.setup(); - renderUsage(); - - await user.click(await screen.findByRole("tab", { name: "Tag Based Usage" })); - - expect(await screen.findByText("Spend Per Tag")).toBeInTheDocument(); - }); - - it("shows the team spend panel when its tab is selected", async () => { - const user = userEvent.setup(); - renderUsage(); - - await user.click(await screen.findByRole("tab", { name: "Team Based Usage" })); - - expect(await screen.findByText("Total Spend Per Team")).toBeInTheDocument(); - expect(screen.getByText("Daily Spend Per Team")).toBeInTheDocument(); - }); - }); - - // org_admin is an organization membership role; those users reach the UI as "Internal User". - describe.each(["Internal User", "Internal Viewer", "internal_user", "internal_user_viewer", "Org Admin"])( - "as %s", - (userRole) => { - it("shows the admin-only notice instead of the usage dashboard", async () => { - renderUsage({ userRole }); - - expect(await screen.findByText(/Proxy-wide usage is only available to admin users/i)).toBeInTheDocument(); - expect(screen.queryByRole("tab", { name: "All Up" })).not.toBeInTheDocument(); - }); - - it("fires no /global/spend or /global/activity request", async () => { - renderUsage({ userRole }); - - await screen.findByText(/Proxy-wide usage is only available to admin users/i); - await flushPendingRequests(); - - expect(networking.getProxyUISettings).not.toHaveBeenCalled(); - expect(networking.adminSpendLogsCall).not.toHaveBeenCalled(); - expect(networking.adminTopKeysCall).not.toHaveBeenCalled(); - expect(networking.adminTopModelsCall).not.toHaveBeenCalled(); - expect(networking.adminTopEndUsersCall).not.toHaveBeenCalled(); - expect(networking.teamSpendLogsCall).not.toHaveBeenCalled(); - expect(networking.tagsSpendLogsCall).not.toHaveBeenCalled(); - expect(networking.allTagNamesCall).not.toHaveBeenCalled(); - expect(networking.adminspendByProvider).not.toHaveBeenCalled(); - expect(networking.adminGlobalActivity).not.toHaveBeenCalled(); - expect(networking.adminGlobalActivityPerModel).not.toHaveBeenCalled(); - }); - }, - ); - - describe("the admin-only gate", () => { - it("proves the flush window is wide enough to catch a leaked request", async () => { - renderUsage({ userRole: "Admin" }); - - await flushPendingRequests(); - - expect(networking.getProxyUISettings).toHaveBeenCalled(); - expect(networking.adminSpendLogsCall).toHaveBeenCalled(); - expect(networking.tagsSpendLogsCall).toHaveBeenCalled(); - expect(networking.adminGlobalActivity).toHaveBeenCalled(); - }); - - it("still lets an admin through, so the notice is a real gate and not a dead branch", async () => { - renderUsage({ userRole: "Admin" }); - - expect(await screen.findByRole("tab", { name: "All Up" })).toBeInTheDocument(); - expect(screen.queryByText(/Proxy-wide usage is only available to admin users/i)).not.toBeInTheDocument(); - await waitFor(() => expect(networking.adminSpendLogsCall).toHaveBeenCalled()); - }); - - it("does not put the session token in the provider spend query", async () => { - renderUsage({ userRole: "Admin", token: "session-jwt-value" }); - - await waitFor(() => expect(networking.adminspendByProvider).toHaveBeenCalled()); - const callArgs = networking.adminspendByProvider.mock.calls[0]; - expect(callArgs).not.toContain("session-jwt-value"); - expect(callArgs[0]).toBe("sk-test"); - }); - }); -}); diff --git a/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx b/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx deleted file mode 100644 index 15b2a30b50c..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/old-usage/_components/usage.tsx +++ /dev/null @@ -1,950 +0,0 @@ -import React, { useState, useEffect } from "react"; - -import ViewUserSpend from "@/components/view_user_spend"; -import { ProxySettings } from "@/app/(dashboard)/hooks/proxySettings/useProxySettings"; -import AdvancedDatePicker from "@/components/shared/advanced_date_picker"; -import { Button } from "@/components/ui/button"; -import { Card, CardContent, CardHeader, CardTitle } from "@/components/ui/card"; -import { - Combobox, - ComboboxChip, - ComboboxChips, - ComboboxChipsInput, - ComboboxContent, - ComboboxEmpty, - ComboboxItem, - ComboboxList, - ComboboxValue, - useComboboxAnchor, -} from "@/components/ui/combobox"; -import { Meter, MeterIndicator, MeterTrack } from "@/components/shared/Meter"; -import { Select, SelectContent, SelectItem, SelectTrigger, SelectValue } from "@/components/ui/select"; -import { - NUMERIC_CELL_CLASS, - Table, - TableBody, - TableCell, - TableHead, - TableHeader, - TableRow, -} from "@/components/ui/table"; -import { Tabs, TabsContent, TabsList, TabsTrigger } from "@/components/ui/tabs"; -import { AreaChart, BarChart, DonutChart } from "@/components/shared/charts"; - -import { - adminSpendLogsCall, - adminTopKeysCall, - adminTopModelsCall, - adminTopEndUsersCall, - teamSpendLogsCall, - tagsSpendLogsCall, - allTagNamesCall, - adminspendByProvider, - adminGlobalActivity, - adminGlobalActivityPerModel, - getProxyUISettings, -} from "@/components/networking"; -import TopKeyView from "@/components/UsagePage/components/EntityUsage/TopKeyView"; -import { MoneyCell } from "@/components/shared/table_cells"; -import { hasCapability } from "@/utils/capabilities"; -import { formatNumberWithCommas } from "@/utils/dataUtils"; - -interface UsagePageProps { - accessToken: string | null; - token: string | null; - userRole: string | null; - userID: string | null; - keys: any[] | null; - premiumUser: boolean; -} - -interface GlobalActivityData { - sum_api_requests: number; - sum_total_tokens: number; - daily_data: { date: string; api_requests: number; total_tokens: number }[]; -} - -const EMPTY_GLOBAL_ACTIVITY: GlobalActivityData = { sum_api_requests: 0, sum_total_tokens: 0, daily_data: [] }; - -type UsageDateRange = { from?: Date; to?: Date }; - -type TeamSpendTotal = { name: string; value: number }; - -type TagOption = { value: string; label: string; disabled: boolean }; - -const ALL_TAGS = "all-tags"; - -const isAdminOrAdminViewer = (role: string | null): boolean => { - if (role === null) return false; - return role === "Admin" || role === "Admin Viewer"; -}; - -const TeamSpendBarList: React.FC<{ data: TeamSpendTotal[] }> = ({ data }) => { - const max = Math.max(0, ...data.map((team) => team.value)); - - return ( -
- {data.map((team) => ( -
-

{team.name}

- - - - - -

- {formatNumberWithCommas(team.value, 2)} -

-
- ))} -
- ); -}; - -const UsagePage: React.FC = ({ accessToken, token, userRole, userID, keys, premiumUser }) => { - const anchor = useComboboxAnchor(); - const canViewGlobalSpend = hasCapability(userRole, "viewGlobalSpend"); - const currentDate = new Date(); - const [keySpendData, setKeySpendData] = useState([]); - const [topKeys, setTopKeys] = useState([]); - const [topModels, setTopModels] = useState([]); - const [topUsers, setTopUsers] = useState([]); - const [teamSpendData, setTeamSpendData] = useState([]); - const [topTagsData, setTopTagsData] = useState([]); - const [allTagNames, setAllTagNames] = useState([]); - const [uniqueTeamIds, setUniqueTeamIds] = useState([]); - const [totalSpendPerTeam, setTotalSpendPerTeam] = useState([]); - const [spendByProvider, setSpendByProvider] = useState([]); - const [globalActivity, setGlobalActivity] = useState(EMPTY_GLOBAL_ACTIVITY); - const [globalActivityPerModel, setGlobalActivityPerModel] = useState([]); - const [selectedKeyToken, setSelectedKeyToken] = useState(null); - const [selectedTags, setSelectedTags] = useState([ALL_TAGS]); - const [dateValue, setDateValue] = useState({ - from: new Date(Date.now() - 7 * 24 * 60 * 60 * 1000), - to: new Date(), - }); - const [proxySettings, setProxySettings] = useState(null); - const [totalMonthlySpend, setTotalMonthlySpend] = useState(0); - - const firstDay = new Date(currentDate.getFullYear(), currentDate.getMonth(), 1); - const lastDay = new Date(currentDate.getFullYear(), currentDate.getMonth() + 1, 0); - - let startTime = formatDate(firstDay); - let endTime = formatDate(lastDay); - - const selectableKeys: { token: string; alias: string }[] = (keys ?? []) - .filter((key: any) => key && typeof key["key_alias"] === "string" && key["key_alias"].length > 0) - .map((key: any) => ({ token: String(key["token"]), alias: String(key["key_alias"]) })); - - const tagOptions: TagOption[] = [ - { value: ALL_TAGS, label: "All Tags", disabled: false }, - ...allTagNames - .filter((tag) => tag !== ALL_TAGS) - .map((tag) => ({ - value: tag, - label: premiumUser ? tag : `✨ ${tag} (Enterprise only Feature)`, - disabled: !premiumUser, - })), - ]; - - function valueFormatterNumbers(number: number) { - const formatter = new Intl.NumberFormat("en-US", { - maximumFractionDigits: 0, - notation: "compact", - compactDisplay: "short", - }); - - return formatter.format(number); - } - - const fetchProxySettings = async () => { - if (accessToken) { - try { - const proxy_settings: ProxySettings = await getProxyUISettings(accessToken); - return proxy_settings; - } catch (error) { - console.error("Error fetching proxy settings:", error); - } - } - }; - - useEffect(() => { - if (!canViewGlobalSpend) { - return; - } - updateTagSpendData(dateValue.from, dateValue.to); - }, [canViewGlobalSpend, dateValue, selectedTags]); - - const updateEndUserData = async ( - startTime: Date | undefined, - endTime: Date | undefined, - uiSelectedKey: string | null, - ) => { - if (!startTime || !endTime || !accessToken) { - return; - } - - let newTopUserData = await adminTopEndUsersCall( - accessToken, - uiSelectedKey, - startTime.toISOString(), - endTime.toISOString(), - ); - setTopUsers(newTopUserData); - }; - - const updateTagSpendData = async (startTime: Date | undefined, endTime: Date | undefined) => { - if (!startTime || !endTime || !accessToken) { - return; - } - - // we refetch because the state variable can be None when the user refreshes the page - const proxy_settings: ProxySettings | undefined = await fetchProxySettings(); - - if (proxy_settings?.DISABLE_EXPENSIVE_DB_QUERIES) { - return; // Don't run expensive DB queries - return out when SpendLogs has more than 1M rows - } - - let top_tags = await tagsSpendLogsCall( - accessToken, - startTime.toISOString(), - endTime.toISOString(), - selectedTags.length === 0 ? undefined : selectedTags, - ); - setTopTagsData(top_tags.spend_per_tag); - }; - - function formatDate(date: Date) { - const year = date.getFullYear(); - let month = date.getMonth() + 1; // JS month index starts from 0 - let day = date.getDate(); - - // Pad with 0 if month or day is less than 10 - const monthStr = month < 10 ? "0" + month : month; - const dayStr = day < 10 ? "0" + day : day; - - return `${year}-${monthStr}-${dayStr}`; - } - - const valueFormatter = (number: number) => `$ ${formatNumberWithCommas(number, 2)}`; - - const fetchAndSetData = async ( - fetchFunction: () => Promise, - setStateFunction: React.Dispatch>, - errorMessage: string, - ) => { - try { - const data = await fetchFunction(); - setStateFunction(data); - } catch (error) { - console.error(errorMessage, error); - // Optionally, update UI to reflect error state for this specific data - } - }; - - // Update the fillMissingDates function to handle different date formats - const fillMissingDates = (data: any[], startDate: Date, endDate: Date, categories: string[]) => { - const filledData = []; - const currentDate = new Date(startDate); - - // Helper function to standardize date format - const standardizeDate = (dateStr: string) => { - if (dateStr.includes("-")) { - // Already in YYYY-MM-DD format - return dateStr; - } else { - // Convert "Jan 06" format - const [month, day] = dateStr.split(" "); - const year = new Date().getFullYear(); - const monthIndex = new Date(`${month} 01 2024`).getMonth(); - const fullDate = new Date(year, monthIndex, parseInt(day)); - return fullDate.toISOString().split("T")[0]; - } - }; - - // Create a map of existing dates for quick lookup - const existingDates = new Map( - data.map((item) => { - const standardizedDate = standardizeDate(item.date); - return [ - standardizedDate, - { - ...item, - date: standardizedDate, // Store standardized date format - }, - ]; - }), - ); - - // Iterate through each date in the range - while (currentDate <= endDate) { - const dateStr = currentDate.toISOString().split("T")[0]; - - if (existingDates.has(dateStr)) { - // Use existing data if we have it - filledData.push(existingDates.get(dateStr)); - } else { - // Create an entry with zero values - const emptyEntry: any = { - date: dateStr, - api_requests: 0, - total_tokens: 0, - }; - - // Add zero values for each model/team if needed - categories.forEach((category) => { - if (!emptyEntry[category]) { - emptyEntry[category] = 0; - } - }); - - filledData.push(emptyEntry); - } - - // Move to next day - currentDate.setDate(currentDate.getDate() + 1); - } - - return filledData; - }; - - // Update the fetchOverallSpend function - const fetchOverallSpend = async () => { - if (!accessToken) { - return; - } - try { - const data = await adminSpendLogsCall(accessToken); - - // Get the first and last day of the current month - const now = new Date(); - const firstDay = new Date(now.getFullYear(), now.getMonth(), 1); - const lastDay = new Date(now.getFullYear(), now.getMonth() + 1, 0); - - // Fill in missing dates - const filledData = fillMissingDates(data, firstDay, lastDay, []); - - // Calculate total spend for the month and round to 2 decimal places - const monthlyTotal = Number(filledData.reduce((sum, day) => sum + (day.spend || 0), 0).toFixed(2)); - setTotalMonthlySpend(monthlyTotal); - - setKeySpendData(filledData); - } catch (error) { - console.error("Error fetching overall spend:", error); - } - }; - - const fetchProviderSpend = () => - fetchAndSetData( - () => (accessToken ? adminspendByProvider(accessToken, startTime, endTime) : Promise.reject("No access token")), - setSpendByProvider, - "Error fetching provider spend", - ); - - const fetchTopKeys = async () => { - if (!accessToken) return; - await fetchAndSetData( - async () => { - const top_keys = await adminTopKeysCall(accessToken); - return top_keys.map((k: any) => ({ - key: k["api_key"].substring(0, 10), - api_key: k["api_key"], - key_alias: k["key_alias"], - spend: Number(k["total_spend"].toFixed(2)), - })); - }, - setTopKeys, - "Error fetching top keys", - ); - }; - - const fetchTopModels = async () => { - if (!accessToken) return; - await fetchAndSetData( - async () => { - const top_models = await adminTopModelsCall(accessToken); - return top_models.map((k: any) => ({ - key: k["model"], - spend: formatNumberWithCommas(k["total_spend"], 2), - })); - }, - setTopModels, - "Error fetching top models", - ); - }; - - // Update the fetchTeamSpend function - const fetchTeamSpend = async () => { - if (!accessToken) return; - await fetchAndSetData( - async () => { - const teamSpend = await teamSpendLogsCall(accessToken); - - // Get the first and last day of the current month - const now = new Date(); - const firstDay = new Date(now.getFullYear(), now.getMonth(), 1); - const lastDay = new Date(now.getFullYear(), now.getMonth() + 1, 0); - - // Fill in missing dates with zero values for all teams - const filledData = fillMissingDates(teamSpend.daily_spend, firstDay, lastDay, teamSpend.teams); - - setTeamSpendData(filledData); - setUniqueTeamIds(teamSpend.teams); - return teamSpend.total_spend_per_team.map((tspt: any) => ({ - name: tspt["team_id"] || "", - value: Number(tspt["total_spend"] || 0), - })); - }, - setTotalSpendPerTeam, - "Error fetching team spend", - ); - }; - - const fetchTagNames = () => { - if (!accessToken) return; - fetchAndSetData( - async () => { - const all_tag_names = await allTagNamesCall(accessToken); - return all_tag_names.tag_names; - }, - setAllTagNames, - "Error fetching tag names", - ); - }; - - const fetchTopTags = () => { - if (!accessToken) return; - fetchAndSetData( - () => tagsSpendLogsCall(accessToken, dateValue.from?.toISOString(), dateValue.to?.toISOString(), undefined), - (data) => setTopTagsData(data.spend_per_tag), - "Error fetching top tags", - ); - }; - - const fetchTopEndUsers = () => { - if (!accessToken) return; - fetchAndSetData( - () => adminTopEndUsersCall(accessToken, null, undefined, undefined), - setTopUsers, - "Error fetching top end users", - ); - }; - - // Update the fetchGlobalActivity function - const fetchGlobalActivity = async () => { - if (!accessToken) return; - try { - const data = await adminGlobalActivity(accessToken, startTime, endTime); - - // Get the date range from the current month - const now = new Date(); - const firstDay = new Date(now.getFullYear(), now.getMonth(), 1); - const lastDay = new Date(now.getFullYear(), now.getMonth() + 1, 0); - - // Fill in missing dates for daily_data - const filledDailyData = fillMissingDates(data.daily_data || [], firstDay, lastDay, [ - "api_requests", - "total_tokens", - ]); - - setGlobalActivity({ - ...data, - daily_data: filledDailyData, - }); - } catch (error) { - console.error("Error fetching global activity:", error); - } - }; - - // Update the fetchGlobalActivityPerModel function - const fetchGlobalActivityPerModel = async () => { - if (!accessToken) return; - try { - const data = await adminGlobalActivityPerModel(accessToken, startTime, endTime); - - // Get the date range from the current month - const now = new Date(); - const firstDay = new Date(now.getFullYear(), now.getMonth(), 1); - const lastDay = new Date(now.getFullYear(), now.getMonth() + 1, 0); - - // Fill in missing dates for each model's daily data - const filledModelData = data.map((modelData: any) => ({ - ...modelData, - daily_data: fillMissingDates(modelData.daily_data || [], firstDay, lastDay, ["api_requests", "total_tokens"]), - })); - - setGlobalActivityPerModel(filledModelData); - } catch (error) { - console.error("Error fetching global activity per model:", error); - } - }; - - useEffect(() => { - const initlizeUsageData = async () => { - if (!canViewGlobalSpend) { - return; - } - if (accessToken && token && userRole && userID) { - const proxy_settings: ProxySettings | undefined = await fetchProxySettings(); - if (proxy_settings) { - setProxySettings(proxy_settings); // saved in state so it can be used when rendering UI - if (proxy_settings?.DISABLE_EXPENSIVE_DB_QUERIES) { - return; // Don't run expensive UI queries - return out of initlizeUsageData at this point - } - } - - fetchOverallSpend(); - fetchProviderSpend(); - fetchTopKeys(); - fetchTopModels(); - fetchGlobalActivity(); - fetchGlobalActivityPerModel(); - - if (isAdminOrAdminViewer(userRole)) { - fetchTeamSpend(); - fetchTagNames(); - fetchTopTags(); - fetchTopEndUsers(); - } - } - }; - - initlizeUsageData(); - }, [canViewGlobalSpend, accessToken, token, userRole, userID, startTime, endTime]); - - if (!canViewGlobalSpend) { - return ( -
- - - Usage - - -

- Proxy-wide usage is only available to admin users. Your own usage is on the Usage page. -

- - -
- ); - } - - if (proxySettings?.DISABLE_EXPENSIVE_DB_QUERIES) { - return ( -
- - - Database Query Limit Reached - - -

- SpendLogs in DB has {proxySettings.NUM_SPEND_LOGS_ROWS} rows. -

- Please follow our guide to view usage when SpendLogs has more than 1M rows. -

-
- ); - } - - return ( -
- - - All Up - - {isAdminOrAdminViewer(userRole) && ( - <> - Team Based Usage - Customer Usage - Tag Based Usage - - )} - - - - - - Cost - Activity - - - -
-
-

- Project Spend {new Date().toLocaleString("default", { month: "long" })} 1 -{" "} - {new Date(new Date().getFullYear(), new Date().getMonth() + 1, 0).getDate()} -

- -
-
- - - Monthly Spend - - - - - -
-
- - - Top Virtual Keys - - - {}} /> - - -
-
- - - Top Models - - - `$${formatNumberWithCommas(value, 2)}`} - /> - - -
-
-
- - - Spend by Provider - - -
-
- `$${formatNumberWithCommas(value, 2)}`} - /> -
-
- - - - Provider - Spend - - - - {spendByProvider.map((provider) => ( - - {provider.provider} - - - - - ))} - -
-
-
- - -
-
- - - -
- - - All Up - - -
-
-

- API Requests {valueFormatterNumbers(globalActivity.sum_api_requests)} -

- -
-
-

- Tokens {valueFormatterNumbers(globalActivity.sum_total_tokens)} -

- -
-
- - - - {globalActivityPerModel.map((globalActivity, index) => ( - - - {globalActivity.model} - - -
-
-

- API Requests {valueFormatterNumbers(globalActivity.sum_api_requests)} -

- -
-
-

- Tokens {valueFormatterNumbers(globalActivity.sum_total_tokens)} -

- -
-
- - - ))} -
- - - - - -
-
- - - Total Spend Per Team - - - - - - - - Daily Spend Per Team - - - - - -
-
- - - -

- Customers of your LLM API calls. Tracked when a `user` param is passed in your LLM calls{" "} - - docs here - -

-
-
- { - setDateValue(value); - updateEndUserData(value.from, value.to, null); - }} - /> -
-
-

Select Key

- -
-
- - - -
- - - - Customer - Spend - Total Events - - - - - {topUsers?.map((user: any, index: number) => ( - - {user.end_user} - - - - {user.total_count} - - ))} - -
-
- - - - - -
-
- { - setDateValue(value); - updateTagSpendData(value.from, value.to); - }} - /> -
- -
- selectedTags.includes(option.value))} - onValueChange={(options: TagOption[]) => setSelectedTags(options.map((option) => option.value))} - isItemEqualToValue={(a: TagOption, b: TagOption) => a.value === b.value} - itemToStringLabel={(option: TagOption) => option.label} - > - }> - - {(options: TagOption[]) => - options.map((option) => ( - - {option.label} - - )) - } - - - - - No tags found - - {(option: TagOption) => ( - - {option.label} - - )} - - - -
-
-
-
- - - Spend Per Tag - - -

- Get Started by Tracking cost per tag{" "} - - here - -

- - - -
-
- - -
- ); -}; - -export default UsagePage; diff --git a/ui/litellm-dashboard/src/app/(dashboard)/old-usage/page.tsx b/ui/litellm-dashboard/src/app/(dashboard)/old-usage/page.tsx deleted file mode 100644 index 138dd97e5e8..00000000000 --- a/ui/litellm-dashboard/src/app/(dashboard)/old-usage/page.tsx +++ /dev/null @@ -1,22 +0,0 @@ -"use client"; - -import Usage from "./_components/usage"; -import { DeprecationBanner } from "@/components/DeprecationBanner"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; - -export default function OldUsagePage() { - const { accessToken, token, userRole, userId: userID, premiumUser } = useAuthorized(); - return ( - <> - - - - ); -} diff --git a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx index 9d39e9c080c..f7ffd6d52cf 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/usage/_components/components/UsagePageView.test.tsx @@ -102,10 +102,6 @@ vi.mock("@/components/activity_metrics", () => ({ processActivityData: vi.fn(), })); -vi.mock("@/components/view_user_spend", () => ({ - default: () =>
View User Spend
, -})); - vi.mock("@/components/UsagePage/components/EntityUsage/TopKeyView", () => ({ default: () =>
Top Keys
, })); diff --git a/ui/litellm-dashboard/src/components/leftnav.test.tsx b/ui/litellm-dashboard/src/components/leftnav.test.tsx index 6cf05af1ed1..f3bff85e2d5 100644 --- a/ui/litellm-dashboard/src/components/leftnav.test.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.test.tsx @@ -444,30 +444,6 @@ describe("Sidebar (leftnav)", () => { }); expect(screen.queryByText("Prompts")).not.toBeInTheDocument(); }); - - it("should hide Old Usage from internal users while keeping other Experimental children", async () => { - mockUseAuthorized.mockReturnValue(internalAuth); - renderWithProviders(); - - act(() => { - fireEvent.click(screen.getByText("Experimental")); - }); - await waitFor(() => { - expect(screen.getByText("API Playground")).toBeInTheDocument(); - }); - expect(screen.queryByText("Old Usage")).not.toBeInTheDocument(); - }); - - it("should show Old Usage to admins", async () => { - renderWithProviders(); - - act(() => { - fireEvent.click(screen.getByText("Experimental")); - }); - await waitFor(() => { - expect(screen.getByText("Old Usage")).toBeInTheDocument(); - }); - }); }); // Workflow Runs, Memory and Guardrails Monitor render a shell and then 401 @@ -704,7 +680,6 @@ describe("Sidebar (leftnav)", () => { expectHref("Models + Endpoints", "/ui/models-and-endpoints"); expectHref("Usage", "/ui/usage"); expectHref("API Reference", "/ui/api-reference"); - expectHref("Old Usage", "/ui/old-usage"); }); it("never links a leaf to the legacy ?page= switch", () => { @@ -760,7 +735,6 @@ describe("getBreadcrumb", () => { it("resolves routes whose segment differs from the sidebar page id", () => { expect(getBreadcrumb("/ui/models-and-endpoints")).toEqual({ section: "AI Gateway", title: "Models + Endpoints" }); expect(getBreadcrumb("/ui/usage")).toEqual({ section: "Observability", title: "Usage" }); - expect(getBreadcrumb("/ui/old-usage")).toEqual({ section: "Developer Tools", title: "Old Usage" }); }); it("titles the dashboard root as Home", () => { diff --git a/ui/litellm-dashboard/src/components/leftnav.tsx b/ui/litellm-dashboard/src/components/leftnav.tsx index 6795cf1c7d1..2cf6f8219db 100644 --- a/ui/litellm-dashboard/src/components/leftnav.tsx +++ b/ui/litellm-dashboard/src/components/leftnav.tsx @@ -414,14 +414,6 @@ const menuGroups: MenuGroup[] = [ icon: , roles: all_admin_roles, }, - { - key: "4", - page: "usage", - route: "old-usage", - label: "Old Usage", - icon: , - roles: rolesWithCapability("viewGlobalSpend"), - }, ], }, ], diff --git a/ui/litellm-dashboard/src/components/page_metadata.ts b/ui/litellm-dashboard/src/components/page_metadata.ts index 18f136c3b28..97220a0ef76 100644 --- a/ui/litellm-dashboard/src/components/page_metadata.ts +++ b/ui/litellm-dashboard/src/components/page_metadata.ts @@ -41,7 +41,6 @@ export const pageDescriptions: Record = { "tag-management": "Organize resources with tags", prompts: "Manage and version prompt templates", skills: "Browse and manage Claude Code skills", - usage: "View legacy usage dashboard", "router-settings": "Configure routing and load balancing settings", telemetry: "Choose which anonymous usage telemetry this proxy collects and see the last report", "logging-and-alerts": "Set up logging and alert configurations", diff --git a/ui/litellm-dashboard/src/components/view_user_spend.tsx b/ui/litellm-dashboard/src/components/view_user_spend.tsx deleted file mode 100644 index 655fc092673..00000000000 --- a/ui/litellm-dashboard/src/components/view_user_spend.tsx +++ /dev/null @@ -1,144 +0,0 @@ -"use client"; -import React, { useEffect, useState } from "react"; -import { modelAvailableCall } from "./networking"; -import { formatNumberWithCommas } from "@/utils/dataUtils"; -import useAuthorized from "@/app/(dashboard)/hooks/useAuthorized"; - -// Define the props type -interface ViewUserSpendProps { - userSpend: number | null; - userMaxBudget: number | null; - selectedTeam: any | null; -} -const ViewUserSpend: React.FC = ({ userSpend, userMaxBudget, selectedTeam }) => { - const { accessToken, userRole, userId: userID } = useAuthorized(); - let [spend, setSpend] = useState(userSpend !== null ? userSpend : 0.0); - const [maxBudget, setMaxBudget] = useState( - selectedTeam ? Number(formatNumberWithCommas(selectedTeam.max_budget, 4)) : null, - ); - useEffect(() => { - if (selectedTeam) { - if (selectedTeam.team_alias === "Default Team") { - setMaxBudget(userMaxBudget); - } else { - let setMaxBudgetFlag = false; - if (selectedTeam.team_memberships) { - /** - * What 'team_memberships' looks like: - * "team_memberships": [ - * { - * "user_id": "2c315de3-e7ce-4269-b73e-b039a06187b1", - * "team_id": "test-team_515e6f42-ded2-4f0d-8919-0a1f43c5a45f", - * "budget_id": "0880769f-716a-4149-ab19-7f7651ad4db5", - * "litellm_budget_table": { - "soft_budget": null, - "max_budget": 20.0, - "max_parallel_requests": null, - "tpm_limit": null, - "rpm_limit": null, - "model_max_budget": null, - "budget_duration": null - } - */ - for (const member of selectedTeam.team_memberships) { - if ( - member.user_id === userID && - "max_budget" in member.litellm_budget_table && - member.litellm_budget_table.max_budget !== null - ) { - setMaxBudget(member.litellm_budget_table.max_budget); - setMaxBudgetFlag = true; - } - } - } - if (!setMaxBudgetFlag) { - setMaxBudget(selectedTeam.max_budget); - } - } - } else { - setMaxBudget(userMaxBudget); - } - }, [selectedTeam, userMaxBudget]); - const [userModels, setUserModels] = useState([]); - useEffect(() => { - const fetchData = async () => { - if (!accessToken || !userID || !userRole) { - return; - } - }; - const fetchUserModels = async () => { - try { - if (userID === null || userRole === null) { - return; - } - - if (accessToken !== null) { - const model_available = await modelAvailableCall(accessToken, userID, userRole); - let available_model_names = model_available["data"].map((element: { id: string }) => element.id); - setUserModels(available_model_names); - } - } catch (error) { - console.error("Error fetching user models:", error); - } - }; - - fetchUserModels(); - fetchData(); - }, [userRole, accessToken, userID]); - - useEffect(() => { - if (userSpend !== null) { - setSpend(userSpend); - } - }, [userSpend]); - - // logic to decide what models to display - let modelsToDisplay = []; - if (selectedTeam && selectedTeam.models) { - modelsToDisplay = selectedTeam.models; - } - - // check if "all-proxy-models" is in modelsToDisplay - if (modelsToDisplay && modelsToDisplay.includes("all-proxy-models")) { - modelsToDisplay = userModels; - } else if (modelsToDisplay && modelsToDisplay.includes("all-team-models")) { - modelsToDisplay = selectedTeam.models; - } else if (modelsToDisplay && modelsToDisplay.length === 0) { - modelsToDisplay = userModels; - } - - const displayMaxBudget = maxBudget !== null ? `$${formatNumberWithCommas(Number(maxBudget), 4)} limit` : "No limit"; - - const roundedSpend = spend !== undefined ? formatNumberWithCommas(spend, 4) : null; - - return ( -
-
-
-

Total Spend

-

${roundedSpend}

-
-
-

Max Budget

-

{displayMaxBudget}

-
-
- {/*
- - Team Models - - - {modelsToDisplay.map((model: string) => ( - - {model} - - ))} - - - -
*/} -
- ); -}; - -export default ViewUserSpend; diff --git a/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx b/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx index 58d7309c078..e5efd4c8e9f 100644 --- a/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx +++ b/ui/litellm-dashboard/tests/CreateKeyPage.expiredToken.test.tsx @@ -127,7 +127,6 @@ vi.mock("@/components/new_usage", () => ({ default: stub("new-usage") })); vi.mock("@/components/api_ref", () => ({ default: stub("api-ref") })); vi.mock("@/components/chat_ui/ChatUI", () => ({ default: stub("chat-ui") })); vi.mock("@/components/leftnav", () => ({ default: stub("sidebar") })); -vi.mock("@/app/(dashboard)/old-usage/_components/usage", () => ({ default: stub("usage") })); vi.mock("@/components/cache_dashboard", () => ({ default: stub("cache-dashboard") })); vi.mock("@/app/(dashboard)/guardrails/_components", () => ({ default: stub("guardrails") })); vi.mock("@/components/prompts", () => ({ default: stub("prompts") })); From e752528ac8f9e1399e0e5bff98b7f404bffda1b4 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 11:53:06 -0700 Subject: [PATCH 06/46] feat(keys): allow assigning a project to an unassigned key on /key/update (#41700) * feat(keys): allow assigning a project to an unassigned key on /key/update Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(keys): suppress LIT002 on project assignment error payloads Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(keys): format project assignment error detail Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(keys): assert accepted project assignment flows into update payload Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(keys): drop stale LIT002 suppressions removed by main Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(keys): mock writer_db team lookup used by check_db_only path Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(keys): write project assignment only while key row is still unassigned Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(keys): satisfy type gate on guarded project assignment write Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(keys): default soft-budget update path to unguarded write Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(keys): run guarded project assignment in one writer transaction Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(keys): include object permission link in assignment guard Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * refactor(keys): return permission-linked payload without mutating the update dict Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(keys): fix type errors and stale noqa after main merge Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(keys): use Sequence in assignment guard model annotation Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(keys): type update fields as a mapping in the key update path Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(keys): place cast reason on the cast call Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(keys): honor explicit null max_budget when assigning a project * test(keys): assert explicit null max_budget flows into update payload * fix(keys): match NULL models column in guarded project assignment * fix(keys): authorize caller before project lookup on /key/update * fix(keys): drop explicit null object_permission before project assignment write --------- Co-authored-by: jesus Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../key_management_endpoints.py | 249 +++++++- .../management_endpoints/key_requests.py | 2 +- .../test_key_management_endpoints.py | 530 +++++++++++++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 +- 4 files changed, 759 insertions(+), 26 deletions(-) diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 3021b21de51..61da5599886 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -162,7 +162,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import ( # noqa: F401 # _is_master_key, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export is_master_key, ) -from litellm.proxy.utils import ( # noqa: F401 # legacy module exports +from litellm.proxy.utils import ( PrismaClient, ProxyLogging, _hash_token_if_needed, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -280,6 +280,10 @@ class _BudgetRowSoftBudgetCreate(TypedDict): updated_by: ReadOnly[str] +class _KeyRowDumpable(Protocol): + def model_dump(self) -> Mapping[str, object]: ... + + class _KeyUpdateTx(Protocol): @property def litellm_verificationtoken(self) -> "TableActions[prisma_models.LiteLLM_VerificationToken]": ... @@ -287,6 +291,9 @@ class _KeyUpdateTx(Protocol): @property def litellm_budgettable(self) -> "TableActions[prisma_models.LiteLLM_BudgetTable]": ... + @property + def litellm_objectpermissiontable(self) -> "TableActions[prisma_models.LiteLLM_ObjectPermissionTable]": ... + class _ConfigTableActions(Protocol): """Config table surface this module needs; the shared repository seam exposes no ``update``.""" @@ -2072,6 +2079,56 @@ async def _check_project_key_limits( ) +async def _validate_project_assignment( + data: UpdateKeyRequest, + existing_key_row: LiteLLM_VerificationToken, + prisma_client: PrismaClient, + user_api_key_cache: UserApiKeyCache, +) -> None: + """Validate assigning a project to a key that does not have one yet.""" + if data.project_id is None or data.project_id == existing_key_row.project_id: + return + if existing_key_row.project_id is not None: + raise HTTPException( + status_code=400, detail="Project reassignment is not supported. Use null to detach the key." + ) + project_obj: Final = await get_project_object( + project_id=data.project_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + if project_obj is None: + raise HTTPException( + status_code=404, + detail={"error": f"Project not found, project_id={data.project_id}"}, + ) + team: Final = data.team_id if "team_id" in data.model_fields_set else existing_key_row.team_id + if team is None or team != project_obj.team_id: + raise HTTPException( + status_code=400, + detail={ + "error": f"Project {data.project_id} belongs to team {project_obj.team_id}. Keys can only be assigned to a project owned by their own team (key team: {team})" + }, + ) + await _check_project_key_limits( + project_id=data.project_id, + data=data.model_copy( + update={ + "models": ( + data.models + if "models" in data.model_fields_set + else cast( # cast-ok: LiteLLM_VerificationToken.models is a bare list + "list[str] | None", existing_key_row.models + ) + ), + "max_budget": data.max_budget if "max_budget" in data.model_fields_set else existing_key_row.max_budget, + } + ), + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + ) + + async def _check_key_project_team( project_id: str, key_team_id: str | None, @@ -2771,6 +2828,115 @@ async def _apply_soft_budget_update( return remaining +async def _apply_object_permission_update( + data_json: Mapping[str, object], + existing_object_permission_id: str | None, + prisma_client: PrismaClient, + table: "TableActions[prisma_models.LiteLLM_ObjectPermissionTable]", +) -> Mapping[str, object]: + """Upsert the requested object permission row inside the caller's writer transaction.""" + new_object_permission: Final = data_json.get("object_permission") + if new_object_permission is None: + return MappingProxyType({k: v for k, v in data_json.items() if k != "object_permission"}) + loaded: Final[object] = ( + json.loads(new_object_permission) if isinstance(new_object_permission, str) else new_object_permission + ) + grants: Final[Mapping[str, object]] = ( + cast( # cast-ok: object_permission payloads are grant-name to grant-list mappings + "Mapping[str, object]", loaded + ) + if isinstance(loaded, dict) + else MappingProxyType({}) + ) + upsert: Final = await prepare_object_permission_upsert( + new_object_permission=grants, + existing_object_permission_id=existing_object_permission_id, + prisma_client=prisma_client, + ) + row: Final = await table.upsert( + where={"object_permission_id": upsert.object_permission_id}, + data={"create": upsert.record, "update": upsert.record}, + ) + return MappingProxyType( + { + **{k: v for k, v in data_json.items() if k != "object_permission"}, + "object_permission_id": row.object_permission_id, + } + ) + + +async def _write_guarded_project_assignment( + table: "TableActions[_KeyRowDumpable]", + hashed_token: str, + existing_key_row: LiteLLM_VerificationToken, + data: Mapping[str, object], +) -> "_KeyRowDumpable | None": + validated_models: Final[Sequence[str]] = ( + cast( # cast-ok: LiteLLM_VerificationToken.models is a bare list + "Sequence[str] | None", existing_key_row.models + ) + or [] + ) + models_guard: Final[Mapping[str, object]] = ( + {"OR": [{"models": {"equals": list[str]()}}, {"models": {"equals": None}}]} + if not validated_models + else {"models": {"equals": validated_models}} + ) + updated_count: Final = await table.update_many( + where={ + "token": hashed_token, + "project_id": None, + "team_id": existing_key_row.team_id, + **models_guard, + "max_budget": existing_key_row.max_budget, + "object_permission_id": existing_key_row.object_permission_id, + }, + data=data, + ) + if updated_count == 0: + raise HTTPException( + status_code=400, + detail="Project assignment failed because the key was updated concurrently. Retry the request.", + ) + return await table.find_unique( + where={"token": hashed_token}, + include={"object_permission": True}, + ) + + +async def _update_key_row_assigning_project( + prisma_client: PrismaClient, + key: str, + non_default_values: Mapping[str, object], + existing_key_row: LiteLLM_VerificationToken, +) -> _KeyUpdateResult: + """Assign a project in one writer transaction, guarded on the key's validated state.""" + hashed_token: Final = _hash_token_if_needed(key) + async with prisma_client.tx() as tx_ctx: + tx: Final[_KeyUpdateTx] = cast( # cast-ok: the transaction object exposes the same table actions + "_KeyUpdateTx", tx_ctx + ) + update_values: Final[Mapping[str, object]] = await _apply_object_permission_update( + data_json=non_default_values, + existing_object_permission_id=existing_key_row.object_permission_id, + prisma_client=prisma_client, + table=tx.litellm_objectpermissiontable, + ) + updated_row: Final = await _write_guarded_project_assignment( + table=tx.litellm_verificationtoken, + hashed_token=hashed_token, + existing_key_row=existing_key_row, + data=with_settings_updated_at( + prisma_client.jsonify_object(MappingProxyType({**update_values, "token": hashed_token})) + ), + ) + updated_data: Final[Mapping[str, object]] = ( + updated_row.model_dump() if updated_row is not None else MappingProxyType({}) + ) + result: Final[_KeyUpdateResult] = {"token": hashed_token, "data": updated_data} + return result + + async def _update_key_row_with_soft_budget( prisma_client: PrismaClient, key: str, @@ -2778,25 +2944,48 @@ async def _update_key_row_with_soft_budget( non_default_values: Mapping[str, object], existing_key_row: LiteLLM_VerificationToken, changed_by: str, + expect_unassigned_project: bool = False, ) -> _KeyUpdateResult: hashed_token: Final = hash_token_if_needed(key) key_where: Final[_KeyRowWhere] = {"token": hashed_token} - tx: _KeyUpdateTx - async with prisma_client.tx() as tx: + async with prisma_client.tx() as tx_ctx: + tx: Final[_KeyUpdateTx] = cast( # cast-ok: the transaction object exposes the same table actions + "_KeyUpdateTx", tx_ctx + ) + update_input: Final[Mapping[str, object]] = ( + await _apply_object_permission_update( + data_json=non_default_values, + existing_object_permission_id=existing_key_row.object_permission_id, + prisma_client=prisma_client, + table=tx.litellm_objectpermissiontable, + ) + if expect_unassigned_project + else non_default_values + ) update_values: Final = await _apply_soft_budget_update( data=data, - non_default_values=non_default_values, + non_default_values=update_input, db=tx, existing_key_row=existing_key_row, changed_by=changed_by, ) include_object_permission: Final[prisma.types.LiteLLM_VerificationTokenInclude] = {"object_permission": True} - updated_row: Final = await tx.litellm_verificationtoken.update( - where=key_where, - data=with_settings_updated_at( - prisma_client.jsonify_object(MappingProxyType({**update_values, "token": hashed_token})) - ), - include=include_object_permission, + update_payload: Final = with_settings_updated_at( + prisma_client.jsonify_object(MappingProxyType({**update_values, "token": hashed_token})) + ) + updated_row: Final = ( + await _write_guarded_project_assignment( + table=tx.litellm_verificationtoken, + hashed_token=hashed_token, + existing_key_row=existing_key_row, + data=update_payload, + ) + if expect_unassigned_project + else await tx.litellm_verificationtoken.update( + where=key_where, + data=update_payload, + include=include_object_permission, + ) ) updated_data: Final[Mapping[str, object]] = ( updated_row.model_dump() if updated_row is not None else MappingProxyType({}) @@ -3563,10 +3752,6 @@ async def _validate_update_key_data( user_api_key_dict=user_api_key_dict, ) - if data.project_id is not None and data.project_id != existing_key_row.project_id: - raise HTTPException( - status_code=400, detail="Project reassignment is not supported. Use null to detach the key." - ) is_project_change: Final = "project_id" in data.model_fields_set and data.project_id != existing_key_row.project_id acting_as_team_admin: Final = await _acting_as_team_admin_for_key_update( @@ -3698,6 +3883,12 @@ async def _validate_update_key_data( user_api_key_cache=user_api_key_cache, route=("/key/update (max_budget/spend)" if _is_budget_change else "/key/update"), ) + await _validate_project_assignment( + data=data, + existing_key_row=existing_key_row, + prisma_client=checked_prisma_client, + user_api_key_cache=user_api_key_cache, + ) check_denied_passthrough_routes_caller_permission( data, user_api_key_dict, @@ -3847,7 +4038,7 @@ async def update_key_fn( - user_id: Optional[str] - User ID associated with key - team_id: Optional[str] - Team ID associated with key - agent_id: Optional[str] - The agent id associated with the key. - - project_id: Optional[str] - Omit to retain the project, or send null to detach. A different project ID is rejected. + - project_id: Optional[str] - Omit to retain the project, send null to detach, or send a project id to assign an unassigned key to a project on the key's team. Moving a key between projects is rejected. - organization_id: Optional[str] - The organization id of the key. - budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`. - end_user_budget_id: Optional[str] - Proxy admin only. Budget id applied to end users first seen through this key that carry no budget of their own. Omit to keep the current value, pass an empty string to clear it. @@ -3995,23 +4186,31 @@ async def update_key_fn( if prisma_client is None: raise Exception("Not connected to DB!") + is_project_assignment: Final = data.project_id is not None and existing_key_row.project_id is None + update_fields: Final[Mapping[str, object]] = cast( # cast-ok: prepare_key_update_data returns a bare dict + "Mapping[str, object]", non_default_values + ) object_permission_upsert: Final = await _prepare_key_update_object_permission( - object_permission_data=non_default_values.get("object_permission"), + object_permission_data=update_fields.get("object_permission"), existing_key_row=existing_key_row, prisma_client=prisma_client, ) key_update_values: Final = MappingProxyType( - {field: value for field, value in non_default_values.items() if field != "object_permission"} + {field: value for field, value in update_fields.items() if field != "object_permission"} ) await _check_key_project_team_on_mutation( data=data, existing_key_row=existing_key_row, prisma_client=prisma_client, ) - update_values: Final = await _write_prepared_key_update_object_permission( - data_json=key_update_values, - upsert=object_permission_upsert, - prisma_client=prisma_client, + update_values: Final = ( + update_fields + if is_project_assignment + else await _write_prepared_key_update_object_permission( + data_json=key_update_values, + upsert=object_permission_upsert, + prisma_client=prisma_client, + ) ) changed_by: Final = user_api_key_dict.user_id or litellm_proxy_admin_name response: Final = ( @@ -4022,8 +4221,16 @@ async def update_key_fn( non_default_values=update_values, existing_key_row=existing_key_row, changed_by=changed_by, + expect_unassigned_project=is_project_assignment, ) if "soft_budget" in data.model_fields_set + else await _update_key_row_assigning_project( + prisma_client=prisma_client, + key=key, + non_default_values=update_fields, + existing_key_row=existing_key_row, + ) + if is_project_assignment else await _update_key_row( prisma_client=prisma_client, key=key, diff --git a/litellm/types/proxy/management_endpoints/key_requests.py b/litellm/types/proxy/management_endpoints/key_requests.py index b49cb63c8e0..f42a412cf2d 100644 --- a/litellm/types/proxy/management_endpoints/key_requests.py +++ b/litellm/types/proxy/management_endpoints/key_requests.py @@ -118,7 +118,7 @@ class UpdateKeyRequest(KeyRequestBase): project_id: str | None = Field( default=None, - description="Omit to retain the project, or send null to detach. Assigning a different project is not supported.", + description="Omit to retain the project, send null to detach, or send a project id to assign an unassigned key to a project on the key's team. Moving a key between projects is not supported.", ) @model_validator(mode="before") diff --git a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py index c26a4bf19b8..0198f582be7 100644 --- a/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_key_management_endpoints.py @@ -67,6 +67,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( _enforce_upperbound_key_params, _execute_virtual_key_regeneration, _get_and_validate_existing_key, + _hash_token_if_needed, _is_team_key, _key_generation_required_param_check, _member_key_request, @@ -77,6 +78,7 @@ from litellm.proxy.management_endpoints.key_management_endpoints import ( _team_admin_requests_member_key, _team_key_generation_check, _transform_verification_tokens_to_deleted_records, + _update_key_row_assigning_project, _validate_end_user_budget_id_change, _validate_max_budget, _validate_reset_spend_value, @@ -20975,12 +20977,20 @@ async def test_regenerate_key_repoints_live_membership_not_the_key_row_it_read( async def _cache_with_project( - project_id: str, project_models: list[str], team_id: str | None = "team-lit-5823" + project_id: str, project_models: list[str], team_id: str | None = "team-lit-5823", + project_max_budget: float | None = None, ) -> UserApiKeyCache: user_api_key_cache = UserApiKeyCache() await user_api_key_cache.async_set_cache( key=project_cache_key(project_id), - value=LiteLLM_ProjectTableCachedObj(project_id=project_id, team_id=team_id, models=project_models), + value=LiteLLM_ProjectTableCachedObj( + project_id=project_id, + team_id=team_id, + models=project_models, + litellm_budget_table=( + LiteLLM_BudgetTable(max_budget=project_max_budget) if project_max_budget is not None else None + ), + ), model_type=LiteLLM_ProjectTableCachedObj, ) return user_api_key_cache @@ -21413,6 +21423,487 @@ async def test_project_detachment_uses_effective_project_for_validation_on_unown assert expected in str(exc.value.detail) +def _project_assignment_prisma(stored_row: Mapping[str, object] | None = None) -> MagicMock: + database = MagicMock() + team_row: Final = LiteLLM_TeamTable(team_id="team-lit-5823", members=[]) + database.db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + database.writer_db.litellm_teamtable.find_unique = AsyncMock(return_value=team_row) + if stored_row is None: + return database + + database.jsonify_object = lambda data: dict(data) + tx: Final = database.tx.return_value.__aenter__.return_value + + def _where_matches(where: Mapping[str, object], row: Mapping[str, object]) -> bool: + for key, clause in where.items(): + if key == "OR": + if not any(_where_matches(part, row) for part in clause): + return False + elif key == "models": + equals = clause["equals"] + if equals is None: + if row[key] is not None: + return False + elif row[key] is None or list(row[key]) != list(equals): + return False + elif row.get(key) != clause: + return False + return True + + async def _update_many(*, where, data): + if "object_permission" in data: + raise AssertionError("object_permission must not reach update_many") + return 1 if _where_matches(where, stored_row) else 0 + + tx.litellm_verificationtoken.update_many = AsyncMock(side_effect=_update_many) + written = MagicMock() + written.model_dump = MagicMock(return_value={**stored_row, "project_id": "project-orbit"}) + tx.litellm_verificationtoken.find_unique = AsyncMock(return_value=written) + return database + + +@pytest.mark.asyncio +async def test_project_assignment_to_unassigned_key_on_same_team(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", models=["model-orbit"] + ) + cache: Final = await _cache_with_project("project-orbit", ["model-orbit"]) + data: Final = UpdateKeyRequest(key=existing.token, project_id="project-orbit") + + await _validate_update_key_data( + data, existing, + UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, _project_assignment_prisma(), cache, + ) + + update: Final = await prepare_key_update_data(data=data, existing_key_row=existing) + assert update["project_id"] == "project-orbit" + + +@pytest.mark.asyncio +async def test_project_assignment_rejects_project_on_another_team(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-other", models=["model-orbit"] + ) + cache: Final = await _cache_with_project("project-orbit", ["model-orbit"]) + + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + UpdateKeyRequest(key=existing.token, project_id="project-orbit"), existing, + UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, _project_assignment_prisma(), cache, + ) + assert exc.value.status_code == 400 + assert "team" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_project_assignment_rejects_key_with_no_team(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id=None, models=["model-orbit"] + ) + cache: Final = await _cache_with_project("project-orbit", ["model-orbit"]) + + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + UpdateKeyRequest(key=existing.token, project_id="project-orbit"), existing, + UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, _project_assignment_prisma(), cache, + ) + assert exc.value.status_code == 400 + assert "team" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_project_assignment_rejects_nonexistent_project(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", models=["model-orbit"] + ) + database: Final = _project_assignment_prisma() + database.db.litellm_projecttable.find_unique = AsyncMock(return_value=None) + + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + UpdateKeyRequest(key=existing.token, project_id="project-ghost"), existing, + UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, database, UserApiKeyCache(), + ) + assert exc.value.status_code == 404 + assert "Project not found" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_project_assignment_validates_existing_models_against_project(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", models=["model-elsewhere"] + ) + cache: Final = await _cache_with_project("project-orbit", ["model-orbit"]) + + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + UpdateKeyRequest(key=existing.token, project_id="project-orbit"), existing, + UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, _project_assignment_prisma(), cache, + ) + assert exc.value.status_code == 400 + assert "not in project's allowed models" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_project_assignment_with_explicit_null_max_budget_ignores_stored_budget(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", models=[], max_budget=100.0 + ) + cache: Final = await _cache_with_project("project-orbit", [], project_max_budget=10.0) + + request: Final = UpdateKeyRequest(key=existing.token, project_id="project-orbit", max_budget=None) + await _validate_update_key_data( + request, existing, + UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, _project_assignment_prisma(), cache, + ) + + result: Final = await prepare_key_update_data(data=request, existing_key_row=existing) + assert result["max_budget"] is None + assert result["project_id"] == "project-orbit" + + +@pytest.mark.asyncio +async def test_project_assignment_without_max_budget_uses_stored_budget(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", models=[], max_budget=100.0 + ) + cache: Final = await _cache_with_project("project-orbit", [], project_max_budget=10.0) + + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + UpdateKeyRequest(key=existing.token, project_id="project-orbit"), existing, + UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN), + None, False, _project_assignment_prisma(), cache, + ) + assert exc.value.status_code == 400 + assert "exceeds project's max_budget" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_prepare_key_update_data_includes_new_project_assignment(): + existing: Final = LiteLLM_VerificationToken(token="project-assign-token", project_id=None) + + result: Final = await prepare_key_update_data( + data=UpdateKeyRequest(key=existing.token, project_id="project-orbit"), existing_key_row=existing, + ) + + assert result["project_id"] == "project-orbit" + + +@pytest.mark.asyncio +async def test_project_assignment_write_requires_row_still_unassigned(): + database = MagicMock() + database.jsonify_object = lambda data: dict(data) + tx: Final = database.tx.return_value.__aenter__.return_value + tx.litellm_verificationtoken.update_many = AsyncMock(return_value=1) + row = MagicMock() + row.model_dump = MagicMock(return_value={"project_id": "project-orbit"}) + tx.litellm_verificationtoken.find_unique = AsyncMock(return_value=row) + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", models=["model-orbit"] + ) + + result: Final = await _update_key_row_assigning_project( + prisma_client=database, + key="sk-assign", + non_default_values={"project_id": "project-orbit"}, + existing_key_row=existing, + ) + + update_where: Final = tx.litellm_verificationtoken.update_many.await_args.kwargs["where"] + assert update_where["token"] == result["token"] + assert result["data"] == {"project_id": "project-orbit"} + + +@pytest.mark.asyncio +async def test_project_assignment_write_rejects_when_row_gained_project(): + database = MagicMock() + database.jsonify_object = lambda data: dict(data) + tx: Final = database.tx.return_value.__aenter__.return_value + tx.litellm_verificationtoken.update_many = AsyncMock(return_value=0) + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", models=["model-orbit"] + ) + + with pytest.raises(HTTPException) as exc: + await _update_key_row_assigning_project( + prisma_client=database, + key="sk-assign", + non_default_values={"project_id": "project-orbit"}, + existing_key_row=existing, + ) + + assert exc.value.status_code == 400 + assert "concurrently" in str(exc.value.detail) + tx.litellm_verificationtoken.find_unique.assert_not_called() + + +@pytest.mark.asyncio +async def test_project_assignment_write_upserts_object_permission_in_same_tx(): + database = MagicMock() + database.jsonify_object = lambda data: dict(data) + database.db.litellm_objectpermissiontable.find_unique = AsyncMock(return_value=None) + tx: Final = database.tx.return_value.__aenter__.return_value + tx.litellm_verificationtoken.update_many = AsyncMock(return_value=1) + permission_row = MagicMock() + permission_row.object_permission_id = "perm-lit-01" + tx.litellm_objectpermissiontable.upsert = AsyncMock(return_value=permission_row) + row = MagicMock() + row.model_dump = MagicMock(return_value={"project_id": "project-orbit"}) + tx.litellm_verificationtoken.find_unique = AsyncMock(return_value=row) + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823" + ) + + result: Final = await _update_key_row_assigning_project( + prisma_client=database, + key="sk-assign", + non_default_values={"project_id": "project-orbit", "object_permission": {"agents": ["agent-a"]}}, + existing_key_row=existing, + ) + + tx.litellm_objectpermissiontable.upsert.assert_awaited_once() + update_data: Final = tx.litellm_verificationtoken.update_many.await_args.kwargs["data"] + assert update_data["object_permission_id"] == "perm-lit-01" + assert "object_permission" not in update_data + assert result["data"] == {"project_id": "project-orbit"} + + +@pytest.mark.asyncio +async def test_project_assignment_write_matches_null_models_column(): + stored: Final = { + "token": _hash_token_if_needed("sk-assign"), + "project_id": None, + "team_id": "team-lit-5823", + "models": None, + "max_budget": None, + "object_permission_id": None, + } + database: Final = _project_assignment_prisma(stored_row=stored) + existing: Final = LiteLLM_VerificationToken( + token=stored["token"], project_id=None, team_id="team-lit-5823", models=[] + ) + + result: Final = await _update_key_row_assigning_project( + prisma_client=database, + key="sk-assign", + non_default_values={"project_id": "project-orbit"}, + existing_key_row=existing, + ) + + assert result["data"]["project_id"] == "project-orbit" + + +@pytest.mark.asyncio +async def test_project_assignment_write_strips_explicit_null_object_permission(): + stored: Final = { + "token": _hash_token_if_needed("sk-assign"), + "project_id": None, + "team_id": "team-lit-5823", + "models": None, + "max_budget": None, + "object_permission_id": None, + } + database: Final = _project_assignment_prisma(stored_row=stored) + existing: Final = LiteLLM_VerificationToken( + token=stored["token"], project_id=None, team_id="team-lit-5823", models=[] + ) + update_values: Final = await prepare_key_update_data( + data=UpdateKeyRequest(key="sk-assign", project_id="project-orbit", object_permission=None), + existing_key_row=existing, + ) + assert "object_permission" in update_values + + result: Final = await _update_key_row_assigning_project( + prisma_client=database, + key="sk-assign", + non_default_values=update_values, + existing_key_row=existing, + ) + + assert result["data"]["project_id"] == "project-orbit" + + +@pytest.mark.asyncio +async def test_project_assignment_write_rejects_when_max_budget_changed(): + stored: Final = { + "token": _hash_token_if_needed("sk-assign"), + "project_id": None, + "team_id": "team-lit-5823", + "models": ["model-orbit"], + "max_budget": 50.0, + "object_permission_id": None, + } + database: Final = _project_assignment_prisma(stored_row=stored) + existing: Final = LiteLLM_VerificationToken( + token=stored["token"], project_id=None, team_id="team-lit-5823", + models=["model-orbit"], max_budget=100.0, + ) + + with pytest.raises(HTTPException) as exc: + await _update_key_row_assigning_project( + prisma_client=database, + key="sk-assign", + non_default_values={"project_id": "project-orbit"}, + existing_key_row=existing, + ) + + assert exc.value.status_code == 400 + assert "concurrently" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_project_assignment_write_rejects_when_object_permission_changed(): + stored: Final = { + "token": _hash_token_if_needed("sk-assign"), + "project_id": None, + "team_id": "team-lit-5823", + "models": ["model-orbit"], + "max_budget": None, + "object_permission_id": "perm-new", + } + database: Final = _project_assignment_prisma(stored_row=stored) + existing: Final = LiteLLM_VerificationToken( + token=stored["token"], project_id=None, team_id="team-lit-5823", models=["model-orbit"] + ) + + with pytest.raises(HTTPException) as exc: + await _update_key_row_assigning_project( + prisma_client=database, + key="sk-assign", + non_default_values={"project_id": "project-orbit"}, + existing_key_row=existing, + ) + + assert exc.value.status_code == 400 + assert "concurrently" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_project_assignment_write_rejects_when_team_changed(): + stored: Final = { + "token": _hash_token_if_needed("sk-assign"), + "project_id": None, + "team_id": "team-moved", + "models": ["model-orbit"], + "max_budget": None, + "object_permission_id": None, + } + database: Final = _project_assignment_prisma(stored_row=stored) + existing: Final = LiteLLM_VerificationToken( + token=stored["token"], project_id=None, team_id="team-lit-5823", models=["model-orbit"] + ) + + with pytest.raises(HTTPException) as exc: + await _update_key_row_assigning_project( + prisma_client=database, + key="sk-assign", + non_default_values={"project_id": "project-orbit"}, + existing_key_row=existing, + ) + + assert exc.value.status_code == 400 + assert "concurrently" in str(exc.value.detail) + + +def _assignment_project_update_prisma(stored_row: Mapping[str, object]) -> MagicMock: + database: Final = _project_assignment_prisma(stored_row=stored_row) + project_row: Final = LiteLLM_ProjectTable(project_id="project-orbit", team_id="team-lit-5823") + database.db.litellm_projecttable.find_unique = AsyncMock(return_value=project_row) + database.writer_db.litellm_projecttable.find_unique = AsyncMock(return_value=project_row) + return database + + +def _setup_assignment_update_mocks(monkeypatch: pytest.MonkeyPatch, database: MagicMock, cache: UserApiKeyCache) -> None: + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", database) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", cache) + proxy_logging_obj: Final = MagicMock() + proxy_logging_obj.service_logging_obj.async_service_success_hook = AsyncMock() + proxy_logging_obj.service_logging_obj.async_service_failure_hook = AsyncMock() + monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", proxy_logging_obj) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + monkeypatch.setattr("litellm.store_audit_logs", False) + + +@pytest.mark.asyncio +async def test_update_key_fn_assigns_project_through_guarded_write(monkeypatch: pytest.MonkeyPatch): + stored: Final = { + "token": _hash_token_if_needed("sk-assign"), + "project_id": None, + "team_id": "team-lit-5823", + "models": ["model-orbit"], + "max_budget": None, + "object_permission_id": None, + } + existing: Final = LiteLLM_VerificationToken( + token=stored["token"], project_id=None, team_id="team-lit-5823", models=["model-orbit"] + ) + database: Final = _assignment_project_update_prisma(stored) + database.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=existing) + cache: Final = await _cache_with_project("project-orbit", ["model-orbit"]) + _setup_assignment_update_mocks(monkeypatch, database, cache) + + result: Final = await update_key_fn( + request=MagicMock(), + data=UpdateKeyRequest(key="sk-assign", project_id="project-orbit"), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + litellm_changed_by=None, + ) + + tx: Final = database.tx.return_value.__aenter__.return_value + tx.litellm_verificationtoken.update_many.assert_awaited_once() + update_where: Final = tx.litellm_verificationtoken.update_many.await_args.kwargs["where"] + assert update_where["project_id"] is None + tx.litellm_verificationtoken.update.assert_not_called() + assert result["project_id"] == "project-orbit" + + +@pytest.mark.asyncio +async def test_update_key_fn_soft_budget_runs_guarded_assignment(monkeypatch: pytest.MonkeyPatch): + stored: Final = { + "token": _hash_token_if_needed("sk-assign"), + "project_id": None, + "team_id": "team-lit-5823", + "models": ["model-orbit"], + "max_budget": None, + "object_permission_id": None, + } + existing: Final = LiteLLM_VerificationToken( + token=stored["token"], project_id=None, team_id="team-lit-5823", models=["model-orbit"] + ) + database: Final = _assignment_project_update_prisma(stored) + database.db.litellm_verificationtoken.find_unique = AsyncMock(return_value=existing) + tx: Final = database.tx.return_value.__aenter__.return_value + tx.litellm_budgettable.create = AsyncMock(return_value=SimpleNamespace(budget_id="budget-new")) + cache: Final = await _cache_with_project("project-orbit", ["model-orbit"]) + _setup_assignment_update_mocks(monkeypatch, database, cache) + + result: Final = await update_key_fn( + request=MagicMock(), + data=UpdateKeyRequest(key="sk-assign", project_id="project-orbit", soft_budget=5.0), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, api_key="sk-admin", user_id="admin-user" + ), + litellm_changed_by=None, + ) + + tx.litellm_budgettable.create.assert_awaited_once() + tx.litellm_verificationtoken.update_many.assert_awaited_once() + update_where: Final = tx.litellm_verificationtoken.update_many.await_args.kwargs["where"] + assert update_where["project_id"] is None + tx.litellm_verificationtoken.update.assert_not_called() + assert result["project_id"] == "project-orbit" + + @pytest.mark.asyncio async def test_key_creator_cannot_detach_project_without_admin_access(): existing: Final = LiteLLM_VerificationToken( @@ -21430,6 +21921,41 @@ async def test_key_creator_cannot_detach_project_without_admin_access(): assert "Only proxy admins, team admins, or org admins" in str(exc.value.detail) +@pytest.mark.asyncio +async def test_unauthorized_caller_gets_403_not_project_lookup_result(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", + user_id="user-orbit", created_by="user-orbit", + ) + database: Final = _project_assignment_prisma() + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + UpdateKeyRequest(key=existing.token, project_id="project-missing"), existing, + UserAPIKeyAuth(user_id="user-other", user_role=LitellmUserRoles.INTERNAL_USER), + None, False, database, UserApiKeyCache(), + ) + assert exc.value.status_code == 403 + assert "Project not found" not in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_unauthorized_caller_gets_403_not_project_team_leak(): + existing: Final = LiteLLM_VerificationToken( + token="project-assign-token", project_id=None, team_id="team-lit-5823", + user_id="user-orbit", created_by="user-orbit", + ) + cache: Final = await _cache_with_project("project-orbit", [], team_id="team-other") + database: Final = _project_assignment_prisma() + with pytest.raises(HTTPException) as exc: + await _validate_update_key_data( + UpdateKeyRequest(key=existing.token, project_id="project-orbit"), existing, + UserAPIKeyAuth(user_id="user-other", user_role=LitellmUserRoles.INTERNAL_USER), + None, False, database, cache, + ) + assert exc.value.status_code == 403 + assert "belongs to team" not in str(exc.value.detail) + + _OWNED_PROJECT: Final = "proj-owned-1" _OWNERSHIP_PROJECT_TEAM: Final = "ownership-project-team" _OWNERSHIP_KEY_TEAM: Final = "ownership-key-team" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index aafa732276b..96c7c34621b 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -8953,7 +8953,7 @@ export interface paths { * - user_id: Optional[str] - User ID associated with key * - team_id: Optional[str] - Team ID associated with key * - agent_id: Optional[str] - The agent id associated with the key. - * - project_id: Optional[str] - Omit to retain the project, or send null to detach. A different project ID is rejected. + * - project_id: Optional[str] - Omit to retain the project, send null to detach, or send a project id to assign an unassigned key to a project on the key's team. Moving a key between projects is rejected. * - organization_id: Optional[str] - The organization id of the key. * - budget_id: Optional[str] - The budget id associated with the key. Created by calling `/budget/new`. * - end_user_budget_id: Optional[str] - Proxy admin only. Budget id applied to end users first seen through this key that carry no budget of their own. Omit to keep the current value, pass an empty string to clear it. @@ -48704,7 +48704,7 @@ export interface components { policies?: string[] | null; /** * Project Id - * @description Omit to retain the project, or send null to detach. Assigning a different project is not supported. + * @description Omit to retain the project, send null to detach, or send a project id to assign an unassigned key to a project on the key's team. Moving a key between projects is not supported. */ project_id?: string | null; /** Prompts */ From b20acb48699e91fbe4d64a8878eaa0c92078a895 Mon Sep 17 00:00:00 2001 From: Maxime Boudier <113669264+MaximeBOUDIER@users.noreply.github.com> Date: Fri, 30 Oct 2026 20:10:13 +0100 Subject: [PATCH 07/46] fix(proxy)!: persist sent model_info fields on POST /model/update (#46902) * fix(proxy): persist sent model_info fields on POST /model/update POST /model/update rebuilt litellm_params from the request but dropped the request's model_info, so changing base_model (or mode, tier, etc.) on an existing DB model was silently ignored. Clients that only use this endpoint, like the Terraform provider, could never fix a deployment's base_model without deleting and recreating it. Persist the model_info fields the caller actually sent, on top of the stored ones. ModelInfo defaults never overwrite stored values, server-owned, team, PTU and pricing fields stay on their dedicated write paths, and values echoed from a /model/info response are not pinned. * fix(proxy): keep echoed and blank model_info values out of POST /model/update writes Blank strings from typed clients no longer overwrite stored values. Values echoed from /model/info (the stored base model's cost map entry and discovered limits) are not pinned, and an echoed catalog value resets a differing stored override like PATCH does. direct_access and access_via_team_ids stay display-only. The tests stop patching the auth check class and give the team lookup a real row instead. * fix(proxy): lock the model row and skip unchanged model_info on POST /model/update * fix(proxy): treat only keyed /model/update bodies as echoes and write the stored model_info as saved A typed model_info value that matched the cost map, or the previous base model's entry, was dropped and could erase a stored override. Echo detection now follows PATCH, only a body carrying the /model/info key counts, compared against the base model the row has after the write. The locked write patches the stored JSON instead of dumping the hydrated Deployment, so litellm_params pricing is no longer copied into model_info. Tests run the real Router with discovery over a mock upstream * docs(terraform): note that model_info updates now take effect and overwrite values set outside Terraform * test: expect POST /model/update to save the cache point flag, and wrap long test lines test_x06 pinned the old contract, where POST /model/update answered 200 and ignored model_info. This PR makes POST save model_info the way PATCH does, so x06 now asserts what its PATCH sibling x01 asserts, the stored flag flips to true and cache points come back on. * fix(proxy): answer 400 when POST /model/update gets an access_groups that is not a list of strings * fix(proxy): read discovered model info through the router registry entry point --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- .../model_management_endpoints.py | 167 ++++++++- terraform/provider/CHANGELOG.md | 3 +- .../provider/litellm/resource_model_crud.go | 2 +- .../test_bedrock_kimi_k3_cache_point_wire.py | 18 +- .../test_model_management_endpoints.py | 345 +++++++++++++++++- 5 files changed, 506 insertions(+), 29 deletions(-) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index 74fa8e3b3c1..dbd77e45d57 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -486,6 +486,7 @@ async def _raise_on_invalid_credential_name( AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY: Final = 5_872_301 _CAPABILITY_LOCK_SQL: Final = "SELECT 1 AS locked FROM pg_advisory_xact_lock($1)" +_MODEL_ROW_LOCK_SQL: Final = 'SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = $1 FOR UPDATE' _STORED_LITELLM_PARAMS_SQL: Final = ( "(CASE jsonb_typeof(litellm_params) WHEN 'string' THEN (litellm_params #>> '{}')::jsonb ELSE litellm_params END)" ) @@ -633,6 +634,7 @@ async def _auto_router_capability_slot( effective_params: Mapping[str, object], model_id: str | None, member_write: MemberAutoRouterWrite | None = None, + lock_model_row: bool = False, ) -> AsyncGenerator[_ProxyModelTable, None]: """Hand out the model table to write through while the row's claim on a licensed capability is settled. @@ -649,6 +651,9 @@ async def _auto_router_capability_slot( A heuristic-v1 router whose tuning has moved off its recorded baseline is judged the same way under the same lock, against the DB rows plus this proxy's config.yaml routers. + + ``lock_model_row`` runs the write in a transaction holding the row's ``FOR UPDATE`` lock, + so a caller that re-reads the row through the yielded table writes onto its committed state. """ from litellm.proxy.proxy_server import ( _license_check, # pyright: ignore[reportPrivateUsage] # existing capability slot reads the proxy license singleton @@ -662,7 +667,8 @@ async def _auto_router_capability_slot( baselines: Final = heuristic_v1_tuning_baselines tuning_candidate: Final = _tuning_candidate(effective_params, model_id=model_id) judges_tuning: Final = baselines is not None and is_mutable_tuned_candidate(tuning_candidate, baselines) - if member_write is None and (limit is None or (not capabilities and not judges_tuning)): + checks_license: Final = member_write is not None or (limit is not None and (bool(capabilities) or judges_tuning)) + if not checks_license and not lock_model_row: yield _proxy_model_table(prisma_client) return transaction_client: Final = _ModelTransactionClient.model_validate(prisma_client.db) @@ -673,14 +679,15 @@ async def _auto_router_capability_slot( ) async with transaction as tx_ctx: tables: Final[_TxModelTables] = tx_ctx - await tx_ctx.query_raw(_CAPABILITY_LOCK_SQL, AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY) + if checks_license: + await tx_ctx.query_raw(_CAPABILITY_LOCK_SQL, AUTO_ROUTER_CAPABILITY_SLOT_LOCK_KEY) + locked_model_id: Final = ( + member_write.model_id if member_write is not None else (model_id if lock_model_row else None) + ) + if locked_model_id is not None: + await tx_ctx.query_raw(_MODEL_ROW_LOCK_SQL, locked_model_id) config_rows: Final = () if llm_router is None else tuple(llm_router.registry().deployments(source="config")) if member_write is not None: - if member_write.model_id is not None: - await tx_ctx.query_raw( - 'SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = $1 FOR UPDATE', - member_write.model_id, - ) pinned_client: Final = _TransactionClient(tx_ctx) team_where: Final[prisma_types.LiteLLM_TeamTableWhereUniqueInput] = {"team_id": member_write.team_id} team_include: Final[prisma_types.LiteLLM_TeamTableInclude] = {"litellm_model_table": True} @@ -752,7 +759,8 @@ async def _auto_router_capability_slot( prisma_client=pinned_client, llm_router=llm_router, ) - for capability in capabilities: + judged_capabilities: Final = capabilities if checks_license else () + for capability in judged_capabilities: if ( violation := await _claimed_capability_violation( tables, capability=capability, config_rows=config_rows, model_id=model_id, limit=limit @@ -761,7 +769,7 @@ async def _auto_router_capability_slot( raise HTTPException( status_code=status.HTTP_403_FORBIDDEN, detail=f"{violation} {AUTO_ROUTER_LICENSE_REMEDY}" ) - if judges_tuning and baselines is not None: + if checks_license and judges_tuning and baselines is not None: model_rows: Final = await ModelRepository(_TransactionClient(tx_ctx)).find_all_except(model_id or "") _raise_on_tuning_quota_violation( candidate=tuning_candidate, @@ -1049,6 +1057,16 @@ def _raise_if_ptu_cost_attribution_disabled(incoming_model_info: Mapping[str, ob ) +_ACCESS_GROUPS: Final[TypeAdapter[tuple[str, ...] | None]] = TypeAdapter(tuple[str, ...] | None) + + +def _raise_if_access_groups_malformed(model_info_writes: Mapping[str, object]) -> None: + try: + _ACCESS_GROUPS.validate_python(model_info_writes.get("access_groups")) + except ValidationError: + raise HTTPException(status_code=400, detail="model_info.access_groups must be a list of strings") from None + + def _validate_ptu_model_info(model_info: Mapping[str, object]) -> None: """Enforce the PTU cross-field invariant on the effective model_info. @@ -1291,6 +1309,109 @@ def _loaded_catalog_entry( return loaded_catalog().get(catalog_key, MappingProxyType({})) +_LEGACY_UPDATE_IGNORED_MODEL_INFO_FIELDS: Final = frozenset( + { + "id", + "db_model", + "created_at", + "created_by", + "updated_at", + "updated_by", + "blocked", + "member_auto_router", + "team_id", + "team_public_model_name", + "direct_access", + "access_via_team_ids", + COST_MAP_LOOKUP_KEY, + *PTU_MODEL_INFO_FIELDS, + } +) + + +def legacy_update_model_info_changes( + db_model: Deployment, + incoming_model_info: ModelInfo, + discovered_model_info: Mapping[str, object] = MappingProxyType({}), + loaded_catalog: LoadedCatalog = GetModelCostMap.loaded_model_cost_map, +) -> Mapping[str, object]: + """The ``model_info`` fields ``POST /model/update`` persists from its request. + + Only fields the caller sent with a value count, so ``ModelInfo`` defaults and the blank + strings typed clients send for unset fields never overwrite stored values. Server-owned, + display-only, team, PTU and pricing fields keep their dedicated write paths. As on + ``update_db_model``, only a body carrying the ``key`` that ``/model/info`` emits is an echo: + its values equal to the cost map entry of the base model the row has after the write, the + entry ``key`` names, or a limit discovered from the upstream stay display-only, and an echoed + value that differs from the stored override resets that override (``None``). + """ + stored: Final = db_model.model_info.model_dump(exclude_none=True) + requested: Final = incoming_model_info.model_dump(exclude_unset=True, exclude_none=True) + sent: Final = {k: v for k, v in requested.items() if k not in _LEGACY_UPDATE_IGNORED_MODEL_INFO_FIELDS and v != ""} + echoed: Final = frozenset( + echoed_cost_map_fields( + requested, + _cost_map_entry(db_model, {**stored, **sent}), + _loaded_catalog_entry(requested, loaded_catalog), + discovered_model_info, + ) + ) + kept: Final = without_server_derived_pricing(sent) + return MappingProxyType( + { + **{k: getattr(incoming_model_info, k) for k in kept if k not in echoed}, + **{k: None for k, v in kept.items() if k in echoed and stored.get(k) not in (None, v)}, + } + ) + + +def _legacy_update_model_info_writes( + current: Deployment, + incoming_model_info: ModelInfo, + discovered_model_info: Mapping[str, object], + member_marker: bool | None, +) -> Mapping[str, object]: + changes: Final = { + **legacy_update_model_info_changes(current, incoming_model_info, discovered_model_info), + **({"member_auto_router": member_marker} if member_marker is not None else {}), + } + return MappingProxyType({k: v for k, v in changes.items() if getattr(current.model_info, k, None) != v}) + + +async def _locked_legacy_update_model_info_json( + table: _ProxyModelTable, + model_id: str, + incoming_model_info: ModelInfo, + discovered_model_info: Mapping[str, object], + member_marker: bool | None, +) -> str | None: + row: Final = await table.find_unique(where={"model_id": model_id}) + if row is None: + return None + payload: Final = row.model_dump() + writes: Final = _legacy_update_model_info_writes( + Deployment(**payload), incoming_model_info, discovered_model_info, member_marker + ) + if not writes: + return None + stored_model_info: Final[object] = payload.get("model_info") + stored: Final = ( + TypeAdapter(dict[str, object]).validate_json(stored_model_info) + if isinstance(stored_model_info, str) + else TypeAdapter(dict[str, object]).validate_python(stored_model_info or {}) + ) + return ( + TypeAdapter(dict[str, object]) + .dump_json( + { + **{k: v for k, v in stored.items() if k not in writes}, + **{k: v for k, v in writes.items() if v is not None}, + } + ) + .decode() + ) + + def _param_keys_to_remove(patch: updateDeployment) -> tuple[str, ...]: return patch.litellm_params_to_remove if isinstance(patch, PatchModelRequest) else () @@ -3279,22 +3400,32 @@ async def update_model( renamed_update: Final[PrismaCompatibleUpdateDBModel] = ( {**base_update, "model_name": renamed_to} if renamed_to is not None else base_update ) - _data: Final[PrismaCompatibleUpdateDBModel] = ( - { - **renamed_update, - "model_info": deployment.model_info.model_copy( - update=MappingProxyType({"member_auto_router": member_marker}) - ).model_dump_json(exclude_none=True), - } - if member_marker is not None - else renamed_update + discovered_model_info: Final[Mapping[str, int]] = ( + llm_router.registry().discovered_model_info(_model_id) + if llm_router is not None + else MappingProxyType({}) ) + requested_model_info: Final = _legacy_update_model_info_writes( + deployment, _model_info, discovered_model_info, member_marker + ) + _raise_if_access_groups_malformed(requested_model_info) async with _auto_router_capability_slot( prisma_client, effective_params=effective_params, model_id=_model_id, member_write=member_write, + lock_model_row=bool(requested_model_info), ) as table: + model_info_json: Final = ( + await _locked_legacy_update_model_info_json( + table, _model_id, _model_info, discovered_model_info, member_marker + ) + if requested_model_info + else None + ) + _data: Final[PrismaCompatibleUpdateDBModel] = ( + {**renamed_update, "model_info": model_info_json} if model_info_json is not None else renamed_update + ) model_response: Final = await table.update( where={"model_id": _model_id}, data=_data, diff --git a/terraform/provider/CHANGELOG.md b/terraform/provider/CHANGELOG.md index 597e0c6d4d8..4a1838a7b60 100644 --- a/terraform/provider/CHANGELOG.md +++ b/terraform/provider/CHANGELOG.md @@ -18,7 +18,7 @@ longer signal it. - **credential, model, mcp_server**: Write-only secret attributes so secrets never land in Terraform state (Terraform 1.11 or later). `litellm_credential` gains `credential_values_wo` (a JSON object, `jsonencode({...})`, since the SDK cannot mark a map write-only) with `credential_values_wo_version`, and `credential_values` becomes optional with exactly one of the two required. `litellm_model` gains `model_api_key_wo`, `aws_secret_access_key_wo` and `vertex_credentials_wo`, each with a `_wo_version` companion and each conflicting with its in-state twin. `litellm_mcp_server` gains `client_id`, `client_secret_wo` and `auth_value_wo` (with `_wo_version` companions), sent to the proxy as `credentials`, and `auth_type` now accepts every value the proxy does (`oauth2`, `api_key`, `bearer_token`, ...). Every `_wo` value is sent on each apply from the configuration and produces no diff, so bump its `_wo_version` to push a rotated secret. `litellm_vector_store.litellm_params` warns at plan time when it carries a secret-looking key such as `api_key` - **team**: Optional `team_member_model_max_budget` argument on `litellm_team`, a JSON string of per-model budgets applied to every team member and sent as `team_member_model_max_budget` on `/team/new` and `/team/update`; read back from `team_member_budget_table.model_max_budget` on `/team/info` so out-of-band changes show as drift, with `budget_limit`/`time_period` treated as equal to the `max_budget`/`budget_duration` spelling the proxy stores; removing it sends `{}`, which clears every per-model member budget; `tpm_limit`/`rpm_limit` are rejected at plan time because the proxy rejects them on team budgets. Needs a proxy that includes [#43887](https://github.com/BerriAI/litellm/pull/43887): v1.105.0 and v1.106.0-rc.1 accept the field but silently drop it, so every plan shows a diff -- **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 +- **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 also persisted through `/model/{id}/update`, since proxies older than this provider drop the `model_info` sent to `/model/update`; 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 @@ -70,6 +70,7 @@ longer signal it. ### Changed +- **model** (breaking): on a proxy that includes [#46902](https://github.com/BerriAI/litellm/pull/46902), `litellm_model` updates now save the `model_info` the provider sends to `/model/update`, so `pricing_base_model`, `tier` and `mode` changes take effect. The first apply after upgrading the proxy overwrites values set outside Terraform with the configuration's, including `tier`, which defaults to `free`, and `base_model`, which falls back to the deployment name when `pricing_base_model` is unset. Run `terraform plan` before that apply and copy any value to keep into the configuration - **model** (breaking): a key named in `additional_litellm_params.additional_drop_params` is no longer deleted from the deployment's `litellm_params`. A configuration that also sets that key, such as `vertex_location`, stores it on the deployment at its next apply, and the proxy applies the drop list to requests instead. For `vertex_location` this changes routing. `/v1/messages` calls now go to the configured region, while chat completions drop a listed `vertex_location` and use the default region. To keep such a key off a deployment, remove its attribute and replace the model with `terraform apply -replace`, which gives it a new model id, since an update keeps a stored value the provider leaves out. `model`, `custom_llm_provider` and `merge_reasoning_content_in_choices` are always sent. A value that is not a JSON list of names, such as a bare `"store"` or an empty string, now fails the plan, and with it `terraform destroy`, instead of being stored as a string the proxy ignores. Rewrite it as `jsonencode([...])`, or remove the model without fixing it through `terraform destroy -refresh=false` - **key** (breaking): `model_max_budget` on `litellm_key` is now a JSON string of per-model budget objects (`jsonencode({"gpt-4o-mini" = {budget_limit = 50, time_period = "30d"}})`), matching `litellm_user`, `litellm_budget` and `litellm_tag`. The old `map(number)` form sent bare numbers to `/key/generate`, which the proxy rejects with a 500 (`'int' object is not iterable`), so every key with a non-empty `model_max_budget` failed to apply. Existing state upgrades automatically (schema version 1) and the attribute is refilled from the proxy on the next read; configurations still using the map form must be rewritten - **Versioning**: the provider is now published at the LiteLLM version, from the same commit as the proxy, on every LiteLLM release (dev, rc, stable). The `0.x` line ends at `0.4.0`; a `~> 0.4` constraint will not receive further releases, so re-pin to the LiteLLM version your proxy runs (for example `~> 1.99.0`). Existing `0.x` versions remain in the registry and keep verifying diff --git a/terraform/provider/litellm/resource_model_crud.go b/terraform/provider/litellm/resource_model_crud.go index 3db6f68760d..1f5aa62560f 100644 --- a/terraform/provider/litellm/resource_model_crud.go +++ b/terraform/provider/litellm/resource_model_crud.go @@ -357,7 +357,7 @@ 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. +// Older proxies drop the model_info sent to /model/update, so display_name also goes 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}, diff --git a/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py b/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py index 0772e928b35..4441582fb6b 100644 --- a/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py +++ b/tests/integration/providers/test_bedrock_kimi_k3_cache_point_wire.py @@ -1325,7 +1325,9 @@ def test_x01_flipping_the_flag_to_true_through_a_patch_update_turns_cache_points @pytest.mark.timeout(600) -def test_x06_the_legacy_post_update_answers_200_and_leaves_the_stored_flag_alone(gateway: Gateway) -> None: +def test_x06_flipping_the_flag_to_true_through_the_legacy_post_update_turns_cache_points_back_on( + gateway: Gateway, +) -> None: arn: Final = _arn("legacy") with wire_server(_peer()) as wire, gateway.scenario() as scenario: name: Final = _deployment( @@ -1333,6 +1335,7 @@ def test_x06_the_legacy_post_update_answers_200_and_leaves_the_stored_flag_alone ) before, before_received = _observe_at(gateway, wire, "chat", _chat_body(name, _prompt(), _SYSTEM_AND_USER)) _assert_answered(before) + assert before_received.cache_points == 0, before_received.body identity: Final = _model_id(name) updated: Final = gateway.request( "POST", @@ -1344,13 +1347,14 @@ def test_x06_the_legacy_post_update_answers_200_and_leaves_the_stored_flag_alone }, ) assert updated.status_code == 200, updated.text - assert _stored_flag(identity) is False - _settled(gateway, name, wire) - after, after_received = _observe_at(gateway, wire, "chat", _chat_body(name, _prompt(), _SYSTEM_AND_USER)) - _assert_answered(after) - assert after_received.cache_points == before_received.cache_points, (before_received.body, after_received.body) + assert _stored_flag(identity) is True + points: Final = eventually( + lambda: _six_cache_point_counts(gateway, wire, name), + lambda counts: len(counts) == 6 and all(count == 2 for count in counts), + seconds=60, + ) + assert points == (2,) * 6, points _success_row(before.response_id) - _success_row(after.response_id) @dataclass(frozen=True, slots=True) diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index 872d1607351..abaa802de2c 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -3,17 +3,19 @@ import contextlib import inspect import json import socket -from collections.abc import Iterator, Mapping -from types import SimpleNamespace +from collections.abc import AsyncIterator, Iterator, Mapping +from types import MappingProxyType, SimpleNamespace from typing import Dict, Final, Optional from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest from fastapi import HTTPException from fastapi.encoders import jsonable_encoder from fastapi.testclient import TestClient from litellm._uuid import uuid +from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.models.credentials import CredentialItem from litellm.models.project import LiteLLM_ProjectTable @@ -1639,11 +1641,350 @@ class TestDeletedAutoRouterAvailability: } +class _ContendedModelDb: + def __init__(self, row: LiteLLM_ProxyModelTable, other_writers: tuple[Mapping[str, object], ...]) -> None: + self.row = row + self.writes: tuple[Mapping[str, object], ...] = () + self._other_writers = other_writers + self._row_locked = False + self.db = self + self.litellm_proxymodeltable = self + self.litellm_teamtable = SimpleNamespace(find_unique=self._find_team) + + async def _find_team(self, *, where: Mapping[str, str]) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable(team_id=where["team_id"]) + + async def find_unique(self, *, where: Mapping[str, object]) -> LiteLLM_ProxyModelTable: + read: Final = self.row + if not self._row_locked and self._other_writers: + self._commit_next_other_writer() + return read + + async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> LiteLLM_ProxyModelTable: + self.writes = (*self.writes, data) + self.row = self.row.model_copy( + update={k: json.loads(v) if k in ("litellm_params", "model_info") else v for k, v in data.items()} + ) + return self.row + + async def query_raw(self, query: str, *args: object) -> tuple[()]: + self._row_locked = self._row_locked or "FOR UPDATE" in query + return () + + @contextlib.asynccontextmanager + async def tx(self, timeout: object = None) -> AsyncIterator["_ContendedModelDb"]: + yield self + self._row_locked = False + while self._other_writers: + self._commit_next_other_writer() + + def _commit_next_other_writer(self) -> None: + self.row = self.row.model_copy( + update={"model_info": {**(self.row.model_info or {}), **self._other_writers[0]}} + ) + self._other_writers = self._other_writers[1:] + + +_AZURE_DEPLOYMENT_PARAMS: Final = MappingProxyType({"model": "azure/my-deployment", "api_key": "sk-existing"}) +_LOCAL_DEPLOYMENT_PARAMS: Final = MappingProxyType( + {"model": "hosted_vllm/local-model", "api_base": "https://local.test/v1", "api_key": "local-key"} +) + + +async def _discover_from_upstream(router: Router, model_card: Mapping[str, object]) -> None: + handler: Final = AsyncHTTPHandler() + await handler.client.aclose() + async with httpx.AsyncClient( + transport=httpx.MockTransport(lambda request: httpx.Response(200, json={"data": [dict(model_card)]})) + ) as client: + handler.client = client + await router.arefresh_model_info(client=handler) + + class TestUpdateModel: """ Tests for the update_model (POST /model/update) handler. """ + @staticmethod + async def _update_model( + stored_model_info: Mapping[str, object], + model_info: ModelInfo, + litellm_params: Mapping[str, object] = _AZURE_DEPLOYMENT_PARAMS, + upstream_model_card: Mapping[str, object] | None = None, + other_writers: tuple[Mapping[str, object], ...] = (), + ) -> "_ContendedModelDb": + db: Final = _ContendedModelDb( + LiteLLM_ProxyModelTable( + model_id="m1", + model_name="my-model", + litellm_params=dict(litellm_params), + model_info=dict(stored_model_info), + ), + other_writers, + ) + router: Final = Router( + model_list=[ + {"model_name": "my-model", "litellm_params": dict(litellm_params), "model_info": {"id": "m1"}} + ] + ) + if upstream_model_card is not None: + await _discover_from_upstream(router, upstream_model_card) + + with ( + patch("litellm.proxy.proxy_server.prisma_client", db), + patch("litellm.proxy.proxy_server.llm_router", router), + patch("litellm.proxy.proxy_server.store_model_in_db", True), + patch("litellm.proxy.proxy_server.premium_user", True), + patch( # test-quality-ok: [TQ008] the route reloads the proxy-global router with no seam to inject + "litellm.proxy.management_endpoints.model_management_endpoints.clear_cache", + new=AsyncMock(return_value=ReconcileOutcome(still_desired=None, live_after=None)), + ), + ): + await update_model( + model_params=updateDeployment(litellm_params=updateLiteLLMParams(), model_info=model_info), + user_api_key_dict=UserAPIKeyAuth(user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN), + ) + + return db + + @pytest.mark.asyncio + async def test_update_model_persists_sent_model_info_over_stored_fields(self): + stored: Final = { + "id": "m1", + "db_model": True, + "base_model": "my-deployment", + "tier": "paid", + "team_id": "team-a", + } + + db: Final = await self._update_model( + stored, + ModelInfo(id="m1", base_model="test-provider/test-model", mode="chat", db_model=False, team_id="team-b"), + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert persisted["base_model"] == "test-provider/test-model" + assert persisted["mode"] == "chat" + assert {k: persisted[k] for k in ("db_model", "tier", "team_id")} == { + "db_model": True, + "tier": "paid", + "team_id": "team-a", + } + + @pytest.mark.asyncio + async def test_update_model_with_only_an_id_leaves_model_info_untouched(self): + db: Final = await self._update_model({"id": "m1", "base_model": "my-deployment"}, ModelInfo(id="m1")) + + assert tuple("model_info" in write for write in db.writes) == (False,) + + @pytest.mark.asyncio + async def test_update_model_resending_the_stored_model_info_does_not_rewrite_it(self): + db: Final = await self._update_model( + {"id": "m1", "db_model": True, "tier": "paid", "mode": "chat"}, + ModelInfo(id="m1", tier="paid", mode="chat"), + ) + + assert tuple("model_info" in write for write in db.writes) == (False,) + + @pytest.mark.asyncio + async def test_update_model_keeps_model_info_fields_a_concurrent_writer_stored(self): + db: Final = await self._update_model( + {"id": "m1", "db_model": True, "mode": "embedding"}, + ModelInfo(id="m1", mode="chat"), + other_writers=(MappingProxyType({"access_groups": ["group-a"]}), MappingProxyType({"tier": "paid"})), + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert {k: persisted.get(k) for k in ("mode", "access_groups", "tier")} == { + "mode": "chat", + "access_groups": ["group-a"], + "tier": "paid", + } + + @pytest.mark.asyncio + async def test_update_model_does_not_pin_values_echoed_from_model_info(self): + import litellm + + catalog_entry: Final = litellm.get_model_info(model="azure/gpt-4o") + echoed: Final = ModelInfo( + id="m1", + base_model="azure/gpt-4o", + key=catalog_entry["key"], + max_tokens=catalog_entry["max_tokens"], + input_cost_per_token=catalog_entry["input_cost_per_token"], + ) + + db: Final = await self._update_model({"id": "m1"}, echoed) + + persisted: Final = db.row.model_info + assert persisted is not None + assert persisted["base_model"] == "azure/gpt-4o" + assert not {"key", "max_tokens", "input_cost_per_token"} & persisted.keys() + + @pytest.mark.asyncio + async def test_update_model_saves_typed_values_that_equal_the_catalog(self): + import litellm + + catalog_entry: Final = litellm.get_model_info(model="azure/gpt-4o") + typed: Final = { + "max_tokens": catalog_entry["max_tokens"], + "max_input_tokens": catalog_entry["max_input_tokens"], + "supports_vision": catalog_entry["supports_vision"], + } + + db: Final = await self._update_model( + {"id": "m1", "max_input_tokens": catalog_entry["max_input_tokens"] // 2}, + ModelInfo(id="m1", base_model="azure/gpt-4o", **typed), + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert {k: persisted.get(k) for k in typed} == typed + + @pytest.mark.asyncio + async def test_update_model_keeps_stored_value_when_a_field_is_sent_blank(self): + db: Final = await self._update_model( + {"id": "m1", "base_model": "my-deployment", "mode": "chat"}, + ModelInfo(id="m1", base_model="test-provider/test-model", mode=""), + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert persisted["base_model"] == "test-provider/test-model" + assert persisted["mode"] == "chat" + + @pytest.mark.asyncio + async def test_update_model_ignores_ptu_blocked_and_pricing_fields(self): + db: Final = await self._update_model( + {"id": "m1"}, + ModelInfo( + id="m1", base_model="test-provider/test-model", ptu_count=5, blocked=True, input_cost_per_token=0.5 + ), + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert persisted["base_model"] == "test-provider/test-model" + assert not {"ptu_count", "blocked", "input_cost_per_token"} & persisted.keys() + + @pytest.mark.asyncio + @pytest.mark.parametrize("access_groups", [5, "group-a", ["group-a", 5], {"group": "group-a"}]) + async def test_update_model_rejects_access_groups_that_are_not_a_list_of_strings(self, access_groups: object): + with pytest.raises(ProxyException) as error: + await self._update_model({"id": "m1"}, ModelInfo.model_validate({"id": "m1", "access_groups": access_groups})) + + assert error.value.code == "400" + + @pytest.mark.asyncio + async def test_update_model_saves_access_groups_sent_as_a_list_of_strings(self): + db: Final = await self._update_model( + {"id": "m1"}, ModelInfo.model_validate({"id": "m1", "access_groups": ["group-a", "group-b"]}) + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert persisted["access_groups"] == ["group-a", "group-b"] + + def test_update_model_does_not_pin_the_previous_base_models_echoed_values(self): + from litellm.proxy.management_endpoints.model_management_endpoints import legacy_update_model_info_changes + + catalog: Final = MappingProxyType( + {"test-provider/previous": MappingProxyType({"supported_openai_params": ["previous-only-param"]})} + ) + stored: Final = Deployment( + model_name="my-model", + litellm_params=LiteLLM_Params(model="test-provider/my-deployment"), + model_info=ModelInfo(id="m1", base_model="test-provider/previous"), + ) + + changes: Final = legacy_update_model_info_changes( + stored, + ModelInfo( + id="m1", + base_model="test-provider/next", + key="test-provider/previous", + supported_openai_params=["previous-only-param"], + ), + loaded_catalog=lambda: catalog, + ) + + assert dict(changes) == {"base_model": "test-provider/next"} + + @pytest.mark.asyncio + async def test_update_model_saves_a_value_the_previous_base_model_has_in_the_catalog(self): + import litellm + + previous: Final = litellm.get_model_info(model="azure/gpt-4o") + + db: Final = await self._update_model( + {"id": "m1", "base_model": "azure/gpt-4o", "max_input_tokens": previous["max_input_tokens"] // 2}, + ModelInfo(id="m1", base_model="azure/gpt-4.1", max_input_tokens=previous["max_input_tokens"]), + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert {k: persisted.get(k) for k in ("base_model", "max_input_tokens")} == { + "base_model": "azure/gpt-4.1", + "max_input_tokens": previous["max_input_tokens"], + } + + @pytest.mark.asyncio + async def test_update_model_does_not_pin_discovered_limits_or_team_access_echoed_from_model_info(self): + db: Final = await self._update_model( + {"id": "m1"}, + ModelInfo( + id="m1", + base_model="test-provider/test-model", + key="test-provider/test-model", + max_input_tokens=128000, + max_output_tokens=64000, + direct_access=True, + access_via_team_ids=["team-a"], + ), + litellm_params=_LOCAL_DEPLOYMENT_PARAMS, + upstream_model_card=MappingProxyType( + {"id": "local-model", "max_model_len": 128000, "max_output_tokens": 32000} + ), + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert persisted["base_model"] == "test-provider/test-model" + assert persisted["max_output_tokens"] == 64000 + assert not {"max_input_tokens", "direct_access", "access_via_team_ids"} & persisted.keys() + + @pytest.mark.asyncio + async def test_update_model_does_not_copy_litellm_params_pricing_into_model_info(self): + db: Final = await self._update_model( + {"id": "m1"}, + ModelInfo(id="m1", tier="paid"), + litellm_params={**_AZURE_DEPLOYMENT_PARAMS, "input_cost_per_token": 0.000003}, + ) + + persisted: Final = db.row.model_info + assert persisted == {"id": "m1", "tier": "paid"} + + @pytest.mark.asyncio + async def test_update_model_resets_a_stored_override_echoed_back_at_its_catalog_value(self): + import litellm + + catalog_entry: Final = litellm.get_model_info(model="azure/gpt-4o") + stored_override: Final = catalog_entry["max_tokens"] // 2 + + db: Final = await self._update_model( + {"id": "m1", "base_model": "azure/gpt-4o", "max_tokens": stored_override}, + ModelInfo( + id="m1", base_model="azure/gpt-4o", key=catalog_entry["key"], max_tokens=catalog_entry["max_tokens"] + ), + ) + + persisted: Final = db.row.model_info + assert persisted is not None + assert "max_tokens" not in persisted + @pytest.mark.asyncio async def test_update_model_clears_cache_after_db_write(self): """ From 6ad4f9cfdf3f2c4a8db3e39acb02d3d7730d2dce Mon Sep 17 00:00:00 2001 From: Nate Armstrong Date: Fri, 30 Oct 2026 12:14:17 -0700 Subject: [PATCH 08/46] ci: cache Rust compilation with sccache (#49117) This should make compilation much faster --- .github/workflows/test-rust-db.yml | 10 ++++++++++ .github/workflows/test-rust.yml | 20 ++++++++++++++++++++ .github/workflows/test-unit.yml | 12 ++++++++++++ 3 files changed, 42 insertions(+) diff --git a/.github/workflows/test-rust-db.yml b/.github/workflows/test-rust-db.yml index bb88b593d5b..80ec256fa6f 100644 --- a/.github/workflows/test-rust-db.yml +++ b/.github/workflows/test-rust-db.yml @@ -42,6 +42,11 @@ concurrency: env: CARGO_TERM_COLOR: always + RUSTC_WRAPPER: sccache + SCCACHE_GHA_ENABLED: "true" + SCCACHE_GHA_VERSION: litellm-ci-v1 + SCCACHE_GHA_RW_MODE: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' && 'READ_WRITE' || 'READ_ONLY' }} + CARGO_INCREMENTAL: "0" jobs: db-queries: @@ -87,6 +92,11 @@ jobs: workspaces: litellm-rust cache-on-failure: true + - name: Set up sccache + uses: mozilla-actions/sccache-action@d651010b8da762cde178750d8eda7b5febfe147a + with: + version: "v0.18.0" + - name: Install Python dependencies working-directory: . run: uv sync --frozen --extra extra-proxy diff --git a/.github/workflows/test-rust.yml b/.github/workflows/test-rust.yml index 04af770b24f..cf9fed192f6 100644 --- a/.github/workflows/test-rust.yml +++ b/.github/workflows/test-rust.yml @@ -78,6 +78,11 @@ concurrency: env: CARGO_TERM_COLOR: always + RUSTC_WRAPPER: sccache + SCCACHE_GHA_ENABLED: "true" + SCCACHE_GHA_VERSION: litellm-ci-v1 + SCCACHE_GHA_RW_MODE: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' && 'READ_WRITE' || 'READ_ONLY' }} + CARGO_INCREMENTAL: "0" jobs: rust-lint: @@ -101,6 +106,11 @@ jobs: cache-on-failure: true save-if: ${{ github.ref == 'refs/heads/main' }} + - name: Set up sccache + uses: mozilla-actions/sccache-action@d651010b8da762cde178750d8eda7b5febfe147a + with: + version: "v0.18.0" + - run: cargo clippy --workspace --all-targets --locked -- -D warnings rust-test: @@ -143,6 +153,11 @@ jobs: cache-on-failure: true save-if: ${{ github.ref == 'refs/heads/main' }} + - name: Set up sccache + uses: mozilla-actions/sccache-action@d651010b8da762cde178750d8eda7b5febfe147a + with: + version: "v0.18.0" + - name: Check generated trace contracts working-directory: . run: uv run scripts/generate_trace_types.py --check @@ -189,6 +204,11 @@ jobs: cache-on-failure: true save-if: ${{ github.ref == 'refs/heads/main' }} + - name: Set up sccache + uses: mozilla-actions/sccache-action@d651010b8da762cde178750d8eda7b5febfe147a + with: + version: "v0.18.0" + - run: uv build --wheel --out-dir dist - run: python .github/scripts/verify_linux_native_wheel.py dist/*.whl diff --git a/.github/workflows/test-unit.yml b/.github/workflows/test-unit.yml index 68823d99519..5ec5975d83b 100644 --- a/.github/workflows/test-unit.yml +++ b/.github/workflows/test-unit.yml @@ -73,9 +73,21 @@ jobs: timeout-minutes: 5 uses: ./.github/actions/cache-cargo-build + - name: Set up sccache + if: steps.changes.outputs.decision != 'skip' && steps.rust-bridge-cache.outputs.cache-hit != 'true' + uses: mozilla-actions/sccache-action@d651010b8da762cde178750d8eda7b5febfe147a + with: + version: "v0.18.0" + - name: Build editable Rust bridge if: steps.changes.outputs.decision != 'skip' && steps.rust-bridge-cache.outputs.cache-hit != 'true' timeout-minutes: 15 + env: + RUSTC_WRAPPER: sccache + SCCACHE_GHA_ENABLED: "true" + SCCACHE_GHA_VERSION: litellm-ci-v1 + SCCACHE_GHA_RW_MODE: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' && 'READ_WRITE' || 'READ_ONLY' }} + CARGO_INCREMENTAL: "0" run: | mkdir -p rust-bridge-dist uv run --no-project --with maturin==1.15.0 python -c 'import maturin; maturin.build_editable("rust-bridge-dist")' From 5898a42f5aac37fc9fa2392f3c7fe16644e3e958 Mon Sep 17 00:00:00 2001 From: norkelbyte <348963847+norkelbyte@users.noreply.github.com> Date: Fri, 30 Oct 2026 12:19:51 -0700 Subject: [PATCH 09/46] fix(openrouter)!: list OpenRouter's live catalog for openrouter/* when check_provider_endpoint is on (#48820) * fix(openrouter): list OpenRouter's live catalog for openrouter/* when check_provider_endpoint is on * fix(openrouter): list models off the event loop and keep only chat-callable catalog ids The model-list builders are sync, and with check_provider_endpoint on they call OpenRouter's catalog with a 10 second timeout and no caching of failures. They ran on the event loop inside /v1/models, /model/info, /model_group/info and the per-user listing, so one slow catalog froze every request on the worker. Those calls now run in a worker thread with asyncio.to_thread File, batch and RAG credential lookups only need each name's deployment, and a wildcard that is itself a deployment sits in the list ahead of its expansions and resolves its own credentials. get_complete_model_list takes expand_deployed_wildcard_routes, which that lookup sets to False, so a deployed wildcard never fetches its provider catalog there. A key wildcard wider than any deployment, like anthropic/* over a deployed anthropic/claude-*, still expands so it finds the matching deployment's credentials. Every other caller keeps the default The pattern router's model lookup tolerates a pattern removed between its two reads of the pattern table, which a deployment delete on the event loop can now do while a listing runs in a thread Listing now reads only OpenRouter's models/user catalog. OpenRouter's own /models list carries no embedding models, the OpenAI-format /v1/models has no mode field so chat pickers cannot tell embeddings apart, /model_group/info labels them mode chat, and one failing catalog emptied the whole list. Ids ending in :batch are dropped because chat completions answers 404 for them, while :free and every other variant stay. An api_base that already ends in /chat/completions, which chat accepts as is, has that suffix stripped before the catalog route is joined * fix(proxy): keep model listing off the shared thread pool and out of credential lookups Listings run on their own thread pool so many slow catalogs cannot starve the default executor that chat and Responses calls share. File and batch credential lookups match deployed names against a key's wildcard grants instead of listing any provider catalog. The token counter keeps OpenRouter deployments on their previous path, and a catch-all "*" over provider/* lists bare ids that route back through it * test(proxy): give the off-loop listing tests a generous deadlock guard and typed ids * fix(proxy): keep routed models in wildcard model info, skip key-only OpenRouter catalogs, cache failed listings briefly * fix(utils): cache failed model listings under their own key so a late failure never hides a good listing * fix(proxy): remember only timed-out provider listings so a provider that failed fast is listed again * fix(proxy): ask each provider catalog once per worker, keep deployed names out of catch-all ids, and match file credentials by name Listings that start together for one deployment now wait for the first catalog request. A catch-all wildcard no longer lists a name a deployment serves by name. A key wildcard only finds file and batch credentials on deployments named under it, so a bare name keeps the default credentials as before. * fix(proxy): keep the listing failure log line the same for timed-out listings * refactor(proxy): fold the cached listing lookup back into the provider listing function * fix(proxy): list each model id once and keep a named deployment's row over a wildcard's * fix(proxy): drop a wildcard's model info row when named deployments serve every name it lists --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/llms/openrouter/common_utils.py | 87 +++++++ litellm/proxy/auth/model_checks.py | 74 +++++- .../openai_files_endpoints/common_utils.py | 53 +++-- litellm/proxy/proxy_server.py | 59 +++-- litellm/proxy/utils.py | 55 +++-- .../router_utils/pattern_match_deployments.py | 4 +- litellm/utils.py | 81 +++++-- .../test_openrouter_common_utils.py | 159 +++++++++++++ tests/unit/proxy/auth/test_model_checks.py | 219 ++++++++++++++++++ .../test_files_common_utils.py | 112 +++++++++ .../proxy_server/test_routes_model_info.py | 12 +- .../proxy/proxy_server/test_routes_models.py | 73 ++++++ tests/unit/proxy/test_proxy_token_counter.py | 33 +++ .../proxy/utils/helpers/test_model_access.py | 67 +++++- tests/unit/test_utils.py | 71 ++++++ 15 files changed, 1055 insertions(+), 104 deletions(-) create mode 100644 tests/unit/llms/openrouter/test_openrouter_common_utils.py diff --git a/litellm/llms/openrouter/common_utils.py b/litellm/llms/openrouter/common_utils.py index 96e53a5aae7..64caf83d745 100644 --- a/litellm/llms/openrouter/common_utils.py +++ b/litellm/llms/openrouter/common_utils.py @@ -1,5 +1,92 @@ +from collections.abc import Mapping, Sequence +from typing import Final + +import httpx +from pydantic import BaseModel, ConfigDict + +import litellm +from litellm.litellm_core_utils.url_utils import join_api_base +from litellm.llms.base_llm.base_utils import BaseLLMModelInfo from litellm.llms.base_llm.chat.transformation import BaseLLMException +from litellm.secret_managers.main import get_secret_str +from litellm.types.llms.openai import AllMessageValues + +_DEFAULT_API_BASE: Final = "https://openrouter.ai/api/v1" +_CHAT_CATALOG_ROUTE: Final = "models/user" +_CHAT_COMPLETIONS_PATH: Final = "/chat/completions" +_BATCH_ONLY_SUFFIX: Final = ":batch" +_CATALOG_TIMEOUT: Final = httpx.Timeout(10.0) class OpenRouterException(BaseLLMException): pass + + +class _CatalogModel(BaseModel): + model_config = ConfigDict(frozen=True) + + id: str + + +class _Catalog(BaseModel): + model_config = ConfigDict(frozen=True) + + data: tuple[_CatalogModel, ...] + + +def _chat_catalog_url(api_base: str) -> str: + return join_api_base(api_base.rstrip("/").removesuffix(_CHAT_COMPLETIONS_PATH), _CHAT_CATALOG_ROUTE) + + +class OpenRouterModelInfo(BaseLLMModelInfo): + @staticmethod + def get_api_key(api_key: str | None = None) -> str | None: + return ( + api_key + or litellm.api_key + or litellm.openrouter_key + or get_secret_str("OPENROUTER_API_KEY") + or get_secret_str("OR_API_KEY") + ) + + @staticmethod + def get_api_base(api_base: str | None = None) -> str: + return api_base or litellm.api_base or get_secret_str("OPENROUTER_API_BASE") or _DEFAULT_API_BASE + + @staticmethod + def get_base_model(model: str) -> str: + return model.removeprefix("openrouter/") + + def validate_environment( + self, + headers: Mapping[str, str], + model: str, + messages: Sequence[AllMessageValues], + optional_params: Mapping[str, object], + litellm_params: Mapping[str, object], + api_key: str | None = None, + api_base: str | None = None, + ) -> dict[str, str]: # mutable-ok: BaseLLMModelInfo.validate_environment returns a dict + resolved_key: Final = self.get_api_key(api_key) + if resolved_key is None: + return {**headers} + return {**headers, "Authorization": f"Bearer {resolved_key}"} + + def get_models( + self, api_key: str | None = None, api_base: str | None = None + ) -> list[str]: # mutable-ok: matches get_models' list[str] contract shared by every provider override + resolved_key: Final = self.get_api_key(api_key) + if resolved_key is None: + raise ValueError("Set OPENROUTER_API_KEY or the deployment's api_key to list OpenRouter models") + headers: Final = self.validate_environment( + headers={}, model="", messages=(), optional_params={}, litellm_params={}, api_key=resolved_key + ) + response: Final = litellm.module_level_client.get( # pyright: ignore[reportUnknownMemberType] # legacy HTTP wrapper has untyped query parameters + url=_chat_catalog_url(self.get_api_base(api_base)), headers=headers, timeout=_CATALOG_TIMEOUT + ) + response.raise_for_status() + return [ + f"openrouter/{model.id}" + for model in _Catalog.model_validate_json(response.content).data + if not model.id.endswith(_BATCH_ONLY_SUFFIX) + ] diff --git a/litellm/proxy/auth/model_checks.py b/litellm/proxy/auth/model_checks.py index b21315529eb..cfc0d2750bd 100644 --- a/litellm/proxy/auth/model_checks.py +++ b/litellm/proxy/auth/model_checks.py @@ -1,12 +1,15 @@ # What is this? ## Common checks for /v1/models and `/model/info` +import asyncio +import contextvars import copy import functools import itertools import re from collections.abc import Callable, Mapping, Sequence +from concurrent.futures import ThreadPoolExecutor from types import MappingProxyType -from typing import Final +from typing import Final, TypeVar import litellm from litellm._logging import verbose_proxy_logger @@ -30,6 +33,17 @@ _CREDENTIAL_LITELLM_PARAM_FIELDS = set(CredentialLiteLLMParams.model_fields) _CREDENTIAL_LITELLM_PARAM_FIELDS = set(CredentialLiteLLMParams.model_fields) +_MODEL_LISTING_POOL: Final = ThreadPoolExecutor(thread_name_prefix="litellm-model-listing") + +_PROVIDERS_LISTED_ONLY_THROUGH_A_DEPLOYMENT: Final = frozenset({LlmProviders.OPENROUTER.value}) + +_Listed: Final = TypeVar("_Listed") + + +async def run_model_listing(listing: Callable[[], _Listed]) -> _Listed: + context: Final = contextvars.copy_context() + return await asyncio.get_running_loop().run_in_executor(_MODEL_LISTING_POOL, context.run, listing) + def _check_wildcard_routing(model: str) -> bool: """ @@ -209,6 +223,7 @@ def get_complete_model_list( include_model_access_groups: bool | None = False, only_model_access_groups: bool | None = False, team_id: str | None = None, + expand_wildcard_routes: bool = True, ) -> list[str]: """Logic for returning complete model list for a given key + team pair""" @@ -254,11 +269,10 @@ def get_complete_model_list( return_wildcard_routes=return_wildcard_routes, llm_router=llm_router, team_id=team_id, + expand_wildcard_routes=expand_wildcard_routes, ) - complete_model_list: Final = unique_models + all_wildcard_models - - return complete_model_list + return list(dict.fromkeys(unique_models + all_wildcard_models)) def _hydrate_litellm_credential_name( @@ -425,6 +439,15 @@ def _cached_wildcard_routing(llm_router: Router | None) -> Callable[[str, str | return functools.cache(wildcard_routed_models) +def _deployed_by_name(llm_router: Router | None) -> Callable[[str, str | None], bool]: + if llm_router is None: + return lambda _model, _team_id: False + registry: Final = llm_router.registry() + deployed_names: Final = frozenset(registry.names(kind="group_set")) + team_deployed_names: Final = frozenset(registry.index("team_model")) + return lambda model, team_id: model in deployed_names or (team_id, model) in team_deployed_names + + def _routes_to_model( routed_models: Callable[[str, str | None], frozenset[str]] | None, team_id: str | None ) -> Callable[[str, str], bool]: @@ -455,7 +478,11 @@ def get_known_models_from_wildcard( wildcard_provider_prefix, wildcard_suffix = wildcard_model_to_expand.split("/", 1) except ValueError: # safely fail return [] - deployment_pattern_naming_listed_ids: Final = _deployment_pattern_naming_listed_ids(wildcard_suffix, litellm_params) + deployment_pattern_naming_listed_ids: Final = ( + wildcard_model_to_expand + if wildcard_model_to_expand != wildcard_model + else _deployment_pattern_naming_listed_ids(wildcard_suffix, litellm_params) + ) # Use provider from litellm_params when available, otherwise from wildcard prefix # (e.g., "openai" from "openai/*" - needed for BYOK where wildcard isn't in router) @@ -514,6 +541,27 @@ def _under_wildcard_prefix(model: str, wildcard_provider_prefix: str, listed_pro return f"{wildcard_provider_prefix}/{model_suffix if leading_is_listed_providers_own_prefix else model}" +@functools.lru_cache(maxsize=256) +def _wildcard_name_regex(name_pattern: str) -> re.Pattern[str]: + return re.compile(PatternMatchRouter().pattern_to_regex(name_pattern)) + + +def _routed_deployment_model(name_pattern: str, listed_name: str, deployment_model: str) -> str: + name_match: Final = _wildcard_name_regex(name_pattern).match(listed_name) + if name_match is None: + return listed_name + return PatternMatchRouter.set_deployment_model_name( + matched_pattern=name_match, litellm_deployment_litellm_model=deployment_model + ) + + +def _names_no_deployment_serves_by_name( + deployment: Mapping[str, object], names: Sequence[str], deployed_by_name: Callable[[str, str | None], bool] +) -> tuple[str, ...]: + team_id: Final = _deployment_team_id(deployment) + return tuple(name for name in names if not deployed_by_name(name, team_id)) + + def expand_wildcard_deployments_for_model_info( deployments: list[dict[str, object]], llm_router: Router | None = None, @@ -527,6 +575,7 @@ def expand_wildcard_deployments_for_model_info( """ expanded: Final[list[dict[str, object]]] = [] routed_models: Final = _cached_wildcard_routing(llm_router) + deployed_by_name: Final = _deployed_by_name(llm_router) for deployment in deployments: model_name = str(deployment.get("model_name") or "") raw_params = deployment.get("litellm_params") @@ -560,22 +609,29 @@ def expand_wildcard_deployments_for_model_info( expanded.append(deployment) continue - for name in expanded_names: + for name in _names_no_deployment_serves_by_name(deployment, expanded_names, deployed_by_name): row = copy.deepcopy(deployment) row["model_name"] = name params = row.get("litellm_params") if isinstance(params, dict): - params["model"] = name + params["model"] = _routed_deployment_model(model_name, name, litellm_model) expanded.append(row) return expanded +def _known_models_without_a_deployment(wildcard_model: str) -> Sequence[str]: + if wildcard_model.partition("/")[0] in _PROVIDERS_LISTED_ONLY_THROUGH_A_DEPLOYMENT: + return () + return get_known_models_from_wildcard(wildcard_model=wildcard_model, litellm_params=None) + + def _get_wildcard_models( unique_models: list[str], return_wildcard_routes: bool | None = False, llm_router: Router | None = None, team_id: str | None = None, + expand_wildcard_routes: bool = True, ) -> list[str]: models_to_remove: Final = set() routes_to_model: Final = _routes_to_model(_cached_wildcard_routing(llm_router), team_id) @@ -586,6 +642,8 @@ def _get_wildcard_models( all_wildcard_models.append(model) models_to_remove.add(model) + if not expand_wildcard_routes: + continue model_list = llm_router.registry().deployments(model, team_id=team_id) if llm_router else None if model_list: @@ -598,7 +656,7 @@ def _get_wildcard_models( ) ) else: - all_wildcard_models.extend(get_known_models_from_wildcard(wildcard_model=model, litellm_params=None)) + all_wildcard_models.extend(_known_models_without_a_deployment(model)) for model in models_to_remove: unique_models.remove(model) diff --git a/litellm/proxy/openai_files_endpoints/common_utils.py b/litellm/proxy/openai_files_endpoints/common_utils.py index 907c2ca303e..c0d81abeba7 100644 --- a/litellm/proxy/openai_files_endpoints/common_utils.py +++ b/litellm/proxy/openai_files_endpoints/common_utils.py @@ -1,4 +1,5 @@ import base64 +import itertools import logging import mimetypes import re @@ -541,7 +542,8 @@ def get_team_provider_credentials( listings on the team's own provider account/key instead of a shared global one. 2. Fallback: any deployment the caller is granted access to for this - provider, expanding wildcard routes and the all-proxy-models sentinel. + provider, matching wildcard grants against deployed model names without + listing any provider's catalog, and expanding the all-proxy-models sentinel. Skipped when ``include_accessible_deployments`` is False. Credential lookup is scoped to both the team's allowlist and the key's own @@ -559,6 +561,7 @@ def get_team_provider_credentials( return None from litellm.proxy._types import SpecialModelNames + from litellm.proxy.auth.auth_checks import is_model_allowed_by_pattern from litellm.proxy.auth.model_checks import get_complete_model_list, get_key_models team_id: Final = user_api_key_dict.team_id @@ -645,28 +648,36 @@ def get_team_provider_credentials( grants_all_models: Final = SpecialModelNames.all_proxy_models.value in team_models effective_team_models: Final = [] if grants_all_models else team_models - models_to_try: Final = list( - dict.fromkeys( - get_complete_model_list( - key_models=list(key_model_allowlist), - team_models=effective_team_models, - proxy_model_list=proxy_model_list, - user_model=None, - infer_model_from_keys=False, - return_wildcard_routes=True, - llm_router=llm_router, - model_access_groups=model_access_groups, - include_model_access_groups=True, - team_id=team_id, - ) + granted_models: Final = get_complete_model_list( + key_models=list(key_model_allowlist), + team_models=effective_team_models, + proxy_model_list=proxy_model_list, + user_model=None, + infer_model_from_keys=False, + return_wildcard_routes=True, + llm_router=llm_router, + model_access_groups=model_access_groups, + include_model_access_groups=True, + team_id=team_id, + expand_wildcard_routes=False, + ) + deployments_granted_by_wildcard: Final = ( + model_name + for model_name in proxy_model_list + if any( + is_model_allowed_by_pattern(model=model_name, allowed_model_pattern=granted) for granted in granted_models ) ) - for model_name in models_to_try: - credentials = _provider_credentials(model_name) - if credentials is not None: - return credentials - - return None + return next( + ( + credentials + for credentials in map( + _provider_credentials, dict.fromkeys(itertools.chain(granted_models, deployments_granted_by_wildcard)) + ) + if credentials is not None + ), + None, + ) HTTP_HEADER_VALUE: Final = re.compile(r"[\x21-\x7e]+(?:[ \t]+[\x21-\x7e]+)*") diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f078267c6b2..7ea1c246da5 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -199,6 +199,7 @@ from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.utils import ( PRICING_OVERRIDES_KEY, CountContextManagement, + LlmProviders, ModelResponse, ModelResponseStream, TextCompletionResponse, @@ -443,6 +444,7 @@ from litellm.proxy.auth.model_checks import ( get_key_models, get_mcp_server_ids, get_team_models, + run_model_listing, ) from litellm.proxy.auth.password_policy import validate_password_not_breached, validate_password_policy from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401, RUF100 # legacy module exports @@ -12487,20 +12489,23 @@ async def model_list( # Get complete model list including wildcard routes if requested from litellm.proxy.auth.model_checks import get_complete_model_list - all_models = get_complete_model_list( - key_models=[], - team_models=[], - proxy_model_list=[ - *proxy_model_list, - *(typesafe_allowlist_model_names() if asks_for_allowlist_names else ()), - ], - user_model=None, - infer_model_from_keys=False, - return_wildcard_routes=include_wildcard_routes, - llm_router=llm_router, - model_access_groups=model_access_groups, - include_model_access_groups=include_model_access_groups or False, - only_model_access_groups=only_model_access_groups or False, + all_models = await run_model_listing( + partial( + get_complete_model_list, + key_models=[], + team_models=[], + proxy_model_list=[ + *proxy_model_list, + *(typesafe_allowlist_model_names() if asks_for_allowlist_names else ()), + ], + user_model=None, + infer_model_from_keys=False, + return_wildcard_routes=include_wildcard_routes, + llm_router=llm_router, + model_access_groups=model_access_groups, + include_model_access_groups=include_model_access_groups or False, + only_model_access_groups=only_model_access_groups or False, + ) ) expanded_undiscoverable_names: Final = undiscoverable_model_names( @@ -14747,6 +14752,8 @@ from litellm.responses.litellm_completion_transformation.count_request import ( with_provider_file_ids, ) +_PROVIDERS_WITH_LISTING_ONLY_MODEL_INFO: Final = frozenset({LlmProviders.OPENROUTER}) + def _with_client_oauth_token( counter: BaseTokenCounter | None, client_oauth_token: str | None @@ -14782,14 +14789,13 @@ def _get_provider_token_counter( ) # Switch case pattern using existing get_provider_model_info - from litellm.types.utils import LlmProviders from litellm.utils import ProviderConfigManager # Convert string provider to LlmProviders enum llm_provider_enum: Final = LlmProviders(custom_llm_provider) # Add more provider mappings as needed - if llm_provider_enum: + if llm_provider_enum not in _PROVIDERS_WITH_LISTING_ONLY_MODEL_INFO: provider_model_info: Final = ProviderConfigManager.get_provider_model_info( model=full_model, provider=llm_provider_enum ) @@ -17941,7 +17947,7 @@ async def model_info_v1( alias_models: Final = copy.deepcopy(registry.deployments(source="alias")) all_models.extend(alias_models) - all_models = expand_wildcard_deployments_for_model_info(all_models, llm_router) + all_models = await run_model_listing(partial(expand_wildcard_deployments_for_model_info, all_models, llm_router)) allowed_model_names: Final = await _get_v1_model_info_allowed_model_names_with_access_groups( user_api_key_dict=user_api_key_dict, @@ -18279,14 +18285,17 @@ async def model_group_info( LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY, ) all_models_str: Final = ( - get_complete_model_list( - key_models=(), - team_models=(), - proxy_model_list=llm_router.registry().names(), - user_model=user_model, - infer_model_from_keys=general_settings.get("infer_model_from_keys", False), - return_wildcard_routes=False, - llm_router=llm_router, + await run_model_listing( + partial( + get_complete_model_list, + key_models=(), + team_models=(), + proxy_model_list=llm_router.registry().names(), + user_model=user_model, + infer_model_from_keys=general_settings.get("infer_model_from_keys", False), + return_wildcard_routes=False, + llm_router=llm_router, + ) ) if is_proxy_admin else await get_available_models_for_user( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index ddaa8d2e69d..00dbcd0f9a8 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -9500,7 +9500,7 @@ async def _agent_access_group_visible_models( ) -> frozenset[str] | None: """Models an agent key may still list once its attached access groups cap it, ``None`` when nothing caps it, so ``/v1/models`` never advertises a model the same key would be denied on.""" - from litellm.proxy.auth.model_checks import get_complete_model_list, get_team_models + from litellm.proxy.auth.model_checks import get_complete_model_list, get_team_models, run_model_listing if not user_api_key_dict.agent_id: return None @@ -9521,17 +9521,20 @@ async def _agent_access_group_visible_models( if not granted: return frozenset() return frozenset( - get_complete_model_list( - key_models=granted, - team_models=(), - proxy_model_list=proxy_model_list, - user_model=None, - infer_model_from_keys=False, - return_wildcard_routes=return_wildcard_routes, - llm_router=llm_router, - model_access_groups=model_access_groups, - include_model_access_groups=include_model_access_groups, - team_id=team_id, + await run_model_listing( + partial( + get_complete_model_list, + key_models=granted, + team_models=(), + proxy_model_list=proxy_model_list, + user_model=None, + infer_model_from_keys=False, + return_wildcard_routes=return_wildcard_routes, + llm_router=llm_router, + model_access_groups=model_access_groups, + include_model_access_groups=include_model_access_groups, + team_id=team_id, + ) ) ) @@ -9576,6 +9579,7 @@ async def get_available_models_for_user( get_complete_model_list, get_key_models, get_team_models, + run_model_listing, ) # Get proxy model list and access groups @@ -9656,18 +9660,21 @@ async def get_available_models_for_user( ) # Get complete model list - complete_models: Final = get_complete_model_list( - key_models=granted_key_models, - team_models=granted_team_models, - proxy_model_list=proxy_model_list, - user_model=user_model, - infer_model_from_keys=general_settings.get("infer_model_from_keys", False), - return_wildcard_routes=return_wildcard_routes, - llm_router=llm_router, - model_access_groups=model_access_groups, - include_model_access_groups=include_model_access_groups, - only_model_access_groups=only_model_access_groups, - team_id=effective_team_id, + complete_models: Final = await run_model_listing( + partial( + get_complete_model_list, + key_models=granted_key_models, + team_models=granted_team_models, + proxy_model_list=proxy_model_list, + user_model=user_model, + infer_model_from_keys=general_settings.get("infer_model_from_keys", False), + return_wildcard_routes=return_wildcard_routes, + llm_router=llm_router, + model_access_groups=model_access_groups, + include_model_access_groups=include_model_access_groups, + only_model_access_groups=only_model_access_groups, + team_id=effective_team_id, + ) ) all_models: Final = ( list( diff --git a/litellm/router_utils/pattern_match_deployments.py b/litellm/router_utils/pattern_match_deployments.py index 4dd217cf59b..2b8d1e9c11b 100644 --- a/litellm/router_utils/pattern_match_deployments.py +++ b/litellm/router_utils/pattern_match_deployments.py @@ -279,7 +279,9 @@ class PatternMatchRouter: PatternMatchRouter.set_deployment_model_name( matched_pattern=pattern_match, litellm_deployment_litellm_model=deployment["litellm_params"]["model"] ) - for deployment in self.patterns[pattern] + for deployment in self.patterns.get( # pyright: ignore[reportUnknownMemberType] # untyped deployment dicts + pattern, () + ) ) diff --git a/litellm/utils.py b/litellm/utils.py index 2ca8a49f968..560d0c31e4a 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -8724,11 +8724,13 @@ class AvailableModelsCache(InMemoryCache): cache_key: Final = self._get_cache_key(custom_llm_provider, litellm_params) - result: Final = cast(list[str] | None, self.get_cache(cache_key)) + result: Final = self._cached(cache_key) - if result is not None: - return copy.deepcopy(result) - return result + if isinstance(result, Sequence): + return list(result) + if self._cached(_failed_listing_key(cache_key)) is not None: + return [] + return None def set_cached_model_info( self, @@ -8738,12 +8740,46 @@ class AvailableModelsCache(InMemoryCache): ): """Set cached model info""" cache_key: Final = self._get_cache_key(custom_llm_provider, litellm_params) - self.set_cache(cache_key, copy.deepcopy(available_models)) + self._store(cache_key, copy.deepcopy(available_models)) + + def set_failed_model_listing( + self, custom_llm_provider: str, litellm_params: LiteLLM_Params | None, ttl_seconds: float + ) -> None: + cache_key: Final = self._get_cache_key(custom_llm_provider, litellm_params) + self._store(_failed_listing_key(cache_key), True, ttl=ttl_seconds) + + def listing_lock(self, custom_llm_provider: str, litellm_params: LiteLLM_Params | None) -> threading.RLock: + return _provider_listing_lock(self._get_cache_key(custom_llm_provider, litellm_params)) + + def _cached(self, key: str) -> Sequence[str] | bool | None: + return cast(Sequence[str] | bool | None, self.get_cache(key)) + + def _store(self, key: str, value: object, ttl: float | None = None) -> None: + self.set_cache(key, value, ttl=ttl) + + +def _failed_listing_key(cache_key: str) -> str: + return f"{cache_key}:failed" + + +_PROVIDER_LISTING_LOCK_GUARD: Final = threading.Lock() + + +@lru_cache(maxsize=1024) +def _cached_provider_listing_lock(cache_key: str) -> threading.RLock: + return threading.RLock() + + +def _provider_listing_lock(cache_key: str) -> threading.RLock: + with _PROVIDER_LISTING_LOCK_GUARD: + return _cached_provider_listing_lock(cache_key) # Global cache instance _model_cache: Final = AvailableModelsCache() +_FAILED_PROVIDER_LISTING_TTL_SECONDS: Final = 30 + def _infer_valid_provider_from_env_vars( custom_llm_provider: str | None = None, @@ -8791,20 +8827,27 @@ def _get_valid_models_from_provider_api( custom_llm_provider: str, litellm_params: LiteLLM_Params | None = None, ) -> list[str]: - try: - cached_result: Final = _model_cache.get_cached_model_info(custom_llm_provider, litellm_params) + with _model_cache.listing_lock(custom_llm_provider, litellm_params): + try: + cached_result: Final = _model_cache.get_cached_model_info(custom_llm_provider, litellm_params) - if cached_result is not None: - return cached_result - models: Final = provider_config.discover_models( - litellm_params=_provider_listing_params(custom_llm_provider, litellm_params) - ) + if cached_result is not None: + return cached_result + models: Final = provider_config.discover_models( + litellm_params=_provider_listing_params(custom_llm_provider, litellm_params) + ) - _model_cache.set_cached_model_info(custom_llm_provider, litellm_params, models) - return models - except Exception as e: - verbose_logger.warning("Error getting valid models: %s", e) - return [] + _model_cache.set_cached_model_info(custom_llm_provider, litellm_params, models) + return models + except httpx.TimeoutException as e: + verbose_logger.warning("Error getting valid models: %s", e) + _model_cache.set_failed_model_listing( + custom_llm_provider, litellm_params, ttl_seconds=_FAILED_PROVIDER_LISTING_TTL_SECONDS + ) + return [] + except Exception as e: + verbose_logger.warning("Error getting valid models: %s", e) + return [] def bundled_provider_models(provider: str) -> tuple[str, ...]: @@ -10470,6 +10513,10 @@ class ProviderConfigManager: from litellm.llms.azure_ai.common_utils import AzureFoundryModelInfo return AzureFoundryModelInfo(model=model) + elif LlmProviders.OPENROUTER == provider: + from litellm.llms.openrouter.common_utils import OpenRouterModelInfo + + return OpenRouterModelInfo() return None @staticmethod diff --git a/tests/unit/llms/openrouter/test_openrouter_common_utils.py b/tests/unit/llms/openrouter/test_openrouter_common_utils.py new file mode 100644 index 00000000000..087c05561b5 --- /dev/null +++ b/tests/unit/llms/openrouter/test_openrouter_common_utils.py @@ -0,0 +1,159 @@ +from typing import Final + +import httpx +import pytest +import respx + +import litellm +from litellm.llms.openrouter.common_utils import OpenRouterModelInfo + +DEFAULT_BASE: Final = "https://openrouter.ai/api/v1" +KEY_SOURCES: Final = ("api_key", "litellm.api_key", "litellm.openrouter_key", "OPENROUTER_API_KEY", "OR_API_KEY") +BASE_SOURCES: Final = ("api_base", "litellm.api_base", "OPENROUTER_API_BASE") + + +@pytest.fixture(autouse=True) +def _no_ambient_openrouter_settings(monkeypatch: pytest.MonkeyPatch): + for env_var in ("OPENROUTER_API_BASE", "OPENROUTER_API_KEY", "OR_API_KEY"): + monkeypatch.delenv(env_var, raising=False) + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "api_base", None) + monkeypatch.setattr(litellm, "openrouter_key", None) + + +def _catalog(*model_ids: str) -> httpx.Response: + return httpx.Response(200, json={"data": [{"id": model_id, "name": model_id} for model_id in model_ids]}) + + +def _configure(monkeypatch: pytest.MonkeyPatch, source: str, value: str) -> None: + if source.startswith("litellm."): + monkeypatch.setattr(litellm, source.removeprefix("litellm."), value) + else: + monkeypatch.setenv(source, value) + + +def test_get_models_lists_only_the_chat_catalog_under_openrouter_keeping_the_org_namespace(): + with respx.mock() as respx_mock: + respx_mock.get(f"{DEFAULT_BASE}/models/user").mock( + return_value=_catalog("openai/gpt-5.4-mini", "anthropic/claude-haiku-5-5") + ) + + assert OpenRouterModelInfo().get_models(api_key="sk-or-test") == [ + "openrouter/openai/gpt-5.4-mini", + "openrouter/anthropic/claude-haiku-5-5", + ] + assert [call.request.url.path for call in respx_mock.calls] == ["/api/v1/models/user"] + + +@respx.mock +def test_get_models_drops_batch_only_ids_and_keeps_every_other_variant(): + respx.get(f"{DEFAULT_BASE}/models/user").mock( + return_value=_catalog("openai/gpt-5.4-mini:batch", "openai/gpt-5.4-mini", "qwen/qwen3-coder:free") + ) + + assert OpenRouterModelInfo().get_models(api_key="sk-or-test") == [ + "openrouter/openai/gpt-5.4-mini", + "openrouter/qwen/qwen3-coder:free", + ] + + +@pytest.mark.parametrize( + "api_base", + [ + "https://gateway.example/openrouter/api/v1/", + "https://gateway.example/openrouter/api/v1/chat/completions", + "https://gateway.example/openrouter/api/v1/chat/completions/", + ], +) +@respx.mock +def test_get_models_keeps_the_api_base_path_and_sends_the_deployment_key(api_base: str): + chat_route: Final = respx.get("https://gateway.example/openrouter/api/v1/models/user").mock( + return_value=_catalog("openai/gpt-5.4-mini") + ) + + assert OpenRouterModelInfo().get_models(api_key="sk-or-deployment", api_base=api_base) == [ + "openrouter/openai/gpt-5.4-mini" + ] + assert chat_route.calls.last.request.headers["authorization"] == "Bearer sk-or-deployment" + + +@pytest.mark.parametrize("winner", range(len(KEY_SOURCES)), ids=KEY_SOURCES) +@respx.mock +def test_get_models_sends_the_first_configured_key_in_the_chat_paths_order( + monkeypatch: pytest.MonkeyPatch, winner: int +): + for source in KEY_SOURCES[winner + 1 :]: + _configure(monkeypatch, source, f"sk-or-from-{source}") + if winner > 0: + _configure(monkeypatch, KEY_SOURCES[winner], "sk-or-winner") + chat_route: Final = respx.get(f"{DEFAULT_BASE}/models/user").mock(return_value=_catalog("qwen/qwen3-coder")) + + assert OpenRouterModelInfo().get_models(api_key="sk-or-winner" if winner == 0 else None) == [ + "openrouter/qwen/qwen3-coder" + ] + assert chat_route.calls.last.request.headers["authorization"] == "Bearer sk-or-winner" + + +@pytest.mark.parametrize("winner", range(len(BASE_SOURCES) + 1), ids=(*BASE_SOURCES, "default")) +@respx.mock +def test_get_models_calls_the_first_configured_base_in_the_chat_paths_order( + monkeypatch: pytest.MonkeyPatch, winner: int +): + for source in BASE_SOURCES[winner + 1 :]: + _configure(monkeypatch, source, f"https://{source.lower().replace('.', '-')}.example/api/v1") + if 0 < winner < len(BASE_SOURCES): + _configure(monkeypatch, BASE_SOURCES[winner], "https://winner.example/api/v1") + expected_base: Final = "https://winner.example/api/v1" if winner < len(BASE_SOURCES) else DEFAULT_BASE + respx.get(f"{expected_base}/models/user").mock(return_value=_catalog("qwen/qwen3-coder")) + + assert OpenRouterModelInfo().get_models( + api_key="sk-or-test", api_base="https://winner.example/api/v1" if winner == 0 else None + ) == ["openrouter/qwen/qwen3-coder"] + + +def test_get_models_without_any_key_raises_before_calling_openrouter(): + with respx.mock(assert_all_called=False) as respx_mock: + chat_route: Final = respx_mock.get(f"{DEFAULT_BASE}/models/user").mock( + return_value=_catalog("openai/gpt-5.4-mini") + ) + + with pytest.raises(ValueError, match="OPENROUTER_API_KEY"): + OpenRouterModelInfo().get_models() + + assert not chat_route.called + + +@respx.mock +def test_get_models_raises_when_the_catalog_answers_an_error(): + respx.get(f"{DEFAULT_BASE}/models/user").mock(return_value=httpx.Response(401, json={"error": {"code": 401}})) + + with pytest.raises(httpx.HTTPStatusError): + OpenRouterModelInfo().get_models(api_key="sk-or-revoked") + + +def test_validate_environment_adds_the_resolved_key_and_keeps_the_callers_headers(monkeypatch: pytest.MonkeyPatch): + monkeypatch.setenv("OPENROUTER_API_KEY", "sk-or-env") + headers: Final = {"X-Title": "my-app"} + + assert OpenRouterModelInfo().validate_environment( + headers=headers, model="openrouter/openai/gpt-5.4-mini", messages=(), optional_params={}, litellm_params={} + ) == {"X-Title": "my-app", "Authorization": "Bearer sk-or-env"} + assert headers == {"X-Title": "my-app"} + + +def test_validate_environment_without_any_key_sends_no_authorization(): + assert OpenRouterModelInfo().validate_environment( + headers={"X-Title": "my-app"}, + model="openrouter/openai/gpt-5.4-mini", + messages=(), + optional_params={}, + litellm_params={}, + ) == {"X-Title": "my-app"} + + +@pytest.mark.parametrize( + ("model", "base_model"), + [("openrouter/openai/gpt-5.4-mini", "openai/gpt-5.4-mini"), ("openai/gpt-5.4-mini", "openai/gpt-5.4-mini")], +) +def test_get_base_model_drops_only_the_openrouter_prefix(model: str, base_model: str): + assert OpenRouterModelInfo.get_base_model(model) == base_model diff --git a/tests/unit/proxy/auth/test_model_checks.py b/tests/unit/proxy/auth/test_model_checks.py index c67610c24a1..d0156ff0195 100644 --- a/tests/unit/proxy/auth/test_model_checks.py +++ b/tests/unit/proxy/auth/test_model_checks.py @@ -7,6 +7,8 @@ import httpx import pytest import respx +from litellm import Router + def test_get_team_models_for_all_models_and_team_only_models(): from litellm.proxy.auth.model_checks import get_team_models @@ -1261,6 +1263,223 @@ def test_partial_bedrock_wildcard_lists_nothing_when_no_discovered_id_carries_it assert get_known_models_from_wildcard("bedrock/anthropic.*", deployment) == [] +@pytest.mark.parametrize( + ("wildcard", "deployment_model", "listed_ids"), + [ + ( + "openrouter/*", + "openrouter/*", + ["openrouter/openai/gpt-5.4-mini", "openrouter/anthropic/claude-haiku-5-5"], + ), + ("or-team/*", "openrouter/*", ["or-team/openai/gpt-5.4-mini", "or-team/anthropic/claude-haiku-5-5"]), + ("openrouter/openai/*", "openrouter/openai/*", ["openrouter/openai/gpt-5.4-mini"]), + ("*", "openrouter/*", ["openai/gpt-5.4-mini", "anthropic/claude-haiku-5-5"]), + ], +) +@respx.mock +def test_openrouter_wildcard_lists_the_live_catalog_with_each_vendor_namespace( + monkeypatch: pytest.MonkeyPatch, wildcard: str, deployment_model: str, listed_ids: list[str] +) -> None: + import litellm + from litellm.proxy.auth.model_checks import get_known_models_from_wildcard + from litellm.types.router import LiteLLM_Params + + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + monkeypatch.delenv("OPENROUTER_API_BASE", raising=False) + respx.get("https://openrouter.ai/api/v1/models/user").mock( + return_value=httpx.Response( + 200, json={"data": [{"id": "openai/gpt-5.4-mini"}, {"id": "anthropic/claude-haiku-5-5"}]} + ) + ) + + assert ( + get_known_models_from_wildcard(wildcard, LiteLLM_Params(model=deployment_model, api_key=f"sk-or-{wildcard}")) + == listed_ids + ) + + +@respx.mock +def test_key_granted_openrouter_wildcard_without_an_openrouter_deployment_lists_no_environment_catalog( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import litellm + from litellm import Router + from litellm.proxy.auth.model_checks import get_complete_model_list + + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + monkeypatch.setenv("OPENROUTER_API_KEY", "sk-or-environment-only") + catalog: Final = respx.get("https://openrouter.ai/api/v1/models/user").mock( + return_value=httpx.Response(200, json={"data": [{"id": "anthropic/claude-haiku-5-5"}]}) + ) + router: Final = Router( + model_list=[{"model_name": "or-team/*", "litellm_params": {"model": "openrouter/*", "api_key": "sk-or-team"}}] + ) + + listed: Final = get_complete_model_list( + key_models=["openrouter/*"], + team_models=[], + proxy_model_list=router.model_names, + user_model=None, + infer_model_from_keys=False, + llm_router=router, + ) + + assert (listed, catalog.call_count) == ([], 0) + + +@contextmanager +def _openai_catalog_behind_a_catch_all( + monkeypatch: pytest.MonkeyPatch, concrete_name: str, concrete_model_info: Mapping[str, str] +) -> Generator[Router, None, None]: + import litellm + + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + with respx.mock(assert_all_called=False) as respx_mock: + respx_mock.get("https://api.openai.com/v1/models").mock( + return_value=httpx.Response(200, json={"data": [{"id": "gpt-4o"}, {"id": "gpt-4o-mini"}]}) + ) + yield Router( + model_list=[ + {"model_name": "*", "litellm_params": {"model": "openai/*", "api_key": "sk-openai-catch-all"}}, + { + "model_name": concrete_name, + "litellm_params": {"model": "openrouter/openai/gpt-4o", "api_key": "sk-or-gpt-4o"}, + "model_info": dict(concrete_model_info), + }, + ] + ) + + +@pytest.mark.parametrize( + ("concrete_name", "concrete_model_info", "team_id"), + [ + ("gpt-4o", {}, None), + ("model_name_team-x_gpt-4o", {"team_id": "team-x", "team_public_model_name": "gpt-4o"}, "team-x"), + ], +) +def test_catch_all_wildcard_lists_a_name_a_concrete_deployment_serves_once( + monkeypatch: pytest.MonkeyPatch, concrete_name: str, concrete_model_info: Mapping[str, str], team_id: str | None +) -> None: + from litellm.proxy.auth.model_checks import get_complete_model_list + + with _openai_catalog_behind_a_catch_all(monkeypatch, concrete_name, concrete_model_info) as router: + listed: Final = get_complete_model_list( + key_models=[], + team_models=[], + proxy_model_list=router.get_model_names(team_id=team_id), + user_model=None, + infer_model_from_keys=False, + llm_router=router, + team_id=team_id, + ) + + assert sorted(listed) == ["gpt-4o", "gpt-4o-mini"] + + +def test_catch_all_wildcard_model_info_keeps_only_the_concrete_row_for_a_deployed_name( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.proxy.auth.model_checks import expand_wildcard_deployments_for_model_info + + with _openai_catalog_behind_a_catch_all(monkeypatch, "gpt-4o", {}) as router: + rows: Final = expand_wildcard_deployments_for_model_info( + [dict(deployment) for deployment in router.get_model_list() or ()], llm_router=router + ) + + assert sorted((row["model_name"], row["litellm_params"]["model"]) for row in rows) == [ + ("gpt-4o", "openrouter/openai/gpt-4o"), + ("gpt-4o-mini", "openai/gpt-4o-mini"), + ] + + +@contextmanager +def _openrouter_catalog_beside_a_named_deployment( + monkeypatch: pytest.MonkeyPatch, + wildcard_name: str, + named: str, + catalog: tuple[str, ...] = ("openai/gpt-5.4-mini", "anthropic/claude-haiku-5-5"), + wildcard_key: str | None = None, +) -> Generator[Router, None, None]: + import litellm + + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + monkeypatch.delenv("OPENROUTER_API_BASE", raising=False) + with respx.mock(assert_all_called=False) as respx_mock: + respx_mock.get("https://openrouter.ai/api/v1/models/user").mock( + return_value=httpx.Response(200, json={"data": [{"id": catalog_id} for catalog_id in catalog]}) + ) + yield Router( + model_list=[ + { + "model_name": wildcard_name, + "litellm_params": { + "model": "openrouter/*", + "api_key": wildcard_key or f"sk-or-wildcard-{wildcard_name}", + }, + }, + { + "model_name": named, + "litellm_params": {"model": "openrouter/openai/gpt-5.4-mini", "api_key": "sk-or-named"}, + }, + ] + ) + + +@pytest.mark.parametrize( + ("wildcard_name", "named"), + [("openrouter/*", "openrouter/openai/gpt-5.4-mini"), ("or-team/*", "or-team/openai/gpt-5.4-mini")], +) +def test_provider_wildcard_lists_a_name_a_named_deployment_serves_once( + monkeypatch: pytest.MonkeyPatch, wildcard_name: str, named: str +) -> None: + from litellm.proxy.auth.model_checks import expand_wildcard_deployments_for_model_info, get_complete_model_list + + with _openrouter_catalog_beside_a_named_deployment(monkeypatch, wildcard_name, named) as router: + listed: Final = get_complete_model_list( + key_models=[], + team_models=[], + proxy_model_list=list(router.model_names), + user_model=None, + infer_model_from_keys=False, + llm_router=router, + ) + rows: Final = expand_wildcard_deployments_for_model_info( + [dict(deployment) for deployment in router.get_model_list() or ()], llm_router=router + ) + + vendor_prefix: Final = wildcard_name.removesuffix("*") + assert sorted(listed) == [f"{vendor_prefix}anthropic/claude-haiku-5-5", named] + assert sorted((row["model_name"], row["litellm_params"]["api_key"]) for row in rows) == [ + (f"{vendor_prefix}anthropic/claude-haiku-5-5", f"sk-or-wildcard-{wildcard_name}"), + (named, "sk-or-named"), + ] + + +@pytest.mark.parametrize( + ("wildcard_name", "named"), + [("openrouter/*", "openrouter/openai/gpt-5.4-mini"), ("or-team/*", "or-team/openai/gpt-5.4-mini")], +) +def test_provider_wildcard_whose_whole_catalog_is_deployed_by_name_shows_only_the_named_row( + monkeypatch: pytest.MonkeyPatch, wildcard_name: str, named: str +) -> None: + from litellm.proxy.auth.model_checks import expand_wildcard_deployments_for_model_info + + with _openrouter_catalog_beside_a_named_deployment( + monkeypatch, + wildcard_name, + named, + catalog=("openai/gpt-5.4-mini",), + wildcard_key=f"sk-or-one-model-{wildcard_name}", + ) as router: + rows: Final = expand_wildcard_deployments_for_model_info( + [dict(deployment) for deployment in router.get_model_list() or ()], llm_router=router + ) + + assert [(row["model_name"], row["litellm_params"]["api_key"]) for row in rows] == [(named, "sk-or-named")] + + @respx.mock def test_custom_prefix_that_starts_a_discovered_id_still_prefixes_every_listed_model( monkeypatch: pytest.MonkeyPatch, diff --git a/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py b/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py index 82523572b28..238f1fa6495 100644 --- a/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py +++ b/tests/unit/proxy/openai_files_endpoint/test_files_common_utils.py @@ -3,10 +3,13 @@ from types import MappingProxyType from typing import Final from unittest.mock import AsyncMock, MagicMock +import httpx import pytest +import respx from fastapi import HTTPException from litellm_enterprise.proxy.hooks.managed_files import _PROXY_LiteLLMManagedFiles +import litellm from litellm import Router from litellm.caching.caching import DualCache from litellm.proxy._types import UserAPIKeyAuth @@ -852,3 +855,112 @@ def test_a_team_keeps_its_first_config_file_openai_deployment_over_older_db_ones assert credentials is not None assert credentials["api_key"] == "deployment-z-key" + + +@pytest.mark.parametrize( + ("custom_llm_provider", "openrouter_key", "api_key"), + [("openai", "sk-or-files-openai", "openai-key"), ("openrouter", "sk-or-files-openrouter", "sk-or-files-openrouter")], +) +def test_a_provider_credential_lookup_never_lists_an_openrouter_wildcards_live_catalog( + monkeypatch: pytest.MonkeyPatch, custom_llm_provider: str, openrouter_key: str, api_key: str +): + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + router: Final = Router( + model_list=[ + { + "model_name": "openrouter/*", + "litellm_params": { + "model": "openrouter/*", + "api_key": openrouter_key, + "api_base": "https://openrouter.ai/api/v1", + }, + }, + {"model_name": "gpt-5.4-mini", "litellm_params": {"model": "openai/gpt-5.4-mini", "api_key": "openai-key"}}, + ] + ) + + with respx.mock(assert_all_called=False) as respx_mock: + catalog_route: Final = respx_mock.get("https://openrouter.ai/api/v1/models/user").mock( + return_value=httpx.Response(200, json={"data": [{"id": "openai/gpt-5.4-mini"}]}) + ) + credentials: Final = get_team_provider_credentials( + llm_router=router, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", user_id="test-user"), + custom_llm_provider=custom_llm_provider, + include_accessible_deployments=True, + ) + + assert credentials is not None + assert (credentials["custom_llm_provider"], credentials["api_key"]) == (custom_llm_provider, api_key) + assert catalog_route.call_count == 0 + + +@pytest.mark.parametrize( + ("key_models", "deployments", "custom_llm_provider", "api_key"), + [ + ( + ["anthropic/*"], + (("anthropic/claude-*", "anthropic/claude-*", "anthropic-narrow-key"),), + "anthropic", + "anthropic-narrow-key", + ), + (["anthropic/claude-3*"], (("anthropic/*", "anthropic/*", "anthropic-wide-key"),), "anthropic", "anthropic-wide-key"), + (["openrouter/*"], (("openrouter/openai/*", "openrouter/openai/*", "sk-or-narrow"),), "openrouter", "sk-or-narrow"), + ( + ["team/*"], + (("team/*", "openrouter/*", "sk-or-team"), ("team/openai/*", "openai/*", "openai-team-key")), + "openai", + "openai-team-key", + ), + ], +) +def test_a_key_wildcard_finds_the_deployment_it_grants_without_listing_any_provider_catalog( + monkeypatch: pytest.MonkeyPatch, + key_models: list[str], + deployments: tuple[tuple[str, str, str], ...], + custom_llm_provider: str, + api_key: str, +): + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + router: Final = Router( + model_list=[ + {"model_name": model_name, "litellm_params": {"model": model, "api_key": deployment_key}} + for model_name, model, deployment_key in deployments + ] + ) + + with respx.mock(assert_all_called=False) as respx_mock: + catalog_requests: Final = respx_mock.route().mock(return_value=httpx.Response(200, json={"data": []})) + credentials: Final = get_team_provider_credentials( + llm_router=router, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", user_id="test-user", models=key_models), + custom_llm_provider=custom_llm_provider, + include_accessible_deployments=True, + ) + + assert credentials is not None + assert (credentials["custom_llm_provider"], credentials["api_key"]) == (custom_llm_provider, api_key) + assert catalog_requests.call_count == 0 + + +@pytest.mark.parametrize( + ("deployed_name", "api_key"), [("gpt-4o-mini", None), ("openai/gpt-4o-mini", "openai-deployment-key")] +) +def test_a_provider_wildcard_grant_finds_only_deployments_named_under_it(deployed_name: str, api_key: str | None): + router: Final = Router( + model_list=[ + { + "model_name": deployed_name, + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "openai-deployment-key"}, + } + ] + ) + + credentials: Final = get_team_provider_credentials( + llm_router=router, + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", user_id="test-user", models=["openai/*"]), + custom_llm_provider="openai", + include_accessible_deployments=True, + ) + + assert (None if credentials is None else credentials["api_key"]) == api_key diff --git a/tests/unit/proxy/proxy_server/test_routes_model_info.py b/tests/unit/proxy/proxy_server/test_routes_model_info.py index 86364d93ed0..bffa4612964 100644 --- a/tests/unit/proxy/proxy_server/test_routes_model_info.py +++ b/tests/unit/proxy/proxy_server/test_routes_model_info.py @@ -550,15 +550,12 @@ async def test_v1_model_info_star_wildcard_filter_keeps_provider_expansion(monke deployment = { "model_name": "*", - "litellm_params": {"model": "openai/*"}, + "litellm_params": {"model": "openai/*", "api_key": "sk-star"}, } - router = route_entry_points_to_stubs(MagicMock()) - router.get_model_access_groups = MagicMock(return_value={}) - router.get_model_names = MagicMock(return_value=["*"]) - router.get_model_list = MagicMock(return_value=[deployment]) + router = litellm.Router(model_list=[deployment]) monkeypatch.setattr(model_checks, "get_provider_models", fake_get_provider_models) - expanded_deployments = proxy_server.expand_wildcard_deployments_for_model_info([deployment]) + expanded_deployments = proxy_server.expand_wildcard_deployments_for_model_info(router.get_model_list(), router) allowed_model_names = await proxy_server._get_v1_model_info_allowed_model_names_with_access_groups( user_api_key_dict=UserAPIKeyAuth( api_key="sk-test", @@ -575,7 +572,8 @@ async def test_v1_model_info_star_wildcard_filter_keeps_provider_expansion(monke allowed_model_names=allowed_model_names, ) - assert [model["model_name"] for model in result] == ["openai/gpt-4o"] + assert [(model["model_name"], model["litellm_params"]["model"]) for model in result] == [("gpt-4o", "openai/gpt-4o")] + assert router.registry().names("gpt-4o", kind="wildcard") == frozenset({"openai/gpt-4o"}) # --------------------------------------------------------------------------- diff --git a/tests/unit/proxy/proxy_server/test_routes_models.py b/tests/unit/proxy/proxy_server/test_routes_models.py index ffd23fa3a6f..874f471f3e7 100644 --- a/tests/unit/proxy/proxy_server/test_routes_models.py +++ b/tests/unit/proxy/proxy_server/test_routes_models.py @@ -9,10 +9,16 @@ Pins (PR2): from __future__ import annotations +import asyncio +import threading +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field from typing import Final from unittest.mock import MagicMock +import httpx import pytest +import respx import litellm from litellm import Router @@ -541,6 +547,73 @@ def test_codex_service_tiers_resolved_via_internal_team_key(client, auth_as, pat assert (entry["slug"], [tier["id"] for tier in entry["service_tiers"]]) == ("gpt-4-team", ["ultrafast"]) +OPENROUTER_BASE: Final = "https://openrouter.ai/api/v1" + + +@dataclass(frozen=True, slots=True) +class _CatalogHeldUntilTheLoopRuns: + requested: threading.Event = field(default_factory=threading.Event) + released: threading.Event = field(default_factory=threading.Event) + + def answer(self, _request: httpx.Request) -> httpx.Response: + self.requested.set() + if not self.released.wait(timeout=30): + return httpx.Response(504) + return httpx.Response(200, json={"data": [{"id": "openai/gpt-5.4-mini"}]}) + + async def release_from_the_loop(self) -> None: + assert await asyncio.to_thread(self.requested.wait, 30) + self.released.set() + + +@pytest.mark.parametrize("scope", [None, "expand"]) +def test_v1_models_fetches_a_live_openrouter_catalog_without_holding_the_event_loop_or_its_default_thread_pool( + monkeypatch: pytest.MonkeyPatch, scope: str | None +): + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(proxy_server, "user_model", None) + monkeypatch.setattr(proxy_server, "prisma_client", None) + monkeypatch.setattr( + proxy_server, + "llm_router", + litellm.Router( + model_list=[ + { + "model_name": "openrouter/*", + "litellm_params": { + "model": "openrouter/*", + "api_key": f"sk-or-v1-models-off-loop-{scope}", + "api_base": OPENROUTER_BASE, + }, + } + ] + ), + ) + catalog: Final = _CatalogHeldUntilTheLoopRuns() + + async def listed_ids_on_a_one_thread_default_pool() -> list[str]: + asyncio.get_running_loop().set_default_executor(ThreadPoolExecutor(max_workers=1)) + listing: Final = asyncio.ensure_future( + proxy_server.model_list( + user_api_key_dict=UserAPIKeyAuth( + api_key="sk-test-key", user_id="admin", user_role=LitellmUserRoles.PROXY_ADMIN + ), + scope=scope, + ) + ) + await catalog.release_from_the_loop() + response: Final = await listing + return [str(model["id"]) for model in response["data"]] + + with respx.mock() as respx_mock: + catalog_route: Final = respx_mock.get(f"{OPENROUTER_BASE}/models/user").mock(side_effect=catalog.answer) + listed_ids: Final = asyncio.run(listed_ids_on_a_one_thread_default_pool()) + + assert catalog_route.calls.last.response.status_code == 200 + assert listed_ids == ["openrouter/openai/gpt-5.4-mini"] + + def _priced_jev_models() -> frozenset[str]: return frozenset(name.removeprefix("typesafe/") for name in litellm.model_cost if name.startswith("typesafe/")) diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index c182cd1ec75..68eb05e60dd 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -179,6 +179,39 @@ async def test_token_counting_model_not_in_model_list(): assert response.model_used == "special-alias" +@pytest.mark.asyncio +async def test_asking_the_provider_endpoint_leaves_an_openrouter_deployments_count_unchanged( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr( + litellm.proxy.proxy_server, + "llm_router", + Router( + model_list=[ + { + "model_name": "or-gpt", + "litellm_params": { + "model": "openai/gpt-4o", + "custom_llm_provider": "openrouter", + "api_key": "sk-or-token-count", + }, + } + ] + ), + ) + request: Final = TokenCountRequest( + model="or-gpt", messages=[{"role": "user", "content": "Count these tokens, ünïcödé 你好 🙂"}] + ) + + via_provider_endpoint: Final = await token_counter(request=request, call_endpoint=True) + counted_locally: Final = await token_counter(request=request, call_endpoint=False) + + assert (via_provider_endpoint.total_tokens, via_provider_endpoint.model_used) == ( + counted_locally.total_tokens, + counted_locally.model_used, + ) + + @pytest.mark.asyncio async def test_gpt_token_counting(): """ diff --git a/tests/unit/proxy/utils/helpers/test_model_access.py b/tests/unit/proxy/utils/helpers/test_model_access.py index 324360cc2ee..30d24e9eddc 100644 --- a/tests/unit/proxy/utils/helpers/test_model_access.py +++ b/tests/unit/proxy/utils/helpers/test_model_access.py @@ -1,10 +1,17 @@ +import asyncio +import threading +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass, field +from typing import Final from unittest.mock import MagicMock +import httpx import pytest +import respx from fastapi import HTTPException import litellm -from litellm import ModelResponse +from litellm import ModelResponse, Router from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.utils import ( create_model_info_response, @@ -650,3 +657,61 @@ async def test_listing_is_unchanged_without_an_agent_or_without_attached_groups( resolve_agent_ceiling=_agent_ceiling(None), ) assert (plain_key, agent_without_groups) == (["model-a", "model-b"], ["model-a", "model-b"]) + + +OPENROUTER_BASE: Final = "https://openrouter.ai/api/v1" + + +@dataclass(frozen=True, slots=True) +class _CatalogHeldUntilTheLoopRuns: + requested: threading.Event = field(default_factory=threading.Event) + released: threading.Event = field(default_factory=threading.Event) + + def answer(self, _request: httpx.Request) -> httpx.Response: + self.requested.set() + if not self.released.wait(timeout=30): + return httpx.Response(504) + return httpx.Response(200, json={"data": [{"id": "openai/gpt-5.4-mini"}]}) + + async def release_from_the_loop(self) -> None: + assert await asyncio.to_thread(self.requested.wait, 30) + self.released.set() + + +def test_get_available_models_for_user_fetches_a_live_catalog_without_holding_the_event_loop_or_its_default_pool( + monkeypatch: pytest.MonkeyPatch, +): + monkeypatch.setattr(litellm, "check_provider_endpoint", True) + catalog: Final = _CatalogHeldUntilTheLoopRuns() + router: Final = Router( + model_list=[ + { + "model_name": "openrouter/*", + "litellm_params": { + "model": "openrouter/*", + "api_key": "sk-or-available-models-off-loop", + "api_base": OPENROUTER_BASE, + }, + } + ] + ) + + async def list_models_on_a_one_thread_default_pool() -> list[str]: + asyncio.get_running_loop().set_default_executor(ThreadPoolExecutor(max_workers=1)) + listing: Final = asyncio.ensure_future( + get_available_models_for_user( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-test-key", user_id="user-1"), + llm_router=router, + general_settings={}, + user_model=None, + ) + ) + await catalog.release_from_the_loop() + return await listing + + with respx.mock() as respx_mock: + catalog_route: Final = respx_mock.get(f"{OPENROUTER_BASE}/models/user").mock(side_effect=catalog.answer) + models: Final = asyncio.run(list_models_on_a_one_thread_default_pool()) + + assert catalog_route.calls.last.response.status_code == 200 + assert models == ["openrouter/openai/gpt-5.4-mini"] diff --git a/tests/unit/test_utils.py b/tests/unit/test_utils.py index 0f471b1de71..5cf872f4ea7 100644 --- a/tests/unit/test_utils.py +++ b/tests/unit/test_utils.py @@ -9269,6 +9269,77 @@ def reset_mock_cache() -> None: _model_cache.flush_cache() +@respx.mock +def test_get_valid_models_does_not_ask_a_timed_out_provider_again_on_the_next_listing(reset_mock_cache: None) -> None: + catalog: Final = respx.get("https://openrouter.ai/api/v1/models/user").mock( + side_effect=httpx.ReadTimeout("catalog hung") + ) + deployment: Final = LiteLLM_Params(model="openrouter/*", api_key="sk-or-during-an-outage") + + listings: Final = tuple( + get_valid_models(check_provider_endpoint=True, custom_llm_provider="openrouter", litellm_params=deployment) + for _ in range(2) + ) + + assert (listings, catalog.call_count) == (([], []), 1) + + +@respx.mock +def test_get_valid_models_lists_a_provider_again_right_after_a_fast_failure(reset_mock_cache: None) -> None: + respx.get("https://openrouter.ai/api/v1/models/user").mock( + side_effect=(httpx.Response(503), httpx.Response(200, json={"data": [{"id": "anthropic/claude-haiku-5-5"}]})) + ) + deployment: Final = LiteLLM_Params(model="openrouter/*", api_key="sk-or-recovers") + + listings: Final = tuple( + get_valid_models(check_provider_endpoint=True, custom_llm_provider="openrouter", litellm_params=deployment) + for _ in range(2) + ) + + assert listings == ([], ["openrouter/anthropic/claude-haiku-5-5"]) + + +@respx.mock +def test_get_valid_models_keeps_a_good_listing_when_an_overlapping_listing_fails_after_it( + reset_mock_cache: None, +) -> None: + deployment: Final = LiteLLM_Params(model="openrouter/*", api_key="sk-or-overlapping-listings") + first_request: Final = iter((True, False)) + + def answer_after_an_overlapping_listing(request: httpx.Request) -> httpx.Response: + if next(first_request): + get_valid_models(check_provider_endpoint=True, custom_llm_provider="openrouter", litellm_params=deployment) + raise httpx.ReadTimeout("catalog hung", request=request) + return httpx.Response(200, json={"data": [{"id": "anthropic/claude-haiku-5-5"}]}) + + respx.get("https://openrouter.ai/api/v1/models/user").mock(side_effect=answer_after_an_overlapping_listing) + get_valid_models(check_provider_endpoint=True, custom_llm_provider="openrouter", litellm_params=deployment) + + assert get_valid_models( + check_provider_endpoint=True, custom_llm_provider="openrouter", litellm_params=deployment + ) == ["openrouter/anthropic/claude-haiku-5-5"] + + +@respx.mock +def test_get_valid_models_asks_a_provider_once_when_many_listings_start_together(reset_mock_cache: None) -> None: + def slow_catalog(request: httpx.Request) -> httpx.Response: + threading.Event().wait(0.2) + return httpx.Response(200, json={"data": [{"id": "anthropic/claude-haiku-5-5"}]}) + + catalog: Final = respx.get("https://openrouter.ai/api/v1/models/user").mock(side_effect=slow_catalog) + deployment: Final = LiteLLM_Params(model="openrouter/*", api_key="sk-or-many-listings") + + def list_models(_: int) -> tuple[str, ...]: + return tuple( + get_valid_models(check_provider_endpoint=True, custom_llm_provider="openrouter", litellm_params=deployment) + ) + + with ThreadPoolExecutor(max_workers=8) as pool: + listings: Final = frozenset(pool.map(list_models, range(8))) + + assert (listings, catalog.call_count) == (frozenset({("openrouter/anthropic/claude-haiku-5-5",)}), 1) + + @pytest.fixture def isolated_openai_model_sets(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr( From 4f7150f832b3180a6dfcce2be5a052038de45eac Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 30 Oct 2026 12:21:30 -0700 Subject: [PATCH 10/46] fix(ollama): forward assistant tool call ids in ollama_chat requests (#49278) * test(ollama): cover assistant tool call ids in ollama_chat requests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ollama): forward assistant tool call ids in ollama_chat requests Co-authored-by: Kenneth Bingham Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(ollama): format tool call id regression assertion Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ollama): keep tool call ID handling type safe Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(ollama): format tool-call id helper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ollama): preserve typed tool call arguments 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: Kenneth Bingham --- litellm/llms/ollama/chat/transformation.py | 24 +++-- litellm/types/llms/ollama.py | 3 + .../ollama/test_ollama_chat_transformation.py | 93 +++++++++++++++++++ 3 files changed, 110 insertions(+), 10 deletions(-) diff --git a/litellm/llms/ollama/chat/transformation.py b/litellm/llms/ollama/chat/transformation.py index 755d647da1d..168b562b358 100644 --- a/litellm/llms/ollama/chat/transformation.py +++ b/litellm/llms/ollama/chat/transformation.py @@ -235,6 +235,19 @@ class OllamaChatConfig(BaseConfig): litellm_params["function_name"] = function_name tools: Final = optional_params.pop("tools", None) + def transform_tool_call(typed_tool: ChatCompletionAssistantToolCall) -> OllamaToolCall: + function: Final = OllamaToolCallFunction( + name=typed_tool["function"].get("name") or "", + arguments=json.loads(typed_tool["function"]["arguments"]) + if "arguments" in typed_tool["function"] + else {}, + ) + tool_id: Final = typed_tool.get("id") + ollama_tool_call: Final = ( + OllamaToolCall(function=function) if tool_id is None else OllamaToolCall(id=tool_id, function=function) + ) + return ollama_tool_call + new_messages: Final = [] for m in messages: if isinstance( @@ -248,16 +261,7 @@ class OllamaChatConfig(BaseConfig): for tool in tool_calls: typed_tool = ChatCompletionAssistantToolCall(**tool) if typed_tool["type"] == "function": - arguments = {} - if "arguments" in typed_tool["function"]: - arguments = json.loads(typed_tool["function"]["arguments"]) - ollama_tool_call = OllamaToolCall( - function=OllamaToolCallFunction( - name=typed_tool["function"].get("name") or "", - arguments=arguments, - ) - ) - new_tools.append(ollama_tool_call) + new_tools.append(transform_tool_call(typed_tool)) reasoning_content, parsed_content = extract_reasoning_content(cast(dict, m)) content_str = convert_content_list_to_str(cast(AllMessageValues, m)) images = extract_images_from_message(cast(AllMessageValues, m)) diff --git a/litellm/types/llms/ollama.py b/litellm/types/llms/ollama.py index 4783eac9a8d..fb9d9e8c363 100644 --- a/litellm/types/llms/ollama.py +++ b/litellm/types/llms/ollama.py @@ -1,4 +1,6 @@ from typing_extensions import ( + NotRequired, + ReadOnly, Required, TypedDict, ) @@ -12,6 +14,7 @@ class OllamaToolCallFunction( class OllamaToolCall(TypedDict): + id: NotRequired[ReadOnly[str]] function: OllamaToolCallFunction diff --git a/tests/unit/llms/ollama/test_ollama_chat_transformation.py b/tests/unit/llms/ollama/test_ollama_chat_transformation.py index af1b85cf325..6a80f29ac43 100644 --- a/tests/unit/llms/ollama/test_ollama_chat_transformation.py +++ b/tests/unit/llms/ollama/test_ollama_chat_transformation.py @@ -856,6 +856,99 @@ class TestOllamaReasoningContentStreaming: class TestOllamaToolCallTransformation: + @pytest.mark.parametrize("stream", [False, True]) + def test_transform_request_keeps_parallel_call_ids_paired_with_results(self, stream: bool): + config: Final = OllamaChatConfig() + messages: Final = cast( + list[AllMessageValues], + [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_first", + "type": "function", + "function": { + "name": "read_fixture", + "arguments": '{"path": "first.txt"}', + }, + }, + { + "id": "call_second", + "type": "function", + "function": { + "name": "read_fixture", + "arguments": '{"path": "second.txt"}', + }, + }, + ], + }, + { + "role": "tool", + "tool_call_id": "call_second", + "content": "second result", + }, + { + "role": "tool", + "tool_call_id": "call_first", + "content": "first result", + }, + ], + ) + + result: Final = config.transform_request( + model="ollama_chat/test-model", + messages=messages, + optional_params={"stream": stream}, + litellm_params={}, + headers={}, + ) + + assert result["messages"][0]["tool_calls"] == [ + { + "id": "call_first", + "function": {"name": "read_fixture", "arguments": {"path": "first.txt"}}, + }, + { + "id": "call_second", + "function": {"name": "read_fixture", "arguments": {"path": "second.txt"}}, + }, + ] + assert [message.get("tool_call_id") for message in result["messages"][1:]] == [ + "call_second", + "call_first", + ] + + def test_transform_request_does_not_invent_an_id_for_an_unidentified_call(self): + config: Final = OllamaChatConfig() + messages: Final = cast( + list[AllMessageValues], + [ + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": None, + "type": "function", + "function": {"name": "read_fixture", "arguments": "{}"}, + } + ], + } + ], + ) + + result: Final = config.transform_request( + model="ollama_chat/test-model", + messages=messages, + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert result["messages"][0]["tool_calls"] == [{"function": {"name": "read_fixture", "arguments": {}}}] + def test_transform_request_preserves_tool_calls(self): """ tool_calls on assistant messages must survive transform_request. From b2466f8d7b2c783c04d424009b4b5b1066e06d23 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Fri, 30 Oct 2026 12:31:16 -0700 Subject: [PATCH 11/46] chore(cost-map): flag Vertex Gemini computer use and add Cohere Rerank 5 (#49297) Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 33 +++++++++++++++++++ model_prices_and_context_window.json | 33 +++++++++++++++++++ 2 files changed, 66 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index dac7af6ab7a..1d18903bb07 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -29561,6 +29561,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -30167,6 +30168,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_pdf_input": true, @@ -30210,6 +30212,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -30274,6 +30277,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -30334,6 +30338,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -30401,6 +30406,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -32405,6 +32411,7 @@ "text" ], "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_parallel_function_calling": true, @@ -32479,6 +32486,7 @@ ], "supports_audio_output": false, "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -32543,6 +32551,7 @@ ], "supports_audio_output": false, "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -32603,6 +32612,7 @@ ], "supports_audio_output": false, "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -32670,6 +32680,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -47624,6 +47635,24 @@ "output_cost_per_token": 0.0, "source": "https://cohere.com/pricing" }, + "rerank-v5.0-fast": { + "input_cost_per_token": 2e-08, + "litellm_provider": "cohere", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cohere.com/blog/rerank-5" + }, + "rerank-v5.0-pro": { + "input_cost_per_token": 5e-08, + "litellm_provider": "cohere", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cohere.com/blog/rerank-5" + }, "nvidia_nim/nvidia/nv-rerankqa-mistral-4b-v3": { "input_cost_per_query": 0.0, "input_cost_per_token": 0.0, @@ -52713,6 +52742,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -76498,6 +76528,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -76543,6 +76574,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -76654,6 +76686,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index dac7af6ab7a..1d18903bb07 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -29561,6 +29561,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -30167,6 +30168,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_pdf_input": true, @@ -30210,6 +30212,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -30274,6 +30277,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -30334,6 +30338,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -30401,6 +30406,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -32405,6 +32411,7 @@ "text" ], "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_native_streaming": true, "supports_parallel_function_calling": true, @@ -32479,6 +32486,7 @@ ], "supports_audio_output": false, "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -32543,6 +32551,7 @@ ], "supports_audio_output": false, "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -32603,6 +32612,7 @@ ], "supports_audio_output": false, "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -32670,6 +32680,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, @@ -47624,6 +47635,24 @@ "output_cost_per_token": 0.0, "source": "https://cohere.com/pricing" }, + "rerank-v5.0-fast": { + "input_cost_per_token": 2e-08, + "litellm_provider": "cohere", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cohere.com/blog/rerank-5" + }, + "rerank-v5.0-pro": { + "input_cost_per_token": 5e-08, + "litellm_provider": "cohere", + "max_input_tokens": 32768, + "max_output_tokens": 32768, + "mode": "rerank", + "output_cost_per_token": 0.0, + "source": "https://cohere.com/blog/rerank-5" + }, "nvidia_nim/nvidia/nv-rerankqa-mistral-4b-v3": { "input_cost_per_query": 0.0, "input_cost_per_token": 0.0, @@ -52713,6 +52742,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -76498,6 +76528,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -76543,6 +76574,7 @@ ], "supports_audio_input": true, "supports_audio_output": false, + "supports_computer_use": true, "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, @@ -76654,6 +76686,7 @@ "text" ], "supports_audio_input": true, + "supports_computer_use": true, "supports_function_calling": true, "supports_minimal_reasoning_effort": false, "supports_native_streaming": true, From 0d5cf9b2f7d8cba315cd6c550bd98aa9a2fa88e1 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 30 Oct 2026 12:32:03 -0700 Subject: [PATCH 12/46] fix(headroom): forward tenant-scoped session affinity (#49255) --- .../litellm_core_utils/provider_affinity.py | 42 ++++- .../guardrail_hooks/headroom/headroom.py | 91 ++++----- litellm/proxy/litellm_pre_call_utils.py | 1 + litellm/types/litellm_params.py | 2 + .../guardrails/guardrail_hooks/headroom.py | 12 +- .../test_provider_affinity.py | 16 +- .../guardrail_hooks/test_headroom.py | 174 +++++++++--------- tests/unit/types/test_litellm_params.py | 3 + 8 files changed, 209 insertions(+), 132 deletions(-) diff --git a/litellm/litellm_core_utils/provider_affinity.py b/litellm/litellm_core_utils/provider_affinity.py index c831e92c09d..767eab962c8 100644 --- a/litellm/litellm_core_utils/provider_affinity.py +++ b/litellm/litellm_core_utils/provider_affinity.py @@ -2,7 +2,12 @@ import re from collections.abc import Mapping from typing import Final -from litellm.constants import SESSION_ID_FROM_TRACE_ID_PARAM_KEY, SESSION_ID_GENERATED_METADATA_KEY +from litellm.constants import ( + CALLER_SESSION_ID_METADATA_KEY, + SESSION_ID_FROM_TRACE_ID_PARAM_KEY, + SESSION_ID_GENERATED_METADATA_KEY, +) +from litellm.litellm_core_utils.core_helpers import as_str_mapping _HTTP_HEADER_NAME_PATTERN: Final = re.compile(r"^[!#$%&'*+\-.^_`|~0-9A-Za-z]+$") _FORBIDDEN_AFFINITY_HEADERS: Final = frozenset( @@ -98,6 +103,41 @@ def get_stable_session_id(litellm_params: object | None) -> str | None: return _non_generated_litellm_session_id(litellm_params) or _metadata_session_id(litellm_params) +def get_caller_session_id(litellm_params: object | None) -> str | None: + """Resolve a caller conversation without accepting a generated session or a promoted trace id.""" + if litellm_params is None: + return None + metadata_values: Final = _metadata_values(litellm_params) + preserved_session: Final = next( + ( + value + for metadata in metadata_values + if isinstance(value := _get_value(metadata, CALLER_SESSION_ID_METADATA_KEY), str) and value + ), + None, + ) + if preserved_session: + return preserved_session + direct_session: Final = _get_value(litellm_params, "session_id") + if isinstance(direct_session, str) and direct_session: + return direct_session + + proxy_request: Final = _get_value(litellm_params, "proxy_server_request") + headers: Final = as_str_mapping(_get_value(proxy_request, "headers")) + trace_header: Final = next( + (value for key, value in (headers or {}).items() if key.lower() == "x-litellm-trace-id"), None + ) + candidates: Final = ( + get_caller_litellm_session_id(litellm_params), + *( + _get_value(metadata, "session_id") + for metadata in metadata_values + if not _get_value(metadata, SESSION_ID_GENERATED_METADATA_KEY) + ), + ) + return next((value for value in candidates if isinstance(value, str) and value and value != trace_header), None) + + def add_provider_affinity_header( headers: Mapping[str, object], litellm_params: object | None, default_header: str | None = None ) -> dict[str, object]: # mutable-ok: downstream handlers add auth and signing headers diff --git a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py index dac802a1ddc..7f48766bf26 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py +++ b/litellm/proxy/guardrails/guardrail_hooks/headroom/headroom.py @@ -1,10 +1,10 @@ from __future__ import annotations +import hashlib import json import math import re import time -import uuid from collections.abc import Mapping, Sequence from dataclasses import dataclass from typing import TYPE_CHECKING, ClassVar, Final, Literal, Protocol, TypeGuard @@ -28,6 +28,7 @@ from litellm.litellm_core_utils.prompt_templates.factory import ( group_tool_exchanges, has_tool_with_name, ) +from litellm.litellm_core_utils.provider_affinity import get_caller_session_id from litellm.llms.custom_httpx.http_handler import ( get_async_httpx_client, # pyright: ignore[reportUnknownVariableType] httpxSpecialProvider, @@ -39,6 +40,7 @@ from litellm.proxy.guardrails.guardrail_hooks.content_text import ( merge_rewritten_text_parts, ) from litellm.proxy.spend_tracking.compression_savings import HEADROOM_GUARDRAIL_PROVIDER +from litellm.router_utils.common_utils import request_user_api_key_auth from litellm.secret_managers.main import get_secret_str from litellm.types.guardrails import GuardrailEventHooks, Mode from litellm.types.integrations.custom_logger import ( @@ -47,6 +49,11 @@ from litellm.types.integrations.custom_logger import ( AgenticLoopRequestPatch, as_converted_stream, ) +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.types.proxy.guardrails.guardrail_hooks.headroom import ( + HEADROOM_REQUEST_CONTEXT_KEY, + HeadroomRequestContext, +) from litellm.types.utils import CallTypes, GenericGuardrailAPIInputs if TYPE_CHECKING: @@ -357,15 +364,18 @@ def _build_headroom_retrieve_tool() -> dict[str, object]: } -def _resolve_call_id(logging_obj: object, request_state: dict[str, object]) -> str | None: - """Resolve the litellm_call_id shared by a request's pre-call hook and its - agentic-loop hooks, so CCR hash validation can be scoped per call instead - of trusting any hash-shaped string that shows up in message text.""" - logging_call_id: Final = getattr(logging_obj, "litellm_call_id", None) - if isinstance(logging_call_id, str) and logging_call_id: - return logging_call_id - kwargs_call_id: Final = request_state.get("litellm_call_id") - return kwargs_call_id if isinstance(kwargs_call_id, str) else None +def _session_affinity(request_data: Mapping[str, object]) -> str | None: + auth: Final = request_user_api_key_auth(request_data) + session_id: Final = get_caller_session_id(request_data) + if not isinstance(auth, UserAPIKeyAuth) or not session_id: + return None + tenant: Final = ( + ("key", auth.api_key) if auth.api_key else ("principal", auth.org_id, auth.team_id, auth.user_id, auth.agent_id) + ) + if not any(tenant[1:]): + return None + identity: Final = json.dumps(["headroom-affinity-v1", tenant, session_id], separators=(",", ":")) + return hashlib.sha256(identity.encode("utf-8")).hexdigest() def has_headroom_retrieve_tool(tools: object) -> bool: @@ -521,7 +531,6 @@ class HeadroomGuardrail(CustomGuardrail): self.async_handler = get_async_httpx_client( llm_provider=httpxSpecialProvider.GuardrailCallback, ) - self._issued_hashes_by_call_id: dict[str, tuple[frozenset[str], float]] = {} super().__init__( # pyright: ignore[reportUnknownMemberType] guardrail_name=guardrail_name, event_hook=event_hook, @@ -540,10 +549,12 @@ class HeadroomGuardrail(CustomGuardrail): value: Final = headers.get(BYPASS_HEADER) return str(value).lower() == "true" - def _request_headers(self) -> dict[str, str]: + def _request_headers(self, affinity: str | None) -> dict[str, str]: headers: Final[dict[str, str]] = {"Content-Type": "application/json"} if self.headroom_api_key: headers["Authorization"] = f"Bearer {self.headroom_api_key}" + if affinity is not None: + headers["x-headroom-affinity"] = affinity return headers @staticmethod @@ -569,14 +580,6 @@ class HeadroomGuardrail(CustomGuardrail): super().update_in_memory_litellm_params(litellm_params) self.timeout = self._resolve_timeout(litellm_params.timeout) - def _prune_expired_hashes(self) -> None: - now: Final = time.monotonic() - self._issued_hashes_by_call_id = { - call_id: (hashes, expiry) - for call_id, (hashes, expiry) in self._issued_hashes_by_call_id.items() - if expiry > now - } - def _handle_compress_failure( self, messages: list[dict[str, object]], @@ -596,16 +599,19 @@ class HeadroomGuardrail(CustomGuardrail): self, messages: list[dict[str, object]], model: str | None, + affinity: str | None, ) -> _CompressResult: payload: Final[dict[str, object]] = {"messages": messages} if model: payload["model"] = model + if self.ccr_retrieval: + payload["config"] = {"mode": "ccr"} try: raw_response: HttpxResponse = await self.async_handler.post( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler.post is untyped url=f"{self.headroom_api_base}/v1/compress", json=payload, - headers=self._request_headers(), + headers=self._request_headers(affinity), timeout=self.timeout, ) except httpx.HTTPStatusError as e: @@ -734,7 +740,7 @@ class HeadroomGuardrail(CustomGuardrail): stats["tokens_saved"] = tokens_before - tokens_after return _CompressResult(filtered, True, stats, _read_ccr_hashes(body)) - async def _call_retrieve(self, hash_value: str, query: str | None = None) -> str: + async def _call_retrieve(self, hash_value: str, affinity: str | None, query: str | None = None) -> str: params: Final[dict[str, str]] = {} if query: params["query"] = query @@ -743,7 +749,7 @@ class HeadroomGuardrail(CustomGuardrail): raw_response: HttpxResponse = await self.async_handler.get( # pyright: ignore[reportUnknownMemberType] # AsyncHTTPHandler.get is untyped url=f"{self.headroom_api_base}/v1/retrieve/{hash_value}", params=params, - headers=self._request_headers(), + headers=self._request_headers(affinity), timeout=self.timeout, ) except (httpx.ConnectError, httpx.TimeoutException, httpx.TransportError, litellm.Timeout) as e: @@ -777,13 +783,14 @@ class HeadroomGuardrail(CustomGuardrail): async def apply_guardrail( self, inputs: GenericGuardrailAPIInputs, - request_data: dict, + request_data: dict[str, object], input_type: Literal["request", "response"], logging_obj: LiteLLMLoggingObj | None = None, ) -> GenericGuardrailAPIInputs: if input_type != "request": return inputs + request_data.pop(HEADROOM_REQUEST_CONTEXT_KEY, None) if self._should_bypass(request_data): verbose_proxy_logger.debug("Headroom: %s header set; skipping compression", BYPASS_HEADER) return inputs @@ -809,9 +816,7 @@ class HeadroomGuardrail(CustomGuardrail): # /v1/compress grows a field for sending the live turn as the retrieval # query without compressing it: query-aware compression reads the newest # user message, so it is withheld here at some cost to history ranking. - # request_data is a bare dict on the base signature; narrow it before - # reading the untranslated messages so long tool names can be recovered. - raw_messages: Final = _REQUEST_DATA_ADAPTER.validate_python(request_data).get("messages") + raw_messages: Final = request_data.get("messages") raw_retrieve_call_ids: Final = _raw_retrieve_call_ids(raw_messages) protected_indices: Final = _protected_indices(messages, raw_retrieve_call_ids) compressible: Final = [m for i, m in enumerate(messages) if i not in protected_indices] @@ -820,9 +825,11 @@ class HeadroomGuardrail(CustomGuardrail): model: Final = self.headroom_model or request_data.get("model") start_time: Final = time.time() + affinity: Final = _session_affinity(request_data) result: Final = await self._call_compress( messages=_flatten_messages_for_compression(compressible), model=model if isinstance(model, str) else None, + affinity=affinity, ) end_time: Final = time.time() @@ -868,12 +875,9 @@ class HeadroomGuardrail(CustomGuardrail): if not hashes: return {**inputs, "structured_messages": compressed} # pyright: ignore[reportReturnType] - self._prune_expired_hashes() - call_id = _resolve_call_id(logging_obj, request_data) - if not call_id: - call_id = str(uuid.uuid4()) - request_data["litellm_call_id"] = call_id - self._issued_hashes_by_call_id[call_id] = (frozenset(hashes), time.monotonic() + _HASH_CACHE_TTL_SECONDS) + request_data[HEADROOM_REQUEST_CONTEXT_KEY] = HeadroomRequestContext( + frozenset(hashes), affinity, time.monotonic() + _HASH_CACHE_TTL_SECONDS + ) existing_tools: Final = inputs.get("tools") retrieve_tool: Final = _build_headroom_retrieve_tool() @@ -934,9 +938,11 @@ class HeadroomGuardrail(CustomGuardrail): ) -> AgenticLoopPlan: tool_calls: Final[list[dict[str, object]]] = tools.get("tool_calls", []) - self._prune_expired_hashes() - call_id: Final = _resolve_call_id(logging_obj, kwargs) - valid_hashes = self._issued_hashes_by_call_id.get(call_id, (frozenset(), 0.0))[0] if call_id else frozenset() + context: Final = _REQUEST_DATA_ADAPTER.validate_python(kwargs).get(HEADROOM_REQUEST_CONTEXT_KEY) + valid_context: Final = ( + context if isinstance(context, HeadroomRequestContext) and context.expires_at > time.monotonic() else None + ) + valid_hashes: Final = valid_context.issued_hashes if valid_context is not None else frozenset() retrieved: Final[list[tuple[dict[str, object], str]]] = [] for tc in tool_calls: @@ -944,11 +950,6 @@ class HeadroomGuardrail(CustomGuardrail): raw_hash = arguments.get("hash", "") if isinstance(arguments, dict) else "" hash_value = str(raw_hash).lower() query = arguments.get("query") if isinstance(arguments, dict) else None - # A hash is only honored if it was issued by *this request's own* - # Headroom /v1/compress call, scoped by litellm_call_id. Scoping by - # message text alone is forgeable -- an attacker can plant a - # hash-shaped string in their own prompt, and a hash issued for one - # request would validate for any other request that echoes it back. if hash_value not in valid_hashes: verbose_proxy_logger.warning( "Headroom CCR: rejecting hash=%s not produced by current request compression", @@ -958,6 +959,7 @@ class HeadroomGuardrail(CustomGuardrail): else: content = await self._call_retrieve( hash_value=hash_value, + affinity=valid_context.affinity if valid_context is not None else None, query=str(query) if query else None, ) verbose_proxy_logger.debug("Headroom CCR: retrieved hash=%s (%d chars)", hash_value, len(content)) @@ -996,7 +998,12 @@ class HeadroomGuardrail(CustomGuardrail): max_tokens=max_tokens, optional_params=optional_params_without_max_tokens, kwargs={ - k: v for k, v in kwargs.items() if not k.startswith("_headroom") and k != "litellm_logging_obj" + **{ + k: v + for k, v in _REQUEST_DATA_ADAPTER.validate_python(kwargs).items() + if not k.startswith("_headroom") and k != "litellm_logging_obj" + }, + **({HEADROOM_REQUEST_CONTEXT_KEY: valid_context} if valid_context is not None else {}), }, ), metadata={"tool_type": "headroom_ccr"}, diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 9b52f2f8a8d..ae9a6d4b959 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -2566,6 +2566,7 @@ def refresh_proxy_server_request_body_snapshot( VECTOR_STORE_CALLER_METADATA_KWARG, STREAM_USAGE_INJECTED_FIELD, "_litellm_strip_stream_usage", + "_headroom_request_context", } ) | _TRANSPORT_ONLY_CREDENTIAL_KEYS diff --git a/litellm/types/litellm_params.py b/litellm/types/litellm_params.py index 2065093f295..3c2c636b5ed 100644 --- a/litellm/types/litellm_params.py +++ b/litellm/types/litellm_params.py @@ -21,6 +21,7 @@ if TYPE_CHECKING: from litellm.router_utils.fallback_event_handlers import AttemptedFallbackTargets from litellm.types.caching import DynamicCacheControl from litellm.types.llms.openai import ChatCompletionAssistantMessage, ChatCompletionUserMessage + from litellm.types.proxy.guardrails.guardrail_hooks.headroom import HeadroomRequestContext from litellm.types.proxy.litellm_pre_call_utils import SecretFields from litellm.types.router import ConfigurableClientsideParamsCustomAuth, DeploymentTypedDict, RetryPolicy from litellm.types.router_weights import RouterWeights @@ -352,6 +353,7 @@ class CallState: @dataclass(frozen=True, slots=True, kw_only=True) class AgenticLoopState: + headroom_context: "HeadroomRequestContext | None" = field(default=None, metadata=wire("_headroom_request_context")) depth: int | None = field(default=None, metadata=wire("_agentic_loop_depth")) fingerprints: Sequence[str] | None = field(default=None, metadata=wire("_agentic_loop_fingerprints")) api_surface: Literal["chat_completions", "responses"] | None = field( diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py b/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py index 0df0f0c80ef..76e806e3899 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/headroom.py @@ -1,9 +1,19 @@ -from typing import Literal +from dataclasses import dataclass +from typing import Final, Literal from pydantic import BaseModel, Field from .base import GuardrailConfigModel +HEADROOM_REQUEST_CONTEXT_KEY: Final = "_headroom_request_context" + + +@dataclass(frozen=True, slots=True) +class HeadroomRequestContext: + issued_hashes: frozenset[str] + affinity: str | None + expires_at: float + class HeadroomGuardrailConfigModel(GuardrailConfigModel[BaseModel]): api_base: str | None = Field( diff --git a/tests/unit/litellm_core_utils/test_provider_affinity.py b/tests/unit/litellm_core_utils/test_provider_affinity.py index edf4f5169b5..6ccd6b17c01 100644 --- a/tests/unit/litellm_core_utils/test_provider_affinity.py +++ b/tests/unit/litellm_core_utils/test_provider_affinity.py @@ -1,12 +1,16 @@ +from collections.abc import Callable + import pytest from litellm.constants import SESSION_ID_GENERATED_METADATA_KEY from litellm.litellm_core_utils.provider_affinity import ( add_provider_affinity_header, + get_caller_session_id, get_stable_session_id, ) +@pytest.mark.parametrize("resolve", [get_stable_session_id, get_caller_session_id]) @pytest.mark.parametrize( ("litellm_params", "expected"), [ @@ -16,18 +20,22 @@ from litellm.litellm_core_utils.provider_affinity import ( ({"litellm_metadata": {"session_id": "litellm-metadata-session"}}, "litellm-metadata-session"), ], ) -def test_get_stable_session_id_uses_explicit_session_sources(litellm_params: dict, expected: str): - assert get_stable_session_id(litellm_params) == expected +def test_session_ids(litellm_params: dict[str, object], expected: str, resolve: Callable[[object], str | None]) -> None: + assert resolve(litellm_params) == expected def test_get_stable_session_id_does_not_use_trace_id(): assert get_stable_session_id({"litellm_trace_id": "per-request-trace"}) is None + assert get_caller_session_id( + {"litellm_session_id": "trace", "proxy_server_request": {"headers": {"x-litellm-trace-id": "trace"}}} + ) is None +@pytest.mark.parametrize("resolve", [get_stable_session_id, get_caller_session_id]) @pytest.mark.parametrize("metadata_key", ["metadata", "litellm_metadata"]) -def test_get_stable_session_id_ignores_proxy_generated_session(metadata_key: str): +def test_generated_sessions_are_ignored(metadata_key: str, resolve: Callable[[object], str | None]) -> None: assert ( - get_stable_session_id( + resolve( { "litellm_session_id": "generated-session", metadata_key: { diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py index 741fea5f2b9..0cdcb2e2d07 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_headroom.py @@ -23,7 +23,6 @@ Tests cover: import datetime import json -import time from typing import Final, Optional from unittest.mock import AsyncMock, MagicMock, patch @@ -33,17 +32,22 @@ import respx from fastapi import HTTPException import litellm - from litellm.litellm_core_utils.litellm_logging import Logging from litellm.proxy.guardrails.guardrail_hooks.headroom.headroom import ( - HeadroomGuardrail, - has_headroom_retrieve_tool, HEADROOM_RETRIEVE_TOOL_NAME, + HeadroomGuardrail, + _session_affinity, + has_headroom_retrieve_tool, ) from litellm.proxy.spend_tracking.compression_savings import ( extract_compression_saved_tokens, ) from litellm.types.integrations.custom_logger import HEADROOM_CONVERTED_STREAM_KEY, HEADROOM_STREAM_OPTIONS_KEY +from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth +from litellm.types.proxy.guardrails.guardrail_hooks.headroom import ( + HEADROOM_REQUEST_CONTEXT_KEY, + HeadroomRequestContext, +) from litellm.types.utils import ( CallTypes, GenericGuardrailAPIInputs, @@ -148,6 +152,10 @@ def guardrail() -> HeadroomGuardrail: return _make_guardrail() +def _issued_context(*hashes: str, affinity: str | None = None) -> HeadroomRequestContext: + return HeadroomRequestContext(frozenset(hashes), affinity, float("inf")) + + def _recorded_guardrail_entries(request_data: dict) -> list: for container_key in ("metadata", "litellm_metadata"): container = request_data.get(container_key) @@ -175,20 +183,22 @@ async def test_apply_guardrail_compresses_and_returns_structured_messages( structured_messages=ORIGINAL_MESSAGES, ) mock_response = _make_compress_response(COMPRESSED_MESSAGES) - request_data = {"model": "gpt-4o"} + auth: Final = UserAPIKeyAuth(api_key="tenant") + request_data: Final = {"session_id": "conversation", "metadata": {"user_api_key_auth": auth}} with patch.object( guardrail.async_handler, "post", new_callable=AsyncMock, return_value=mock_response, - ): + ) as mock_post: result = await guardrail.apply_guardrail( inputs=inputs, request_data=request_data, input_type="request", ) + assert mock_post.call_args.kwargs["headers"]["x-headroom-affinity"] == _session_affinity(request_data) assert result.get("structured_messages") == EXPECTED_MESSAGES entries = _recorded_guardrail_entries(request_data) @@ -510,10 +520,6 @@ async def test_async_build_agentic_loop_plan_calls_retrieve_and_builds_messages( tool_id="call_abc123", ) messages = [{"role": "user", "content": "What does it say? hash=b573993006976af767214fac"}] - guardrail._issued_hashes_by_call_id["call-1"] = ( - frozenset({"b573993006976af767214fac"}), - time.monotonic() + 999, - ) with patch.object( guardrail.async_handler, @@ -530,11 +536,12 @@ async def test_async_build_agentic_loop_plan_calls_retrieve_and_builds_messages( anthropic_messages_optional_request_params={}, logging_obj=None, stream=False, - kwargs={"litellm_call_id": "call-1"}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: _issued_context("b573993006976af767214fac", affinity="original")}, ) assert plan.run_agentic_loop is True assert plan.request_patch is not None + assert HEADROOM_REQUEST_CONTEXT_KEY in plan.request_patch.kwargs follow_up = plan.request_patch.messages assert follow_up is not None @@ -545,6 +552,7 @@ async def test_async_build_agentic_loop_plan_calls_retrieve_and_builds_messages( assert tool_result_message["tool_call_id"] == "call_abc123" mock_get.assert_called_once() + assert mock_get.call_args.kwargs["headers"]["x-headroom-affinity"] == "original" call_url = mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0] assert "b573993006976af767214fac" in call_url @@ -576,10 +584,6 @@ async def test_async_build_agentic_loop_plan_handles_retrieve_404( "content": "Retrieve more: hash=deadbeef000000000000dead", } ] - guardrail._issued_hashes_by_call_id["call-1"] = ( - frozenset({"deadbeef000000000000dead"}), - time.monotonic() + 999, - ) with patch.object( guardrail.async_handler, @@ -596,7 +600,7 @@ async def test_async_build_agentic_loop_plan_handles_retrieve_404( anthropic_messages_optional_request_params={}, logging_obj=None, stream=False, - kwargs={"litellm_call_id": "call-1"}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: _issued_context("deadbeef000000000000dead")}, ) follow_up = plan.request_patch.messages # type: ignore[union-attr] @@ -606,9 +610,13 @@ async def test_async_build_agentic_loop_plan_handles_retrieve_404( @pytest.mark.asyncio +@pytest.mark.parametrize("context", [ + None, {"issued_hashes": ["deadbeef000000000000dead"]}, + HeadroomRequestContext(frozenset({"deadbeef000000000000dead"}), None, 0.0), +]) async def test_async_build_agentic_loop_plan_rejects_hash_with_no_known_call( - guardrail: HeadroomGuardrail, -): + guardrail: HeadroomGuardrail, context: object, +) -> None: """A hash-shaped string planted in message text must not be honored when this guardrail has no record of ever issuing it, even if it's echoed back in the current request's own messages (e.g. via prompt injection).""" @@ -625,8 +633,6 @@ async def test_async_build_agentic_loop_plan_rejects_hash_with_no_known_call( arguments={"hash": "deadbeef000000000000dead"}, tool_id="call_xyz", ) - assert not guardrail._issued_hashes_by_call_id - with patch.object( guardrail.async_handler, "get", @@ -641,7 +647,7 @@ async def test_async_build_agentic_loop_plan_rejects_hash_with_no_known_call( anthropic_messages_optional_request_params={}, logging_obj=None, stream=False, - kwargs={"litellm_call_id": "call-unknown"}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: context}, ) mock_get.assert_not_called() @@ -653,17 +659,11 @@ async def test_async_build_agentic_loop_plan_rejects_hash_with_no_known_call( @pytest.mark.asyncio -async def test_async_build_agentic_loop_plan_rejects_hash_issued_for_different_call( +async def test_async_build_agentic_loop_plan_rejects_hash_issued_for_different_request( guardrail: HeadroomGuardrail, ): - """A hash issued for one request must not be retrievable by a different - request just because the second request echoes that hash-shaped string - back in its own messages -- retrieval must be scoped per litellm_call_id, - not derived by re-scanning attacker-controlled message text.""" - guardrail._issued_hashes_by_call_id["call-A"] = ( - frozenset({"b573993006976af767214fac"}), - time.monotonic() + 999, - ) + issuer_data: Final[dict[str, object]] = {"model": "gpt-4o", "litellm_call_id": "call-A"} + await _apply(guardrail, ORIGINAL_MESSAGES, ["b573993006976af767214fac"], request_data=issuer_data) tool_calls = [ { @@ -693,7 +693,7 @@ async def test_async_build_agentic_loop_plan_rejects_hash_issued_for_different_c anthropic_messages_optional_request_params={}, logging_obj=None, stream=False, - kwargs={"litellm_call_id": "call-B"}, + kwargs={"litellm_call_id": "call-A"}, ) mock_get.assert_not_called() @@ -736,10 +736,6 @@ async def test_async_build_agentic_loop_plan_builds_responses_api_function_call_ } ] messages = [{"role": "user", "content": "What does it say? hash=b573993006976af767214fac"}] - guardrail._issued_hashes_by_call_id["call-1"] = ( - frozenset({"b573993006976af767214fac"}), - time.monotonic() + 999, - ) with patch.object( guardrail.async_handler, @@ -756,7 +752,7 @@ async def test_async_build_agentic_loop_plan_builds_responses_api_function_call_ anthropic_messages_optional_request_params={}, logging_obj=None, stream=False, - kwargs={"litellm_call_id": "call-1"}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: _issued_context("b573993006976af767214fac")}, ) follow_up = plan.request_patch.messages # type: ignore[union-attr] @@ -806,7 +802,6 @@ async def test_responses_retrieval_followup_reuses_the_model_the_caller_requeste "arguments": json.dumps({"hash": retrieve_hash}), } ] - guardrail._issued_hashes_by_call_id["call-1"] = (frozenset({retrieve_hash}), time.monotonic() + 999) with patch.object( guardrail.async_handler, @@ -832,7 +827,7 @@ async def test_responses_retrieval_followup_reuses_the_model_the_caller_requeste anthropic_messages_optional_request_params={}, logging_obj=logging_obj, stream=False, - kwargs={"litellm_call_id": "call-1"}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: _issued_context(retrieve_hash)}, ) assert plan.request_patch is not None @@ -876,10 +871,6 @@ async def test_async_build_agentic_loop_plan_builds_anthropic_tool_result_messag } ] messages = [{"role": "user", "content": "What does it say? hash=b573993006976af767214fac"}] - guardrail._issued_hashes_by_call_id["call-1"] = ( - frozenset({"b573993006976af767214fac"}), - time.monotonic() + 999, - ) with patch.object( guardrail.async_handler, @@ -896,7 +887,7 @@ async def test_async_build_agentic_loop_plan_builds_anthropic_tool_result_messag anthropic_messages_optional_request_params={}, logging_obj=None, stream=False, - kwargs={"litellm_call_id": "call-1"}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: _issued_context("b573993006976af767214fac")}, ) follow_up = plan.request_patch.messages # type: ignore[union-attr] @@ -929,16 +920,20 @@ HASH_SHAPED_HISTORY = [ ] -async def _apply(guardrail: HeadroomGuardrail, messages: list, ccr_hashes: list | None = None) -> dict: - request_data = {"model": "claude-sonnet-5"} +async def _apply( + guardrail: HeadroomGuardrail, messages: list, ccr_hashes: list | None = None, + request_data: dict[str, object] | None = None, +) -> dict: + data: Final = request_data if request_data is not None else {"model": "claude-sonnet-5"} def _echo(**kwargs): + assert kwargs["json"].get("config") == ({"mode": "ccr"} if guardrail.ccr_retrieval else None) return _make_compress_response(json.loads(json.dumps(kwargs["json"]["messages"])), ccr_hashes=ccr_hashes) with patch.object(guardrail.async_handler, "post", new_callable=AsyncMock, side_effect=_echo): return await guardrail.apply_guardrail( inputs=GenericGuardrailAPIInputs(texts=["x"], structured_messages=json.loads(json.dumps(messages))), - request_data=request_data, + request_data=data, input_type="request", ) @@ -950,10 +945,11 @@ async def test_hash_shaped_text_in_history_never_injects_retrieve_tool(guardrail tool must not appear and no hash may be registered as issued. Covers a service that omits ccr_hashes, returns it empty, or returns a non-list.""" for ccr_hashes in (None, [], "b573993006976af767214fac"): - result = await _apply(guardrail, HASH_SHAPED_HISTORY, ccr_hashes=ccr_hashes) + request_data: Final[dict[str, object]] = {} + result = await _apply(guardrail, HASH_SHAPED_HISTORY, ccr_hashes=ccr_hashes, request_data=request_data) assert not has_headroom_retrieve_tool(result.get("tools") or []) - assert not guardrail._issued_hashes_by_call_id + assert HEADROOM_REQUEST_CONTEXT_KEY not in request_data @pytest.mark.asyncio @@ -961,15 +957,18 @@ async def test_service_declared_ccr_hashes_drive_injection_and_validation(guardr """Only the hashes the service reports in ccr_hashes are honored, in the service's own 12 to 24 hex grammar; anything else is dropped because each entry is interpolated into the /v1/retrieve URL.""" + request_data: Final[dict[str, object]] = {} result = await _apply( guardrail, HASH_SHAPED_HISTORY, ccr_hashes=["98CA69107318", "b573993006976af767214fac", "../../etc/passwd", "tooshort", 42], + request_data=request_data, ) assert has_headroom_retrieve_tool(result.get("tools") or []) - ((issued, _expiry),) = guardrail._issued_hashes_by_call_id.values() - assert issued == frozenset({"98ca69107318", "b573993006976af767214fac"}) + context: Final = request_data[HEADROOM_REQUEST_CONTEXT_KEY] + assert isinstance(context, HeadroomRequestContext) + assert context.issued_hashes == frozenset({"98ca69107318", "b573993006976af767214fac"}) @pytest.mark.asyncio @@ -987,7 +986,6 @@ async def test_ccr_retrieval_disabled_ignores_service_declared_hashes(monkeypatc assert result["structured_messages"][-1] == HASH_SHAPED_HISTORY[-1] assert not has_headroom_retrieve_tool(result.get("tools") or []) - assert not guardrail._issued_hashes_by_call_id @pytest.mark.asyncio @@ -1498,6 +1496,7 @@ async def test_apply_guardrail_fail_open_does_not_register_hashes_from_original_ ] guardrail = _make_guardrail(unreachable_fallback="fail_open") + request_data: Final = {HEADROOM_REQUEST_CONTEXT_KEY: _issued_context("deadbeef000000000000dead")} inputs = GenericGuardrailAPIInputs( texts=["hello"], structured_messages=messages_with_fake_hash, @@ -1511,13 +1510,13 @@ async def test_apply_guardrail_fail_open_does_not_register_hashes_from_original_ ): result = await guardrail.apply_guardrail( inputs=inputs, - request_data={}, + request_data=request_data, input_type="request", ) assert result["structured_messages"] == messages_with_fake_hash assert not has_headroom_retrieve_tool(result.get("tools") or []) - assert not guardrail._issued_hashes_by_call_id + assert HEADROOM_REQUEST_CONTEXT_KEY not in request_data @pytest.mark.asyncio @@ -2164,18 +2163,18 @@ def _openai_completion_payload(message: dict, finish_reason: str) -> dict: } -def _openai_tool_call_payload() -> dict: +def _openai_tool_call_payload(hash_value: str = CCR_HASH) -> dict: return _openai_completion_payload( { "role": "assistant", "content": None, "tool_calls": [ { - "id": "call_ccr", + "id": f"call_{hash_value}", "type": "function", "function": { "name": HEADROOM_RETRIEVE_TOOL_NAME, - "arguments": json.dumps({"hash": CCR_HASH}), + "arguments": json.dumps({"hash": hash_value}), }, } ], @@ -2215,14 +2214,14 @@ def _openai_responses_payload(output_item: dict) -> dict: } -def _openai_responses_retrieve_call_payload() -> dict: +def _openai_responses_retrieve_call_payload(hash_value: str = CCR_HASH) -> dict: return _openai_responses_payload( { "type": "function_call", "id": "fc_ccr", - "call_id": "call_ccr", + "call_id": f"call_{hash_value}", "name": HEADROOM_RETRIEVE_TOOL_NAME, - "arguments": json.dumps({"hash": CCR_HASH}), + "arguments": json.dumps({"hash": hash_value}), "status": "completed", } ) @@ -2381,10 +2380,6 @@ async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end( to a client that never declared the tool.""" original_content = "the full uncompressed document" final_answer = "the document says hello" - guardrail._issued_hashes_by_call_id["ccr-call-id"] = ( - frozenset({CCR_HASH}), - time.monotonic() + 999, - ) monkeypatch.setenv("OPENAI_API_KEY", "sk-test") monkeypatch.setattr(litellm, "callbacks", [guardrail]) @@ -2392,6 +2387,7 @@ async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end( upstream = respx_mock.post("https://api.openai.com/v1/chat/completions").mock( side_effect=[ httpx.Response(200, json=_openai_tool_call_payload()), + httpx.Response(200, json=_openai_tool_call_payload("a" * 24)), httpx.Response(200, json=_openai_text_payload(final_answer)), ] ) @@ -2408,20 +2404,22 @@ async def test_streaming_chat_completion_resolves_ccr_retrieval_end_to_end( tools=[_retrieve_tool_definition()], stream=True, litellm_call_id="ccr-call-id", + **{HEADROOM_REQUEST_CONTEXT_KEY: _issued_context(CCR_HASH, "a" * 24, affinity="original-affinity")}, ) chunks = [chunk async for chunk in response] streamed_text = "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) assert streamed_text == final_answer assert not any(chunk.choices and chunk.choices[0].delta.tool_calls for chunk in chunks) - mock_get.assert_called_once() - assert CCR_HASH in (mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0]) + assert mock_get.call_count == 2 + assert mock_get.call_args.kwargs["url"].endswith("a" * 24) + assert all(call.kwargs["headers"]["x-headroom-affinity"] == "original-affinity" for call in mock_get.call_args_list) - assert len(upstream.calls) == 2 - followup_body = json.loads(upstream.calls[1].request.content) + assert len(upstream.calls) == 3 + followup_body = json.loads(upstream.calls[2].request.content) assert not followup_body.get("stream") assert original_content in json.dumps(followup_body["messages"]) - assert not any(key.startswith("_headroom_interception") for key in followup_body) + assert not any(key.startswith("_headroom") for key in followup_body) @pytest.mark.asyncio @@ -2435,10 +2433,6 @@ async def test_streaming_responses_resolves_ccr_retrieval_end_to_end( instead of streaming a headroom_retrieve function_call to the client.""" original_content = "the full uncompressed document" final_answer = "the document says hello" - guardrail._issued_hashes_by_call_id["ccr-call-id"] = ( - frozenset({CCR_HASH}), - time.monotonic() + 999, - ) monkeypatch.setenv("OPENAI_API_KEY", "sk-test") monkeypatch.setattr(litellm, "callbacks", [guardrail]) @@ -2446,6 +2440,7 @@ async def test_streaming_responses_resolves_ccr_retrieval_end_to_end( upstream = respx_mock.post("https://api.openai.com/v1/responses").mock( side_effect=[ httpx.Response(200, json=_openai_responses_retrieve_call_payload()), + httpx.Response(200, json=_openai_responses_retrieve_call_payload("a" * 24)), httpx.Response(200, json=_openai_responses_text_payload(final_answer)), ] ) @@ -2462,6 +2457,7 @@ async def test_streaming_responses_resolves_ccr_retrieval_end_to_end( tools=[_responses_retrieve_tool_definition()], stream=True, litellm_call_id="ccr-call-id", + **{HEADROOM_REQUEST_CONTEXT_KEY: _issued_context(CCR_HASH, "a" * 24, affinity="original-affinity")}, ) events = [event async for event in response] @@ -2471,14 +2467,15 @@ async def test_streaming_responses_resolves_ccr_retrieval_end_to_end( assert streamed_text == final_answer assert not any("function_call" in str(getattr(event, "type", "")) for event in events) assert not any(getattr(getattr(event, "item", None), "type", None) == "function_call" for event in events) - mock_get.assert_called_once() - assert CCR_HASH in (mock_get.call_args.kwargs.get("url") or mock_get.call_args.args[0]) + assert mock_get.call_count == 2 + assert mock_get.call_args.kwargs["url"].endswith("a" * 24) + assert all(call.kwargs["headers"]["x-headroom-affinity"] == "original-affinity" for call in mock_get.call_args_list) - assert len(upstream.calls) == 2 - followup_body = json.loads(upstream.calls[1].request.content) + assert len(upstream.calls) == 3 + followup_body = json.loads(upstream.calls[2].request.content) assert not followup_body.get("stream") assert original_content in json.dumps(followup_body["input"]) - assert not any(key.startswith("_headroom_interception") for key in followup_body) + assert not any(key.startswith("_headroom") for key in followup_body) def test_sync_streaming_responses_resolves_ccr_retrieval_end_to_end( @@ -2491,10 +2488,6 @@ def test_sync_streaming_responses_resolves_ccr_retrieval_end_to_end( completed response object.""" original_content = "the full uncompressed document" final_answer = "the document says hello" - guardrail._issued_hashes_by_call_id["ccr-call-id"] = ( - frozenset({CCR_HASH}), - time.monotonic() + 999, - ) monkeypatch.setenv("OPENAI_API_KEY", "sk-test") monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) @@ -2518,6 +2511,7 @@ def test_sync_streaming_responses_resolves_ccr_retrieval_end_to_end( tools=[_responses_retrieve_tool_definition()], stream=True, litellm_call_id="ccr-call-id", + **{HEADROOM_REQUEST_CONTEXT_KEY: _issued_context(CCR_HASH)}, ) events = list(response) @@ -2917,7 +2911,6 @@ def _anthropic_response_with_text_and_tool_call() -> dict: async def _plan_for(guardrail: HeadroomGuardrail, response, messages: list): - guardrail._issued_hashes_by_call_id["call-1"] = (frozenset({"h" * 24}), time.monotonic() + 60) logging_obj = MagicMock() logging_obj.litellm_call_id = "call-1" logging_obj.model_call_details = {} @@ -2938,7 +2931,7 @@ async def _plan_for(guardrail: HeadroomGuardrail, response, messages: list): anthropic_messages_optional_request_params={}, logging_obj=logging_obj, stream=False, - kwargs={}, + kwargs={HEADROOM_REQUEST_CONTEXT_KEY: _issued_context("h" * 24)}, ) @@ -3033,7 +3026,7 @@ async def test_retrieve_call_passes_bounded_timeout(guardrail: HeadroomGuardrail new_callable=AsyncMock, return_value=_make_retrieve_response("original"), ) as mock_get: - result = await guardrail._call_retrieve("a" * 24) + result = await guardrail._call_retrieve("a" * 24, affinity=None) assert result == "original" timeout = _timeout_of(mock_get.call_args) @@ -3143,3 +3136,16 @@ def test_unusable_timeout_falls_back_to_the_default(configured: float): assert guardrail.timeout.read == 60.0 assert guardrail.timeout.connect == 5.0 + + +@pytest.mark.parametrize("session", [None, "conversation", "other"]) +@pytest.mark.parametrize("scope", [None, {}, {"api_key": "key"}, {"api_key": "other"}, {"user_id": "user"}]) +def test_affinity_requires_trusted_identity_and_caller_session( + scope: dict[str, str] | None, session: str | None +) -> None: + baseline: Final = {"session_id": "conversation", "metadata": {"user_api_key_auth": UserAPIKeyAuth(api_key="key")}} + auth: Final = UserAPIKeyAuth(**scope) if scope else scope + candidate: Final = {"session_id": session, "metadata": {"user_api_key_auth": auth}} + affinity: Final = _session_affinity(candidate) + assert (affinity is None) is (session is None or not isinstance(auth, UserAPIKeyAuth)) + assert (affinity == _session_affinity(baseline)) is (candidate == baseline) diff --git a/tests/unit/types/test_litellm_params.py b/tests/unit/types/test_litellm_params.py index e369b7e93ab..22eb9929e7f 100644 --- a/tests/unit/types/test_litellm_params.py +++ b/tests/unit/types/test_litellm_params.py @@ -33,6 +33,7 @@ from litellm.types.litellm_params import ( wire_names, ) from litellm.types.llms.openai import ChatCompletionAssistantMessage, ChatCompletionUserMessage +from litellm.types.proxy.guardrails.guardrail_hooks.headroom import HeadroomRequestContext from litellm.types.proxy.litellm_pre_call_utils import SecretFields from litellm.types.router import ( ConfigurableClientsideParamsCustomAuth, @@ -231,6 +232,7 @@ OPTION_NAMES: Final = ( ) AGENTIC_LOOP_STATE_NAMES: Final = ( + "_headroom_request_context", "_agentic_loop_depth", "_agentic_loop_fingerprints", "_agentic_loop_api_surface", @@ -529,6 +531,7 @@ TYPE_HINT_NAMESPACE: Final[Mapping[str, object]] = { "Logging": Logging, "SecretFields": SecretFields, "CompactionState": CompactionState, + "HeadroomRequestContext": HeadroomRequestContext, "RouterWeights": RouterWeights, "AttemptedFallbackTargets": AttemptedFallbackTargets, } From 7b259690e7eac9720ee21067a71152912f1080ef Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 30 Oct 2026 12:36:11 -0700 Subject: [PATCH 13/46] fix(proxy): don't send llm_exceptions alert for BudgetExceededError (#49290) * fix(proxy): don't send llm_exceptions alert for BudgetExceededError * test(proxy): verify budget alert suppression through failure hooks --------- Co-authored-by: Arvind Kumar --- litellm/proxy/utils.py | 2 + ..._utils_model_creation_and_error_logging.py | 49 +++++++++++++++++++ 2 files changed, 51 insertions(+) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 00dbcd0f9a8..cc72dd03a40 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -633,6 +633,8 @@ def _is_client_error_exception(exc: Exception) -> bool: return exc.status_code < 500 if isinstance(exc, ProxyException): return not (exc.code.isdigit() and int(exc.code) >= 500) + if isinstance(exc, litellm.BudgetExceededError): + return exc.status_code < 500 return False diff --git a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py index 3cdd2edb2fb..0a483840336 100644 --- a/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py +++ b/tests/unit/proxy/test_proxy_utils_model_creation_and_error_logging.py @@ -1844,6 +1844,26 @@ class TestPostCallFailureHookLLMExceptionAlerting: alerting_handler = await self._alerted(HTTPException(status_code=429, detail="rate limited")) assert alerting_handler.called is False + @pytest.mark.asyncio + @pytest.mark.parametrize("status_code, should_alert", ((400, False), (422, False), (429, False), (500, True))) + async def test_budget_exceeded_alerts_only_for_server_status( + self, status_code: int, should_alert: bool, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "budget_exceeded_status_code", status_code) + exc: Final = litellm.BudgetExceededError(current_cost=1.0, max_budget=0.5) + assert (await self._alerted(exc)).called is should_alert + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "exc", + ( + litellm.AuthenticationError(message="invalid provider key", model="test-model", llm_provider="openai"), + litellm.RateLimitError(message="provider quota exceeded", model="test-model", llm_provider="openai"), + ), + ) + async def test_provider_client_status_still_alerts(self, exc: Exception) -> None: + assert (await self._alerted(exc)).called is True + class TestPostCallFailureHookProxyExceptionLogging: """A guardrail block raises a ProxyException; on an LLM route it must still @@ -4109,3 +4129,32 @@ def test_handle_exception_on_proxy_logs_bug_report_only_for_unmapped_500(caplog) assert provider_result.code == internal_result.code == "500" assert ISSUE_URL_BASE in caplog.text assert ISSUE_URL_BASE not in internal_result.message + + +def test_is_client_error_exception_budget_exceeded_is_client_error() -> None: + from litellm.exceptions import BudgetExceededError + from litellm.proxy.utils import _is_client_error_exception + + exc: Final = BudgetExceededError(current_cost=100.56, max_budget=100.0) + assert _is_client_error_exception(exc) is True + + +def test_is_client_error_exception_budget_exceeded_5xx_is_not_client_error() -> None: + from litellm.exceptions import BudgetExceededError + from litellm.proxy.utils import _is_client_error_exception + + exc: Final = BudgetExceededError(current_cost=100.56, max_budget=100.0) + exc.status_code = 500 + assert _is_client_error_exception(exc) is False + + +def test_is_client_error_exception_existing_behaviour_unchanged() -> None: + from litellm.proxy._types import ProxyException + from litellm.proxy.utils import _is_client_error_exception + + assert _is_client_error_exception(HTTPException(status_code=400, detail="x")) is True + assert _is_client_error_exception(HTTPException(status_code=500, detail="x")) is False + assert _is_client_error_exception( + ProxyException(message="x", type="x", param=None, code="429") + ) is True + assert _is_client_error_exception(ValueError("x")) is False From 67954b5c3a4b00a1da3c86b18bca8452194002ac Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 30 Oct 2026 12:38:21 -0700 Subject: [PATCH 14/46] fix(proxy): apply key and team router_settings to /v1/images/generations (#49286) * fix(proxy): apply key and team router_settings to /v1/images/generations Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * style(proxy): format router settings helper Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): resolve key model_group_alias once on image generations 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/proxy/common_request_processing.py | 93 ++++++---- litellm/proxy/image_endpoints/endpoints.py | 18 ++ ...st_key_router_settings_image_generation.py | 162 ++++++++++++++++++ .../proxy/image_endpoints/test_endpoints.py | 148 +++++++++++++++- 4 files changed, 389 insertions(+), 32 deletions(-) create mode 100644 tests/integration/routing/test_key_router_settings_image_generation.py diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 762f552dadc..9b2813a8cd8 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -151,7 +151,7 @@ from litellm.router_utils.common_utils import resolve_model_group_alias, resolve from litellm.types.extraction import ExtractionResponse from litellm.types.guardrails import GuardrailEventHooks from litellm.types.router import RouterRateLimitError -from litellm.types.router_weights import validate_router_weights +from litellm.types.router_weights import RouterWeights, validate_router_weights _LateResponseT = TypeVar("_LateResponseT", bound=Response) _LlmCallT = TypeVar("_LlmCallT") @@ -894,6 +894,61 @@ async def _resolve_per_request_model_group_alias( return target +async def hierarchical_router_settings_updates( + requested_model: object, + user_api_key_dict: UserAPIKeyAuth, + llm_router: Router, + proxy_config: ProxyConfig, + proxy_logging_obj: ProxyLogging, + apply_model_group_alias: bool, +) -> Mapping[str, object]: + # Apply hierarchical router_settings (Key > Team) + # Global router_settings are already on the Router object itself. + # This sits with the other alias rewrites, and ahead of the guardrail + # merge and the pre-call hooks, so everything that keys off the model + # group -- model-level guardrails, per-model budgets and rate limits, + # the logging object -- sees the group that will actually serve. + from litellm.proxy.proxy_server import prisma_client + + router_settings: Final[Mapping[str, object] | None] = _ROUTER_SETTINGS_ADAPTER.validate_python( + await proxy_config._get_hierarchical_router_settings( + user_api_key_dict=user_api_key_dict, + prisma_client=prisma_client, + proxy_logging_obj=proxy_logging_obj, + ) + ) + if router_settings is None: + return MappingProxyType({}) + + # If router_settings found (from key or team), apply them + # Pass settings as per-request overrides instead of creating a new Router + # This avoids expensive Router instantiation on each request + def validated_router_weights() -> RouterWeights | None: + try: + return validate_router_weights(router_settings.get("weights")) + except ValidationError: + verbose_proxy_logger.warning("Ignoring invalid saved router weights; update team/key router_settings") + return None + + router_weights: Final = validated_router_weights() + alias_target: Final = ( + await _resolve_per_request_model_group_alias( + requested_model=requested_model, + router_settings=router_settings, + user_api_key_dict=user_api_key_dict, + llm_router=llm_router, + ) + if apply_model_group_alias + else None + ) + settings_updates: Final[Mapping[str, object]] = { + "router_settings_override": router_settings, + "_router_weights": router_weights, + **({"model": alias_target} if alias_target is not None else {}), + } + return MappingProxyType(settings_updates) + + _REQUEST_MODEL: Final[TypeAdapter[str | list[str] | None]] = TypeAdapter(str | list[str] | None) _DECODED_JSON: Final = TypeAdapter(object) _ROUTER_SETTINGS_ADAPTER: Final[TypeAdapter[Mapping[str, object] | None]] = TypeAdapter(Mapping[str, object] | None) @@ -2495,41 +2550,17 @@ class ProxyBaseLLMRequestProcessing: ): self.data["model"] = user_api_key_dict.aliases[self.data["model"]] - # Apply hierarchical router_settings (Key > Team) - # Global router_settings are already on the Router object itself. - # This sits with the other alias rewrites, and ahead of the guardrail - # merge and the pre-call hooks, so everything that keys off the model - # group -- model-level guardrails, per-model budgets and rate limits, - # the logging object -- sees the group that will actually serve. if llm_router is not None and proxy_config is not None: - from litellm.proxy.proxy_server import prisma_client - - router_settings: Final = await proxy_config._get_hierarchical_router_settings( - user_api_key_dict=user_api_key_dict, - prisma_client=prisma_client, - proxy_logging_obj=proxy_logging_obj, - ) - - # If router_settings found (from key or team), apply them - # Pass settings as per-request overrides instead of creating a new Router - # This avoids expensive Router instantiation on each request - if router_settings is not None: - self.data["router_settings_override"] = router_settings - try: - self.data["_router_weights"] = validate_router_weights(router_settings.get("weights")) - except ValidationError: - self.data["_router_weights"] = None - verbose_proxy_logger.warning( - "Ignoring invalid saved router weights; update team/key router_settings" - ) - alias_target: Final = await _resolve_per_request_model_group_alias( + self.data.update( + await hierarchical_router_settings_updates( requested_model=self.data.get("model"), - router_settings=router_settings, user_api_key_dict=user_api_key_dict, llm_router=llm_router, + proxy_config=proxy_config, + proxy_logging_obj=proxy_logging_obj, + apply_model_group_alias=True, ) - if alias_target is not None: - self.data["model"] = alias_target + ) aliased_model: Final = self.data.get("model") if llm_router is not None and route_type in ROUTER_SCHEDULED_ROUTE_TYPES and isinstance(aliased_model, str): diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 28fb9baf2aa..6244e5025ce 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -8,8 +8,10 @@ from typing import Final, get_type_hints import orjson from fastapi import APIRouter, Depends, File, HTTPException, Request, Response, UploadFile, status from fastapi.responses import ORJSONResponse +from pydantic import TypeAdapter import litellm +from litellm.constants import MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY from litellm.litellm_core_utils.prompt_templates.common_utils import ( get_str_from_messages, ) @@ -17,6 +19,7 @@ from litellm.proxy._types import * from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, + hierarchical_router_settings_updates, log_llm_api_exception, resolve_litellm_call_id, ) @@ -36,6 +39,7 @@ from litellm.types.images.main import ImageEditRequestParams from litellm.types.llms.openai import ChatCompletionUserMessage router: Final = APIRouter() +_IMAGE_GENERATION_MODEL_ADAPTER: Final[TypeAdapter[object]] = TypeAdapter(object) IMAGE_EDIT_NUMERIC_FORM_FIELDS: Final = numeric_form_fields(get_type_hints(ImageEditRequestParams)) IMAGE_EDIT_OPTIONAL_FIELD_DEFAULTS: Final = MappingProxyType({"prompt": None, "image": None}) @@ -132,6 +136,19 @@ async def image_generation( if data["model"] in litellm.model_alias_map: data["model"] = litellm.model_alias_map[data["model"]] + if llm_router is not None: + requested_model: Final = _IMAGE_GENERATION_MODEL_ADAPTER.validate_python(data.get("model")) + data.update( + await hierarchical_router_settings_updates( + requested_model=requested_model, + user_api_key_dict=user_api_key_dict, + llm_router=llm_router, + proxy_config=proxy_config, + proxy_logging_obj=proxy_logging_obj, + apply_model_group_alias=request.scope.get(MODEL_GROUP_ALIAS_RESOLVED_SCOPE_KEY) is not True, + ) + ) + ### CALL HOOKS ### - modify incoming data / reject request before calling the model prompt_value: Final = data.get("prompt") if prompt_value is not None: @@ -156,6 +173,7 @@ async def image_generation( route_type="aimage_generation", llm_router=llm_router, user_model=user_model, + user_api_key_dict=user_api_key_dict, ) response = await llm_call diff --git a/tests/integration/routing/test_key_router_settings_image_generation.py b/tests/integration/routing/test_key_router_settings_image_generation.py new file mode 100644 index 00000000000..84415c72f2f --- /dev/null +++ b/tests/integration/routing/test_key_router_settings_image_generation.py @@ -0,0 +1,162 @@ +from __future__ import annotations + +import json +import uuid +from types import MappingProxyType +from typing import Final + +from integration._support.client import Gateway, Scenario, gateway_from_environment +from integration._support.wire import Reply, Request, Wire, wire_server + +FRESH_CONNECTION: Final = MappingProxyType({"Connection": "close"}) +IMAGE_RESPONSE: Final = { + "created": 1, + "background": None, + "data": [ + { + "b64_json": None, + "revised_prompt": None, + "url": "https://images.example/generated.png", + "provider_specific_fields": None, + } + ], + "output_format": None, + "quality": None, + "size": None, + "usage": { + "total_tokens": 0, + "input_tokens": 0, + "input_tokens_details": {"image_tokens": 0, "text_tokens": 0}, + "output_tokens": 0, + }, +} + + +def _primary_unavailable(_: Request) -> Reply: + return Reply( + status=503, + body=json.dumps({"error": {"message": "primary unavailable", "type": "server_error"}}).encode(), + ) + + +def _image_reply(_: Request) -> Reply: + return Reply( + body=json.dumps( + { + "created": 1, + "data": [{"url": "https://images.example/generated.png"}], + } + ).encode() + ) + + +def _image_model(scenario: Scenario, group: str, wire: Wire) -> str: + return scenario.model( + name=group, + model="openai/dall-e-3", + api_base=f"{wire.url}/v1", + api_key="integration-provider-key", + num_retries=0, + ) + + +def _assert_image_fallback( + gateway: Gateway, + key: str, + primary: Wire, + fallback: Wire, + primary_model: str, +) -> None: + response: Final = gateway.request( + "POST", + "/v1/images/generations", + {"model": primary_model, "prompt": "draw a tree"}, + key=key, + headers=FRESH_CONNECTION, + ) + + assert response.status_code == 200, response.text + assert response.json() == IMAGE_RESPONSE, response.text + assert tuple((request.method, request.target) for request in primary.drain()) == ( + ("POST", "/v1/images/generations"), + ), response.text + assert tuple((request.method, request.target) for request in fallback.drain()) == ( + ("POST", "/v1/images/generations"), + ), response.text + + +def test_key_router_settings_fallback_routes_image_generation() -> None: + with ( + wire_server(_primary_unavailable) as primary_wire, + wire_server(_image_reply) as fallback_wire, + gateway_from_environment() as gateway, + gateway.scenario() as scenario, + ): + primary_model: Final = _image_model(scenario, f"image-primary-{uuid.uuid4().hex}", primary_wire) + fallback_model: Final = _image_model(scenario, f"image-fallback-{uuid.uuid4().hex}", fallback_wire) + key: Final = scenario.key(router_settings={"fallbacks": [{primary_model: [fallback_model]}], "num_retries": 0}) + + _assert_image_fallback(gateway, key, primary_wire, fallback_wire, primary_model) + + +def test_team_router_settings_fallback_routes_image_generation() -> None: + with ( + wire_server(_primary_unavailable) as primary_wire, + wire_server(_image_reply) as fallback_wire, + gateway_from_environment() as gateway, + gateway.scenario() as scenario, + ): + primary_model: Final = _image_model(scenario, f"image-primary-{uuid.uuid4().hex}", primary_wire) + fallback_model: Final = _image_model(scenario, f"image-fallback-{uuid.uuid4().hex}", fallback_wire) + team: Final = scenario.team( + router_settings={"fallbacks": [{primary_model: [fallback_model]}], "num_retries": 0} + ) + key: Final = scenario.key(team_id=team) + + _assert_image_fallback(gateway, key, primary_wire, fallback_wire, primary_model) + + +def test_key_model_group_alias_resolves_once_for_image_generation() -> None: + with ( + wire_server(_image_reply) as primary_wire, + wire_server(_image_reply) as target_wire, + wire_server(_image_reply) as chained_target_wire, + gateway_from_environment() as gateway, + gateway.scenario() as scenario, + ): + primary_model: Final = _image_model(scenario, f"image-primary-{uuid.uuid4().hex}", primary_wire) + target_model: Final = _image_model(scenario, f"image-target-{uuid.uuid4().hex}", target_wire) + chained_target_model: Final = _image_model( + scenario, f"image-chained-target-{uuid.uuid4().hex}", chained_target_wire + ) + alias_name: Final = f"image-alias-{uuid.uuid4().hex}" + key: Final = scenario.key( + router_settings={ + "model_group_alias": { + alias_name: target_model, + target_model: chained_target_model, + } + } + ) + + response: Final = gateway.request( + "POST", + "/v1/images/generations", + {"model": alias_name, "prompt": "draw a tree"}, + key=key, + headers=FRESH_CONNECTION, + ) + primary_requests: Final = tuple((request.method, request.target) for request in primary_wire.drain()) + target_requests: Final = tuple((request.method, request.target) for request in target_wire.drain()) + chained_target_requests: Final = tuple( + (request.method, request.target) for request in chained_target_wire.drain() + ) + + assert alias_name not in {primary_model, target_model, chained_target_model}, response.text + assert response.status_code == 200, response.text + assert response.json() == IMAGE_RESPONSE, response.text + assert primary_requests == (), response.text + assert target_requests == (("POST", "/v1/images/generations"),), ( + f"{response.text}; chained target requests: {chained_target_requests}" + ) + assert chained_target_requests == (), response.text diff --git a/tests/unit/proxy/image_endpoints/test_endpoints.py b/tests/unit/proxy/image_endpoints/test_endpoints.py index f4aebecc11e..6a2c7e3ed41 100644 --- a/tests/unit/proxy/image_endpoints/test_endpoints.py +++ b/tests/unit/proxy/image_endpoints/test_endpoints.py @@ -3,7 +3,7 @@ import copy import logging from collections.abc import Iterator, Mapping from types import SimpleNamespace -from typing import Any, Dict +from typing import Any, Dict, Final, cast import orjson import pytest @@ -12,10 +12,17 @@ from fastapi.testclient import TestClient from starlette.requests import Request from starlette.responses import Response +import litellm from litellm._logging import verbose_proxy_logger +from litellm.models.team import LiteLLM_TeamTableCachedObj from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.image_endpoints import endpoints +from litellm.proxy.proxy_server import PrismaClient, ProxyConfig +from litellm.proxy.utils import ProxyLogging +from litellm.types.utils import ImageObject, ImageResponse +from litellm.utils import custom_llm_setup @pytest.mark.asyncio @@ -410,3 +417,142 @@ async def test_failure_before_the_provider_call_bills_the_callers_litellm_call_i assert raised.value.headers["x-litellm-call-id"] == call_id assert [data["litellm_call_id"] for data in hook_request_data] == [call_id] + + +class _ImageRouterHandler(litellm.CustomLLM): + def __init__(self) -> None: + self.calls: list[str] = [] + + async def aimage_generation( + self, + model: str, + prompt: str, + model_response: ImageResponse, + api_key: str | None, + api_base: str | None, + optional_params: Mapping[str, object], + logging_obj: object, + timeout: object | None = None, + client: object | None = None, + ) -> ImageResponse: + self.calls.append(model) + if model.endswith("failing"): + raise litellm.RateLimitError(message="primary failed", llm_provider="image-regression", model=model) + return ImageResponse(created=1, data=[ImageObject(url="https://images.example/generated.png")]) + + +def _image_generation_client( + monkeypatch: pytest.MonkeyPatch, + user_api_key_dict: UserAPIKeyAuth, + team_settings: Mapping[str, object] | None = None, +) -> tuple[TestClient, _ImageRouterHandler]: + from litellm.proxy import proxy_server + + handler: Final = _ImageRouterHandler() + cache: Final = UserApiKeyCache() + monkeypatch.setattr(litellm, "provider_list", [*litellm.provider_list]) + monkeypatch.setattr(litellm, "_custom_providers", [*litellm._custom_providers]) + monkeypatch.setattr(litellm, "custom_provider_map", [{"provider": "image-regression", "custom_handler": handler}]) + custom_llm_setup() + if team_settings is not None: + team_id: Final = user_api_key_dict.team_id + assert team_id is not None + cache.set_cache( + key=f"team_id:{team_id}", + value=LiteLLM_TeamTableCachedObj(team_id=team_id, router_settings=dict(team_settings)), + local_only=True, + model_type=LiteLLM_TeamTableCachedObj, + ) + monkeypatch.setattr(proxy_server, "prisma_client", cast(PrismaClient, object())) + + router: Final = litellm.Router( + model_list=[ + {"model_name": "A", "litellm_params": {"model": "image-regression/failing"}}, + {"model_name": "B", "litellm_params": {"model": "image-regression/working"}}, + ], + num_retries=0, + ) + monkeypatch.setattr(litellm, "model_alias_map", {}) + monkeypatch.setattr(proxy_server, "general_settings", {}) + monkeypatch.setattr(proxy_server, "llm_router", router) + monkeypatch.setattr(proxy_server, "proxy_config", ProxyConfig()) + monkeypatch.setattr(proxy_server, "proxy_logging_obj", ProxyLogging(user_api_key_cache=cache)) + monkeypatch.setattr(proxy_server, "user_api_key_cache", cache) + monkeypatch.setattr(proxy_server, "user_model", None) + monkeypatch.setattr(proxy_server, "version", "test-version") + + app: Final = FastAPI() + app.include_router(endpoints.router) + app.dependency_overrides[user_api_key_auth] = lambda: user_api_key_dict + return TestClient(app, raise_server_exceptions=False), handler + + +def test_image_generation_uses_key_router_settings_fallbacks(monkeypatch: pytest.MonkeyPatch) -> None: + client, handler = _image_generation_client( + monkeypatch, + UserAPIKeyAuth(router_settings={"fallbacks": [{"A": ["B"]}]}), + ) + + response: Final = client.post("/v1/images/generations", json={"model": "A", "prompt": "draw a tree"}) + + assert response.status_code == 200, response.text + assert response.json() == ImageResponse( + created=1, + data=[ImageObject(url="https://images.example/generated.png")], + ).model_dump(), response.text + assert handler.calls == ["failing", "working"], response.text + + +def test_image_generation_uses_team_router_settings_fallbacks(monkeypatch: pytest.MonkeyPatch) -> None: + team_id: Final = "image-regression-team" + client, handler = _image_generation_client( + monkeypatch, + UserAPIKeyAuth(team_id=team_id), + {"fallbacks": [{"A": ["B"]}]}, + ) + + response: Final = client.post("/v1/images/generations", json={"model": "A", "prompt": "draw a tree"}) + + assert response.status_code == 200, response.text + assert response.json() == ImageResponse( + created=1, + data=[ImageObject(url="https://images.example/generated.png")], + ).model_dump(), response.text + assert handler.calls == ["failing", "working"], response.text + + +def test_image_generation_without_router_settings_does_not_use_fallback(monkeypatch: pytest.MonkeyPatch) -> None: + client, handler = _image_generation_client(monkeypatch, UserAPIKeyAuth()) + + response: Final = client.post("/v1/images/generations", json={"model": "A", "prompt": "draw a tree"}) + + assert response.status_code == 500, response.text + assert handler.calls == ["failing"], response.text + + +def test_image_generation_uses_key_model_group_alias(monkeypatch: pytest.MonkeyPatch) -> None: + client, handler = _image_generation_client( + monkeypatch, + UserAPIKeyAuth(router_settings={"model_group_alias": {"alias": "B"}}), + ) + + response: Final = client.post("/v1/images/generations", json={"model": "alias", "prompt": "draw a tree"}) + + assert response.status_code == 200, response.text + assert response.json() == ImageResponse( + created=1, + data=[ImageObject(url="https://images.example/generated.png")], + ).model_dump(), response.text + assert handler.calls == ["working"], response.text + + +def test_image_generation_uses_key_num_retries(monkeypatch: pytest.MonkeyPatch) -> None: + client, handler = _image_generation_client( + monkeypatch, + UserAPIKeyAuth(router_settings={"num_retries": 2}), + ) + + response: Final = client.post("/v1/images/generations", json={"model": "A", "prompt": "draw a tree"}) + + assert response.status_code == 500, response.text + assert handler.calls == ["failing", "failing", "failing"], response.text From 8f968b6c8e2ed2a7c3a1bdbc4d75dc733308e14e Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 30 Oct 2026 12:51:48 -0700 Subject: [PATCH 15/46] feat(openrouter): record upstream provider in logging payload and spend logs (#49287) * test: cover OpenRouter upstream provider logging Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test: scope streaming provider metadata guard Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * feat(openrouter): preserve upstream provider in logs Co-authored-by: BeWayllar <124474371+BeWayllar@users.noreply.github.com> Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(proxy): omit upstream_provider from spend metadata when unset * refactor(openrouter): set upstream provider hidden params without mutation --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Co-authored-by: BeWayllar <124474371+BeWayllar@users.noreply.github.com> --- .../streaming_chunk_builder_utils.py | 34 ++- .../litellm_core_utils/streaming_handler.py | 23 +- .../llms/openrouter/chat/transformation.py | 30 ++- litellm/proxy/_types.py | 1 + .../spend_tracking/spend_tracking_utils.py | 21 ++ litellm/types/utils.py | 1 + .../test_openrouter_upstream_provider.py | 216 ++++++++++++++++++ .../test_litellm_logging.py | 7 +- .../test_streaming_chunk_builder_utils.py | 65 ++++++ .../test_openrouter_chat_transformation.py | 117 ++++++++++ .../test_spend_tracking_utils.py | 64 +++++- 11 files changed, 564 insertions(+), 15 deletions(-) create mode 100644 tests/integration/spend/test_openrouter_upstream_provider.py diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 0ddfa1210c2..4857973d51d 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -329,6 +329,20 @@ class ChunkProcessor: models[0] if models else None, ) + @staticmethod + def _get_upstream_provider(chunks: Sequence[object]) -> str | None: + def _provider_from_chunk(chunk: object) -> str | None: + hidden_params: Final = get_hidden_params(chunk) + if hidden_params is None: + return None + provider: Final[object] = hidden_params.get("upstream_provider") + return provider if isinstance(provider, str) and provider else None + + return next( + (provider for chunk in chunks if (provider := _provider_from_chunk(chunk)) is not None), + None, + ) + @staticmethod def apply_provider_assembled_streaming_metadata( response: ModelResponse, @@ -506,10 +520,22 @@ class ChunkProcessor: chunks, first_chunk_model, ) - if provider_response_model is not None: - response._hidden_params = dict( # pyright: ignore[reportPrivateUsage] # ModelResponse exposes no public hidden-params setter - response._hidden_params, # pyright: ignore[reportPrivateUsage] # ModelResponse exposes no public hidden-params getter - provider_response_model=provider_response_model, + upstream_provider: Final = self._get_upstream_provider(chunks) + provider_hidden_params: Final[Mapping[str, str]] = MappingProxyType( + { + key: value + for key, value in ( + ("provider_response_model", provider_response_model), + ("upstream_provider", upstream_provider), + ) + if value is not None + } + ) + if provider_hidden_params: + response_hidden_params: Final[Mapping[str, builtins.object]] = get_hidden_params(response) or {} + set_hidden_params( + response, + dict(response_hidden_params, **provider_hidden_params), ) return response diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 9710956c97e..484edd869cb 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -277,6 +277,7 @@ class _VertexChunkLike(Protocol): class _ParsedChunkHiddenParams(LiteLLMBaseModel): provider_specific_fields: Mapping[str, object] | None = None + upstream_provider: str | None = None def _provider_response_model(chunk: object) -> str | None: @@ -293,9 +294,17 @@ def _parsed_provider_hidden_params(hidden: object) -> _ParsedChunkHiddenParams | return None +def _provider_upstream_provider(chunk: object) -> str | None: + parsed: Final = _parsed_provider_hidden_params(get_hidden_params(chunk)) + if parsed is None or not parsed.upstream_provider: + return None + return parsed.upstream_provider + + def _provider_hidden_params( chunk: object, provider_response_model: str | None, + upstream_provider: str | None, ) -> Mapping[str, object] | None: hidden: Final = get_hidden_params(chunk) parsed: Final = _parsed_provider_hidden_params(hidden) @@ -308,6 +317,10 @@ def _provider_hidden_params( for key, value in ( ("provider_response_model", provider_response_model), ("provider_specific_fields", provider_specific_fields), + ( + "upstream_provider", + parsed.upstream_provider if parsed is not None and parsed.upstream_provider else upstream_provider, + ), ) if value is not None } @@ -360,6 +373,7 @@ class CustomStreamWrapper: self.system_fingerprint: str | None = None self._provider_response_model: str | None = None + self._upstream_provider: str | None = None self.received_finish_reason: str | None = None self.intermittent_finish_reason: str | None = None # finish reasons that show up mid-stream self.special_tokens = [ @@ -1653,10 +1667,17 @@ class CustomStreamWrapper: provider_response_model: Final = _provider_response_model(chunk) if provider_response_model is not None: self._provider_response_model = provider_response_model + upstream_provider: Final = _provider_upstream_provider(chunk) + if self._upstream_provider is None and upstream_provider is not None: + self._upstream_provider = upstream_provider prebuilt_delta: Final = self._prebuilt_plain_delta(chunk) plain_delta_created: bool = prebuilt_delta is not None # rebind-ok: the first chunk builds its delta below model_response = self.model_response_creator( - hidden_params=_provider_hidden_params(chunk, self._provider_response_model), + hidden_params=_provider_hidden_params( + chunk, + self._provider_response_model, + self._upstream_provider, + ), delta=prebuilt_delta, ) response_obj: dict[str, Any] = {} diff --git a/litellm/llms/openrouter/chat/transformation.py b/litellm/llms/openrouter/chat/transformation.py index cca6815399f..69f3d76c034 100644 --- a/litellm/llms/openrouter/chat/transformation.py +++ b/litellm/llms/openrouter/chat/transformation.py @@ -14,7 +14,7 @@ import httpx from pydantic import ConfigDict, TypeAdapter from litellm.litellm_core_utils.core_helpers import finite_cost -from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, set_hidden_params +from litellm.litellm_core_utils.hidden_params import HIDDEN_PARAMS_ATTR, get_hidden_params, set_hidden_params from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolParam @@ -212,19 +212,27 @@ class OpenrouterConfig(OpenAIGPTConfig): # OpenRouter returns cost information in the usage object when usage.include=true try: response_json: Final = _JSON_OBJECT.validate_python(raw_response.json()) + response_provider: Final = response_json.get("provider") + if isinstance(response_provider, str) and response_provider: + set_hidden_params( + model_response, + {**(get_hidden_params(model_response) or {}), "upstream_provider": response_provider}, + ) if "usage" in response_json and response_json["usage"]: response_cost: Final = _JSON_OBJECT.validate_python(response_json["usage"]).get("cost") if response_cost is not None: # Store cost in hidden params for the cost calculator to use if not hasattr(model_response, HIDDEN_PARAMS_ATTR): set_hidden_params(model_response, {}) - hidden_params: Final = cast( # cast-ok: preserve mapping operations on dynamic response metadata - dict[str, object], getattr(model_response, HIDDEN_PARAMS_ATTR) + cost_hidden_params: Final = ( + cast( # cast-ok: preserve mapping operations on dynamic response metadata + dict[str, object], getattr(model_response, HIDDEN_PARAMS_ATTR) + ) ) - if "additional_headers" not in hidden_params: - hidden_params["additional_headers"] = {} + if "additional_headers" not in cost_hidden_params: + cost_hidden_params["additional_headers"] = {} additional_headers: Final = cast( # cast-ok: preserve mapping operations on response metadata - dict[str, object], hidden_params["additional_headers"] + dict[str, object], cost_hidden_params["additional_headers"] ) reported_cost: Final = ( finite_cost(float(response_cost)) if isinstance(response_cost, str | int | float) else None @@ -282,7 +290,7 @@ class OpenRouterChatCompletionStreamingHandler(BaseModelResponseIterator): for choice in chunk["choices"]: choice["delta"]["reasoning_content"] = choice["delta"].get("reasoning") new_choices.append(choice) - return ModelResponseStream( + model_response: Final = ModelResponseStream( id=chunk["id"], object="chat.completion.chunk", created=chunk["created"], @@ -290,6 +298,12 @@ class OpenRouterChatCompletionStreamingHandler(BaseModelResponseIterator): model=chunk["model"], choices=new_choices, ) + typed_chunk: Final[Mapping[str, object]] = cast( # cast-ok: stream chunks are decoded JSON objects + Mapping[str, object], chunk + ) + upstream_provider: Final[object] = typed_chunk.get("provider") + if isinstance(upstream_provider, str) and upstream_provider: + set_hidden_params(model_response, {"upstream_provider": upstream_provider}) except KeyError as e: raise OpenRouterException( message=f"KeyError: {e}, Got unexpected response from OpenRouter: {chunk}", @@ -298,3 +312,5 @@ class OpenRouterChatCompletionStreamingHandler(BaseModelResponseIterator): ) except Exception as e: raise e + else: + return model_response diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index f165355da5a..73b12399af4 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3758,6 +3758,7 @@ class SpendLogsMetadata(TypedDict): kubernetes_pod_routing: ReadOnly[StandardLoggingKubernetesPodRouting | None] internal_call_origin: InternalCallOrigin | None billed_retry_attempt: ReadOnly[NotRequired[bool | None]] + upstream_provider: ReadOnly[NotRequired[str | None]] litellm_roi_estimator: ReadOnly[NotRequired[bool | None]] is_silent_experiment: ReadOnly[NotRequired[bool | None]] guardrail_information: list[StandardLoggingGuardrailInformation] | None diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index f8d5cba4229..306e70f7029 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -11,6 +11,7 @@ from types import MappingProxyType from typing import TYPE_CHECKING, Final, Literal, Protocol, cast, runtime_checkable from pydantic import BaseModel, JsonValue, TypeAdapter, ValidationError +from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger @@ -194,6 +195,7 @@ _STAMPED_METADATA_KEYS: Final = frozenset( "autorouter_savings_estimate", "autorouter_baseline_observation", "used_client_oauth_token", + "upstream_provider", KUBERNETES_POD_ROUTING_KEY, "litellm_roi_estimator", "is_silent_experiment", @@ -203,6 +205,10 @@ _STAMPED_METADATA_KEYS: Final = frozenset( ) +class _UpstreamProviderMetadata(TypedDict, total=False): + upstream_provider: ReadOnly[str] + + def _stored_user_email(metadata: Mapping[str, object]) -> str | None: email: Final = metadata.get("user_api_key_user_email") return email if isinstance(email, str) and email.strip() else None @@ -221,6 +227,7 @@ def _get_spend_logs_metadata( model_map_information: StandardLoggingModelInformation | None = None, cold_storage_object_key: str | None = None, litellm_overhead_time_ms: float | None = None, + upstream_provider: str | None = None, cost_breakdown: CostBreakdown | None = None, litellm_call_id: str | None = None, autorouter_savings: float | None = None, @@ -232,6 +239,9 @@ def _get_spend_logs_metadata( kubernetes_pod_routing: StandardLoggingKubernetesPodRouting | None = None, is_silent_experiment: bool = False, ) -> SpendLogsMetadata: + upstream_provider_metadata: Final[_UpstreamProviderMetadata] = ( + {"upstream_provider": upstream_provider} if upstream_provider is not None else {} + ) if metadata is None: return SpendLogsMetadata( user_api_key=None, @@ -265,6 +275,7 @@ def _get_spend_logs_metadata( eval_information=None, cold_storage_object_key=cold_storage_object_key, litellm_overhead_time_ms=None, + **upstream_provider_metadata, attempted_retries=None, max_retries=None, attempted_fallbacks=None, @@ -290,6 +301,7 @@ def _get_spend_logs_metadata( **MappingProxyType( {key: metadata.get(key) for key in SpendLogsMetadata.__annotations__ if key not in _STAMPED_METADATA_KEYS} ), + **upstream_provider_metadata, autorouter_savings=autorouter_savings, autorouter_savings_estimate=autorouter_savings_estimate, autorouter_baseline_observation=autorouter_baseline_observation, @@ -706,6 +718,14 @@ def get_logging_payload( if standard_logging_payload is not None: hidden_params: Final = standard_logging_payload.get("hidden_params", {}) litellm_overhead_time_ms = hidden_params.get("litellm_overhead_time_ms") + upstream_provider_value: Final[object] = ( + standard_logging_payload["hidden_params"]["upstream_provider"] + if standard_logging_payload is not None + and "hidden_params" in standard_logging_payload + and "upstream_provider" in standard_logging_payload["hidden_params"] + else None + ) + upstream_provider: Final = upstream_provider_value if isinstance(upstream_provider_value, str) else None logged_provider: Final = ( kwargs.get("custom_llm_provider") @@ -807,6 +827,7 @@ def get_logging_payload( else None ), litellm_overhead_time_ms=litellm_overhead_time_ms, + upstream_provider=upstream_provider, cost_breakdown=( standard_logging_payload.get("cost_breakdown", None) if standard_logging_payload is not None else None ), diff --git a/litellm/types/utils.py b/litellm/types/utils.py index f8115571dfd..db865c7488a 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3460,6 +3460,7 @@ class StandardLoggingHiddenParams(TypedDict): batch_models: list[str] | None batch_successful_requests: ReadOnly[int | None] batch_failed_requests: ReadOnly[int | None] + upstream_provider: ReadOnly[NotRequired[str | None]] litellm_model_name: str | None # the model name sent to the provider by litellm usage_object: dict | None diff --git a/tests/integration/spend/test_openrouter_upstream_provider.py b/tests/integration/spend/test_openrouter_upstream_provider.py new file mode 100644 index 00000000000..fa17a364aea --- /dev/null +++ b/tests/integration/spend/test_openrouter_upstream_provider.py @@ -0,0 +1,216 @@ +import json +import uuid +from collections.abc import Callable +from typing import Final, Literal + +import httpx +import pytest +from integration._support.client import JSON_OBJECT, Gateway, Scenario, eventually, object_value, string_value +from integration._support.database import read_rows +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.shares_worker_state + +_BACKEND: Final = "openai/gpt-4o-mini" +_PROMPT: Final = "upstream provider trace" +_PROVIDER: Final = "serving-site-a" +_ANSWER: Final = "hello" +_USAGE: Final[dict[str, JsonValue]] = { + "prompt_tokens": 5, + "completion_tokens": 1, + "total_tokens": 6, + "cost": 0.00015, +} + + +def _reply( + request_id: str, + *, + response_provider: bool, + stream_provider_mode: Literal["all", "last", "none"], +) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + assert request.method == "POST" + assert request.target.split("?", 1)[0].endswith("/chat/completions"), request.target + body: Final = JSON_OBJECT.validate_json(request.body) + assert body["messages"] == [{"role": "user", "content": _PROMPT}], body + if body.get("stream") is True: + common: Final = { + "id": request_id, + "object": "chat.completion.chunk", + "created": 1, + "model": _BACKEND, + } + all_provider: Final[dict[str, JsonValue]] = {"provider": _PROVIDER} + first_provider: Final = all_provider if stream_provider_mode == "all" else {} + finish_provider: Final = ( + all_provider if stream_provider_mode in ("all", "last") else {} + ) + usage_provider: Final = all_provider if stream_provider_mode == "all" else {} + chunks: Final[tuple[dict[str, JsonValue], ...]] = ( + { + **common, + **first_provider, + "choices": [{"index": 0, "delta": {"role": "assistant", "content": _ANSWER}, "finish_reason": None}], + }, + { + **common, + **finish_provider, + "choices": [{"index": 0, "delta": {}, "finish_reason": "stop"}], + }, + {**common, **usage_provider, "choices": [], "usage": _USAGE}, + ) + return Reply( + content_type="text/event-stream", + chunks=tuple(f"data: {json.dumps(chunk)}\n\n".encode() for chunk in chunks) + + (b"data: [DONE]\n\n",), + ) + + provider: Final[dict[str, JsonValue]] = {"provider": _PROVIDER} if response_provider else {} + return Reply( + body=json.dumps( + { + "id": request_id, + "object": "chat.completion", + "created": 1, + "model": _BACKEND, + **provider, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": _ANSWER}, + "finish_reason": "stop", + } + ], + "usage": _USAGE, + } + ).encode() + ) + + return respond + + +def _stream_chunks(response: httpx.Response) -> tuple[dict[str, JsonValue], ...]: + return tuple( + JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in response.text.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + + +def _send( + gateway: Gateway, + scenario: Scenario, + wire_url: str, + request_id: str, + *, + openrouter: bool, + stream: bool, +) -> httpx.Response: + model: Final = scenario.model( + model=f"{'openrouter' if openrouter else 'openai'}/{_BACKEND}", + api_base=wire_url, + api_key="synthetic-provider-key", + ) + key: Final = scenario.key(models=[model]) + response: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [{"role": "user", "content": _PROMPT}], + "stream": stream, + **({"stream_options": {"include_usage": True}} if stream else {}), + }, + key=key, + ) + assert response.status_code == 200, response.text + returned_id: Final = ( + string_value(_stream_chunks(response)[0]["id"]) + if stream + else string_value(JSON_OBJECT.validate_json(response.content)["id"]) + ) + assert returned_id == request_id, response.text + return response + + +def _spend_metadata(request_id: str) -> dict[str, JsonValue]: + rows: Final = eventually( + lambda: read_rows( + 'SELECT metadata FROM "LiteLLM_SpendLogs" WHERE request_id=%s', + (request_id,), + ), + lambda values: len(values) == 1, + seconds=70, + ) + return object_value(rows[0]["metadata"]) + + +@pytest.mark.parametrize( + ("stream", "stream_provider_mode"), + [ + pytest.param(False, "none", id="nonstream"), + pytest.param(True, "all", id="stream-every-chunk"), + pytest.param(True, "last", id="stream-later-chunk"), + ], +) +def test_openrouter_provider_is_persisted_in_spend_metadata_regression( + gateway: Gateway, stream: bool, stream_provider_mode: Literal["all", "last", "none"] +) -> None: + request_id: Final = f"chatcmpl-{uuid.uuid4().hex}" + with ( + wire_server( + _reply( + request_id, + response_provider=True, + stream_provider_mode=stream_provider_mode, + ) + ) as wire, + gateway.scenario() as scenario, + ): + _send(gateway, scenario, wire.url, request_id, openrouter=True, stream=stream) + metadata: Final = _spend_metadata(request_id) + + assert metadata.get("upstream_provider") == _PROVIDER, metadata + + +def test_openrouter_client_response_keeps_provider_in_body_guard(gateway: Gateway) -> None: + request_id: Final = f"chatcmpl-{uuid.uuid4().hex}" + with ( + wire_server( + _reply( + request_id, + response_provider=True, + stream_provider_mode="none", + ) + ) as wire, + gateway.scenario() as scenario, + ): + response: Final = _send(gateway, scenario, wire.url, request_id, openrouter=True, stream=False) + + body: Final = JSON_OBJECT.validate_json(response.content) + choices: Final = body["choices"] + assert isinstance(choices, list) + message: Final = object_value(object_value(choices[0])["message"]) + assert body["provider"] == _PROVIDER + assert "upstream_provider" not in body + assert message["content"] == _ANSWER + + +def test_non_openrouter_spend_metadata_has_no_provider_guard(gateway: Gateway) -> None: + request_id: Final = f"chatcmpl-{uuid.uuid4().hex}" + with ( + wire_server( + _reply( + request_id, + response_provider=False, + stream_provider_mode="none", + ) + ) as wire, + gateway.scenario() as scenario, + ): + _send(gateway, scenario, wire.url, request_id, openrouter=False, stream=False) + metadata: Final = _spend_metadata(request_id) + + assert "upstream_provider" not in metadata, metadata diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index d0e31cdd57a..f7c2082b683 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -12574,9 +12574,10 @@ def test_get_hidden_params(): assert result["api_base"] is None assert result["response_cost"] is None assert result["additional_headers"] is None + assert result.get("upstream_provider") is None # assert all fields in StandardLoggingHiddenParams are present - assert all(field in result for field in StandardLoggingHiddenParams.__annotations__) + assert all(field in result for field in StandardLoggingHiddenParams.__annotations__ if field != "upstream_provider") # Test with valid params hidden_params = { @@ -12584,6 +12585,7 @@ def test_get_hidden_params(): "cache_key": "test-cache", "api_base": "https://api.test.com", "response_cost": 0.001, + "upstream_provider": "serving-site-a", "additional_headers": { "x-ratelimit-limit-requests": "2000", "x-ratelimit-remaining-requests": "1999", @@ -12596,8 +12598,9 @@ def test_get_hidden_params(): assert result["response_cost"] == 0.001 assert result["additional_headers"] is not None assert result["additional_headers"]["x_ratelimit_limit_requests"] == 2000 + assert result.get("upstream_provider") == "serving-site-a" # assert all fields in StandardLoggingHiddenParams are present - assert all(field in result for field in StandardLoggingHiddenParams.__annotations__) + assert all(field in result for field in StandardLoggingHiddenParams.__annotations__ if field != "upstream_provider") @pytest.mark.usefixtures("_vcr_outcome_gate", "drain_logging_worker", "isolate_litellm_state", "setup_and_teardown") diff --git a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index 9df45b92686..bbf1897b3a7 100644 --- a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -910,6 +910,71 @@ def test_stream_chunk_builder_dict_snapshot_preserves_hidden_provider_fields(): assert response._hidden_params["provider_specific_fields"]["traffic_type"] == "default" +def test_stream_chunk_builder_finds_upstream_provider_in_later_chunk_regression() -> None: + first: Final = ModelResponseStream( + id="chatcmpl-provider", + created=1, + model="gpt-4.1-mini", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content="hello", role="assistant"), + ) + ], + ) + later: Final = ModelResponseStream( + id="chatcmpl-provider", + created=1, + model="gpt-4.1-mini", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content=" world", role=None), + ) + ], + ) + first_hidden_params: Final = {"origin": "first"} + later_hidden_params: Final = {"upstream_provider": "serving-site-a"} + chunks: Final = [ + {**first.model_dump(), "_hidden_params": first_hidden_params}, + {**later.model_dump(), "_hidden_params": later_hidden_params}, + ] + + response: Final = stream_chunk_builder(chunks=chunks) + + assert response is not None + assert response._hidden_params.get("upstream_provider") == "serving-site-a" + assert first_hidden_params.get("upstream_provider") is None + assert later_hidden_params.get("upstream_provider") == "serving-site-a" + + +def test_stream_chunk_builder_preserves_source_provider_metadata_guard() -> None: + chunk: Final = ModelResponseStream( + id="chatcmpl-provider", + created=1, + model="gpt-4.1-mini", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content="hello", role="assistant"), + ) + ], + ) + hidden_params: Final = {"upstream_provider": "serving-site-a"} + chunk_dict: Final = {**chunk.model_dump(), "_hidden_params": hidden_params} + + response: Final = stream_chunk_builder(chunks=[chunk_dict]) + + assert response is not None + assert chunk_dict["_hidden_params"]["upstream_provider"] == "serving-site-a" + + def test_stream_chunk_builder_propagates_vertex_ai_metadata_from_chunks(): """Vertex AI metadata on streaming chunks must appear on assembled response.""" grounding_metadata = [{"webSearchQueries": ["weather in SF"]}] diff --git a/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py b/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py index 888c47505a4..3d659641f54 100644 --- a/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py +++ b/tests/unit/llms/openrouter/chat/test_openrouter_chat_transformation.py @@ -872,3 +872,120 @@ def test_openrouter_streaming_reasoning_is_preserved() -> None: assembled: Final = litellm.stream_chunk_builder(chunks=list(stream_chunks), messages=messages) assert assembled is not None assert assembled.choices[0].message.reasoning_content == "Consider the evidence. Apply the rule." + + +def _transform_openrouter_response(body: dict[str, object]) -> litellm.ModelResponse: + return OpenrouterConfig().transform_response( + model="openrouter/openai/gpt-4o-mini", + raw_response=httpx.Response(200, json=body), + model_response=litellm.ModelResponse(), + logging_obj=MagicMock(), + request_data={}, + messages=[{"role": "user", "content": "hi"}], + optional_params={}, + litellm_params={}, + encoding=None, + ) + + +def test_openrouter_nonstream_provider_without_usage_is_preserved_regression() -> None: + response: Final = _transform_openrouter_response( + { + "id": "gen-provider-no-usage", + "object": "chat.completion", + "created": 1, + "model": "openai/gpt-4o-mini", + "provider": "serving-site-a", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ], + } + ) + + assert response._hidden_params.get("upstream_provider") == "serving-site-a" + + +@pytest.mark.parametrize( + ("provider_fields", "case"), + [ + pytest.param({}, "missing", id="missing"), + pytest.param({"provider": None}, "null", id="null"), + pytest.param({"provider": ""}, "empty", id="empty"), + pytest.param({"provider": 123}, "non-string", id="non-string"), + ], +) +def test_openrouter_invalid_provider_keeps_reported_cost_guard( + provider_fields: dict[str, object], case: str +) -> None: + response: Final = _transform_openrouter_response( + { + "id": f"gen-invalid-provider-{case}", + "object": "chat.completion", + "created": 1, + "model": "openai/gpt-4o-mini", + **provider_fields, + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "hi"}, + "finish_reason": "stop", + } + ], + "usage": { + "prompt_tokens": 10, + "completion_tokens": 20, + "total_tokens": 30, + "cost": 0.00015, + }, + } + ) + + assert response._hidden_params.get("upstream_provider") is None + assert response._hidden_params["additional_headers"][OPENROUTER_COST_HEADER] == 0.00015 + + +def test_openrouter_streaming_chunk_preserves_provider_regression() -> None: + chunk: Final = { + "id": "gen-provider-stream", + "object": "chat.completion.chunk", + "created": 1, + "model": "openai/gpt-4o-mini", + "provider": "serving-site-a", + "choices": [{"index": 0, "delta": {"content": "hi"}, "finish_reason": None}], + } + + response: Final = OpenRouterChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ).chunk_parser(chunk) + + assert response._hidden_params.get("upstream_provider") == "serving-site-a" + + +@pytest.mark.parametrize( + ("provider_fields", "case"), + [ + pytest.param({}, "missing", id="missing"), + pytest.param({"provider": None}, "null", id="null"), + pytest.param({"provider": ""}, "empty", id="empty"), + pytest.param({"provider": 123}, "non-string", id="non-string"), + ], +) +def test_openrouter_streaming_chunk_ignores_invalid_provider_guard(provider_fields: dict[str, object], case: str) -> None: + chunk: Final = { + "id": f"gen-invalid-provider-{case}", + "object": "chat.completion.chunk", + "created": 1, + "model": "openai/gpt-4o-mini", + **provider_fields, + "choices": [{"index": 0, "delta": {"content": "hi"}, "finish_reason": None}], + } + + response: Final = OpenRouterChatCompletionStreamingHandler( + streaming_response=None, sync_stream=True + ).chunk_parser(chunk) + + assert response._hidden_params.get("upstream_provider") is None diff --git a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py index 1ea445fda0d..a91cdd512fe 100644 --- a/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/unit/proxy/spend_tracking/test_spend_tracking_utils.py @@ -2394,6 +2394,65 @@ def test_get_logging_payload_includes_overhead_in_spend_logs_metadata(): ) +@patch("litellm.proxy.proxy_server.master_key", None) +@patch("litellm.proxy.proxy_server.general_settings", {}) +def test_get_logging_payload_carries_upstream_provider_from_hidden_params_regression() -> None: + standard_logging_payload: Final = _make_standard_logging_payload_with_usage_object(usage_object={}) + payload_with_provider: Final = { + **dict(standard_logging_payload), + "hidden_params": { + **dict(standard_logging_payload["hidden_params"]), + "upstream_provider": "serving-site-a", + }, + } + payload: Final = get_logging_payload( + kwargs={ + "model": "gpt-5.6", + "call_type": "responses", + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + "standard_logging_object": payload_with_provider, + }, + response_obj={ + "id": "resp-test", + "usage": { + "input_tokens": 1000, + "output_tokens": 10, + "total_tokens": 1010, + }, + }, + start_time=datetime.datetime(2025, 1, 1, tzinfo=timezone.utc), + end_time=datetime.datetime(2025, 1, 1, 0, 0, 1, tzinfo=timezone.utc), + ) + + assert json.loads(payload["metadata"]).get("upstream_provider") == "serving-site-a" + + +@patch("litellm.proxy.proxy_server.master_key", None) +@patch("litellm.proxy.proxy_server.general_settings", {}) +def test_get_logging_payload_omits_upstream_provider_when_hidden_param_absent_guard() -> None: + standard_logging_payload: Final = _make_standard_logging_payload_with_usage_object(usage_object={}) + payload: Final = get_logging_payload( + kwargs={ + "model": "gpt-5.6", + "call_type": "responses", + "litellm_params": {"metadata": {"user_api_key": "test-key"}}, + "standard_logging_object": standard_logging_payload, + }, + response_obj={ + "id": "resp-test", + "usage": { + "input_tokens": 1000, + "output_tokens": 10, + "total_tokens": 1010, + }, + }, + start_time=datetime.datetime(2025, 1, 1, tzinfo=timezone.utc), + end_time=datetime.datetime(2025, 1, 1, 0, 0, 1, tzinfo=timezone.utc), + ) + + assert "upstream_provider" not in json.loads(payload["metadata"]) + + @patch("litellm.proxy.proxy_server.master_key", None) @patch("litellm.proxy.proxy_server.general_settings", {}) def test_get_logging_payload_strips_null_bytes_from_request_tags(): @@ -6911,7 +6970,10 @@ def test_spend_logs_payload(model_id: str | None): ) assert len(payload["request_id"]) > 0 - expected_metadata_keys: Final = SpendLogsMetadata.__annotations__.keys() - {"is_silent_experiment"} + expected_metadata_keys: Final = SpendLogsMetadata.__annotations__.keys() - { + "is_silent_experiment", + "upstream_provider", + } assert "metadata" in payload assert isinstance(payload["metadata"], str) From f4d1f45d87dcba757edd1feae71300146fbb1570 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Fri, 30 Oct 2026 13:05:12 -0700 Subject: [PATCH 16/46] fix(proxy): raise live budget counters for failure spend rows (#47741) * fix(proxy): raise live budget counters for failure spend rows A stream that breaks mid-flight writes its recovered cost to the spend log, but the failure hook had already released the reservation and never touched the live counters, so an over-budget key kept passing until the counter was next repaired from the database Every charged failure row with a cost now raises the same counters the success path raises. The error waits at most a second for them, so a hung Redis cannot hold it back Fixes #35568 * test(proxy): cover failure spend counters across budgets, endpoints and outages Integration cells for team, user, end-user, tag and window budgets, concurrent streams, billed retries, native /v1/messages, Bedrock guardrail cost, a Redis stall on two workers and disable_spend_updates * test(proxy): check stalled failure counters and use the real spend writer - The Redis stall cell now reads the burst key's live counter and requires at least the full failure spend - Unit tests drop the charging writer double, bare dicts, list mutation and wall-clock timestamps - Partial-stream test models are frozen * test(proxy): move the stalled counter bound reason into the assert message --------- Co-authored-by: yucheng-berri <293926326+yucheng-berri@users.noreply.github.com> --- .../proxy/hooks/proxy_track_cost_callback.py | 124 ++++++- litellm/proxy/proxy_server.py | 2 +- tests/integration/spend/_partial_stream.py | 173 +++++++++ .../test_failure_spend_live_counters_chaos.py | 239 ++++++++++++ .../spend/test_partial_stream_live_budget.py | 350 ++++++++++++++++++ .../hooks/test_proxy_track_cost_callback.py | 225 +++++++++++ 6 files changed, 1103 insertions(+), 10 deletions(-) create mode 100644 tests/integration/spend/_partial_stream.py create mode 100644 tests/integration/spend/test_failure_spend_live_counters_chaos.py create mode 100644 tests/integration/spend/test_partial_stream_live_budget.py diff --git a/litellm/proxy/hooks/proxy_track_cost_callback.py b/litellm/proxy/hooks/proxy_track_cost_callback.py index 0d333d4fd08..2951bc02d8c 100644 --- a/litellm/proxy/hooks/proxy_track_cost_callback.py +++ b/litellm/proxy/hooks/proxy_track_cost_callback.py @@ -12,6 +12,7 @@ from pydantic import TypeAdapter, ValidationError import litellm from litellm._logging import verbose_proxy_logger from litellm.batches.batch_utils import batch_cost_is_final +from litellm.caching.redis_batch import flush_post_call_redis_batches from litellm.constants import BACKGROUND_INTERACTION_COST_POLLING_ENABLED, BILLED_RETRY_ATTEMPT_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.core_helpers import ( @@ -134,6 +135,8 @@ _CAPTURED_IDENTITY_CALL_TYPES: Final[frozenset[str]] = frozenset( ) ) +FAILURE_SPEND_COUNTER_WAIT_SECONDS: Final = 1.0 + _OBJECT_MAPPING: Final = TypeAdapter(ObjectMapping) @@ -235,6 +238,12 @@ def _proxy_spend_writer() -> DBSpendUpdateWriter: return proxy_logging_obj.db_spend_update_writer +def _proxy_spend_counters() -> "_IncrementSpendCounters": + from litellm.proxy.proxy_server import increment_spend_counters + + return increment_spend_counters + + DeploymentLookup: TypeAlias = Callable[[str], Deployment | None] @@ -361,11 +370,16 @@ class _ProxyDBLogger(CustomLogger): message_logging: bool = True, spend_writer: Callable[[], DBSpendUpdateWriter] = _proxy_spend_writer, deployment_lookup: DeploymentLookup = _proxy_deployment, + spend_counters: Callable[[], "_IncrementSpendCounters"] = _proxy_spend_counters, + failure_counter_wait_seconds: float = FAILURE_SPEND_COUNTER_WAIT_SECONDS, ) -> None: super().__init__(turn_off_message_logging=turn_off_message_logging, message_logging=message_logging) self.spend_event_producer = spend_event_producer self._spend_writer: Final = spend_writer self._deployment_lookup: Final = deployment_lookup + self._spend_counters: Final = spend_counters + self._failure_counter_wait_seconds: Final = failure_counter_wait_seconds + self._pending_increments: Final[set[asyncio.Task[None]]] = set() # mutable-ok: tasks leave when done async def async_log_success_event( self, kwargs: ObjectMapping, response_obj: object, start_time: datetime, end_time: datetime @@ -590,9 +604,11 @@ class _ProxyDBLogger(CustomLogger): completed_cost: bool = False, request_tags: list[str] | None = None, ) -> None: - """One spend row (and one key/team/user spend deduction) per ``(response_cost, kwargs)`` pair.""" - for row_cost, row_kwargs in spend_rows: - if completed_cost: + """One spend row (and one key/team/user spend deduction) per ``(response_cost, kwargs)`` pair, then the + live spend counters raised by every charged row that cost something. A completed cost settles through the + success path's write, which raises the counters against the kept reservation.""" + if completed_cost: + for row_cost, row_kwargs in spend_rows: from litellm.proxy.proxy_server import increment_spend_counters, proxy_logging_obj await _update_database_and_spend_counters( @@ -623,12 +639,16 @@ class _ProxyDBLogger(CustomLogger): router=get_llm_router(), ), ) - continue + return + + token: Final = LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict) + end_user_id: Final = None if litellm.disable_end_user_cost_tracking else user_api_key_dict.end_user_id + charged: Final = [ await self._spend_writer().update_database( - token=LiteLLMProxyRequestSetup.get_logged_api_key(user_api_key_dict), + token=token, response_cost=row_cost, user_id=user_api_key_dict.user_id, - end_user_id=None if litellm.disable_end_user_cost_tracking else user_api_key_dict.end_user_id, + end_user_id=end_user_id, team_id=user_api_key_dict.team_id, kwargs=dict(row_kwargs), completion_response=original_exception, @@ -637,6 +657,91 @@ class _ProxyDBLogger(CustomLogger): org_id=user_api_key_dict.org_id, project_id=user_api_key_dict.project_id, ) + for row_cost, row_kwargs in spend_rows + ] + counted_rows: Final = tuple( + row for row, row_charged in zip(spend_rows, charged, strict=True) if row_charged and row[0] > 0 + ) + if counted_rows and not ProxyUpdateSpend.disable_spend_updates(): + await self._increment_failure_spend_counters( + user_api_key_dict, token, end_user_id, counted_rows, start_time + ) + + async def _increment_failure_spend_counters( + self, + user_api_key_dict: UserAPIKeyAuth, + token: str | None, + end_user_id: str | None, + counted_rows: Sequence[tuple[float, Mapping[str, object]]], + start_time: datetime, + ) -> None: + """The counters the success path raises for the same spend, sent before the client sees the error. A + counter store slower than ``failure_counter_wait_seconds`` finishes them in the background instead of + holding the error back. The reservation was already released, so every counter takes the full row cost.""" + increment: Final = asyncio.create_task( + self._raise_failure_spend_counters(user_api_key_dict, token, end_user_id, counted_rows, start_time) + ) + self._pending_increments.add(increment) + increment.add_done_callback(self._pending_increments.discard) + finished, _ = await asyncio.wait((increment,), timeout=self._failure_counter_wait_seconds) + if not finished: + verbose_proxy_logger.warning( + "Spend counters for failure spend rows took over %ss, finishing them in the background", + self._failure_counter_wait_seconds, + ) + + async def _raise_failure_spend_counters( + self, + user_api_key_dict: UserAPIKeyAuth, + token: str | None, + end_user_id: str | None, + counted_rows: Sequence[tuple[float, Mapping[str, object]]], + start_time: datetime, + ) -> None: + for row_cost, row_kwargs in counted_rows: + await self._raise_row_spend_counters( + user_api_key_dict, token, end_user_id, row_cost, row_kwargs, start_time + ) + await flush_post_call_redis_batches() + + async def _raise_row_spend_counters( + self, + user_api_key_dict: UserAPIKeyAuth, + token: str | None, + end_user_id: str | None, + row_cost: float, + row_kwargs: Mapping[str, object], + start_time: datetime, + ) -> None: + standard_logging_object: Final = row_kwargs.get("standard_logging_object") + sl_object: Final = ( + cast("StandardLoggingPayload", standard_logging_object) # cast-ok: the failure snapshot dict + if isinstance(standard_logging_object, dict) + else None + ) + metadata: Final = _str_keyed_fields(_str_keyed_fields(row_kwargs.get("litellm_params")).get("metadata")) + served_model_id: Final = _str_keyed_fields(metadata.get("model_info")).get("id") or _str_keyed_fields( + standard_logging_object + ).get("model_id") + try: + await self._spend_counters()( + token=token, + team_id=user_api_key_dict.team_id, + user_id=user_api_key_dict.user_id, + response_cost=row_cost, + org_id=user_api_key_dict.org_id, + end_user_id=end_user_id, + tags=_get_request_tags_for_cost_tracking(sl_object=sl_object, metadata=dict(metadata)), + request_started_at=start_time, + model_access_groups=debitable_model_access_groups( + attributed=get_request_model_access_groups(row_kwargs), + served_model_id=served_model_id if isinstance(served_model_id, str) else None, + router=get_llm_router(), + ), + project_id=user_api_key_dict.project_id, + ) + except Exception: # noqa: BLE001 # one row's counter failure must not keep the other rows' counters from landing + verbose_proxy_logger.exception("Failed to raise spend counters for a failure spend row") async def _PROXY_track_cost_callback( self, @@ -1060,7 +1165,7 @@ def _metadata_keys(metadata: object) -> tuple[str, ...]: return tuple(sorted(str(key) for key in metadata)) -def _get_budget_reservation_from_metadata(metadata: dict) -> dict | None: +def _get_budget_reservation_from_metadata(metadata: dict) -> dict[str, object] | None: return budget_reservation_from_metadata(metadata) @@ -1095,6 +1200,7 @@ class _IncrementSpendCounters(Protocol): tags: list[str] | None = None, request_started_at: datetime | None = None, model_access_groups: Sequence[str] | None = None, + project_id: str | None = None, ) -> None: ... @@ -1111,7 +1217,7 @@ async def _update_database_and_spend_counters( start_time: datetime | None, end_time: datetime | None, response_cost: float, - budget_reservation: dict | None, + budget_reservation: dict[str, object] | None, request_tags: list[str] | None = None, model_access_groups: Sequence[str] | None = None, project_id: str | None = None, @@ -1174,7 +1280,7 @@ async def _update_database_and_spend_counters_in_batch( start_time: datetime | None, end_time: datetime | None, response_cost: float, - budget_reservation: dict | None, + budget_reservation: dict[str, object] | None, request_tags: list[str] | None, model_access_groups: Sequence[str] | None, project_id: str | None, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 7ea1c246da5..ea6b51b228b 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -3510,7 +3510,7 @@ async def increment_spend_counters( user_id: str | None, response_cost: float | None, org_id: str | None = None, - budget_reservation: dict | None = None, + budget_reservation: dict[str, object] | None = None, end_user_id: str | None = None, tags: list[str] | None = None, request_started_at: datetime | None = None, diff --git a/tests/integration/spend/_partial_stream.py b/tests/integration/spend/_partial_stream.py new file mode 100644 index 00000000000..e35320a8d8e --- /dev/null +++ b/tests/integration/spend/_partial_stream.py @@ -0,0 +1,173 @@ +import json +import threading +import uuid +from collections.abc import Callable, Iterable, Iterator +from hashlib import sha256 +from typing import Final + +import httpx +from integration._support.client import Gateway, Scenario, eventually, held +from integration._support.wire import Reply, Request, Wire +from openai.types.chat import ChatCompletionChunk +from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter + +JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +BACKEND: Final = "integration-partial-stream" +PIECE: Final = "partial stream billed tokens " +DELIVERED_PIECES: Final = 12 +_FAILURE_ROWS: Final = 'SELECT spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s AND status=%s ORDER BY "startTime"' +ROWS: Final = 'SELECT status, spend FROM "LiteLLM_SpendLogs" WHERE api_key=%s ORDER BY "startTime"' + + +class _SpendRow(BaseModel): + model_config = ConfigDict(frozen=True) + + spend: float + + +class StatusRow(BaseModel): + model_config = ConfigDict(frozen=True) + + status: str + spend: float + + +class _BudgetError(BaseModel): + model_config = ConfigDict(frozen=True) + + type: str + + +class _ErrorResponse(BaseModel): + model_config = ConfigDict(frozen=True) + + error: _BudgetError + + +def _chunk(stream_id: str, delta: dict[str, JsonValue], finish_reason: str | None = None) -> bytes: + body: Final = { + "id": stream_id, + "object": "chat.completion.chunk", + "created": 1700000000, + "model": BACKEND, + "choices": [{"index": 0, "delta": delta, "finish_reason": finish_reason}], + } + return f"data: {json.dumps(body)}\n\n".encode() + + +def completion() -> bytes: + return json.dumps( + { + "id": f"chatcmpl-{uuid.uuid4().hex}", + "object": "chat.completion", + "created": 1700000000, + "model": BACKEND, + "choices": [{"index": 0, "message": {"role": "assistant", "content": "ok"}, "finish_reason": "stop"}], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + } + ).encode() + + +def peer(delivered: int | None, gate: threading.Event | None = None) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method != "POST": + return Reply(status=404) + if JSON_OBJECT.validate_json(request.body).get("stream") is True: + stream_id: Final = f"chatcmpl-{uuid.uuid4().hex}" + return Reply( + content_type="text/event-stream", + chunks=( + *(_chunk(stream_id, {"role": "assistant", "content": PIECE}) for _ in range(DELIVERED_PIECES)), + _chunk(stream_id, {}, "stop"), + b"data: [DONE]\n\n", + ), + abort_after=delivered, + pause_between_chunks=0.02, + gate_before_first=gate, + ) + return Reply(body=completion()) + + return respond + + +def priced_model(scenario: Scenario, wire: Wire) -> str: + return scenario.model( + model=f"openai/{BACKEND}", + api_base=f"{wire.url}/v1", + api_key="synthetic-openai-key", + input_cost_per_token=0.001, + output_cost_per_token=0.002, + ) + + +def stream(gateway: Gateway, model: str, key: str, **fields: JsonValue) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "max_tokens": 5, + "stream": True, + "messages": [{"role": "user", "content": "partial stream"}], + **fields, + }, + key=key, + ) + + +def broken_stream(gateway: Gateway, model: str, key: str, **fields: JsonValue) -> None: + response: Final = stream(gateway, model, key, **fields) + assert response.status_code == 200, response.text + events: Final = tuple( + line.removeprefix("data: ") + for line in response.text.splitlines() + if line.startswith("data: ") and line != "data: [DONE]" + ) + assert "error" in JSON_OBJECT.validate_json(events[-1]), response.text + delivered: Final = "".join(_streamed_text(ChatCompletionChunk.model_validate_json(event) for event in events[:-1])) + assert delivered == PIECE * DELIVERED_PIECES, response.text + + +def _streamed_text(chunks: Iterable[ChatCompletionChunk]) -> Iterator[str]: + for chunk in chunks: + yield from (choice.delta.content or "" for choice in chunk.choices) + + +def chat(gateway: Gateway, model: str, key: str, **fields: JsonValue) -> httpx.Response: + return gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "max_tokens": 5, + "messages": [{"role": "user", "content": f"after {uuid.uuid4().hex}"}], + **fields, + }, + key=key, + ) + + +def failure_spend(gateway: Gateway, key: str, rows: int = 1) -> tuple[float, ...]: + digest: Final = sha256(key.encode()).hexdigest() + eventually(lambda: gateway.rows(_FAILURE_ROWS, (digest, "failure")), lambda found: len(found) >= rows, seconds=70) + found: Final = held( + lambda: gateway.rows(_FAILURE_ROWS, (digest, "failure")), lambda latest: len(latest) == rows, holding=2 + ) + return tuple(_SpendRow.model_validate(row).spend for row in found) + + +def chats_sent(wire: Wire) -> int: + return sum(1 for request in wire.drain() if request.method == "POST") + + +def recovered_cost(gateway: Gateway, scenario: Scenario, model: str) -> float: + key: Final = scenario.key(models=[model], max_budget=100) + broken_stream(gateway, model, key) + (spend,) = failure_spend(gateway, key) + assert spend > 0 + return spend + + +def assert_budget_refusal(response: httpx.Response) -> None: + assert response.status_code == 422, response.text + assert _ErrorResponse.model_validate_json(response.content).error.type == "budget_exceeded", response.text diff --git a/tests/integration/spend/test_failure_spend_live_counters_chaos.py b/tests/integration/spend/test_failure_spend_live_counters_chaos.py new file mode 100644 index 00000000000..ecdd5aadb8e --- /dev/null +++ b/tests/integration/spend/test_failure_spend_live_counters_chaos.py @@ -0,0 +1,239 @@ +import json +import threading +import uuid +from concurrent.futures import ThreadPoolExecutor +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 +from integration._support.process import owned_proxy_process +from integration._support.redis_process import owned_redis +from integration._support.timing import assert_not_delayed_by, timed +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.spend._partial_stream import ( + DELIVERED_PIECES, + assert_budget_refusal, + broken_stream, + chat, + chats_sent, + failure_spend, + peer, + priced_model, + recovered_cost, +) +from pydantic import JsonValue +from redis import Redis + +from litellm.constants import REDIS_SOCKET_TIMEOUT + +_BURST: Final = 12 +_SPREAD: Final = 6 +_ATTACK_MARKER: Final = "synthetic-attack-marker" +_POLICY_ID: Final = "synthetic-priced-policy" +_POLICY_APPLY_PATH: Final = f"/guardrail/{_POLICY_ID}/version/DRAFT/apply" + + +def _proxy_config(directory: Path, name: str, sections: dict[str, JsonValue]) -> Path: + base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + path: Final = directory / f"{name}.yaml" + path.write_text(yaml.safe_dump({**base, **sections})) + return path + + +def _redis_counter_config(directory: Path) -> Path: + base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + return _proxy_config( + directory, + "redis-counters", + { + "router_settings": { + **base["router_settings"], + "redis_host": "os.environ/REDIS_HOST", + "redis_port": "os.environ/REDIS_PORT", + } + }, + ) + + +def _spread_requests(gateway: Gateway, model: str, key: str) -> tuple[httpx.Response, ...]: + return tuple( + gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "max_tokens": 5, "messages": [{"role": "user", "content": f"spread {uuid.uuid4().hex}"}]}, + key=key, + headers={"connection": "close"}, + ) + for _ in range(_SPREAD) + ) + + +@pytest.mark.timeout(600) +def test_broken_streams_through_a_redis_stall_log_once_and_the_budget_holds_on_every_worker( + gateway: Gateway, tmp_path: Path +) -> None: + release: Final = threading.Event() + with ( + owned_redis(tmp_path) as counters, + wire_server(peer(DELIVERED_PIECES)) as wire, + wire_server(peer(DELIVERED_PIECES, gate=release)) as held_wire, + owned_proxy_process( + gateway, + tmp_path, + {"REDIS_HOST": counters.host, "REDIS_PORT": str(counters.port)}, + config=_redis_counter_config(tmp_path), + workers=2, + ) as owned, + owned.gateway.scenario() as scenario, + ): + scenario.cleanups.callback(release.set) + candidate: Final = owned.gateway + model: Final = priced_model(scenario, wire) + stalled_model: Final = priced_model(scenario, held_wire) + cost: Final = recovered_cost(candidate, scenario, model) + control: Final = timed(lambda: broken_stream(candidate, model, scenario.key(models=[model]))) + burst_key: Final = scenario.key(models=[stalled_model]) + chats_sent(held_wire) + + with ThreadPoolExecutor(max_workers=_BURST) as pool: + pending: Final = tuple( + pool.submit(broken_stream, candidate, stalled_model, burst_key) for _ in range(_BURST) + ) + eventually(held_wire.received.qsize, lambda held: held >= _BURST, seconds=30) + with counters.paused(): + release.set() + drained: Final = timed(lambda: tuple(future.result() for future in pending)) + + assert_not_delayed_by((drained.seconds,), control=(control.seconds,), injected_delay=REDIS_SOCKET_TIMEOUT) + assert failure_spend(candidate, burst_key, rows=_BURST) == pytest.approx((cost,) * _BURST) + spent: Final = cost * _BURST + with Redis(host=counters.host, port=counters.port) as live_counters: + landed: Final = eventually( + lambda: float(live_counters.get(f"spend:key:{sha256(burst_key.encode()).hexdigest()}") or 0), + lambda counted: counted >= spent - 1e-9, + seconds=30, + return_last_on_timeout=True, + ) + assert spent - 1e-9 <= landed <= 2 * spent + 1e-9, ( + f"one broken stream costs {cost}. The cold counter is seeded from the database, and a spend flush" + " during the stall can land these rows there first, so up to twice their cost is expected, as for" + " completed streams held past a flush" + ) + + budgeted: Final = scenario.key(models=[model], max_budget=cost / 2) + broken_stream(candidate, model, budgeted) + before_kill: Final = _spread_requests(candidate, model, budgeted) + owned.kill_worker(min(owned.worker_pids())) + after_kill: Final = _spread_requests(candidate, model, budgeted) + + for refused in (*before_kill, *after_kill): + assert_budget_refusal(refused) + + +def _policy(request: Request) -> Reply: + if request.method != "POST" or request.target != _POLICY_APPLY_PATH: + return Reply(status=404) + if _ATTACK_MARKER.encode() not in request.body: + return Reply(body=json.dumps({"action": "NONE", "outputs": [], "assessments": []}).encode()) + return Reply( + body=json.dumps( + { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "blocked by synthetic policy"}], + "assessments": [ + {"contentPolicy": {"filters": [{"type": "VIOLENCE", "confidence": "HIGH", "action": "BLOCKED"}]}} + ], + "usage": {"contentPolicyUnits": 1}, + } + ).encode() + ) + + +def _policy_config(directory: Path, policy: Wire) -> Path: + return _proxy_config( + directory, + "priced-policy", + { + "guardrails": [ + { + "guardrail_name": "priced-bedrock-policy", + "litellm_params": { + "guardrail": "bedrock", + "mode": "pre_call", + "default_on": True, + "guardrailIdentifier": _POLICY_ID, + "guardrailVersion": "DRAFT", + "aws_region_name": "us-east-1", + "aws_access_key_id": "AKIASYNTHETICGUARDRAIL", + "aws_secret_access_key": "synthetic-guardrail-secret", + "aws_bedrock_runtime_endpoint": policy.url, + }, + } + ] + }, + ) + + +def _attack(gateway: Gateway, model: str, key: str) -> httpx.Response: + blocked: Final = gateway.request( + "POST", + "/v1/chat/completions", + {"model": model, "messages": [{"role": "user", "content": f"{_ATTACK_MARKER} {uuid.uuid4().hex}"}]}, + key=key, + ) + assert blocked.status_code == 400, blocked.text + return blocked + + +@pytest.mark.timeout(480) +def test_guardrail_cost_on_blocked_requests_counts_toward_the_budget(gateway: Gateway, tmp_path: Path) -> None: + with ( + wire_server(_policy) as policy, + wire_server(peer(None)) as wire, + owned_proxy_process(gateway, tmp_path, {}, config=_policy_config(tmp_path, policy), workers=2) as owned, + owned.gateway.scenario() as scenario, + ): + candidate: Final = owned.gateway + model: Final = priced_model(scenario, wire) + meter: Final = scenario.key(models=[model], max_budget=100) + _attack(candidate, model, meter) + (cost,) = failure_spend(candidate, meter) + assert cost > 0 + key: Final = scenario.key(models=[model], max_budget=cost * 1.5) + chats_sent(wire) + + _attack(candidate, model, key) + _attack(candidate, model, key) + after: Final = chat(candidate, model, key) + + assert_budget_refusal(after) + assert chats_sent(wire) == 0 + + +@pytest.mark.timeout(480) +def test_broken_streams_leave_the_budget_alone_while_spend_updates_are_disabled( + gateway: Gateway, tmp_path: Path +) -> None: + base: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config: Final = _proxy_config( + tmp_path, "spend-updates-off", {"general_settings": {**base["general_settings"], "disable_spend_updates": True}} + ) + with ( + wire_server(peer(DELIVERED_PIECES)) as wire, + gateway.scenario() as metering, + owned_proxy_process(gateway, tmp_path, {}, config=config) as owned, + owned.gateway.scenario() as scenario, + ): + cost: Final = recovered_cost(gateway, metering, priced_model(metering, wire)) + candidate: Final = owned.gateway + model: Final = priced_model(scenario, wire) + key: Final = scenario.key(models=[model], max_budget=cost / 2) + + broken_stream(candidate, model, key) + after: Final = chat(candidate, model, key) + + assert after.status_code == 200, after.text diff --git a/tests/integration/spend/test_partial_stream_live_budget.py b/tests/integration/spend/test_partial_stream_live_budget.py new file mode 100644 index 00000000000..89fb66c9bbf --- /dev/null +++ b/tests/integration/spend/test_partial_stream_live_budget.py @@ -0,0 +1,350 @@ +import json +import threading +import uuid +from collections.abc import Callable +from concurrent.futures import ThreadPoolExecutor +from hashlib import sha256 +from typing import Final + +import httpx +import pytest +from integration._support.client import Gateway, eventually, held +from integration._support.wire import Reply, Request, wire_server +from integration.spend._partial_stream import ( + BACKEND, + DELIVERED_PIECES, + JSON_OBJECT, + PIECE, + ROWS, + StatusRow, + assert_budget_refusal, + broken_stream, + chat, + chats_sent, + completion, + failure_spend, + peer, + priced_model, + recovered_cost, + stream, +) +from pydantic import JsonValue + +_CONCURRENT_STREAMS: Final = 6 +_MESSAGES_BACKEND: Final = "claude-integration-partial-stream" +_UNCACHED: Final[dict[str, JsonValue]] = {"cache": {"no-cache": True}} + + +@pytest.mark.timeout(240) +def test_a_broken_stream_that_spends_past_the_budget_blocks_the_next_request(gateway: Gateway) -> None: + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + key: Final = scenario.key(models=[model], max_budget=cost / 2) + chats_sent(wire) + + broken_stream(gateway, model, key) + after: Final = chat(gateway, model, key) + + assert_budget_refusal(after) + assert failure_spend(gateway, key, rows=2) == pytest.approx((cost, 0.0)) + assert chats_sent(wire) == 1 + + +@pytest.mark.timeout(240) +def test_broken_stream_spend_counts_once_toward_the_budget(gateway: Gateway) -> None: + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + key: Final = scenario.key(models=[model], max_budget=cost * 1.5) + + broken_stream(gateway, model, key) + broken_stream(gateway, model, key) + after: Final = chat(gateway, model, key) + + assert_budget_refusal(after) + assert failure_spend(gateway, key, rows=3) == pytest.approx((cost, cost, 0.0)) + + +@pytest.mark.timeout(240) +def test_a_stream_that_breaks_before_any_chunk_leaves_the_budget_untouched(gateway: Gateway) -> None: + with wire_server(peer(0)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + key: Final = scenario.key(models=[model], max_budget=0.05) + + refused: Final = stream(gateway, model, key) + after: Final = chat(gateway, model, key) + + assert refused.status_code == 500, refused.text + + assert after.status_code == 200, after.text + assert failure_spend(gateway, key) == (0.0,) + + +def _broken_stream_then_refusal( + gateway: Gateway, model: str, spender: str, other: str, **fields: JsonValue +) -> httpx.Response: + broken_stream(gateway, model, spender, **fields) + return chat(gateway, model, other, **fields) + + +@pytest.mark.timeout(240) +def test_a_broken_stream_on_one_team_key_blocks_the_teams_other_keys(gateway: Gateway) -> None: + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + team: Final = scenario.team(models=[model], max_budget=cost / 2) + + after: Final = _broken_stream_then_refusal( + gateway, model, scenario.key(team_id=team, models=[model]), scenario.key(team_id=team, models=[model]) + ) + + assert_budget_refusal(after) + + +@pytest.mark.timeout(240) +def test_a_broken_stream_on_one_of_a_users_keys_blocks_their_other_keys(gateway: Gateway) -> None: + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + user: Final = scenario.user(max_budget=cost / 2) + + after: Final = _broken_stream_then_refusal( + gateway, model, scenario.key(user_id=user, models=[model]), scenario.key(user_id=user, models=[model]) + ) + + assert_budget_refusal(after) + + +@pytest.mark.timeout(240) +def test_a_broken_stream_for_an_end_user_blocks_their_next_request_on_any_key(gateway: Gateway) -> None: + end_user: Final = f"integration-end-user-{uuid.uuid4().hex}" + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + gateway.post("/customer/new", {"user_id": end_user, "max_budget": cost / 2}) + scenario.cleanups.callback(gateway.post, "/customer/delete", {"user_ids": [end_user]}) + + after: Final = _broken_stream_then_refusal( + gateway, model, scenario.key(models=[model]), scenario.key(models=[model]), user=end_user + ) + + assert_budget_refusal(after) + + +@pytest.mark.timeout(240) +def test_a_broken_stream_under_a_tag_blocks_the_next_request_with_that_tag(gateway: Gateway) -> None: + tag: Final = f"integration-partial-stream-{uuid.uuid4().hex}" + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + gateway.post("/tag/new", {"name": tag, "max_budget": cost / 2}) + scenario.cleanups.callback(gateway.post, "/tag/delete", {"name": tag}) + key: Final = scenario.key(models=[model]) + + after: Final = _broken_stream_then_refusal(gateway, model, key, key, metadata={"tags": [tag]}) + untagged: Final = chat(gateway, model, key) + + assert_budget_refusal(after) + assert untagged.status_code == 200, untagged.text + + +@pytest.mark.timeout(240) +def test_a_broken_stream_fills_the_keys_budget_window(gateway: Gateway) -> None: + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + key: Final = scenario.key(models=[model], budget_limits=[{"budget_duration": "1d", "max_budget": cost / 2}]) + + after: Final = _broken_stream_then_refusal(gateway, model, key, key) + + assert_budget_refusal(after) + + +@pytest.mark.timeout(240) +def test_concurrent_broken_streams_each_count_once_toward_the_budget(gateway: Gateway) -> None: + with wire_server(peer(DELIVERED_PIECES)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + cost: Final = recovered_cost(gateway, scenario, model) + key: Final = scenario.key(models=[model], max_budget=cost * (_CONCURRENT_STREAMS + 0.5)) + + with ThreadPoolExecutor(max_workers=_CONCURRENT_STREAMS) as pool: + tuple(pool.map(lambda _: broken_stream(gateway, model, key), range(_CONCURRENT_STREAMS))) + within_budget: Final = chat(gateway, model, key) + broken_stream(gateway, model, key) + over_budget: Final = chat(gateway, model, key) + + assert within_budget.status_code == 200, within_budget.text + assert_budget_refusal(over_budget) + assert sorted(failure_spend(gateway, key, rows=_CONCURRENT_STREAMS + 2)) == pytest.approx( + [0.0] + [cost] * (_CONCURRENT_STREAMS + 1) + ) + + +@pytest.mark.timeout(240) +def test_completed_streams_count_once_toward_the_budget(gateway: Gateway) -> None: + with wire_server(peer(None)) as wire, gateway.scenario() as scenario: + model: Final = priced_model(scenario, wire) + meter: Final = scenario.key(models=[model], max_budget=100) + assert stream(gateway, model, meter, **_UNCACHED).status_code == 200 + (metered,) = eventually( + lambda: gateway.rows(ROWS, (sha256(meter.encode()).hexdigest(),)), + lambda found: len(found) == 1, + seconds=70, + ) + cost: Final = StatusRow.model_validate(metered).spend + key: Final = scenario.key(models=[model], max_budget=cost * 1.5) + + first: Final = stream(gateway, model, key, **_UNCACHED) + second: Final = stream(gateway, model, key, **_UNCACHED) + assert (first.status_code, second.status_code) == (200, 200), (first.text, second.text) + assert _rows(gateway, key, 2) == (StatusRow(status="success", spend=cost),) * 2 + third: Final = stream(gateway, model, key, **_UNCACHED) + + assert_budget_refusal(third) + + +def _hung_then_served(release: threading.Event, hung_model: str) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if request.method != "POST": + return Reply(status=404) + if JSON_OBJECT.validate_json(request.body).get("model") == hung_model: + return Reply(chunks=(completion(),), gate_before_first=release) + return Reply(body=completion()) + + return respond + + +def _timed_out_spend(gateway: Gateway, model: str, key: str) -> float: + timed_out: Final = chat(gateway, model, key, messages=[{"role": "user", "content": "hung"}]) + assert timed_out.status_code in (408, 500), timed_out.text + digest: Final = sha256(key.encode()).hexdigest() + landed: Final = eventually(lambda: gateway.rows(ROWS, (digest,)), lambda found: len(found) >= 3, seconds=70) + rows: Final = tuple( + StatusRow.model_validate(row) + for row in held(lambda: gateway.rows(ROWS, (digest,)), lambda latest: len(latest) == len(landed), holding=2) + ) + assert {row.status for row in rows} == {"failure"}, rows + return sum(row.spend for row in rows) + + +@pytest.mark.timeout(240) +def test_retries_the_provider_billed_before_a_timeout_count_toward_the_budget(gateway: Gateway) -> None: + hung_backend: Final = f"{BACKEND}-hung" + release: Final = threading.Event() + with wire_server(_hung_then_served(release, hung_backend)) as wire, gateway.scenario() as scenario: + scenario.cleanups.callback(release.set) + hung: Final = scenario.model( + model=f"openai/{hung_backend}", + api_base=f"{wire.url}/v1", + api_key="synthetic-openai-key", + input_cost_per_token=0.001, + output_cost_per_token=0.002, + timeout=1, + num_retries=2, + ) + served: Final = priced_model(scenario, wire) + billed: Final = _timed_out_spend(gateway, hung, scenario.key(models=[hung, served], max_budget=100)) + assert billed > 0 + key: Final = scenario.key(models=[hung, served], max_budget=billed / 2) + + assert _timed_out_spend(gateway, hung, key) == pytest.approx(billed) + after: Final = chat(gateway, served, key) + + assert_budget_refusal(after) + + +def _messages_event(kind: str, fields: dict[str, JsonValue]) -> bytes: + return f"event: {kind}\ndata: {json.dumps({'type': kind, **fields})}\n\n".encode() + + +def _messages_peer(request: Request) -> Reply: + if request.method != "POST" or request.target != "/v1/messages": + return Reply(status=404) + message: Final[dict[str, JsonValue]] = { + "id": f"msg_{uuid.uuid4().hex}", + "type": "message", + "role": "assistant", + "model": _MESSAGES_BACKEND, + "stop_reason": None, + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 1}, + } + if JSON_OBJECT.validate_json(request.body).get("stream") is not True: + return Reply( + body=json.dumps( + {**message, "content": [{"type": "text", "text": "ok"}], "stop_reason": "end_turn"} + ).encode() + ) + return Reply( + content_type="text/event-stream", + chunks=( + _messages_event("message_start", {"message": {**message, "content": []}}), + _messages_event("content_block_start", {"index": 0, "content_block": {"type": "text", "text": ""}}), + *( + _messages_event("content_block_delta", {"index": 0, "delta": {"type": "text_delta", "text": PIECE}}) + for _ in range(DELIVERED_PIECES) + ), + _messages_event("content_block_stop", {"index": 0}), + _messages_event( + "message_delta", + {"delta": {"stop_reason": "end_turn", "stop_sequence": None}, "usage": {"output_tokens": 60}}, + ), + _messages_event("message_stop", {}), + ), + abort_after=DELIVERED_PIECES + 2, + pause_between_chunks=0.02, + ) + + +def _messages(gateway: Gateway, model: str, key: str, *, streamed: bool) -> httpx.Response: + return gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 5, + "stream": streamed, + "messages": [{"role": "user", "content": f"partial stream {uuid.uuid4().hex}"}], + }, + key=key, + ) + + +def _broken_messages_stream(gateway: Gateway, model: str, key: str) -> None: + response: Final = _messages(gateway, model, key, streamed=True) + assert response.status_code == 200, response.text + assert response.text.count(PIECE) == DELIVERED_PIECES, response.text + assert "message_stop" not in response.text, response.text + + +def _rows(gateway: Gateway, key: str, count: int) -> tuple[StatusRow, ...]: + digest: Final = sha256(key.encode()).hexdigest() + eventually(lambda: gateway.rows(ROWS, (digest,)), lambda found: len(found) >= count, seconds=70) + return tuple( + StatusRow.model_validate(row) + for row in held(lambda: gateway.rows(ROWS, (digest,)), lambda latest: len(latest) == count, holding=2) + ) + + +@pytest.mark.timeout(240) +def test_a_broken_messages_stream_counts_toward_the_budget(gateway: Gateway) -> None: + with wire_server(_messages_peer) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"anthropic/{_MESSAGES_BACKEND}", + api_base=wire.url, + api_key="synthetic-anthropic-key", + input_cost_per_token=0.001, + output_cost_per_token=0.002, + ) + meter: Final = scenario.key(models=[model], max_budget=100) + _broken_messages_stream(gateway, model, meter) + (metered,) = _rows(gateway, meter, 1) + assert metered.spend > 0, metered + key: Final = scenario.key(models=[model], max_budget=metered.spend / 2) + + _broken_messages_stream(gateway, model, key) + after: Final = _messages(gateway, model, key, streamed=False) + + assert_budget_refusal(after) + assert _rows(gateway, key, 2) == (metered, StatusRow(status="failure", spend=0.0)) diff --git a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py index 8d4f8cd87ad..e482b6552a4 100644 --- a/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py +++ b/tests/unit/proxy/hooks/test_proxy_track_cost_callback.py @@ -14,6 +14,7 @@ from fastapi import HTTPException import litellm from litellm._logging import verbose_proxy_logger +from litellm.caching.dual_cache import DualCache from litellm.litellm_core_utils.internal_call_metadata import MODEL_ACCESS_GROUP_METADATA_KEY from litellm.proxy._types import ProxyException, SpendLogsPayload, UserAPIKeyAuth from litellm.proxy.collector import SpendEventConsumer @@ -36,6 +37,7 @@ from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payloa from litellm.proxy.utils import ( REJECTED_BEFORE_REQUEST_SETUP, REJECTED_BEFORE_REQUEST_SETUP_KEY, + BilledFailedAttempt, ProxyUpdateSpend, ) from litellm.types.router import Deployment, LiteLLM_Params @@ -2676,6 +2678,229 @@ async def test_async_post_call_failure_hook_records_recovered_partial_spend(): assert mock_update_database.call_args[1]["response_cost"] == 3.5e-05 +@pytest.fixture() +def live_spend_counters(monkeypatch: pytest.MonkeyPatch) -> DualCache: + counters: Final = DualCache() + monkeypatch.setattr("litellm.proxy.proxy_server.spend_counter_cache", counters) + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", DualCache()) + monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", None) + monkeypatch.setattr("litellm.proxy.proxy_server.llm_router", None) + return counters + + +def _failure_with_billed_attempts() -> Mapping[str, object]: + usage: Final = Usage(prompt_tokens=24, completion_tokens=0, total_tokens=24) + return { + "combined_usage_object": usage, + "response_cost": 0.5, + "billed_failed_attempts": ( + BilledFailedAttempt(attempted_retries=0, deployment_id="dep-a", usage=usage, response_cost=0.25), + BilledFailedAttempt(attempted_retries=1, deployment_id="dep-b", usage=usage, response_cost=0.5), + ), + } + + +_LIVE_COUNTERS: Final = { + "spend:key:hashed-key": 0.2, + "spend:team:team-1": 0.3, + "spend:team_member:user-1:team-1": 0.35, + "spend:user:user-1": 0.4, + "spend:end_user:end-user-1": 0.45, + "spend:org:org-1": 0.5, + "spend:project:project-1": 0.55, + "spend:tag:batch-jobs": 0.6, + "spend:model_access_group:premium": 0.65, +} + + +async def _fail_request(logger: ProxyDBLogger, failure: Mapping[str, object]) -> None: + await logger.async_post_call_failure_hook( + request_data={ + "model": "gpt-4.1", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {"tags": ["batch-jobs"], MODEL_ACCESS_GROUP_METADATA_KEY: ["premium"]}, + "proxy_server_request": {"request_id": "rid"}, + **failure, + }, + original_exception=Exception("stream broke mid-flight"), + user_api_key_dict=UserAPIKeyAuth( + api_key="hashed-key", + team_id="team-1", + user_id="user-1", + end_user_id="end-user-1", + org_id="org-1", + project_id="project-1", + ), + ) + + +async def _counter_values(counters: DualCache) -> dict[str, float]: + return {key: await counters.async_get_cache(key) for key in _LIVE_COUNTERS} + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("failure", "charged"), + [ + ( + { + "combined_usage_object": Usage(prompt_tokens=30, completion_tokens=40, total_tokens=70), + "response_cost": 0.125, + }, + 0.125, + ), + (_failure_with_billed_attempts(), 1.25), + ({}, 0.0), + ], +) +async def test_failure_spend_rows_raise_the_live_spend_counters_by_their_cost( + live_spend_counters: DualCache, failure: Mapping[str, object], charged: float +): + """#35568: budget checks read the live counters, so a failure row's cost has to land there as well as in the + spend log, or an over-budget key keeps passing until the counter is next repaired from the database.""" + for key, spend in _LIVE_COUNTERS.items(): + await live_spend_counters.async_set_cache(key, spend) + + await _fail_request(ProxyDBLogger(), failure) + + assert await _counter_values(live_spend_counters) == pytest.approx( + {key: spend + charged for key, spend in _LIVE_COUNTERS.items()} + ) + + +@pytest.mark.asyncio +async def test_a_completed_cost_failure_raises_the_live_spend_counters_once(live_spend_counters: DualCache): + """A failure whose cost was already finalized settles through the success path's write, which raises the + counters itself.""" + for key, spend in _LIVE_COUNTERS.items(): + await live_spend_counters.async_set_cache(key, spend) + + await _fail_request( + ProxyDBLogger(), {"standard_logging_object": {"cost_finalized": True, "response_cost": 0.125}} + ) + + assert await _counter_values(live_spend_counters) == pytest.approx( + {key: spend + 0.125 for key, spend in _LIVE_COUNTERS.items()} + ) + + +@pytest.mark.asyncio +async def test_failure_spend_rows_leave_the_live_spend_counters_alone_while_spend_updates_are_disabled( + live_spend_counters: DualCache, monkeypatch: pytest.MonkeyPatch +): + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", {"disable_spend_updates": True}) + for key, spend in _LIVE_COUNTERS.items(): + await live_spend_counters.async_set_cache(key, spend) + + await _fail_request(ProxyDBLogger(), _failure_with_billed_attempts()) + + assert await _counter_values(live_spend_counters) == pytest.approx(_LIVE_COUNTERS) + + +@pytest.mark.asyncio +async def test_a_failing_spend_counter_store_is_still_sent_every_failure_spend_row(): + counters: Final = AsyncMock(side_effect=RuntimeError("redis unavailable")) + + await _fail_request(ProxyDBLogger(spend_counters=lambda: counters), _failure_with_billed_attempts()) + + assert [call.kwargs["response_cost"] for call in counters.call_args_list] == [0.5, 0.25, 0.5] + + +@pytest.mark.asyncio +@pytest.mark.timeout(30) +async def test_a_slow_spend_counter_store_lets_the_failure_out_and_raises_the_counters_once_it_answers(): + store_answers: Final = asyncio.Event() + raised_cost: Final[asyncio.Future[object]] = asyncio.get_running_loop().create_future() + + async def slow_counter_store(**counter_args: object) -> None: + await store_answers.wait() + raised_cost.set_result(counter_args["response_cost"]) + + logger: Final = ProxyDBLogger(spend_counters=lambda: slow_counter_store, failure_counter_wait_seconds=0) + + await _fail_request( + logger, + { + "combined_usage_object": Usage(prompt_tokens=30, completion_tokens=40, total_tokens=70), + "response_cost": 0.125, + }, + ) + assert not raised_cost.done() + + store_answers.set() + assert await raised_cost == 0.125 + + +@pytest.mark.asyncio +async def test_a_failure_spend_row_lands_in_the_key_budget_window_as_of_the_request_start( + live_spend_counters: DualCache, monkeypatch: pytest.MonkeyPatch +): + """The window row's one-time seed from the spend log stops at the request start, which keeps the failure row + it already logged from being counted twice.""" + from litellm.proxy.proxy_server import proxy_logging_obj + + window_queue: Final = proxy_logging_obj.db_spend_update_writer.window_spend_update_queue + await window_queue.flush_and_get_aggregated_window_spend_transactions() + key_cache: Final = DualCache() + monkeypatch.setattr("litellm.proxy.proxy_server.user_api_key_cache", key_cache) + request_start: Final = datetime(2026, 10, 20, 12, 0, tzinfo=timezone.utc) + reset_at: Final = request_start + timedelta(days=10) + await key_cache.async_set_cache( + "hashed-key", + UserAPIKeyAuth( + token="hashed-key", + budget_limits=[{"budget_duration": "30d", "max_budget": 100.0, "reset_at": reset_at.isoformat()}], + ), + ) + + await _fail_request( + ProxyDBLogger(), + { + "combined_usage_object": Usage(prompt_tokens=30, completion_tokens=40, total_tokens=70), + "response_cost": 0.125, + "start_time": request_start, + }, + ) + + assert [ + (row["entity_type"], row["entity_id"], row["spend"], row["started_at"]) + for row in await window_queue.flush_and_get_aggregated_window_spend_transactions() + ] == [("key", "hashed-key", pytest.approx(0.125), "2026-10-20T12:00:00.000000")] + + +@pytest.mark.asyncio +async def test_a_billed_attempt_on_a_deployment_the_router_dropped_debits_every_attributed_access_group( + live_spend_counters: DualCache, monkeypatch: pytest.MonkeyPatch +): + """Its spend row names the dropped deployment, which narrows nothing, so the counters must not narrow it to the + final attempt's groups either.""" + monkeypatch.setattr( + "litellm.proxy.proxy_server.llm_router", + litellm.Router( + model_list=[ + { + "model_name": "gpt-4.1", + "litellm_params": {"model": "openai/gpt-4.1", "api_key": "sk-test"}, + "model_info": {"id": "dep-b", "access_groups": ["pool"]}, + } + ] + ), + ) + + await _fail_request( + ProxyDBLogger(), + { + **_failure_with_billed_attempts(), + "metadata": {MODEL_ACCESS_GROUP_METADATA_KEY: ["pool", "premium"], "model_info": {"id": "dep-b"}}, + }, + ) + + assert { + group: await live_spend_counters.async_get_cache(f"spend:model_access_group:{group}") + for group in ("premium", "pool") + } == pytest.approx({"premium": 0.25, "pool": 1.25}) + + @pytest.mark.asyncio async def test_track_cost_callback_enriches_user_id_for_mcp_style_metadata(): """MCP tool calls may only carry user_api_key; user/team rollups still need user_id.""" From 632f71b850d157efbfa75aa4c1f7a2066576446c Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 30 Oct 2026 13:19:58 -0700 Subject: [PATCH 17/46] fix(router): recover streamed context window errors before output (#49293) --- litellm/responses/streaming_iterator.py | 2 +- litellm/router.py | 81 ++++++- ..._mantle_context_overflow_fallbacks_wire.py | 10 +- .../test_streaming_iterator_error_events.py | 19 +- tests/unit/test_router/test_router.py | 15 +- .../unit/test_router/test_router_fallbacks.py | 201 ++++++++++++++++++ 6 files changed, 303 insertions(+), 25 deletions(-) diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index fb91ba6a618..d42f1af380e 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -312,7 +312,7 @@ def map_stream_error_to_exception( error_body: Final = {"message": error_message, "type": error_type, "code": error_code} provider_exception: Final = BaseLLMException( status_code=status_code, - message=f"Error code: {status_code} - {{'error': {error_body}}}", + message=f"Error code: {status_code} - {json.dumps({'error': error_body})}", body=error_body, headers=headers, ) diff --git a/litellm/router.py b/litellm/router.py index 233e7581a40..538cd751fb7 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -897,6 +897,23 @@ def _stream_pre_content_error( ) +def _context_window_stream_fallback_error( + error: "MidStreamFallbackError | litellm.ContextWindowExceededError", + model: object, + llm_provider: str, + has_output: bool, +) -> "MidStreamFallbackError": + if isinstance(error, litellm.ContextWindowExceededError): + if has_output: + raise error + return _stream_pre_content_error(error, model if isinstance(model, str) else "", llm_provider) + if isinstance(error.original_exception, litellm.ContextWindowExceededError) and ( + has_output or error.generated_content or error.emitted_disqualifying_content + ): + raise error.original_exception + return error + + def _deployment_num_retries(deployment: "Deployment | None") -> int | None: """The deployment's own num_retries litellm_param, an int or a digit string the way the config loader leaves it.""" configured: Final = getattr(deployment.litellm_params, "num_retries", None) if deployment is not None else None @@ -3845,7 +3862,19 @@ class Router(DeprecatedRouterMembers): try: async for item in _achunks_with_first(model_response, first_chunk): yield item - except MidStreamFallbackError as e: + except (MidStreamFallbackError, litellm.ContextWindowExceededError) as stream_error: + e: Final = ( + _context_window_stream_fallback_error( + stream_error, + getattr(model_response, "model", ""), + model_response.custom_llm_provider or "", + bool(getattr(model_response, "response_uptil_now", "")) + or _stream_chunks_have_generated_content(model_response.chunks), + ) + if isinstance(stream_error, litellm.ContextWindowExceededError) + or isinstance(stream_error.original_exception, litellm.ContextWindowExceededError) + else stream_error + ) with anyio.CancelScope(shield=True): await close_model_response() await held_slot.aclose() @@ -3917,7 +3946,9 @@ class Router(DeprecatedRouterMembers): initial_kwargs["messages"] = messages self._update_kwargs_before_fallbacks(model=model_group, kwargs=initial_kwargs) fallback_response = await self.async_function_with_fallbacks_common_utils( - e=e, + e=e.original_exception + if isinstance(e.original_exception, litellm.ContextWindowExceededError) + else e, disable_fallbacks=fallbacks_disabled_for_request(initial_kwargs), fallbacks=fallbacks, context_window_fallbacks=context_window_fallbacks, @@ -3985,10 +4016,10 @@ class Router(DeprecatedRouterMembers): if fallback_response is not None and hasattr(fallback_response, "aclose"): try: await fallback_response.aclose() - except BaseException as e: + except BaseException as close_error: verbose_router_logger.debug( "stream_with_fallbacks: error closing fallback_response: %s", - e, + close_error, ) wrapped_response: Final = FallbackStreamWrapper(stream_with_fallbacks()) @@ -4348,6 +4379,10 @@ class Router(DeprecatedRouterMembers): def __aiter__(self): return self + @property + def yielded_first_chunk(self) -> bool: + return self._yielded_first_chunk + async def __anext__(self): try: chunk: Final = await self._async_generator.__anext__() @@ -4375,6 +4410,7 @@ class Router(DeprecatedRouterMembers): # the very key that created the container (#30210). if self.completed_response is None and getattr(chunk, "type", None) in _RESPONSES_TERMINAL_EVENT_TYPES: self.completed_response = chunk + self._yielded_first_chunk = True return chunk async def aclose(self): @@ -4401,7 +4437,18 @@ class Router(DeprecatedRouterMembers): yield item for held_event in held_lifecycle_events: yield held_event - except MidStreamFallbackError as e: + except (MidStreamFallbackError, litellm.ContextWindowExceededError) as stream_error: + e: Final = ( + _context_window_stream_fallback_error( + stream_error, + source_iterator.model or "", + source_iterator.custom_llm_provider or "", + wrapper.yielded_first_chunk, + ) + if isinstance(stream_error, litellm.ContextWindowExceededError) + or isinstance(stream_error.original_exception, litellm.ContextWindowExceededError) + else stream_error + ) if response_id_exposed: self._raise_released_stream_error(e, initial_kwargs) async with contextlib.aclosing( @@ -4572,7 +4619,11 @@ class Router(DeprecatedRouterMembers): # The content-policy dispatch branch matches on the trigger's own type, so a refusal's # MidStreamFallbackError envelope is unwrapped here or the wrong fallback list is consulted. fallback_trigger: Final[Exception] = ( - e.original_exception if isinstance(e.original_exception, litellm.ContentPolicyViolationError) else e + e.original_exception + if isinstance( + e.original_exception, (litellm.ContentPolicyViolationError, litellm.ContextWindowExceededError) + ) + else e ) fallback_response = await self.async_function_with_fallbacks_common_utils( # rebind-ok: set on success e=fallback_trigger, @@ -4658,7 +4709,19 @@ class Router(DeprecatedRouterMembers): try: for item in _chunks_with_first(model_response, first_chunk): yield item - except MidStreamFallbackError as e: + except (MidStreamFallbackError, litellm.ContextWindowExceededError) as stream_error: + e: Final = ( + _context_window_stream_fallback_error( + stream_error, + getattr(model_response, "model", ""), + model_response.custom_llm_provider or "", + bool(getattr(model_response, "response_uptil_now", "")) + or _stream_chunks_have_generated_content(model_response.chunks), + ) + if isinstance(stream_error, litellm.ContextWindowExceededError) + or isinstance(stream_error.original_exception, litellm.ContextWindowExceededError) + else stream_error + ) if fallbacks_disabled_for_request(initial_kwargs) or ( not e.is_pre_first_chunk and (e.generated_content or _stream_chunks_have_generated_content(model_response.chunks)) @@ -4683,7 +4746,9 @@ class Router(DeprecatedRouterMembers): router_self._update_kwargs_before_fallbacks(model=model_group, kwargs=initial_kwargs) fallback_response = run_async_function( router_self.async_function_with_fallbacks_common_utils, - e=e, + e=e.original_exception + if isinstance(e.original_exception, litellm.ContextWindowExceededError) + else e, disable_fallbacks=fallbacks_disabled_for_request(initial_kwargs), fallbacks=fallbacks, context_window_fallbacks=context_window_fallbacks, diff --git a/tests/integration/providers/test_bedrock_mantle_context_overflow_fallbacks_wire.py b/tests/integration/providers/test_bedrock_mantle_context_overflow_fallbacks_wire.py index 37aa96c2467..23b75a5a8cf 100644 --- a/tests/integration/providers/test_bedrock_mantle_context_overflow_fallbacks_wire.py +++ b/tests/integration/providers/test_bedrock_mantle_context_overflow_fallbacks_wire.py @@ -419,12 +419,11 @@ def test_context_window_fallback_is_not_consulted_on_streamed_messages(cwf_rig: _assert_no_fallback(cwf_rig, "overflow", prompt) -def test_context_window_fallback_does_not_fire_on_streamed_responses(cwf_rig: Rig) -> None: +def test_context_window_fallback_fires_on_streamed_responses(cwf_rig: Rig) -> None: prompt: Final = _prompt("cwf responses stream") response, text = _consumed(cwf_rig.gateway, "/v1/responses", _responses(cwf_rig.models["cwf"], prompt, True)) assert response.status_code == 200, text - assert _GENERIC in str(_only_failed(text)["message"]), text - _assert_no_fallback(cwf_rig, "overflow", prompt) + _assert_fell_back(cwf_rig, "overflow", prompt, text) def test_plain_fallback_fires_for_the_overflow_on_chat_completions(fb_rig: Rig) -> None: @@ -467,12 +466,11 @@ def test_streamed_messages_overflow_on_a_plain_fallback_deployment_returns_400_w _assert_no_fallback(fb_rig, "overflow", prompt) -def test_plain_fallback_does_not_fire_on_streamed_responses(fb_rig: Rig) -> None: +def test_plain_fallback_fires_on_streamed_responses(fb_rig: Rig) -> None: prompt: Final = _prompt("fb responses stream") response, text = _consumed(fb_rig.gateway, "/v1/responses", _responses(fb_rig.models["fb"], prompt, True)) assert response.status_code == 200, text - assert _GENERIC in str(_only_failed(text)["message"]), text - _assert_no_fallback(fb_rig, "overflow", prompt) + _assert_fell_back(fb_rig, "overflow", prompt, text) def test_streamed_messages_overflow_on_an_openai_fallback_deployment_returns_the_error_event_without_falling_back( diff --git a/tests/unit/responses/test_streaming_iterator_error_events.py b/tests/unit/responses/test_streaming_iterator_error_events.py index 7445070bd3a..14ecb185f06 100644 --- a/tests/unit/responses/test_streaming_iterator_error_events.py +++ b/tests/unit/responses/test_streaming_iterator_error_events.py @@ -166,21 +166,24 @@ def test_maybe_raise_for_error_event_maps_context_length_code_to_400(): assert not isinstance(exc_info.value, MidStreamFallbackError) -def test_maybe_raise_for_error_event_raises_context_window_exceeded_directly(): +@pytest.mark.parametrize( + "message", + ( + "This model's maximum context length is 128000 tokens. However, your messages resulted in 130000 tokens.", + "Your input exceeds the context window of this model. Please adjust your input and try again.", + ), +) +def test_maybe_raise_for_error_event_raises_context_window_exceeded_directly(message: str) -> None: """A context-length error event maps to ContextWindowExceededError exactly like the non-streaming path and, being a non-retriable client error, is raised directly rather than wrapped for mid-stream fallback, preserving the direct-SDK 400 contract from issue #15785.""" - iterator = _make_iterator() - chunk = _make_error_chunk( - "invalid_request_error", - "context_length_exceeded", - "This model's maximum context length is 128000 tokens. However, your messages resulted in 130000 tokens.", - ) + iterator: Final = _make_iterator() + chunk: Final = _make_error_chunk("invalid_request_error", "context_length_exceeded", message) with pytest.raises(litellm.ContextWindowExceededError) as exc_info: iterator._maybe_raise_for_error_event(chunk) assert exc_info.value.status_code == 400 assert not isinstance(exc_info.value, MidStreamFallbackError) - assert "maximum context length" in str(exc_info.value) + assert message in str(exc_info.value) CONTENT_POLICY_MESSAGE = "This content was flagged for possible cybersecurity risk. The response was halted mid-stream." diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index a6436e47507..96df8d3e1f0 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -5140,7 +5140,10 @@ async def test_set_response_headers_is_the_only_complexity_header_source_for_pro @pytest.mark.asyncio -async def test_acompletion_streaming_iterator_adopts_fallback_response_headers(): +@pytest.mark.parametrize("close_fails", [False, True]) +async def test_acompletion_streaming_iterator_adopts_fallback_response_headers( + close_fails: bool, +) -> None: """LIT-6767: after a successful pre-first-chunk fallback, the wrapper must describe the deployment that served the stream, with no value left over from the attempt that failed.""" @@ -5187,6 +5190,7 @@ async def test_acompletion_streaming_iterator_adopts_fallback_response_headers() class FallbackStream: def __init__(self): + self.close_attempted: bool = False self._response_headers = {"x-request-id": "req-FALLBACK"} self._hidden_params = { "model_id": "fallback-deployment", @@ -5204,6 +5208,11 @@ async def test_acompletion_streaming_iterator_adopts_fallback_response_headers() except StopIteration: raise StopAsyncIteration from None + async def aclose(self) -> None: + self.close_attempted = True + if close_fails: + raise RuntimeError("fallback transport failed during close") + fallback_stream = FallbackStream() with patch.object( router, @@ -5220,12 +5229,14 @@ async def test_acompletion_streaming_iterator_adopts_fallback_response_headers() # the very first chunk the fallback produces must already be published under # the fallback's identity: the proxy commits response headers once that chunk # is buffered, so adopting any later is adopting too late - await result.__anext__() + first_chunk: Final = await result.__anext__() + assert first_chunk.choices[0].delta.content == "OK" assert result._response_headers == {"x-request-id": "req-FALLBACK"} assert result._hidden_params["model_id"] == "fallback-deployment" async for _ in result: pass + assert fallback_stream.close_attempted assert result._response_headers == {"x-request-id": "req-FALLBACK"} assert result._hidden_params["model_id"] == "fallback-deployment" assert result._hidden_params["api_base"] == "https://fallback.example" diff --git a/tests/unit/test_router/test_router_fallbacks.py b/tests/unit/test_router/test_router_fallbacks.py index 88d89696297..bccd6290f44 100644 --- a/tests/unit/test_router/test_router_fallbacks.py +++ b/tests/unit/test_router/test_router_fallbacks.py @@ -1006,3 +1006,204 @@ async def test_unknown_default_fallback_raises_after_primary_failure( assert route.call_count == 1 assert "missing" in str(exc_info.value) + + +def _context_error_response_stream( + *, event_type: str = "error", preceding_event: dict[str, object] | None = None, completed: bool = False +) -> httpx.Response: + response: Final = { + "id": "resp_backup" if completed else "resp_primary", + "object": "response", + "created_at": 1700000000, + "model": "backup" if completed else "primary", + "status": "completed" if completed else "in_progress", + "output": [], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "temperature": 1.0, + "top_p": 1.0, + } + error: Final = { + "type": "invalid_request_error", + "code": "context_length_exceeded", + "message": "Your input exceeds the context window of this model. Please adjust your input and try again.", + } + terminal: Final = ( + {"type": "response.completed", "response": response, "sequence_number": 2} + if completed + else {"type": event_type, "error": error, "sequence_number": 2} + if event_type == "error" + else {"type": event_type, "response": {**response, "status": "failed", "error": error}, "sequence_number": 2} + ) + events: Final = ( + {"type": "response.created", "response": response, "sequence_number": 0}, + *((preceding_event,) if preceding_event is not None else ()), + terminal, + ) + return httpx.Response( + 200, + headers={"content-type": "text/event-stream"}, + content="".join(f"data: {json.dumps(event)}\n\n" for event in events).encode(), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("event_type", ("error", "response.failed")) +@pytest.mark.parametrize("context_fallback", (False, True)) +async def test_responses_context_error_uses_the_configured_fallback_before_output( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, event_type: str, context_fallback: bool +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = _fallback_router(fallbacks=[{"primary": ["missing"]}] if context_fallback else None) + route: Final = respx_mock.post("https://fallback-migration.local/v1/responses").mock( + side_effect=( + _context_error_response_stream(event_type=event_type), + _context_error_response_stream(completed=True), + ) + ) + stream: Final = await router.aresponses( + model="primary", + input="oversized prompt", + stream=True, + context_window_fallbacks=[{"primary": ["backup"]}] if context_fallback else None, + ) + events: Final = [event async for event in stream] + assert tuple(json.loads(call.request.content)["model"] for call in route.calls) == ("primary", "backup") + assert tuple(event.type for event in events) == ("response.created", "response.completed") + assert events[0].response.id == events[-1].response.id + assert events[-1].response.model == "backup" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("barrier", ("text", "tool", "disabled", "unconfigured")) +async def test_responses_context_error_does_not_restart_committed_or_disabled_streams( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, barrier: str +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = _fallback_router(fallbacks=[] if barrier == "unconfigured" else None) + preceding: Final = { + "text": { + "type": "response.output_text.delta", + "delta": "already visible", + "item_id": "msg_primary", + "output_index": 0, + "content_index": 0, + "sequence_number": 1, + }, + "tool": { + "type": "response.output_item.added", + "item": { + "type": "function_call", "id": "fc_primary", "call_id": "call_primary", + "name": "lookup", "arguments": "", + }, + "output_index": 0, + "sequence_number": 1, + }, + }.get(barrier) + route: Final = respx_mock.post("https://fallback-migration.local/v1/responses").mock( + return_value=_context_error_response_stream(preceding_event=preceding) + ) + stream: Final = await router.aresponses( + model="primary", input="oversized prompt", stream=True, disable_fallbacks=barrier == "disabled" + ) + with pytest.raises(litellm.ContextWindowExceededError): + async for _ in stream: + pass + assert route.call_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.parametrize("sync_mode", (False, True)) +@pytest.mark.parametrize("emission", ("role", "text", "tool")) +async def test_chat_context_error_falls_back_after_role_only_but_not_after_output( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, sync_mode: bool, emission: str +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + router: Final = _fallback_router(fallbacks=[{"primary": ["missing"]}]) + delta: Final = ( + {"role": "assistant", "tool_calls": [{ + "index": 0, "id": "call_primary", "type": "function", "function": {"name": "lookup", "arguments": ""} + }]} + if emission == "tool" + else {"role": "assistant", "content": "visible" if emission == "text" else ""} + ) + first_chunk: Final = { + "id": "chatcmpl-primary", "object": "chat.completion.chunk", "created": 1700000000, "model": "primary", + "choices": [{"index": 0, "delta": delta}], + } + error: Final = {"error": { + "type": "invalid_request_error", "code": "context_length_exceeded", + "message": "This model's maximum context length is 128000 tokens. Your messages resulted in 130000 tokens.", + }} + backup: Final = { + "id": "chatcmpl-backup", "object": "chat.completion.chunk", "created": 1700000000, "model": "backup", + "choices": [{"index": 0, "delta": {"role": "assistant", "content": "backup answer"}, "finish_reason": "stop"}], + } + route: Final = respx_mock.post("https://fallback-migration.local/v1/chat/completions").mock( + side_effect=( + httpx.Response(200, content=f"data: {json.dumps(first_chunk)}\n\ndata: {json.dumps(error)}\n\n".encode()), + httpx.Response(200, content=f"data: {json.dumps(backup)}\n\ndata: [DONE]\n\n".encode()), + ) + ) + request: Final = { + "model": "primary", "messages": [{"role": "user", "content": "oversized prompt"}], "stream": True, + "context_window_fallbacks": [{"primary": ["backup"]}], + } + stream: Final = router.completion(**request) if sync_mode else await router.acompletion(**request) + if emission != "role": + with pytest.raises(litellm.ContextWindowExceededError): + tuple(stream) if sync_mode else tuple([chunk async for chunk in stream]) + assert route.call_count == 1 + return + chunks: Final = tuple(stream) if sync_mode else tuple([chunk async for chunk in stream]) + assert tuple(json.loads(call.request.content)["model"] for call in route.calls) == ("primary", "backup") + assert "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) == "backup answer" + + +@pytest.mark.parametrize( + "wrapped,has_output,generated_content,disqualifying,recoverable", + ( + (False, False, "", False, True), + (False, True, "", False, False), + (True, False, "", False, True), + (True, True, "", False, False), + (True, False, "already visible", False, False), + (True, False, "", True, False), + ), +) +def test_context_stream_fallback_obeys_all_output_receipts( + wrapped: bool, has_output: bool, generated_content: str, disqualifying: bool, recoverable: bool +) -> None: + from litellm.exceptions import MidStreamFallbackError + from litellm.router import _context_window_stream_fallback_error + + original: Final = litellm.ContextWindowExceededError(message="too long", model="primary", llm_provider="openai") + error: Final = ( + MidStreamFallbackError( + message="too long", model="primary", llm_provider="openai", original_exception=original, + is_pre_first_chunk=True, generated_content=generated_content, emitted_disqualifying_content=disqualifying, + ) + if wrapped + else original + ) + if not recoverable: + with pytest.raises(litellm.ContextWindowExceededError) as raised: + _context_window_stream_fallback_error(error, "primary", "openai", has_output) + assert raised.value is original + return + recovered: Final = _context_window_stream_fallback_error(error, "primary", "openai", has_output) + assert recovered.original_exception is original + assert recovered.is_pre_first_chunk is True + if wrapped: + assert recovered is error + + +def test_context_stream_fallback_preserves_other_errors_for_existing_continuation_policy() -> None: + from litellm.exceptions import MidStreamFallbackError + from litellm.router import _context_window_stream_fallback_error + + error: Final = MidStreamFallbackError( + message="provider unavailable", model="primary", llm_provider="openai", generated_content="already visible" + ) + assert _context_window_stream_fallback_error(error, "primary", "openai", True) is error From fbd126e5182040a5ed8b97d54d41b5febc94452b Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 30 Oct 2026 13:25:47 -0700 Subject: [PATCH 18/46] test: fix integration tests that fail on test order, races and intended changes (#49310) Seven integration tests fail on main for reasons that are not product bugs. Two are order or race bugs in the tests, one reads rows other tests wrote, one passed on an error body, and three still expect behavior that #48644, #39864 and #39540 changed on purpose --- .../compatibility/test_missing_body_param_status.py | 2 +- .../database/migrations/test_request_log_indexes.py | 9 ++++++--- .../integration/observability/test_guardrail_effects.py | 7 +++++-- .../providers/test_gemini_tts_voice_mapping_wire.py | 4 ++-- .../providers/test_hosted_vllm_stream_stall_wire.py | 6 +++++- tests/integration/routing/test_observed_routing.py | 2 +- .../security/test_hide_secrets_tool_compatibility.py | 8 +++++++- 7 files changed, 27 insertions(+), 11 deletions(-) diff --git a/tests/integration/compatibility/test_missing_body_param_status.py b/tests/integration/compatibility/test_missing_body_param_status.py index 005e581c643..d47503e7f08 100644 --- a/tests/integration/compatibility/test_missing_body_param_status.py +++ b/tests/integration/compatibility/test_missing_body_param_status.py @@ -1199,7 +1199,7 @@ def test_anthropic_messages_explicit_null_reaches_upstream(gateway: Gateway) -> "anthropic_messages", deployment_params={"max_tokens": 32}, provider_model="openai/gpt-4o-mini", - response_route="acompletion", + response_route="aresponses", ) response: Final = _post( gateway, diff --git a/tests/integration/database/migrations/test_request_log_indexes.py b/tests/integration/database/migrations/test_request_log_indexes.py index d06f0b13cff..5741c85470b 100644 --- a/tests/integration/database/migrations/test_request_log_indexes.py +++ b/tests/integration/database/migrations/test_request_log_indexes.py @@ -1033,13 +1033,14 @@ def test_the_entry_size_check_draws_the_line_exactly_where_postgres_does( create_index: Final = sql.SQL('CREATE INDEX "entry_probe_idx" ON "entry_probe" ({})').format( sql.SQL(", ").join(sql.Identifier(column) for column in index_columns) ) + fitting_key: Final = secrets.token_hex(longest_fitting_key)[:longest_fitting_key] with psycopg.connect(scratch_database, autocommit=True) as conn: conn.execute('CREATE TABLE "entry_probe" ("s" TEXT, "k" TEXT, "t" TIMESTAMP(3))') conn.execute( 'INSERT INTO "entry_probe" ("s", "k", "t") VALUES (%s, %s, %s), (%s, %s, %s)', ( s_value, - secrets.token_hex(longest_fitting_key)[:longest_fitting_key], + fitting_key, t_value, s_value, secrets.token_hex(longest_fitting_key)[: longest_fitting_key + 1], @@ -1051,7 +1052,8 @@ def test_the_entry_size_check_draws_the_line_exactly_where_postgres_does( ) with pytest.raises(psycopg.errors.ProgramLimitExceeded, match="exceeds btree version 4 maximum 2704"): conn.execute(create_index) - conn.execute('DELETE FROM "entry_probe" WHERE octet_length("k") > %s', (longest_fitting_key,)) + conn.execute('TRUNCATE "entry_probe"') + conn.execute('INSERT INTO "entry_probe" ("s", "k", "t") VALUES (%s, %s, %s)', (s_value, fitting_key, t_value)) conn.execute(create_index) (finding,) = findings @@ -1075,7 +1077,8 @@ def test_a_value_postgres_compresses_under_the_limit_is_not_counted_while_random ) with pytest.raises(psycopg.errors.ProgramLimitExceeded, match="index row size 3016 exceeds"): conn.execute(create_index) - conn.execute('DELETE FROM "compress_probe" WHERE "k" <> %s', ("a" * 3000,)) + conn.execute('TRUNCATE "compress_probe"') + conn.execute('INSERT INTO "compress_probe" ("k") VALUES (%s)', ("a" * 3000,)) conn.execute(create_index) stored: Final = conn.execute('SELECT octet_length("k"), pg_column_size("k") FROM "compress_probe"').fetchone() diff --git a/tests/integration/observability/test_guardrail_effects.py b/tests/integration/observability/test_guardrail_effects.py index 4f5c04bddcc..1220a65528c 100644 --- a/tests/integration/observability/test_guardrail_effects.py +++ b/tests/integration/observability/test_guardrail_effects.py @@ -23,7 +23,7 @@ import yaml from integration._support.client import Gateway, eventually, object_value from integration._support.database import read_rows from integration._support.mcp import mcp_peer, register_mcp, tool_names -from integration._support.process import OwnedProxy, owned_proxy, owned_proxy_process +from integration._support.process import OwnedProxy, OwnedState, owned_proxy, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server from integration._support.ports import allocate_port from openai import APIError, AsyncOpenAI, OpenAI @@ -3263,7 +3263,10 @@ def test_client_disconnect_during_the_body_upload_runs_no_endpoint_and_writes_no with ( wire_server(_openai_upstream) as upstream, owned_proxy_process( - gateway, tmp_path, {"OPENAI_BASE_URL": upstream.url + "/v1", "OPENAI_API_KEY": "synthetic-openai-key"} + gateway, + tmp_path, + {"OPENAI_BASE_URL": upstream.url + "/v1", "OPENAI_API_KEY": "synthetic-openai-key"}, + state=OwnedState.PRIVATE, ) as owned, owned.gateway.scenario() as scenario, ): diff --git a/tests/integration/providers/test_gemini_tts_voice_mapping_wire.py b/tests/integration/providers/test_gemini_tts_voice_mapping_wire.py index 90ff32e01dd..20472dc2399 100644 --- a/tests/integration/providers/test_gemini_tts_voice_mapping_wire.py +++ b/tests/integration/providers/test_gemini_tts_voice_mapping_wire.py @@ -501,7 +501,7 @@ def test_vertex_batch_file_rows_carry_the_mapped_voice(gateway: Gateway) -> None assert tuple(text_in(row) for row in rows) == (first, second) -def test_vertex_live_default_mode_sends_one_setup_without_the_client_voice(gateway: Gateway) -> None: +def test_vertex_live_sends_one_setup_carrying_the_mapped_session_update_voice(gateway: Gateway) -> None: with gateway.scenario() as scenario: handle: Final = live_scenario(scenario) key: Final = scenario.key() @@ -511,7 +511,7 @@ def test_vertex_live_default_mode_sends_one_setup_without_the_client_voice(gatew assert types[0] == "session.created" and types[-1] == "response.done", frames setups: Final = live_setups(gateway, handle.scenario_id) assert len(setups) == 1, setups - assert voice_in(setups[0]) == NO_VOICE, setups[0] + assert voice_in(setups[0]) == "Umbriel", setups[0] @dataclass(frozen=True, slots=True) diff --git a/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py b/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py index fad7898871d..018c6a653d1 100644 --- a/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py +++ b/tests/integration/providers/test_hosted_vllm_stream_stall_wire.py @@ -279,7 +279,11 @@ def cooldown_rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Cooldown upstream, directory, {"PROXY_CONFIG_RELOAD_INTERVAL_SECONDS": "2"}, - config=cooldown_config(directory, router_settings(cache, allowed_fails=0), prometheus=False), + config=cooldown_config( + directory, + router_settings(cache, allowed_fails=0, allowed_fails_policy={"TimeoutErrorAllowedFails": 0}), + prometheus=False, + ), ) ) yield _CooldownRig(owned.gateway, cache) diff --git a/tests/integration/routing/test_observed_routing.py b/tests/integration/routing/test_observed_routing.py index ca5f248e0de..64ab799fe3a 100644 --- a/tests/integration/routing/test_observed_routing.py +++ b/tests/integration/routing/test_observed_routing.py @@ -593,7 +593,7 @@ def test_health_routing_fallbacks_preserve_error_specific_targets_and_request_co "choices": [{"index": 0, "delta": {"role": "assistant", "content": ""}}], }, response_text assert {key: events[1]["error"][key] for key in events[1]["error"] if key != "message"} == { - "type": None, + "type": "internal_server_error", "param": None, "code": "500", }, response_text diff --git a/tests/integration/security/test_hide_secrets_tool_compatibility.py b/tests/integration/security/test_hide_secrets_tool_compatibility.py index d8b900de6d0..a89678aef8e 100644 --- a/tests/integration/security/test_hide_secrets_tool_compatibility.py +++ b/tests/integration/security/test_hide_secrets_tool_compatibility.py @@ -140,9 +140,15 @@ def _key(rig: Rig, *, permissions: dict[str, JsonValue] | None = None) -> Iterat ) +def _without_cooldowns(configuration: dict[str, JsonValue]) -> None: + settings: Final = configuration["router_settings"] + assert isinstance(settings, dict) + settings["disable_cooldowns"] = True + + @pytest.fixture(scope="module") def caller(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Caller]: - with tool_rig(tmp_path_factory.mktemp("hide-secrets-compat"), _peer) as rig: + with tool_rig(tmp_path_factory.mktemp("hide-secrets-compat"), _peer, configure=_without_cooldowns) as rig: with _key(rig) as key: yield Caller(rig, key) From 0382f134e74e027d923c0cde7ad013174990b9b4 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 30 Oct 2026 13:39:01 -0700 Subject: [PATCH 19/46] fix(ui): add Meta Model API to the Add Model provider dropdown (#49314) * test(ui): cover separate Meta Model API provider onboarding Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): add Meta Model API provider onboarding Co-authored-by: lostmartian Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * test(ui): align Meta coverage with provider keys and listed registry Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): show Meta model placeholder 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: lostmartian --- .../provider_create_fields.json | 28 +++++++++++++++++++ .../public_endpoints/test_public_endpoints.py | 23 ++++++++++++++- .../components/provider_info_helpers.test.tsx | 16 +++++++++++ .../src/components/provider_info_helpers.tsx | 6 +++- 4 files changed, 71 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index df3af4ad487..d849b9fbaac 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -2377,6 +2377,34 @@ ], "default_model_placeholder": "gpt-3.5-turbo" }, + { + "provider": "META", + "provider_display_name": "Meta Model API", + "litellm_provider": "meta", + "credential_fields": [ + { + "key": "api_base", + "label": "API Base", + "placeholder": null, + "tooltip": null, + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "api_key", + "label": "API Key", + "placeholder": null, + "tooltip": null, + "required": false, + "field_type": "password", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "meta/muse-spark-1.1" + }, { "provider": "MILVUS", "provider_display_name": "Milvus", diff --git a/tests/unit/proxy/public_endpoints/test_public_endpoints.py b/tests/unit/proxy/public_endpoints/test_public_endpoints.py index 48672c7a305..3fe5345d008 100644 --- a/tests/unit/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/unit/proxy/public_endpoints/test_public_endpoints.py @@ -28,6 +28,28 @@ from litellm.types.proxy.public_endpoints.public_endpoints import ProviderCreate from litellm.types.utils import LlmProviders +def test_meta_model_api_fields_are_separate_from_meta_llama() -> None: + app: Final = FastAPI() + app.include_router(router) + response: Final = TestClient(app).get("/public/providers/fields") + assert response.status_code == 200 + providers: Final = response.json() + meta: Final = next((p for p in providers if p["litellm_provider"] == "meta"), None) + llama: Final = next(p for p in providers if p["litellm_provider"] == "meta_llama") + assert meta is not None + assert meta["provider"] == "META" + assert meta["provider_display_name"] == "Meta Model API" + assert meta["default_model_placeholder"] == "meta/muse-spark-1.1" + fields: Final = {field["key"]: field for field in meta["credential_fields"]} + assert set(fields) == {"api_base", "api_key"} + assert fields["api_base"]["field_type"] == "text" + assert fields["api_key"]["field_type"] == "password" + assert fields["api_base"]["required"] is False + assert fields["api_key"]["required"] is False + assert llama["provider"] == "LLAMA" + assert llama["provider_display_name"] == "Meta Llama" + + def test_bigmodel_fields_require_private_key_and_share_default_base_with_adapter() -> None: from litellm.llms.bigmodel.common_utils import BIGMODEL_API_BASE @@ -568,7 +590,6 @@ ADD_MODEL_UNLISTED_PROVIDERS: Final = frozenset( "libertai", "litellm_agent", "manus", - "meta", "modelscope", "mongodb", "nano-gpt", diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx index 02c31ecc082..4c483434831 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.test.tsx @@ -20,6 +20,22 @@ const BUNDLED_MODEL_MAP: Record = JSON.pa ); describe("provider_info_helpers", () => { + it("resolves Meta Model API separately from Meta Llama and isolates their bundled models", () => { + expect(resolveLitellmProviderSlug("META")).toBe("meta"); + expect(getPlaceholder(Providers.META)).toBe("meta/muse-spark-1.1"); + expect(getProviderLogoAndName("meta").displayName).toBe("Meta Model API"); + const metaModels = Object.keys(BUNDLED_MODEL_MAP).filter( + (key) => BUNDLED_MODEL_MAP[key]?.litellm_provider === "meta", + ); + const llamaModels = Object.keys(BUNDLED_MODEL_MAP).filter( + (key) => BUNDLED_MODEL_MAP[key]?.litellm_provider === "meta_llama", + ); + expect(metaModels).toContain("meta/muse-spark-1.1"); + expect(llamaModels.length).toBeGreaterThan(0); + expect(getProviderModels("META", BUNDLED_MODEL_MAP).sort()).toEqual(metaModels.sort()); + expect(getProviderModels("LLAMA", BUNDLED_MODEL_MAP).sort()).toEqual(llamaModels.sort()); + }); + it("resolves BigModel metadata, name, logo and model placeholder", () => { const metadata = PROVIDER_CREATE_FIELDS.find((provider) => provider.provider === "BIGMODEL"); expect(provider_map.BIGMODEL).toBe(metadata?.litellm_provider); diff --git a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx index 6448bd936cc..a8491f656a0 100644 --- a/ui/litellm-dashboard/src/components/provider_info_helpers.tsx +++ b/ui/litellm-dashboard/src/components/provider_info_helpers.tsx @@ -140,6 +140,7 @@ export enum Providers { LLAMAFILE = "Llamafile", LM_STUDIO = "Lm Studio", LLAMA = "Meta Llama", + META = "Meta Model API", MARITALK = "Maritalk", MICROSOFT_365_COPILOT = "Microsoft 365 Copilot", MiniMax = "MiniMax", @@ -270,6 +271,7 @@ export const provider_map: Record = { LEMONADE: "lemonade", LLAMAFILE: "llamafile", LLAMA: "meta_llama", + META: "meta", LM_STUDIO: "lm_studio", MICROSOFT_365_COPILOT: "microsoft_365_copilot", MARITALK: "maritalk", @@ -337,7 +339,7 @@ export const provider_map: Record = { ZeroGPU: "zerogpu", }; -const standaloneSubproviderSlugs = new Set(["bedrock_mantle"]); +const standaloneSubproviderSlugs = new Set(["bedrock_mantle", "meta_llama"]); export const resolveLitellmProviderSlug = (provider: string): string => provider_map[provider] ?? provider.toLowerCase(); @@ -393,6 +395,7 @@ export const providerLogoMap: Partial> = { [Providers.LAMBDA_AI]: lambdaLogo.src, [Providers.LM_STUDIO]: lmstudioLogo.src, [Providers.LLAMA]: metaLlamaLogo.src, + [Providers.META]: metaLlamaLogo.src, [Providers.MiniMax]: minimaxLogo.src, [Providers.MICROSOFT_365_COPILOT]: microsoft365CopilotLogo.src, [Providers.MistralAI]: mistralLogo.src, @@ -491,6 +494,7 @@ const providerPlaceholderMap: Partial> = { [Providers.FalAI]: "fal_ai/fal-ai/flux-pro/v1.1-ultra", [Providers.Google_AI_Studio]: "gemini-pro", [Providers.JinaAI]: "jina_ai/", + [Providers.META]: "meta/muse-spark-1.1", [Providers.MICROSOFT_365_COPILOT]: "microsoft_365_copilot/chat", [Providers.NVIDIA_RIVA]: "nvidia_riva/nvidia/parakeet-ctc-1_1b-asr", [Providers.Oracle]: "oci/xai.grok-4", From 516c796e86a5b1028b2073cd74ca341a174fa289 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 30 Oct 2026 13:46:27 -0700 Subject: [PATCH 20/46] fix(rerank): preserve token usage in spend logs (#49308) --- litellm/litellm_core_utils/litellm_logging.py | 11 +- litellm/llms/bedrock/rerank/transformation.py | 6 +- .../llms/hosted_vllm/rerank/transformation.py | 4 +- .../llms/together_ai/rerank/transformation.py | 4 +- litellm/types/rerank.py | 46 +++++++- .../test_litellm_logging.py | 100 ++++++++++++++++++ .../test_hosted_vllm_rerank_transformation.py | 35 ++++++ .../jina_ai/rerank/test_transformation.py | 5 +- 8 files changed, 200 insertions(+), 11 deletions(-) diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 26949d9bbfb..924603dc7e8 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -165,7 +165,7 @@ from litellm.types.llms.openai import ( ) from litellm.types.mcp import MCPPostCallResponseObject from litellm.types.prompts.init_prompts import PromptSpec -from litellm.types.rerank import RerankResponse +from litellm.types.rerank import RerankResponse, get_rerank_usage from litellm.types.router import RouterRateLimitError, SilentExperimentFlag from litellm.types.utils import ( DEPLOYMENT_SCOPED_PRICING_FIELDS, @@ -355,6 +355,7 @@ def _get_provider_request_id(original_exception: Exception) -> str | None: # Cache custom pricing keys as frozenset for O(1) lookups instead of looping through 49 keys _CUSTOM_PRICING_KEYS: Final[frozenset[str]] = frozenset(CustomPricingLiteLLMParams.model_fields.keys()) _METADATA_MAPPING: Final = TypeAdapter(Mapping[str, object]) +_RESPONSE_MAPPING: Final = TypeAdapter(Mapping[str, object]) _MODEL_INFO_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = _CUSTOM_PRICING_KEYS | DEPLOYMENT_SCOPED_PRICING_FIELDS _UNSERIALIZABLE_METADATA_KEYS: Final[frozenset[str]] = frozenset( ("user_api_key_auth", "user_api_key_budget_reservation") @@ -6817,7 +6818,11 @@ class StandardLoggingPayloadSetup: total_tokens=0, ) - usage: Final = response_obj.get("usage", None) or {} + usage: Final = response_obj.get("usage", None) + if usage is None: + rerank_usage: Final = get_rerank_usage(_RESPONSE_MAPPING.validate_python(response_obj)) + if rerank_usage is not None: + return Usage(**rerank_usage) if usage is None or (not isinstance(usage, dict) and not isinstance(usage, Usage)): return Usage( prompt_tokens=0, @@ -6853,7 +6858,7 @@ class StandardLoggingPayloadSetup: return _empty _raw: Final = response_obj.get("usage", None) if _raw is None: - return _empty + return dict(get_rerank_usage(_RESPONSE_MAPPING.validate_python(response_obj)) or _empty) if isinstance(_raw, ResponseAPIUsage): return ResponseAPILoggingUtils.transform_response_api_usage_to_chat_usage(_raw).model_dump() if isinstance(_raw, dict): diff --git a/litellm/llms/bedrock/rerank/transformation.py b/litellm/llms/bedrock/rerank/transformation.py index 442903fc05e..c2f5035528d 100644 --- a/litellm/llms/bedrock/rerank/transformation.py +++ b/litellm/llms/bedrock/rerank/transformation.py @@ -6,6 +6,8 @@ Why separate file? Make it easy to see how transformation works from typing import Final +from pydantic import TypeAdapter + from litellm._uuid import uuid from litellm.types.llms.bedrock import ( BedrockRerankBedrockRerankingConfiguration, @@ -90,7 +92,9 @@ class BedrockRerankConfig: example input: {"results":[{"index":0,"relevanceScore":0.6847912669181824},{"index":1,"relevanceScore":0.5980774760246277}]} """ - _billed_units = RerankBilledUnits(**response.get("usage", {"search_units": 1})) # by default 1 search unit + _billed_units: Final = TypeAdapter(RerankBilledUnits).validate_python( + response.get("usage", {"search_units": 1}) + ) _tokens: Final = RerankTokens(**response.get("usage", {})) rerank_meta: Final = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens) diff --git a/litellm/llms/hosted_vllm/rerank/transformation.py b/litellm/llms/hosted_vllm/rerank/transformation.py index 0b77a04c69a..1b7da6f6635 100644 --- a/litellm/llms/hosted_vllm/rerank/transformation.py +++ b/litellm/llms/hosted_vllm/rerank/transformation.py @@ -244,5 +244,7 @@ class HostedVLLMRerankConfig(BaseRerankConfig): return RerankResponse( id=response.get("id") or str(uuid.uuid4()), results=rerank_results, - meta=rerank_meta, + meta=TypeAdapter(RerankResponseMeta).validate_python(response["meta"]) + if "meta" in response and response["meta"] is not None + else rerank_meta, ) diff --git a/litellm/llms/together_ai/rerank/transformation.py b/litellm/llms/together_ai/rerank/transformation.py index f1c3ca9e87d..eb7ec53416d 100644 --- a/litellm/llms/together_ai/rerank/transformation.py +++ b/litellm/llms/together_ai/rerank/transformation.py @@ -6,6 +6,8 @@ Why separate file? Make it easy to see how transformation works from typing import Final +from pydantic import TypeAdapter + from litellm._uuid import uuid from litellm.types.rerank import ( RerankBilledUnits, @@ -19,7 +21,7 @@ from litellm.types.rerank import ( class TogetherAIRerankConfig: def _transform_response(self, response: dict) -> RerankResponse: - _billed_units: Final = RerankBilledUnits(**response.get("usage", {})) + _billed_units: Final = TypeAdapter(RerankBilledUnits).validate_python(response.get("usage", {})) _tokens: Final = RerankTokens(**response.get("usage", {})) rerank_meta: Final = RerankResponseMeta(billed_units=_billed_units, tokens=_tokens) diff --git a/litellm/types/rerank.py b/litellm/types/rerank.py index ffdcf53fd7e..ca2767a7174 100644 --- a/litellm/types/rerank.py +++ b/litellm/types/rerank.py @@ -4,10 +4,10 @@ https://docs.cohere.com/reference/rerank """ -from collections.abc import Sequence -from typing import Literal, overload +from collections.abc import Mapping, Sequence +from typing import Final, Literal, overload -from pydantic import ConfigDict, PrivateAttr +from pydantic import ConfigDict, PrivateAttr, TypeAdapter, ValidationError from typing_extensions import ReadOnly, Required, TypedDict from litellm.types.llms.base import LiteLLMBaseModel @@ -57,6 +57,8 @@ class OptionalRerankParams(TypedDict, total=False): class RerankBilledUnits(TypedDict, total=False): search_units: int | None total_tokens: int | None + input_tokens: ReadOnly[int | None] + output_tokens: ReadOnly[int | None] class RerankTokens(TypedDict, total=False): @@ -65,11 +67,47 @@ class RerankTokens(TypedDict, total=False): class RerankResponseMeta(TypedDict, total=False): - api_version: dict | None + api_version: dict[str, object] | None billed_units: RerankBilledUnits | None tokens: RerankTokens | None +class RerankUsage(TypedDict): + prompt_tokens: ReadOnly[int] + completion_tokens: ReadOnly[int] + total_tokens: ReadOnly[int] + + +_RERANK_META: Final = TypeAdapter(RerankResponseMeta) + + +def get_rerank_usage(response: Mapping[str, object]) -> RerankUsage | None: + if "results" not in response or response.get("meta") is None: + return None + try: + meta: Final = _RERANK_META.validate_python(response.get("meta")) + except ValidationError: + return None + tokens: Final = meta.get("tokens") or RerankTokens() + billed: Final = meta.get("billed_units") or RerankBilledUnits() + input_tokens: Final = next( + (count for count in (tokens.get("input_tokens"), billed.get("input_tokens")) if count is not None), None + ) + output_tokens: Final = next( + (count for count in (tokens.get("output_tokens"), billed.get("output_tokens")) if count is not None), None + ) + total_tokens: Final = billed.get("total_tokens") + if input_tokens is None and output_tokens is None and total_tokens is None: + return None + completion_tokens: Final = output_tokens or 0 + prompt_tokens: Final = input_tokens if input_tokens is not None else max((total_tokens or 0) - completion_tokens, 0) + return RerankUsage( + prompt_tokens=prompt_tokens, + completion_tokens=completion_tokens, + total_tokens=total_tokens if total_tokens is not None else prompt_tokens + completion_tokens, + ) + + class RerankResponseDocument(TypedDict): text: str diff --git a/tests/unit/litellm_core_utils/test_litellm_logging.py b/tests/unit/litellm_core_utils/test_litellm_logging.py index f7c2082b683..e952127a32c 100644 --- a/tests/unit/litellm_core_utils/test_litellm_logging.py +++ b/tests/unit/litellm_core_utils/test_litellm_logging.py @@ -59,6 +59,7 @@ from litellm.proxy.hooks.max_iterations_limiter import PROXY_MaxIterationsHandle from litellm.proxy.hooks.parallel_request_limiter_v3 import PROXY_MaxParallelRequestsHandler_v3 from litellm.proxy.utils import InternalUsageCache from litellm.types.llms.openai import ResponseAPIUsage, ResponseCompletedEvent, ResponsesAPIResponse +from litellm.types.rerank import RerankResponse from litellm.types.utils import ( CallTypes, CostBreakdown, @@ -4514,6 +4515,105 @@ def test_logging_preserves_missing_usage_without_accepting_request_metadata( assert (payload["prompt_tokens"], payload["completion_tokens"], payload["total_tokens"]) == (0, 0, 0) +@pytest.mark.parametrize( + ("meta", "expected"), + [ + ({"tokens": {"input_tokens": 7, "output_tokens": 3}}, (7, 3, 10)), + ({"tokens": {}, "billed_units": {"total_tokens": 42}}, (42, 0, 42)), + ({"tokens": {"output_tokens": 3}, "billed_units": {"total_tokens": 12}}, (9, 3, 12)), + ({"tokens": {"input_tokens": 7, "output_tokens": 3}, "billed_units": {"total_tokens": 0}}, (7, 3, 0)), + ({"billed_units": {"input_tokens": 7, "output_tokens": 3}}, (7, 3, 10)), + ( + { + "tokens": {"input_tokens": 7, "output_tokens": 3}, + "billed_units": {"input_tokens": 90, "output_tokens": 10, "total_tokens": 12}, + }, + (7, 3, 12), + ), + ( + { + "tokens": {"input_tokens": 0, "output_tokens": 0}, + "billed_units": {"input_tokens": 90, "output_tokens": 10, "total_tokens": 12}, + }, + (0, 0, 12), + ), + ( + { + "tokens": {"input_tokens": None, "output_tokens": None}, + "billed_units": {"input_tokens": 7, "output_tokens": 3, "total_tokens": None}, + }, + (7, 3, 10), + ), + ({"tokens": {"input_tokens": 0}}, (0, 0, 0)), + (None, None), + ({}, None), + ({"tokens": None, "billed_units": None}, None), + ({"billed_units": {"search_units": 3}}, None), + ], +) +def test_rerank_metadata_reaches_standard_and_spend_logs( + logging_obj: Logging, meta: dict[str, object] | None, expected: tuple[int, int, int] | None +) -> None: + from litellm.proxy.spend_tracking.spend_tracking_utils import get_logging_payload + + now: Final = datetime_unit_test(2026, 1, 1, 12, 0, 0) + native: Final = {"id": "rerank-usage", "results": [{"index": 0, "relevance_score": 0.9}], "meta": meta} + response: Final = RerankResponse.model_validate(native) + counts: Final = expected if expected is not None else (0, 0, 0) + expected_usage: Final = dict(zip(("prompt_tokens", "completion_tokens", "total_tokens"), counts)) + assert response.model_dump() == native + usage: Final = StandardLoggingPayloadSetup.get_usage_from_response_obj(response.model_dump()) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == counts + assert StandardLoggingPayloadSetup.get_usage_as_dict(response.model_dump()) == ( + expected_usage if expected is not None else {} + ) + kwargs: Final = {"model": "rerank-model", "call_type": "arerank", "litellm_params": {"metadata": {}}} + standard: Final = get_standard_logging_object_payload( + kwargs=kwargs, + init_response_obj=response, + start_time=now, + end_time=now, + logging_obj=logging_obj, + status="success", + ) + assert standard is not None + assert (standard["prompt_tokens"], standard["completion_tokens"], standard["total_tokens"]) == counts + assert standard["metadata"]["usage_object"] == (expected_usage if expected is not None else {}) + spend: Final = get_logging_payload( + kwargs={**kwargs, "standard_logging_object": standard}, response_obj=response, start_time=now, end_time=now + ) + assert (spend["prompt_tokens"], spend["completion_tokens"], spend["total_tokens"]) == counts + + +@pytest.mark.parametrize("combined", [None, Usage(prompt_tokens=11, completion_tokens=2, total_tokens=13)]) +@pytest.mark.parametrize("explicit_usage", [{}, {"prompt_tokens": 0, "completion_tokens": 0, "total_tokens": 0}]) +def test_rerank_metadata_does_not_override_explicit_or_combined_usage( + combined: Usage | None, explicit_usage: dict[str, int] +) -> None: + response: Final = {"results": [], "meta": {"tokens": {"input_tokens": 42}}, "usage": explicit_usage} + expected: Final = combined.model_dump() if combined is not None else explicit_usage + assert StandardLoggingPayloadSetup.get_usage_as_dict(response, combined) == expected + usage: Final = StandardLoggingPayloadSetup.get_usage_from_response_obj(response, combined) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == ( + expected.get("prompt_tokens", 0), + expected.get("completion_tokens", 0), + expected.get("total_tokens", 0), + ) + + +@pytest.mark.parametrize( + "response", + [ + {"meta": {"tokens": {"input_tokens": 42}}}, + {"results": [], "meta": {"tokens": {"input_tokens": "invalid"}}}, + ], +) +def test_rerank_metadata_ignores_unrelated_and_invalid_responses(response: dict[str, object]) -> None: + assert StandardLoggingPayloadSetup.get_usage_as_dict(response) == {} + usage: Final = StandardLoggingPayloadSetup.get_usage_from_response_obj(response) + assert (usage.prompt_tokens, usage.completion_tokens, usage.total_tokens) == (0, 0, 0) + + def test_append_system_prompt_messages(): """ Test append_system_prompt_messages prepends system message from kwargs to messages list. diff --git a/tests/unit/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py b/tests/unit/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py index 9a62fcf6f0f..7461611e0fa 100644 --- a/tests/unit/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py +++ b/tests/unit/llms/hosted_vllm/test_hosted_vllm_rerank_transformation.py @@ -148,6 +148,41 @@ class TestHostedVLLMRerankTransform: self.config._transform_response(response_dict) +@pytest.mark.parametrize( + "meta", + [ + { + "api_version": {"version": "2"}, + "tokens": {"input_tokens": 7, "output_tokens": 3}, + "billed_units": {"total_tokens": 12}, + }, + {"tokens": {"input_tokens": 0}, "billed_units": {"input_tokens": 0, "output_tokens": 0}}, + {"tokens": None, "billed_units": None}, + {}, + ], +) +@pytest.mark.parametrize("legacy_usage", [{}, {"usage": {"total_tokens": 99}}]) +def test_hosted_vllm_rerank_preserves_native_metadata_over_legacy_usage( + meta: dict[str, object], legacy_usage: dict[str, object] +) -> None: + response: Final = HostedVLLMRerankConfig()._transform_response( + {"id": "rerank-native", "results": [{"index": 0, "relevance_score": 0.9}], "meta": meta, **legacy_usage} + ) + assert response.meta == meta + assert response.model_dump() == { + "id": "rerank-native", + "results": [{"index": 0, "relevance_score": 0.9}], + "meta": meta, + } + + +def test_hosted_vllm_rerank_null_metadata_uses_legacy_usage() -> None: + response: Final = HostedVLLMRerankConfig()._transform_response( + {"results": [], "meta": None, "usage": {"total_tokens": 42}} + ) + assert response.meta == {"billed_units": {"total_tokens": 42}, "tokens": {"input_tokens": 42}} + + class TestGetOptionalRerankParamsInstruction: """`instruction` is threaded through get_optional_rerank_params only when set.""" diff --git a/tests/unit/llms/jina_ai/rerank/test_transformation.py b/tests/unit/llms/jina_ai/rerank/test_transformation.py index 9ec02a4548b..b752f7b4649 100644 --- a/tests/unit/llms/jina_ai/rerank/test_transformation.py +++ b/tests/unit/llms/jina_ai/rerank/test_transformation.py @@ -46,7 +46,10 @@ def test_transform_rerank_response_maps_results_and_usage(): ({"results": [], "usage": {}}, {"billed_units": {}, "tokens": {}}), ( {"results": [], "usage": {"total_tokens": 12, "input_tokens": 7, "output_tokens": 3, "unknown": 1}}, - {"billed_units": {"total_tokens": 12}, "tokens": {"input_tokens": 7, "output_tokens": 3}}, + { + "billed_units": {"total_tokens": 12, "input_tokens": 7, "output_tokens": 3}, + "tokens": {"input_tokens": 7, "output_tokens": 3}, + }, ), ], ) From a42c1c138a83cb15cf4586f967f105ff90c5c3e0 Mon Sep 17 00:00:00 2001 From: Graham Neubig Date: Fri, 30 Oct 2026 16:46:37 -0400 Subject: [PATCH 21/46] fix(exceptions): classify SGLang context window overflow (#48837) * fix(exceptions): classify SGLang context window overflow Co-authored-by: openhands * test(exceptions): use SGLang's real top-level error body in overflow regression Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: openhands Co-authored-by: Moe Khalil Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../exception_mapping_utils.py | 1 + .../test_exception_mapping_utils.py | 61 +++++++++++++++++++ 2 files changed, 62 insertions(+) diff --git a/litellm/litellm_core_utils/exception_mapping_utils.py b/litellm/litellm_core_utils/exception_mapping_utils.py index 33b6d4e1857..85c232ad3f4 100644 --- a/litellm/litellm_core_utils/exception_mapping_utils.py +++ b/litellm/litellm_core_utils/exception_mapping_utils.py @@ -98,6 +98,7 @@ class ExceptionCheckers: known_exception_substrings: Final = [ "exceed context limit", "this model's maximum context length is", + "requested token count exceeds the model's maximum context length", "string too long. expected a string with maximum length", "model's maximum context limit", "is longer than the model's context length", diff --git a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py index efb9648cf13..2b1e4a5eed3 100644 --- a/tests/unit/litellm_core_utils/test_exception_mapping_utils.py +++ b/tests/unit/litellm_core_utils/test_exception_mapping_utils.py @@ -440,6 +440,67 @@ def test_gemini_context_window_error_mapping(error_message, should_raise_context ) +@pytest.mark.parametrize("sync_mode", [True, False]) +@pytest.mark.parametrize("streaming", [True, False]) +@pytest.mark.parametrize( + ("error_message", "expected_type"), + [ + ( + ( + "Requested token count exceeds the model's maximum context length of 131072 tokens. " + "You requested a total of 140000 tokens: 130000 tokens from the input messages and " + "10000 tokens for the completion. Please reduce the number of tokens in the input " + "messages or the completion to fit within the limit." + ), + litellm.ContextWindowExceededError, + ), + ( + ( + "requested token count exceeds the model's maximum context length of 8192 tokens. " + "you requested a total of 9000 tokens" + ), + litellm.ContextWindowExceededError, + ), + ("max_tokens must be a positive integer", litellm.BadRequestError), + ], + ids=["sglang-context-overflow", "lowercase-context-overflow", "unrelated-bad-request"], +) +@pytest.mark.asyncio +async def test_sglang_context_window_error_mapping( + sync_mode: bool, + streaming: bool, + error_message: str, + expected_type: type[litellm.BadRequestError], +) -> None: + def reject_request(request: httpx.Request) -> httpx.Response: + return httpx.Response( + 400, + request=request, + json={"object": "error", "message": error_message, "type": "BadRequestError", "param": None, "code": 400}, + ) + + transport: Final = httpx.MockTransport(reject_request) + if sync_mode: + with ( + httpx.Client(transport=transport) as http_client, + OpenAI(api_key="sk-test", max_retries=0, http_client=http_client) as client, + pytest.raises(litellm.BadRequestError) as caught, + ): + _call_openai_api_sync("openai/sglang-model", "chat_completion", streaming, client) + else: + async with ( + httpx.AsyncClient(transport=transport) as async_http_client, + AsyncOpenAI(api_key="sk-test", max_retries=0, http_client=async_http_client) as async_client, + ): + with pytest.raises(litellm.BadRequestError) as caught: + await _call_openai_api_async("openai/sglang-model", "chat_completion", streaming, async_client) + + assert type(caught.value) is expected_type + assert caught.value.status_code == 400 + assert error_message in str(caught.value) + assert caught.value.llm_provider == "openai" + + def test_lemonade_context_window_error_mapping(): """Lemonade's llama.cpp backend should map context overflows to LiteLLM's standard error.""" From 467847b2e9da67dd1c4ef43ff7d95ca56c013fa3 Mon Sep 17 00:00:00 2001 From: moe-berri Date: Fri, 30 Oct 2026 13:58:04 -0700 Subject: [PATCH 22/46] fix(ui): restore Z.AI credential fields in provider catalog (#49325) Carry the Z.AI catalog fix from #40088 and move its regression into the mapped unit test file. Co-authored-by: Chaitanya Laxman Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../provider_create_fields.json | 28 ++++++++++++++++++ .../public_endpoints/test_public_endpoints.py | 29 ++++++++++++++++++- 2 files changed, 56 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/public_endpoints/provider_create_fields.json b/litellm/proxy/public_endpoints/provider_create_fields.json index d849b9fbaac..01a822c5d93 100644 --- a/litellm/proxy/public_endpoints/provider_create_fields.json +++ b/litellm/proxy/public_endpoints/provider_create_fields.json @@ -4285,6 +4285,34 @@ ], "default_model_placeholder": "gpt-3.5-turbo" }, + { + "provider": "ZAI", + "provider_display_name": "Z.AI (Zhipu AI)", + "litellm_provider": "zai", + "credential_fields": [ + { + "key": "api_base", + "label": "API Base", + "placeholder": "https://api.z.ai/api/paas/v4", + "tooltip": null, + "required": false, + "field_type": "text", + "options": null, + "default_value": null + }, + { + "key": "api_key", + "label": "API Key", + "placeholder": null, + "tooltip": null, + "required": true, + "field_type": "password", + "options": null, + "default_value": null + } + ], + "default_model_placeholder": "zai/glm-4.5" + }, { "provider": "ZeroGPU", "provider_display_name": "ZeroGPU", diff --git a/tests/unit/proxy/public_endpoints/test_public_endpoints.py b/tests/unit/proxy/public_endpoints/test_public_endpoints.py index 3fe5345d008..72755633376 100644 --- a/tests/unit/proxy/public_endpoints/test_public_endpoints.py +++ b/tests/unit/proxy/public_endpoints/test_public_endpoints.py @@ -408,6 +408,34 @@ def test_cognition_provider_fields(): assert fields_by_key["api_base"]["required"] is False +def test_zai_provider_fields(): + from litellm.llms.zai.chat.transformation import ZAI_API_BASE + + app_instance: Final = FastAPI() + app_instance.include_router(router) + test_client: Final = TestClient(app_instance) + + response: Final = test_client.get("/public/providers/fields") + assert response.status_code == 200 + providers: Final = response.json() + + zai: Final = next((p for p in providers if p["provider"] == "ZAI"), None) + assert zai is not None, "Z.AI provider entry not found" + + assert zai["provider_display_name"] == "Z.AI (Zhipu AI)" + assert zai["litellm_provider"] == "zai" + assert zai["default_model_placeholder"].startswith("zai/") + + fields_by_key: Final = {f["key"]: f for f in zai["credential_fields"]} + + assert fields_by_key["api_key"]["required"] is True + assert fields_by_key["api_key"]["field_type"] == "password" + + assert fields_by_key["api_base"]["required"] is False + assert fields_by_key["api_base"]["field_type"] == "text" + assert fields_by_key["api_base"]["placeholder"] == ZAI_API_BASE + + def test_zerogpu_add_model_form_defaults_resolve_to_a_priced_zerogpu_model(): import litellm from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider @@ -611,7 +639,6 @@ ADD_MODEL_UNLISTED_PROVIDERS: Final = frozenset( "transcribe", "valkey", "xiaomi_mimo", - "zai", } ) From 7e37c207482dde6b17c19783e9d8aeafadaecd11 Mon Sep 17 00:00:00 2001 From: tavren27 <348969968+tavren27@users.noreply.github.com> Date: Fri, 30 Oct 2026 14:12:51 -0700 Subject: [PATCH 23/46] fix(responses): run the guardrails a Responses WebSocket frame selects (#48725) * fix(responses): run the guardrails a Responses WebSocket frame selects * fix(responses): name the field a malformed frame guardrail selection got wrong A frame whose guardrails or guardrail_config has the wrong shape now gets an error event naming the field, instead of pydantic's raw validation dump * fix(responses): accept every guardrail selection shape a Responses WebSocket frame sends, and run during_call guardrails A frame can pick guardrails as a list of names or name-to-config dicts, inside metadata or litellm_metadata, with policies validated as a list of strings. during_call guardrails now run before the turn goes upstream, and selection fields sent next to a nested response are rejected instead of dropped * fix(proxy): keep a key's configured guardrails when policies add guardrails A policy's guardrails were merged with dict.fromkeys, so a configured guardrail entry ({name: {extra_body}}) in key metadata raised unhashable type: 'dict' and every request naming policies failed with a 500 on HTTP and on the Responses WebSocket. The merge now keeps existing entries and skips policy guardrails already selected by name. Also types the WebSocket test guardrail's hooks. * fix(responses): run streaming-scoped guardrails on socket turns, keep guardrail fields on blocks, merge guardrail lists by name * fix(responses): check frame guardrail picks against the team permission, honor picks beside a nested response, prefer a guardrail entry with settings * style(responses): format the frame guard and give the pass-through guardrail test an enterprise license * fix(proxy): keep pass-through guardrail settings, gate frame policy picks, skip moderation-only hooks on socket turns * fix(responses): keep string metadata guardrail values as client metadata on socket turns * fix(responses): ask no team permission for string metadata named like a guardrail pick on socket turns * fix(proxy): socket guardrail picks read the key's team metadata when the proxy has no database --------- Co-authored-by: mateo-berri <277851410+mateo-berri@users.noreply.github.com> --- litellm/integrations/custom_guardrail.py | 43 +- litellm/proxy/auth/auth_checks.py | 81 +-- litellm/proxy/litellm_pre_call_utils.py | 67 +- .../passthrough_guardrails.py | 30 +- .../proxy/response_api_endpoints/endpoints.py | 294 ++++++++- litellm/proxy/utils.py | 6 +- litellm/responses/streaming_iterator.py | 36 +- .../integrations/test_custom_guardrail.py | 13 + ...st_auth_checks_object_access_and_lookup.py | 6 +- .../test_passthrough_guardrails.py | 19 + .../response_api_endpoints/test_endpoints.py | 615 +++++++++++++++++- .../unit/proxy/test_litellm_pre_call_utils.py | 82 +++ .../utils/helpers/test_guardrail_merge.py | 18 +- 13 files changed, 1134 insertions(+), 176 deletions(-) diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index a9fb5c29bad..736ebd58130 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -3,8 +3,9 @@ import copy import hashlib import os import secrets -from collections.abc import Mapping, Sequence +from collections.abc import Iterable, Mapping, Sequence from datetime import datetime +from itertools import chain from types import MappingProxyType from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, cast, get_args @@ -208,6 +209,46 @@ def guardrail_request_data_with_streaming( } +_GUARDRAIL_ENTRY_WITH_SETTINGS: Final = TypeAdapter(Mapping[str, object]) + + +def merged_guardrail_entries(*entry_lists: Iterable[object]) -> list[object]: + """Merges guardrail lists so each guardrail is named once: the first ``{name: settings}`` entry naming it wins over + any plain name in either order, else its first plain name stays""" + entries: Final = tuple(chain.from_iterable(entry_lists)) + settings_names: Final = tuple(_settings_entry_names(entry) for entry in entries) + named_with_settings: Final = frozenset[str]().union(*settings_names) + return [ + entry + for index, entry in enumerate(entries) + if not _superseded_guardrail_entry(entries, settings_names, named_with_settings, index) + ] + + +def guardrail_entry_names(entry: object) -> tuple[str, ...]: + if isinstance(entry, str): + return (entry,) + if isinstance(entry, dict): + return tuple(_GUARDRAIL_ENTRY_WITH_SETTINGS.validate_python(entry)) + return () + + +def _settings_entry_names(entry: object) -> frozenset[str]: + return frozenset(guardrail_entry_names(entry)) if isinstance(entry, dict) else frozenset() + + +def _superseded_guardrail_entry( + entries: tuple[object, ...], + settings_names: tuple[frozenset[str], ...], + named_with_settings: frozenset[str], + index: int, +) -> bool: + entry: Final = entries[index] + if isinstance(entry, str): + return entry in named_with_settings or entry in entries[:index] + return bool(settings_names[index]) and settings_names[index] <= frozenset[str]().union(*settings_names[:index]) + + def _request_is_streaming(data: object, event_type: GuardrailEventHooks | None = None) -> bool: if event_type in _REALTIME_STREAMING_HOOKS: return True diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index 79a46ef355e..7d379a6f042 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -937,53 +937,58 @@ _GUARDRAIL_MODIFICATION_KEYS: Final[tuple] = ( ) -def _guardrail_modification_check(request_body: Mapping[str, object], team_object: LiteLLM_TeamTable | None) -> None: +def requests_guardrail_modification(request_body: Mapping[str, object]) -> bool: + """Whether a request body sets the plural ``guardrails`` list or a per-request toggle that influences whether + default-on guardrails run (``disable_global_guardrails``, ``disable_global_guardrail`` singular, and + ``opted_out_global_guardrails``), at the top level or in either metadata key""" + return ( + _requests_guardrail_modification_in(request_body.get("metadata")) + or _requests_guardrail_modification_in(request_body.get("litellm_metadata")) + or _requests_guardrail_modification_in(request_body) + ) + + +def _requests_guardrail_modification_in(container: object) -> bool: + coerced: Final = _guardrail_container_as_dict(container) + if coerced is None: + return False + return any(key in coerced for key in _GUARDRAIL_MODIFICATION_KEYS) + + +def _guardrail_container_as_dict(container: object) -> Mapping[str, object] | None: + """Accept dict or JSON-string (from multipart/form-data or extra_body). + + Without this, an attacker can smuggle guardrail keys past the check by + sending ``{"metadata": "{\\"disable_global_guardrails\\": true}"}`` — + ``isinstance(dict)`` on the string returns False, the check returns + no-modification, and ``add_litellm_data_to_request`` parses the string + to a dict downstream. + """ + if isinstance(container, Mapping): + return container + if isinstance(container, str): + parsed: Final = _safe_json_loads_obj(container) + return parsed if isinstance(parsed, dict) else None + return None + + +def guardrail_modification_check(request_body: Mapping[str, object], team_object: LiteLLM_TeamTable | None) -> None: """ Reject user-supplied metadata flags that would modify guardrail behavior - unless the team has explicit permission. Checked keys include the plural - ``guardrails`` list plus the per-request toggles that influence whether - default-on guardrails run (``disable_global_guardrails``, - ``disable_global_guardrail`` singular, and ``opted_out_global_guardrails``). + unless the team has explicit permission. User-supplied values for the bypass toggles are also silently ignored by ``_get_admin_metadata`` at read time; this check adds defense in depth by failing loudly at the auth layer so operators see an explicit 403 instead of a confusing silent-ignore. """ + if requests_guardrail_modification(request_body): + require_guardrail_modification_permission(team_object) + + +def require_guardrail_modification_permission(team_object: LiteLLM_TeamTable | None) -> None: from litellm.proxy.guardrails.guardrail_helpers import can_modify_guardrails - def _coerce_to_dict(container: object) -> dict | None: - """Accept dict or JSON-string (from multipart/form-data or extra_body). - - Without this, an attacker can smuggle guardrail keys past the check by - sending ``{"metadata": "{\\"disable_global_guardrails\\": true}"}`` — - ``isinstance(dict)`` on the string returns False, the check returns - no-modification, and ``add_litellm_data_to_request`` parses the string - to a dict downstream. - """ - if isinstance(container, dict): - return container - if isinstance(container, str): - parsed: Final = _safe_json_loads_obj(container) - return parsed if isinstance(parsed, dict) else None - return None - - def _user_requested_modification(container: object) -> bool: - coerced: Final = _coerce_to_dict(container) - if coerced is None: - return False - return any(key in coerced for key in _GUARDRAIL_MODIFICATION_KEYS) - - # Check both metadata keys — callers can populate either depending on the - # endpoint. Cover the top-level too so root-level injection is rejected. - modifies: Final = ( - _user_requested_modification(request_body.get("metadata")) - or _user_requested_modification(request_body.get("litellm_metadata")) - or _user_requested_modification(request_body) - ) - if not modifies: - return - if not can_modify_guardrails(team_object): raise HTTPException( status_code=403, @@ -1485,7 +1490,7 @@ async def common_checks( _enforce_user_param_check(general_settings, request, request_body, route) _global_proxy_budget_check(global_proxy_spend, skip_all_budget_checks, route) - _guardrail_modification_check(_typed_request_body(request_body), team_object) + guardrail_modification_check(_typed_request_body(request_body), team_object) # 10 [OPTIONAL] Organization RBAC checks organization_role_based_access_check(user_object=user_object, route=route, request_body=request_body) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index ae9a6d4b959..a9cffcc6566 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -48,7 +48,7 @@ from litellm.constants import ( VECTOR_STORE_CALLER_METADATA_KWARG, X_LITELLM_DISABLE_CALLBACKS, ) -from litellm.integrations.custom_guardrail import without_server_streaming_classification +from litellm.integrations.custom_guardrail import merged_guardrail_entries, without_server_streaming_classification from litellm.litellm_core_utils.core_helpers import is_codex_user_agent, normalize_drop_params from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.initialize_dynamic_callback_params import ( @@ -3948,30 +3948,14 @@ def add_guardrails_from_key_or_team_metadata( """ from litellm.proxy.utils import premium_user_check - # Initialize guardrails set (avoiding duplicates) - combined_guardrails: Final = set() - - # Add key-level guardrails first - if key_metadata and "guardrails" in key_metadata: - if isinstance(key_metadata["guardrails"], list) and len(key_metadata["guardrails"]) > 0: - premium_user_check() - combined_guardrails.update(key_metadata["guardrails"]) - - # Add team-level guardrails (set automatically handles duplicates) - if team_metadata and "guardrails" in team_metadata: - if isinstance(team_metadata["guardrails"], list) and len(team_metadata["guardrails"]) > 0: - premium_user_check() - combined_guardrails.update(team_metadata["guardrails"]) - - # Add project-level guardrails (set automatically handles duplicates) - if project_metadata and "guardrails" in project_metadata: - if isinstance(project_metadata["guardrails"], list) and len(project_metadata["guardrails"]) > 0: - premium_user_check() - combined_guardrails.update(project_metadata["guardrails"]) - - # Set combined guardrails in metadata as list - if combined_guardrails: - data[metadata_variable_name]["guardrails"] = list(combined_guardrails) + configured_guardrails: Final = tuple( + metadata["guardrails"] + for metadata in (key_metadata, team_metadata, project_metadata) + if metadata and isinstance(metadata.get("guardrails"), list) and len(metadata["guardrails"]) > 0 + ) + if configured_guardrails: + premium_user_check() + data[metadata_variable_name]["guardrails"] = merged_guardrail_entries(*configured_guardrails) _add_guardrails_from_key_or_team_metadata: Final = add_guardrails_from_key_or_team_metadata @@ -4066,14 +4050,10 @@ def _add_guardrails_from_policies_in_metadata( if metadata_variable_name not in data: data[metadata_variable_name] = {} - existing_guardrails = data[metadata_variable_name].get("guardrails", []) - if not isinstance(existing_guardrails, list): - existing_guardrails = [] - - # Combine existing guardrails with policy-resolved guardrails (no duplicates) - combined: Final = set(existing_guardrails) - combined.update(resolved_guardrails) - data[metadata_variable_name]["guardrails"] = list(combined) + existing_guardrails: Final = data[metadata_variable_name].get("guardrails", []) + data[metadata_variable_name]["guardrails"] = merged_guardrail_entries( + existing_guardrails if isinstance(existing_guardrails, list) else (), resolved_guardrails + ) # Store applied policies in metadata for tracking if "applied_policies" not in data[metadata_variable_name]: @@ -4308,10 +4288,11 @@ def _apply_resolved_guardrails_to_metadata( return existing_guardrails: Final = data[metadata_variable_name].get("guardrails", []) - existing_guardrails_list: Final = existing_guardrails if isinstance(existing_guardrails, list) else [] + existing_guardrails_list: Final = ( + _REQUEST_ENTRIES.validate_python(existing_guardrails) if isinstance(existing_guardrails, list) else () + ) - # Combine existing guardrails with policy-resolved guardrails (no duplicates) - combined: Final = list(dict.fromkeys((*existing_guardrails_list, *resolved_guardrails))) + combined: Final = merged_guardrail_entries(existing_guardrails_list, resolved_guardrails) data[metadata_variable_name]["guardrails"] = combined verbose_proxy_logger.debug("Policy engine: added guardrails to request metadata: %s", combined) @@ -4328,7 +4309,7 @@ def _policy_match_tags(request_data: Mapping[str, object], metadata_variable_nam async def add_guardrails_from_policy_engine( - data: dict, + data: dict[str, object], # mutable-ok: pops the request's policies and records the resolved guardrails in place metadata_variable_name: str, user_api_key_dict: UserAPIKeyAuth, ) -> None: @@ -4371,11 +4352,13 @@ async def add_guardrails_from_policy_engine( all_tags: Final = _policy_match_tags(data, metadata_variable_name) _team_alias: Final = user_api_key_dict.team_alias _key_alias: Final = user_api_key_dict.key_alias - context: Final = PolicyMatchContext( - team_alias=_team_alias if isinstance(_team_alias, str) else None, - key_alias=_key_alias if isinstance(_key_alias, str) else None, - model=data.get("model"), - tags=all_tags, + context: Final = PolicyMatchContext.model_validate( + { + "team_alias": _team_alias if isinstance(_team_alias, str) else None, + "key_alias": _key_alias if isinstance(_key_alias, str) else None, + "model": data.get("model"), + "tags": all_tags, + } ) verbose_proxy_logger.debug( diff --git a/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py b/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py index 1c5d270542f..bd8390d704e 100644 --- a/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py +++ b/litellm/proxy/pass_through_endpoints/passthrough_guardrails.py @@ -10,6 +10,7 @@ Handles guardrail execution for passthrough endpoints with: from typing import Any, Final from litellm._logging import verbose_proxy_logger +from litellm.integrations.custom_guardrail import merged_guardrail_entries from litellm.proxy._types import ( PassThroughGuardrailsConfig, PassThroughGuardrailSettings, @@ -205,7 +206,7 @@ class PassthroughGuardrailHandler: def collect_guardrails( user_api_key_dict: UserAPIKeyAuth, passthrough_guardrails_config: PassThroughGuardrailsConfigInput | None, - ) -> dict[str, bool] | None: + ) -> list[object] | None: """ Collect guardrails for a passthrough endpoint. @@ -225,7 +226,7 @@ class PassthroughGuardrailHandler: passthrough_guardrails_config: List or Dict of guardrail names/settings Returns: - Dict of guardrail names to run (format: {guardrail_name: True}), or None + Guardrail entries to run, each a name or a ``{name: settings}`` entry, or None """ from litellm.proxy.litellm_pre_call_utils import ( add_guardrails_from_key_or_team_metadata, @@ -242,15 +243,6 @@ class PassthroughGuardrailHandler: verbose_proxy_logger.debug("Passthrough guardrails config is empty, skipping") return None - # Passthrough is enabled - collect guardrails - guardrails_to_run: Final[dict[str, bool]] = {} - - # Add passthrough-specific guardrails - for guardrail_name in normalized_config: - guardrails_to_run[guardrail_name] = True - verbose_proxy_logger.debug("Added passthrough-specific guardrail: %s", guardrail_name) - - # Add org/team/key level guardrails using shared helper temp_data: Final[dict[str, Any]] = {"metadata": {}} add_guardrails_from_key_or_team_metadata( key_metadata=user_api_key_dict.metadata, @@ -258,20 +250,12 @@ class PassthroughGuardrailHandler: data=temp_data, metadata_variable_name="metadata", ) + inherited_guardrails: Final[list[object]] = temp_data["metadata"].get("guardrails", []) + guardrails_to_run: Final = merged_guardrail_entries(normalized_config, inherited_guardrails) - # Merge inherited guardrails into guardrails_to_run - inherited_guardrails: Final = temp_data["metadata"].get("guardrails", []) - for guardrail_name in inherited_guardrails: - if guardrail_name not in guardrails_to_run: - guardrails_to_run[guardrail_name] = True - verbose_proxy_logger.debug("Added inherited guardrail (key/team level): %s", guardrail_name) + verbose_proxy_logger.debug("Collected guardrails for passthrough endpoint: %s", guardrails_to_run) - verbose_proxy_logger.debug( - "Collected guardrails for passthrough endpoint: %s", - list(guardrails_to_run.keys()), - ) - - return guardrails_to_run if guardrails_to_run else None + return guardrails_to_run or None @staticmethod def get_field_targeted_text( diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index ef59deb2e5e..089cb75a795 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1,11 +1,13 @@ import asyncio import contextlib +import copy import json import uuid from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Callable, Mapping, Sequence from dataclasses import dataclass from enum import Enum from functools import partial +from itertools import chain from types import MappingProxyType from typing import ( TYPE_CHECKING, @@ -35,7 +37,7 @@ from litellm.anthropic_interface.exceptions.exception_mapping_utils import Anthr from litellm.caching.dual_cache import DualCache from litellm.constants import EMPTY_MAPPING, MANAGED_BACKGROUND_RESPONSE_METADATA_KEY, TERMINAL_RESPONSE_STATUSES from litellm.exceptions import EndpointNotSupportedByModelError -from litellm.integrations.custom_guardrail import ModifyResponseException +from litellm.integrations.custom_guardrail import ModifyResponseException, guardrail_request_data_with_streaming from litellm.litellm_core_utils.core_helpers import as_str_mapping, get_metadata_variable_name_from_kwargs from litellm.litellm_core_utils.internal_call_metadata import is_background_response, requests_background from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj @@ -47,6 +49,11 @@ from litellm.llms.openai.responses.guardrail_translation.handler import ( build_blocked_response, ) from litellm.proxy._types import * +from litellm.proxy.auth.auth_checks import ( + get_team_object, + requests_guardrail_modification, + require_guardrail_modification_permission, +) from litellm.proxy.auth.auth_utils import carries_client_side_overrides from litellm.proxy.auth.user_api_key_auth import ( UserAPIKeyAuth, @@ -57,6 +64,7 @@ from litellm.proxy.common_request_processing import ( ProxyBaseLLMRequestProcessing, cost_headers, create_response, + proxy_exception_from_route_error, resolve_counted_tool_fields, resolve_litellm_call_id, ) @@ -79,6 +87,7 @@ from litellm.proxy.hooks.responses_id_security import ResponsesIDSecurity, decry from litellm.proxy.litellm_pre_call_utils import ( _CLIENT_PRICING_METADATA_FIELDS, # pyright: ignore[reportPrivateUsage] # the HTTP path's own strip list _UNTRUSTED_METADATA_CONTROL_FIELDS, # pyright: ignore[reportPrivateUsage] # the HTTP path's own strip list + add_guardrails_from_policy_engine, is_chat_completions_body, ) from litellm.proxy.pass_through_endpoints.llm_provider_handlers.batch_attribution import ( @@ -89,7 +98,12 @@ from litellm.proxy.utils import count_request_input_tokens, handle_exception_on_ from litellm.responses.litellm_completion_transformation.compaction_items import messages_since_newest_compaction from litellm.responses.litellm_completion_transformation.count_request import MODEL_FILE_ID_MAPPING_FIELD from litellm.responses.retrieve_stream_query import RETRIEVE_INCLUDE_QUERY_KEYS, RETRIEVE_STREAM_QUERY_KEYS -from litellm.responses.streaming_iterator import GuardedResponsesWsFrame, ResponsesWsFrameGuard, spawn_logging_task +from litellm.responses.streaming_iterator import ( + GuardedResponsesWsFrame, + ResponsesWsFrameGuard, + ResponsesWsFrameRejection, + spawn_logging_task, +) from litellm.router import Router from litellm.types.llms.base import LiteLLMBaseModel from litellm.types.llms.openai import ( @@ -2149,11 +2163,15 @@ _PII_TOKENS: Final = TypeAdapter(dict[str, str]) _GUARDRAIL_INFORMATION_KEY: Final = "standard_logging_guardrail_information" _GUARDED_FRAME_FIELDS: Final = ("input", "instructions", "tools") _CLIENT_METADATA_KEYS: Final = ("metadata", "litellm_metadata") -# the guardrail list is merged from key, team and the admission frame; a later frame cannot shrink it _CLIENT_GUARDRAIL_SELECTORS: Final = ("guardrails", "guardrail_config", "policies") +_METADATA_GUARDRAIL_SELECTORS: Final = ("guardrails", "guardrail_config") +_PROXY_ONLY_FRAME_FIELDS: Final = (*_CLIENT_GUARDRAIL_SELECTORS, "include_guardrail_response") +_UNFORWARDED_FRAME_FIELDS: Final = (*_GUARDED_FRAME_FIELDS, *_PROXY_ONLY_FRAME_FIELDS, "litellm_metadata") +_OBJECTS: Final = TypeAdapter(list[object]) +_TURN_COPIED_CONTAINERS: Final = (list, dict) -class _PreCallGuardrailWalk(Protocol): +class _GuardrailWalk(Protocol): async def pre_call_hook( self, *, @@ -2163,15 +2181,49 @@ class _PreCallGuardrailWalk(Protocol): guardrails_only: bool, ) -> dict[str, object]: ... # mutable-ok: ProxyLogging.pre_call_hook returns the guardrails' own dict + async def during_call_hook( + self, + *, + data: dict[str, object], # mutable-ok: ProxyLogging.during_call_hook hands the guardrails a dict they edit + user_api_key_dict: UserAPIKeyAuth | None, + call_type: Literal["aresponses"], + guardrails_only: bool, + ) -> object: ... + + +async def _key_team_object(user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTable | None: + from litellm.proxy.proxy_server import prisma_client, proxy_logging_obj, user_api_key_cache + + team_id: Final = user_api_key_dict.team_id + if team_id is None or team_id == UI_TEAM_ID: + return None + team_from_key: Final = LiteLLM_TeamTable(team_id=team_id, metadata=user_api_key_dict.team_metadata) + if prisma_client is None: + return team_from_key + try: + return await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + parent_otel_span=user_api_key_dict.parent_otel_span, + proxy_logging_obj=proxy_logging_obj, + ) + except HTTPException: + return team_from_key + def _responses_ws_frame_guard( admitted: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, - proxy_logging_obj: _PreCallGuardrailWalk, + proxy_logging_obj: _GuardrailWalk, record_refused_turn: Callable[[tuple[Mapping[str, object], ...], Exception], None], + *, + team_object: Callable[[UserAPIKeyAuth], Awaitable[LiteLLM_TeamTable | None]] = _key_team_object, ) -> ResponsesWsFrameGuard: """Socket admission skips guardrails (the first frame's ``input`` is not part of the admitted request), so every - ``response.create`` frame gets the pre_call guardrail walk the HTTP ``/v1/responses`` path runs on its body.""" + ``response.create`` frame gets the pre_call and during_call guardrail walks the HTTP ``/v1/responses`` path runs on + its body. during_call runs before the frame is forwarded, since a forwarded frame cannot be recalled, and like HTTP, + which sends the model call beside during_call, the frame carries the pre_call output and not during_call rewrites.""" metadata_key: Final = get_metadata_variable_name_from_kwargs(admitted) admitted_bucket: Final = admitted.get(metadata_key) bucket: Final = { @@ -2186,32 +2238,57 @@ def _responses_ws_frame_guard( # a blocked frame is billed as its own failed request; only the first turn settles the admission reservation rejected_frame_auth: Final = user_api_key_dict.model_copy(update={"budget_reservation": None}) + async def rejection( + exception: Exception, + data: dict[str, object], # mutable-ok: billed + ) -> ResponsesWsFrameRejection: + await _bill_rejected_frame(rejected_frame_auth, exception, data) + proxy_exception: Final = proxy_exception_from_route_error(exception) + return ResponsesWsFrameRejection(proxy_exception.message, proxy_exception.provider_specific_fields) + async def guard(frame: str) -> GuardedResponsesWsFrame: frame_obj: Final = _WS_FRAME.validate_json(frame) nested: Final = frame_obj.get("response") - if isinstance(nested, dict) and any(key in frame_obj for key in _GUARDED_FRAME_FIELDS): - raise ValueError( - "response.create carries input, instructions or tools both at the top level and under response" - ) body: Final = _WS_FRAME.validate_python(nested) if isinstance(nested, dict) else frame_obj - frame_data: Final = { - **{ - key: _without_untrusted_metadata(key, value) - for key, value in body.items() - if key != "type" and key not in _CLIENT_GUARDRAIL_SELECTORS + selection_sources: Final = (frame_obj, body) if isinstance(nested, dict) else (body,) + turn_data: Final = guardrail_request_data_with_streaming( + { + **{ + key: _without_untrusted_metadata(key, value) + for key, value in body.items() + if key != "type" and key not in _PROXY_ONLY_FRAME_FIELDS + }, + "model": model, + "litellm_call_id": str(uuid.uuid4()), + metadata_key: {**bucket}, }, - "model": model, - "litellm_call_id": str(uuid.uuid4()), - # shallow: the bucket carries proxy-owned objects (the OTEL parent span) that cannot be deep-copied - metadata_key: {**bucket}, - } + is_streaming=True, + ) + try: + _reject_guarded_field_beside_nested_response(frame_obj, nested) + await _check_frame_may_modify_guardrails(selection_sources, user_api_key_dict, team_object) + frame_data: Final = { + **turn_data, + metadata_key: await _metadata_with_frame_guardrails( + bucket, selection_sources, model, metadata_key, user_api_key_dict + ), + } + except Exception as e: + raise await rejection(e, turn_data) from e try: guarded_data: Final = await proxy_logging_obj.pre_call_hook( user_api_key_dict=user_api_key_dict, data=frame_data, call_type="aresponses", guardrails_only=True ) + # a guardrail may delete a field (the Responses translation pops ``instructions``), so start without them + guarded_fields: Final = { + key: copy.deepcopy(guarded_data[key]) for key in _GUARDED_FRAME_FIELDS if key in guarded_data + } + pii_tokens: Final = _recorded_pii_tokens(guarded_data.get("metadata")) + await proxy_logging_obj.during_call_hook( + data=guarded_data, user_api_key_dict=user_api_key_dict, call_type="aresponses", guardrails_only=True + ) except Exception as e: - await _bill_rejected_frame(rejected_frame_auth, e, frame_data) - raise + raise await rejection(e, frame_data) from e if guarded_data.get("model") != model: rerouted: Final = ValueError( f"a guardrail rerouted this turn to {guarded_data.get('model')!r}; an open Responses WebSocket" @@ -2219,25 +2296,186 @@ def _responses_ws_frame_guard( ) await _bill_rejected_frame(rejected_frame_auth, rerouted, frame_data) raise rerouted - # a guardrail may delete a field (the Responses translation pops ``instructions``), so start without them - guarded_fields: Final = {key: guarded_data[key] for key in _GUARDED_FRAME_FIELDS if key in guarded_data} - kept: Final = {key: value for key, value in body.items() if key not in _GUARDED_FRAME_FIELDS} rewritten: Final = ( - {**frame_obj, "response": {**kept, **guarded_fields}} + {**_forwarded_frame_fields(frame_obj), "response": {**_forwarded_frame_fields(body), **guarded_fields}} if isinstance(nested, dict) - else {**kept, **guarded_fields} + else {**_forwarded_frame_fields(body), **guarded_fields} ) guardrail_information: Final = _recorded_guardrail_information(guarded_data.get(metadata_key)) return GuardedResponsesWsFrame( frame=json.dumps(rewritten), guardrail_information=guardrail_information, - pii_tokens=_recorded_pii_tokens(guarded_data.get("metadata")), + pii_tokens=pii_tokens, bill_refusal=partial(record_refused_turn, guardrail_information), ) return guard +def _reject_guarded_field_beside_nested_response(frame: Mapping[str, object], nested: object) -> None: + """Only the nested ``response`` body runs through the guardrails, so a guarded field beside it would go upstream + unguarded""" + stray: Final = ( + next((key for key in _GUARDED_FRAME_FIELDS if key in frame), None) if isinstance(nested, dict) else None + ) + if stray is not None: + raise ValueError(f"response.create sends {stray} at the top level next to response; send it inside response") + + +async def _check_frame_may_modify_guardrails( + sources: tuple[Mapping[str, object], ...], + user_api_key_dict: UserAPIKeyAuth, + team_object: Callable[[UserAPIKeyAuth], Awaitable[LiteLLM_TeamTable | None]], +) -> None: + """A frame that picks guardrails or policies needs the team's permission to modify guardrails""" + if any( + "policies" in source or requests_guardrail_modification(_without_client_metadata(source)) for source in sources + ): + require_guardrail_modification_permission(await team_object(user_api_key_dict)) + + +def _without_client_metadata(source: Mapping[str, object]) -> Mapping[str, object]: + metadata: Final = source.get("metadata") + return {**source, "metadata": _metadata_picks("metadata", metadata)} if isinstance(metadata, dict) else source + + +def _forwarded_frame_fields(source: Mapping[str, object]) -> dict[str, object]: + return { + key: _without_guardrail_selection(value) if key == "metadata" else value + for key, value in source.items() + if key not in _UNFORWARDED_FRAME_FIELDS + } + + +class _FrameGuardrailSelection(TypedDict, total=False): + guardrails: ReadOnly[Sequence[str | Mapping[str, Mapping[str, object]]]] + guardrail_config: ReadOnly[Mapping[str, object]] + policies: ReadOnly[Sequence[str]] + + +_FRAME_GUARDRAIL_SELECTION: Final = TypeAdapter(_FrameGuardrailSelection) + + +async def _metadata_with_frame_guardrails( + bucket: Mapping[str, object], + sources: tuple[Mapping[str, object], ...], + model: object, + metadata_key: str, + user_api_key_dict: UserAPIKeyAuth, +) -> Mapping[str, object]: + """A frame selects guardrails the way an HTTP ``/v1/responses`` body does, adding to the key, team and policy + guardrails the socket was admitted with, so a frame can never shrink that list. A nested frame selects at its top + level and in its ``response`` body, the body last so its ``guardrail_config`` wins""" + selections: Final = tuple(chain.from_iterable(_frame_guardrail_selections(source) for source in sources)) + picked_policies: Final = tuple(selection["policies"] for selection in selections if "policies" in selection) + metadata: Final = ( + await _metadata_with_frame_policies( + bucket, tuple(chain.from_iterable(picked_policies)), model, metadata_key, user_api_key_dict + ) + if picked_policies + else bucket + ) + selected_guardrails: Final = tuple(chain.from_iterable(selection.get("guardrails", ()) for selection in selections)) + selected_config: Final = dict( + chain.from_iterable( + selection["guardrail_config"].items() for selection in selections if "guardrail_config" in selection + ) + ) + existing_guardrails: Final = metadata.get("guardrails") + existing_config: Final = metadata.get("guardrail_config") + return { + **metadata, + **( + { + "guardrails": [ + *(_OBJECTS.validate_python(existing_guardrails) if isinstance(existing_guardrails, list) else ()), + *selected_guardrails, + ] + } + if selected_guardrails + else {} + ), + **( + { + "guardrail_config": { + **(_WS_FRAME.validate_python(existing_config) if isinstance(existing_config, dict) else {}), + **selected_config, + } + } + if selected_config + else {} + ), + } + + +def _frame_guardrail_selections(body: Mapping[str, object]) -> tuple[_FrameGuardrailSelection, ...]: + """Like HTTP ``/v1/responses``, a frame selects guardrails at its top level and in its ``metadata`` and + ``litellm_metadata``, the top level last so its ``guardrail_config`` wins""" + return ( + *( + _frame_guardrail_selection((key,), _metadata_picks(key, value), _METADATA_GUARDRAIL_SELECTORS) + for key, value in body.items() + if key in _CLIENT_METADATA_KEYS and isinstance(value, dict) + ), + _frame_guardrail_selection((), body, _CLIENT_GUARDRAIL_SELECTORS), + ) + + +def _metadata_picks(key: str, metadata: object) -> Mapping[str, object]: + """OpenAI ``metadata`` values are strings, so a string there is the client's own metadata and never a pick""" + validated: Final = _WS_FRAME.validate_python(metadata) + if key != "metadata": + return validated + return { + name: item + for name, item in validated.items() + if name not in _METADATA_GUARDRAIL_SELECTORS or not isinstance(item, str) + } + + +def _frame_guardrail_selection( + location: tuple[str, ...], source: Mapping[str, object], selectors: tuple[str, ...] +) -> _FrameGuardrailSelection: + try: + return _FRAME_GUARDRAIL_SELECTION.validate_python({key: source[key] for key in selectors if key in source}) + except ValidationError as e: + raise ValueError(f"response.create {first_validation_error(e, prefix=location)}") from e + + +def _without_guardrail_selection(metadata: object) -> object: + if not isinstance(metadata, dict): + return metadata + return { + name: item + for name, item in _WS_FRAME.validate_python(metadata).items() + if name not in _METADATA_GUARDRAIL_SELECTORS or isinstance(item, str) + } + + +async def _metadata_with_frame_policies( + bucket: Mapping[str, object], + policies: Sequence[str], + model: object, + metadata_key: str, + user_api_key_dict: UserAPIKeyAuth, +) -> Mapping[str, object]: + selection: Final = { + "model": model, + "policies": list(policies), + metadata_key: {key: _turn_copy(value) for key, value in bucket.items()}, + } + await add_guardrails_from_policy_engine( + data=selection, metadata_variable_name=metadata_key, user_api_key_dict=user_api_key_dict + ) + return _WS_FRAME.validate_python(selection[metadata_key]) + + +def _turn_copy(value: object) -> object: + """Shallow, since the bucket carries proxy-owned objects (the OTEL parent span) that cannot be copied, and the + policy engine writes its lists in place""" + return copy.copy(value) if type(value) in _TURN_COPIED_CONTAINERS else value + + def _without_untrusted_metadata(key: str, value: object) -> object: """Like the HTTP path, drop the proxy-owned ``user_api_key_*`` and guardrail control keys a client may forge.""" if key not in _CLIENT_METADATA_KEYS or not isinstance(value, dict): diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index cc72dd03a40..0b0698faa9a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -147,6 +147,7 @@ from litellm.exceptions import ( from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, + merged_guardrail_entries, ) from litellm.integrations.custom_logger import CustomLogger from litellm.integrations.prometheus import PrometheusLogger @@ -3506,6 +3507,7 @@ class ProxyLogging: data: dict, user_api_key_dict: UserAPIKeyAuth | None, call_type: CallTypesLiteral, + guardrails_only: bool = False, ): caps: Final = ProxyLogging.callback_capabilities() if not caps.has_guardrail and not caps.has_moderation_override: @@ -3517,6 +3519,7 @@ class ProxyLogging: if ( isinstance(callback, CustomLogger) and not isinstance(callback, CustomGuardrail) + and (not guardrails_only or callback.enforces_request_content) and _overrides_moderation_hook(callback) and user_api_key_dict is not None ): @@ -9100,8 +9103,7 @@ def _merge_guardrails_with_existing(data: dict, model_level_guardrails: object) if not isinstance(model_level_guardrails, list): model_level_guardrails = [model_level_guardrails] if model_level_guardrails else [] - # Combine existing and model-level guardrails - metadata["guardrails"] = list(set(existing_guardrails + model_level_guardrails)) + metadata["guardrails"] = merged_guardrail_entries(existing_guardrails, model_level_guardrails) return modified_data diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index d42f1af380e..a8e55a5e94d 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -2277,6 +2277,29 @@ class GuardedResponsesWsFrame: ResponsesWsFrameGuard: TypeAlias = Callable[[str], Awaitable[GuardedResponsesWsFrame]] +class ResponsesWsFrameRejection(ValueError): + def __init__(self, message: str, provider_specific_fields: Mapping[str, object] | None = None) -> None: + super().__init__(message) + self.provider_specific_fields: Final = provider_specific_fields + + +def _frame_rejection_event(rejection: Exception) -> str: + provider_fields: Final = ( + rejection.provider_specific_fields if isinstance(rejection, ResponsesWsFrameRejection) else None + ) + return json.dumps( + { + "type": "error", + "error": { + "type": "invalid_request_error", + "message": str(rejection), + **({"provider_specific_fields": dict(provider_fields)} if provider_fields else {}), + }, + }, + default=str, + ) + + async def _reject_guarded_turn( websocket: ResponsesClientWebSocket, guarded: GuardedResponsesWsFrame, rejection: Exception ) -> None: @@ -2562,7 +2585,8 @@ class ResponsesWebSocketStreaming: return await self.frame_guard(message) except Exception as e: # noqa: BLE001 # any guardrail rejection maps to one client-facing error event verbose_logger.info("Responses WebSocket: pre_call guardrail rejected a response.create frame: %s", e) - await self._send_error_event("invalid_request_error", str(e)) + with contextlib.suppress(Exception): + await self.websocket.send_text(_frame_rejection_event(e)) return None async def _forward_guarded_frame( @@ -2995,14 +3019,6 @@ class ResponsesWebSocketStreaming: except Exception: # noqa: BLE001, S110 # client may already be gone pass - async def _send_error_event(self, error_type: str, message: str) -> None: - try: - await self.websocket.send_text( - json.dumps({"type": "error", "error": {"type": error_type, "message": message}}) - ) - except Exception: # noqa: BLE001, S110 # client may already be gone - pass - async def client_to_backend(self) -> None: """Forward response.create events from client to backend.""" try: @@ -3416,7 +3432,7 @@ class ManagedResponsesWebSocketHandler: return await self.frame_guard(raw_message) except Exception as e: # noqa: BLE001 # any guardrail rejection maps to one client-facing error event verbose_logger.info("ManagedResponsesWS: pre_call guardrail rejected a response.create frame: %s", e) - await self._send_error(str(e), error_type="invalid_request_error") + await self._send_frame(_frame_rejection_event(e)) return None @staticmethod diff --git a/tests/unit/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py index 0483888533e..c32c1b71f75 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/tests/unit/integrations/test_custom_guardrail.py @@ -3800,3 +3800,16 @@ def test_default_on_guardrail() -> None: ) is True ) + + +def test_merged_guardrail_entries_names_each_guardrail_once_preferring_its_first_entry_with_settings(): + from litellm.integrations.custom_guardrail import merged_guardrail_entries + + assert merged_guardrail_entries(["a"], [{"a": {}, "b": {}}, {"a": {"x": 1}}], ["b", "c", "c"]) == [ + {"a": {}, "b": {}}, + "c", + ] + assert merged_guardrail_entries(["pii", {"pii": {"extra_body": {"threshold": 1}}}], ["model-guard"]) == [ + {"pii": {"extra_body": {"threshold": 1}}}, + "model-guard", + ] diff --git a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py index 6b59f6d40dc..dc7bd8cb312 100644 --- a/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py +++ b/tests/unit/proxy/auth/test_auth_checks_object_access_and_lookup.py @@ -5884,7 +5884,7 @@ async def test_team_member_budget_check_reads_from_spend_counter(): class TestGuardrailModificationCheck: - """Defense-in-depth: `_guardrail_modification_check` must 403 when the + """Defense-in-depth: `guardrail_modification_check` must 403 when the caller's metadata attempts to modify any guardrail-related key and the team lacks the `modify_guardrails` permission. Checks both the historically-covered `guardrails` list and the bypass toggles that @@ -5892,11 +5892,11 @@ class TestGuardrailModificationCheck: """ def _call(self, request_body): - from litellm.proxy.auth.auth_checks import _guardrail_modification_check + from litellm.proxy.auth.auth_checks import guardrail_modification_check team_object = MagicMock() team_object.metadata = {} # no permission - return _guardrail_modification_check(request_body=request_body, team_object=team_object) + return guardrail_modification_check(request_body=request_body, team_object=team_object) def test_noop_when_no_guardrail_keys_present(self): # no-op — should return silently diff --git a/tests/unit/proxy/pass_through_endpoints/test_passthrough_guardrails.py b/tests/unit/proxy/pass_through_endpoints/test_passthrough_guardrails.py index 40403de846f..ef19482b30a 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_passthrough_guardrails.py +++ b/tests/unit/proxy/pass_through_endpoints/test_passthrough_guardrails.py @@ -73,6 +73,25 @@ class TestPassthroughGuardrailHandlerGetGuardrailNames: assert set(result) == {"pii-detection", "content-moderation"} +class TestPassthroughGuardrailHandlerCollectGuardrails: + def test_inherits_a_key_guardrail_with_its_settings(self, monkeypatch: pytest.MonkeyPatch): + from litellm.integrations.custom_guardrail import CustomGuardrail + from litellm.proxy._types import UserAPIKeyAuth + + monkeypatch.setattr("litellm.proxy.proxy_server.premium_user", True) + key = UserAPIKeyAuth( + api_key="hashed", metadata={"guardrails": [{"pii": {"extra_body": {"threshold": 1}}}, "toxic"]} + ) + + guardrails = PassthroughGuardrailHandler.collect_guardrails(key, ["passthrough-guard", "toxic"]) + + assert guardrails == ["passthrough-guard", "toxic", {"pii": {"extra_body": {"threshold": 1}}}] + request_data = {"metadata": {"guardrails": guardrails}} + assert CustomGuardrail(guardrail_name="pii").get_guardrail_dynamic_request_body_params(request_data) == { + "threshold": 1 + } + + class TestPassthroughGuardrailHandlerNormalizeConfig: """Tests for PassthroughGuardrailHandler.normalize_config method.""" diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index 1e4dded801e..07f85cb9628 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -7,6 +7,7 @@ import contextlib import dataclasses import datetime import json +import re import unittest from collections.abc import Callable, Mapping from typing import Any, Final, Literal, cast @@ -22,12 +23,17 @@ from fastapi.testclient import TestClient from httpx import Response import litellm +from litellm.caching.caching import DualCache from litellm.constants import EMPTY_MAPPING -from litellm.proxy._types import ProxyException +from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.proxy._types import ProxyException, UserAPIKeyAuth from litellm.proxy.proxy_server import app +from litellm.proxy.utils import ProxyLogging from litellm.responses.streaming_iterator import ResponsesAPIStreamingIterator +from litellm.types.guardrails import GuardrailEventHooks from litellm.types.llms.openai import ResponseAPIUsage, ResponsesAPIResponse from litellm.types.router import RouterErrors +from litellm.types.utils import CallTypesLiteral from tests._master_key import MASTER_KEY @@ -2213,6 +2219,59 @@ class TestResponsesWSFirstFrameValidation: ws.close.assert_awaited_once_with(code=1011, reason="Internal server error") +class _FrameGuardrail(CustomGuardrail): + def __init__( + self, + name: str, + event_hook: GuardrailEventHooks = GuardrailEventHooks.pre_call, + rewrites: bool = False, + blocks: bool = False, + ): + super().__init__(guardrail_name=name, event_hook=event_hook, default_on=False) + self.inputs: list[object] = [] + self.configs: list[object] = [] + self.rewrites = rewrites + self.blocks = blocks + + async def async_pre_call_hook( + self, + user_api_key_dict: UserAPIKeyAuth, + cache: DualCache, + data: dict[str, object], + call_type: CallTypesLiteral, + ) -> dict[str, object]: + return self._scan(data) + + async def async_moderation_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, call_type: CallTypesLiteral + ) -> None: + data["input"] = self._scan(data)["input"] + + def _scan(self, data: dict[str, object]) -> dict[str, object]: + self.inputs.append(data["input"]) + self.configs.append(cast(dict[str, object], data["litellm_metadata"]).get("guardrail_config")) + if self.blocks: + raise HTTPException( + status_code=400, + detail={"error": f"Content blocked by {self.guardrail_name}", "guardrail_name": self.guardrail_name}, + ) + return {**data, "input": f"{data['input']} [{self.guardrail_name}]"} if self.rewrites else data + + +@contextlib.contextmanager +def _guardrails_on(*guardrails: CustomGuardrail): + with patch.object(litellm, "callbacks", list(guardrails)): + ProxyLogging._callback_capabilities_cache.clear() + try: + yield ProxyLogging(user_api_key_cache=DualCache()) + finally: + ProxyLogging._callback_capabilities_cache.clear() + + +def _key() -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="hashed") + + class TestResponsesWSFirstFrameModelAuth: @pytest.mark.asyncio async def test_endpoint_enforces_auth_after_model_from_first_frame(self): @@ -2344,7 +2403,7 @@ class TestResponsesWSFirstFrameModelAuth: data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_name": "turn"}] return {**data, "input": data["input"] + " [guarded]"} - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) admitted = { "model": "gpt-4o-mini", @@ -2375,7 +2434,7 @@ class TestResponsesWSFirstFrameModelAuth: data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}] raise ValueError("blocked") - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) proxy_logging_obj.post_call_failure_hook = AsyncMock() spend_logger = MagicMock(spec=ProxyDBLogger) @@ -2413,7 +2472,7 @@ class TestResponsesWSFirstFrameModelAuth: async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): return {**data, "input": data["input"] + " [guarded]", "tools": data["tools"][:1]} - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) guard = _responses_ws_frame_guard( {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() @@ -2432,10 +2491,10 @@ class TestResponsesWSFirstFrameModelAuth: async def test_frame_guard_strips_the_proxy_owned_metadata_keys_a_client_sends(self): from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) guard = _responses_ws_frame_guard( - {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() + {"model": "gpt-4o-mini", "litellm_metadata": {}}, _key(), proxy_logging_obj, MagicMock() ) metadata = { "user_api_key_metadata": {"disable_global_guardrails": True}, @@ -2454,17 +2513,18 @@ class TestResponsesWSFirstFrameModelAuth: [ {"type": "response.create", "input": "forbidden", "guardrails": []}, {"type": "response.create", "input": "forbidden", "metadata": {"guardrails": []}}, + {"type": "response.create", "input": "forbidden", "litellm_metadata": {"guardrails": []}}, {"type": "response.create", "response": {"input": "forbidden", "guardrails": [], "policies": []}}, ], ) async def test_frame_guard_keeps_the_admitted_guardrail_list_when_a_frame_tries_to_replace_it(self, frame): from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) guard = _responses_ws_frame_guard( {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}}, - MagicMock(), + _key(), proxy_logging_obj, MagicMock(), ) @@ -2476,13 +2536,258 @@ class TestResponsesWSFirstFrameModelAuth: assert "guardrails" not in (scanned.get("metadata") or {}) assert scanned["litellm_metadata"]["guardrails"] == ["key-required"] + @pytest.mark.asyncio + @pytest.mark.parametrize( + "selected", + [ + {"guardrails": ["frame-pick"]}, + {"guardrails": [{"frame-pick": {"extra_body": {"threshold": 0.8}}}]}, + {"metadata": {"guardrails": ["frame-pick"], "tag": "kept"}}, + {"litellm_metadata": {"guardrails": ["frame-pick"]}}, + {"response": {"guardrails": ["frame-pick"]}}, + ], + ids=["top-level", "with-extra-body", "metadata", "litellm-metadata", "nested"], + ) + async def test_frame_guard_runs_the_guardrails_a_frame_selects_on_top_of_the_admitted_ones(self, selected: dict[str, object]): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + key_required = _FrameGuardrail("key-required") + frame_pick = _FrameGuardrail("frame-pick", rewrites=True) + unselected = _FrameGuardrail("unselected") + admitted = {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}} + nested = "response" in selected + body = {"input": "hi", "store": False, **(selected["response"] if nested else selected)} + frame = {"type": "response.create", "response": body} if nested else {"type": "response.create", **body} + + with _guardrails_on(key_required, frame_pick, unselected) as proxy_logging_obj: + guard = _responses_ws_frame_guard(admitted, _key(), proxy_logging_obj, MagicMock()) + selected_turn = await guard(json.dumps(frame)) + plain_turn = await guard(json.dumps({"type": "response.create", "input": "again"})) + + kept_metadata = {"metadata": {"tag": "kept"}} if "metadata" in selected else {} + forwarded = {"input": "hi [frame-pick]", "store": False, **kept_metadata} + assert json.loads(selected_turn.frame) == ( + {"type": "response.create", "response": forwarded} if nested else {"type": "response.create", **forwarded} + ) + assert json.loads(plain_turn.frame) == {"type": "response.create", "input": "again"} + assert key_required.inputs == ["hi", "again"] + assert frame_pick.inputs == ["hi"] + assert unselected.inputs == [] + assert admitted["litellm_metadata"] == {"guardrails": ["key-required"]} + + @pytest.mark.asyncio + async def test_frame_guard_forwards_string_metadata_named_like_a_pick_without_running_it(self): + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + key_required = _FrameGuardrail("key-required") + frame_pick = _FrameGuardrail("frame-pick") + metadata = {"guardrails": "frame-pick", "guardrail_config": "strict", "trace": "t1"} + + async def locked_team(user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable(team_id="team-1", metadata={"guardrails": {"modify_guardrails": False}}) + + with _guardrails_on(key_required, frame_pick) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}}, + _key(), + proxy_logging_obj, + MagicMock(), + team_object=locked_team, + ) + guarded = await guard(json.dumps({"type": "response.create", "input": "hi", "metadata": metadata})) + + assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi", "metadata": metadata} + assert key_required.inputs == ["hi"] + assert frame_pick.inputs == [] + + @pytest.mark.asyncio + async def test_frame_guard_hands_a_selected_guardrail_its_config_with_the_top_level_winning(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + frame_pick = _FrameGuardrail("frame-pick") + frame = { + "type": "response.create", + "input": "hi", + "metadata": {"guardrail_config": {"frame-pick": {"threshold": 1}, "metadata-only": {"threshold": 3}}}, + "guardrails": ["frame-pick"], + "guardrail_config": {"frame-pick": {"threshold": 2}}, + } + + with _guardrails_on(frame_pick) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"guardrail_config": {"key-required": {"threshold": 0}}}}, + _key(), + proxy_logging_obj, + MagicMock(), + ) + guarded = await guard(json.dumps(frame)) + + assert frame_pick.configs == [ + {"key-required": {"threshold": 0}, "frame-pick": {"threshold": 2}, "metadata-only": {"threshold": 3}} + ] + assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi", "metadata": {}} + + @pytest.mark.asyncio + async def test_frame_guard_runs_streaming_scoped_guardrails_and_skips_non_streaming_ones_on_a_socket_turn(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + streaming_only = _FrameGuardrail("streaming-only") + streaming_only.apply_stream_scope("streaming") + streaming_during_call = _FrameGuardrail("streaming-during-call", event_hook=GuardrailEventHooks.during_call) + streaming_during_call.apply_stream_scope("streaming") + non_streaming_only = _FrameGuardrail("non-streaming-only") + non_streaming_only.apply_stream_scope("non_streaming") + frame = { + "type": "response.create", + "input": "hi", + "guardrails": ["streaming-only", "streaming-during-call", "non-streaming-only"], + } + + with _guardrails_on(streaming_only, streaming_during_call, non_streaming_only) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, _key(), proxy_logging_obj, MagicMock() + ) + guarded = await guard(json.dumps(frame)) + + assert streaming_only.inputs == ["hi"] + assert streaming_during_call.inputs == ["hi"] + assert non_streaming_only.inputs == [] + assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi"} + + @pytest.mark.asyncio + async def test_frame_guard_blocks_a_turn_with_the_message_a_selected_during_call_guardrail_raises(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + during_pick = _FrameGuardrail("during-pick", event_hook=GuardrailEventHooks.during_call, blocks=True) + + with _guardrails_on(during_pick) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, _key(), proxy_logging_obj, MagicMock() + ) + unselected_turn = await guard(json.dumps({"type": "response.create", "input": "plain"})) + with pytest.raises(ValueError, match=r"^Content blocked by during-pick$"): + await guard(json.dumps({"type": "response.create", "input": "hi", "guardrails": ["during-pick"]})) + + assert json.loads(unselected_turn.frame) == {"type": "response.create", "input": "plain"} + assert during_pick.inputs == ["hi"] + + @pytest.mark.asyncio + async def test_frame_guard_runs_only_the_moderation_hooks_that_enforce_request_content(self): + from litellm.integrations.custom_logger import CustomLogger + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + class _ModerationHook(CustomLogger): + def __init__(self, enforces_request_content: bool) -> None: + super().__init__() + self.enforces_request_content = enforces_request_content + self.inputs: list[object] = [] + + async def async_moderation_hook( + self, data: dict[str, object], user_api_key_dict: UserAPIKeyAuth, call_type: CallTypesLiteral + ) -> None: + self.inputs.append(data.get("input")) + + vendor_moderation = _ModerationHook(enforces_request_content=False) + content_enforcer = _ModerationHook(enforces_request_content=True) + + with _guardrails_on(vendor_moderation, content_enforcer) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, _key(), proxy_logging_obj, MagicMock() + ) + await guard(json.dumps({"type": "response.create", "input": "hi"})) + + assert vendor_moderation.inputs == [] + assert content_enforcer.inputs == ["hi"] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("selection", "message"), + [ + ( + {"guardrails": "frame-pick"}, + "response.create guardrails: 'str' instances are not allowed as a Sequence value", + ), + ({"guardrails": [123]}, "response.create guardrails.0.str: Input should be a valid string"), + ({"guardrail_config": []}, "response.create guardrail_config: Input should be a valid dictionary"), + ( + {"policies": "frame-policy"}, + "response.create policies: 'str' instances are not allowed as a Sequence value", + ), + ({"policies": [1]}, "response.create policies.0: Input should be a valid string"), + ( + {"litellm_metadata": {"guardrails": "frame-pick"}}, + "response.create litellm_metadata.guardrails: 'str' instances are not allowed as a Sequence value", + ), + ( + {"litellm_metadata": {"guardrail_config": []}}, + "response.create litellm_metadata.guardrail_config: Input should be a valid dictionary", + ), + ], + ) + async def test_frame_guard_refuses_and_bills_a_frame_whose_guardrail_selection_is_malformed( + self, selection: dict[str, object], message: str + ): + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + key_required = _FrameGuardrail("key-required") + spend_logger = MagicMock(spec=ProxyDBLogger) + spend_logger.async_post_call_failure_hook = AsyncMock() + + with _guardrails_on(key_required, spend_logger) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}}, + _key(), + proxy_logging_obj, + MagicMock(), + ) + with pytest.raises(ValueError, match=f"^{re.escape(message)}$"): + await guard(json.dumps({"type": "response.create", "input": "hi", **selection})) + + assert key_required.inputs == [] + spend_logger.async_post_call_failure_hook.assert_awaited_once() + + @pytest.mark.asyncio + async def test_frame_guard_runs_the_guardrails_of_the_policies_a_frame_selects(self): + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + from litellm.types.proxy.policy_engine import Policy, PolicyGuardrails + + key_required = _FrameGuardrail("key-required") + policy_pick = _FrameGuardrail("policy-pick") + policy_registry = get_policy_registry() + policy_registry._policies = {} + policy_registry._policies_by_id = {} + policy_registry._initialized = False + try: + policy_registry.add_policy("frame-policy", Policy(guardrails=PolicyGuardrails(add=["policy-pick"]))) + with _guardrails_on(key_required, policy_pick) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}}, + _key(), + proxy_logging_obj, + MagicMock(), + ) + guarded = await guard( + json.dumps({"type": "response.create", "input": "hi", "policies": ["frame-policy"]}) + ) + finally: + policy_registry._policies = {} + policy_registry._policies_by_id = {} + policy_registry._initialized = False + + assert key_required.inputs == ["hi"] + assert policy_pick.inputs == ["hi"] + assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi"} + @pytest.mark.asyncio async def test_frame_guard_copies_an_admitted_bucket_that_holds_the_otel_parent_span(self): import threading from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) span = threading.RLock() guard = _responses_ws_frame_guard( @@ -2510,7 +2815,7 @@ class TestResponsesWSFirstFrameModelAuth: data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}] return {**data, "model": "onprem-llama"} - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) spend_logger = MagicMock(spec=ProxyDBLogger) spend_logger.async_post_call_failure_hook = AsyncMock() @@ -2543,7 +2848,7 @@ class TestResponsesWSFirstFrameModelAuth: data["litellm_metadata"]["standard_logging_guardrail_information"] = [{"guardrail_cost": 0.5}] raise ValueError("blocked") - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) spend_logger = MagicMock(spec=ProxyDBLogger) spend_logger.async_post_call_failure_hook = AsyncMock() @@ -2571,7 +2876,7 @@ class TestResponsesWSFirstFrameModelAuth: async def pre_call_hook(user_api_key_dict, data, call_type, guardrails_only): return {key: value for key, value in data.items() if key != "instructions"} - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) guard = _responses_ws_frame_guard( {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() @@ -2596,7 +2901,7 @@ class TestResponsesWSFirstFrameModelAuth: data["metadata"]["pii_tokens"] = {"": "alice@example.com"} return {**data, "input": "email please"} - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) guard = _responses_ws_frame_guard( {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() @@ -2617,19 +2922,279 @@ class TestResponsesWSFirstFrameModelAuth: assert dict(guarded.pii_tokens) == {"": "alice@example.com"} @pytest.mark.asyncio - async def test_frame_guard_rejects_a_frame_that_carries_input_at_both_levels(self): + @pytest.mark.parametrize("field", ["input", "instructions", "tools"]) + async def test_frame_guard_rejects_and_bills_a_guarded_field_beside_a_nested_response(self, field: str): + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) + proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) + spend_logger = MagicMock(spec=ProxyDBLogger) + spend_logger.async_post_call_failure_hook = AsyncMock() + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() + ) + message = f"response.create sends {field} at the top level next to response; send it inside response" + + with ( + patch.object(litellm, "callbacks", [spend_logger]), + pytest.raises(ValueError, match=f"^{re.escape(message)}$"), + ): + await guard(json.dumps({"type": "response.create", field: "unguarded", "response": {"input": "hi"}})) + + proxy_logging_obj.pre_call_hook.assert_not_awaited() + assert spend_logger.async_post_call_failure_hook.await_args.kwargs["request_data"]["input"] == "hi" + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "top_level", + [ + {"guardrails": ["frame-pick"]}, + {"metadata": {"guardrails": ["frame-pick"], "tag": "kept"}}, + {"litellm_metadata": {"guardrails": ["frame-pick"]}}, + ], + ids=["top-level", "metadata", "litellm-metadata"], + ) + async def test_frame_guard_runs_the_guardrails_a_nested_frame_selects_beside_its_response(self, top_level: dict[str, object]): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + key_required = _FrameGuardrail("key-required") + frame_pick = _FrameGuardrail("frame-pick", rewrites=True) + unselected = _FrameGuardrail("unselected") + + with _guardrails_on(key_required, frame_pick, unselected) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}}, + _key(), + proxy_logging_obj, + MagicMock(), + ) + guarded = await guard(json.dumps({"type": "response.create", **top_level, "response": {"input": "hi"}})) + + kept_metadata = {"metadata": {"tag": "kept"}} if "metadata" in top_level else {} + assert json.loads(guarded.frame) == { + "type": "response.create", + **kept_metadata, + "response": {"input": "hi [frame-pick]"}, + } + assert key_required.inputs == ["hi"] + assert frame_pick.inputs == ["hi"] + assert unselected.inputs == [] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "frame", + [ + {"type": "response.create", "input": "hi", "guardrails": ["frame-pick"]}, + {"type": "response.create", "input": "hi", "metadata": {"guardrails": ["frame-pick"]}}, + {"type": "response.create", "input": "hi", "metadata": {"disable_global_guardrails": True}}, + {"type": "response.create", "input": "hi", "metadata": {"disable_global_guardrails": "true"}}, + {"type": "response.create", "response": {"input": "hi", "guardrails": ["frame-pick"]}}, + {"type": "response.create", "litellm_metadata": {"guardrails": ["frame-pick"]}, "response": {"input": "hi"}}, + {"type": "response.create", "input": "hi", "policies": ["frame-policy"]}, + {"type": "response.create", "response": {"input": "hi", "policies": ["frame-policy"]}}, + ], + ids=[ + "top-level", + "metadata", + "global-toggle", + "string-global-toggle", + "nested", + "beside-nested", + "policies", + "nested-policies", + ], + ) + async def test_frame_guard_refuses_guardrail_and_policy_picks_from_a_team_without_permission(self, frame: dict[str, object]): + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + key_required = _FrameGuardrail("key-required") + frame_pick = _FrameGuardrail("frame-pick") + spend_logger = MagicMock(spec=ProxyDBLogger) + spend_logger.async_post_call_failure_hook = AsyncMock() + team_lookups: list[UserAPIKeyAuth] = [] + key = _key() + + async def locked_team(user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTable: + team_lookups.append(user_api_key_dict) + return LiteLLM_TeamTable(team_id="team-1", metadata={"guardrails": {"modify_guardrails": False}}) + + with _guardrails_on(key_required, frame_pick, spend_logger) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {"guardrails": ["key-required"]}}, + key, + proxy_logging_obj, + MagicMock(), + team_object=locked_team, + ) + plain_turn = await guard(json.dumps({"type": "response.create", "input": "plain"})) + with pytest.raises(ValueError, match=r"^Your team does not have permission to modify guardrails\.$"): + await guard(json.dumps(frame)) + + assert json.loads(plain_turn.frame) == {"type": "response.create", "input": "plain"} + assert key_required.inputs == ["plain"] + assert frame_pick.inputs == [] + assert team_lookups == [key] + spend_logger.async_post_call_failure_hook.assert_awaited_once() + + @pytest.mark.asyncio + async def test_frame_guard_runs_a_frame_s_picks_for_a_team_allowed_to_modify_guardrails(self): + from litellm.proxy._types import LiteLLM_TeamTable + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + frame_pick = _FrameGuardrail("frame-pick") + + async def open_team(user_api_key_dict: UserAPIKeyAuth) -> LiteLLM_TeamTable: + return LiteLLM_TeamTable(team_id="team-1", metadata={"guardrails": {"modify_guardrails": True}}) + + with _guardrails_on(frame_pick) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, + _key(), + proxy_logging_obj, + MagicMock(), + team_object=open_team, + ) + await guard(json.dumps({"type": "response.create", "input": "hi", "guardrails": ["frame-pick"]})) + + assert frame_pick.inputs == ["hi"] + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("modify_guardrails", "picked_inputs"), + [(True, ["hi"]), (False, [])], + ids=["open", "locked"], + ) + @pytest.mark.parametrize( + "team_lookup", + [ + lambda: patch( + "litellm.proxy.response_api_endpoints.endpoints.get_team_object", + AsyncMock(side_effect=HTTPException(status_code=404, detail={"error": "Team doesn't exist"})), + ), + lambda: patch("litellm.proxy.proxy_server.prisma_client", None), + ], + ids=["team-row-gone", "no-database"], + ) + async def test_frame_guard_reads_the_key_s_team_metadata_when_the_team_cannot_be_loaded( + self, + modify_guardrails: bool, + picked_inputs: list[str], + team_lookup: Callable[[], contextlib.AbstractContextManager[object]], + ): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + frame_pick = _FrameGuardrail("frame-pick") + key = UserAPIKeyAuth( + api_key="hashed", + team_id="deleted-team", + team_metadata={"guardrails": {"modify_guardrails": modify_guardrails}}, + ) + frame = json.dumps({"type": "response.create", "input": "hi", "guardrails": ["frame-pick"]}) + + with _guardrails_on(frame_pick) as proxy_logging_obj, team_lookup(): + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, key, proxy_logging_obj, MagicMock() + ) + if modify_guardrails: + await guard(frame) + else: + with pytest.raises(ValueError, match=r"^Your team does not have permission to modify guardrails\.$"): + await guard(frame) + + assert frame_pick.inputs == picked_inputs + + @pytest.mark.asyncio + async def test_frame_guard_forwards_the_pre_call_output_and_not_a_during_call_rewrite_like_http(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + pre_pick = _FrameGuardrail("pre-pick", rewrites=True) + during_pick = _FrameGuardrail("during-pick", event_hook=GuardrailEventHooks.during_call, rewrites=True) + + with _guardrails_on(pre_pick, during_pick) as proxy_logging_obj: + guard = _responses_ws_frame_guard( + {"model": "gpt-4o-mini", "litellm_metadata": {}}, _key(), proxy_logging_obj, MagicMock() + ) + guarded = await guard( + json.dumps({"type": "response.create", "input": "hi", "guardrails": ["pre-pick", "during-pick"]}) + ) + + assert during_pick.inputs == ["hi [pre-pick]"] + assert json.loads(guarded.frame) == {"type": "response.create", "input": "hi [pre-pick]"} + + @pytest.mark.asyncio + async def test_frame_guard_forwards_a_nested_turn_without_top_level_fields_that_pick_no_guardrail(self): + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=lambda **kwargs: kwargs["data"]) guard = _responses_ws_frame_guard( {"model": "gpt-4o-mini", "litellm_metadata": {}}, MagicMock(), proxy_logging_obj, MagicMock() ) - with patch.object(litellm, "callbacks", []), pytest.raises(ValueError, match="both at the top level"): - await guard(json.dumps({"type": "response.create", "input": "forbidden", "response": {}})) + with patch.object(litellm, "callbacks", []): + guarded = await guard( + json.dumps( + { + "type": "response.create", + "include_guardrail_response": True, + "guardrail_config": {"frame-pick": {"threshold": 1}}, + "metadata": {"trace": "t1"}, + "response": {"input": "hi"}, + } + ) + ) - proxy_logging_obj.pre_call_hook.assert_not_awaited() + assert json.loads(guarded.frame) == { + "type": "response.create", + "metadata": {"trace": "t1"}, + "response": {"input": "hi"}, + } + + @pytest.mark.asyncio + async def test_a_blocked_socket_turn_answers_with_the_guardrail_fields_http_returns(self): + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard + from litellm.responses.streaming_iterator import ManagedResponsesWebSocketHandler + + frame_pick = _FrameGuardrail("frame-pick", blocks=True) + websocket = MagicMock(send_text=AsyncMock()) + + with _guardrails_on(frame_pick) as proxy_logging_obj: + handler = ManagedResponsesWebSocketHandler( + websocket=websocket, + model="anthropic/claude-haiku-5-5", + logging_obj=Logging( + model="anthropic/claude-haiku-5-5", + messages=[], + stream=True, + call_type="aresponses", + start_time=0, + litellm_call_id="test-id", + function_id="test-func", + ), + frame_guard=_responses_ws_frame_guard( + {"model": "claude-managed", "litellm_metadata": {}}, _key(), proxy_logging_obj, MagicMock() + ), + ) + await handler._process_response_create( + json.dumps({"type": "response.create", "input": "hi", "guardrails": ["frame-pick"]}) + ) + + assert json.loads(websocket.send_text.await_args.args[0]) == { + "type": "error", + "error": { + "type": "invalid_request_error", + "message": "Content blocked by frame-pick", + "provider_specific_fields": { + "error": "Content blocked by frame-pick", + "guardrail_name": "frame-pick", + "guardrail_mode": "pre_call", + }, + }, + } @pytest.mark.asyncio async def test_frame_guard_keeps_the_guardrail_error_when_billing_the_rejected_frame_fails(self): @@ -2637,7 +3202,7 @@ class TestResponsesWSFirstFrameModelAuth: from litellm.proxy.hooks.proxy_track_cost_callback import ProxyDBLogger from litellm.proxy.response_api_endpoints.endpoints import _responses_ws_frame_guard - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=ValueError("blocked")) spend_logger = MagicMock(spec=ProxyDBLogger) spend_logger.async_post_call_failure_hook = AsyncMock(side_effect=RuntimeError("db down")) @@ -2679,7 +3244,7 @@ class TestResponsesWSFirstFrameModelAuth: async def fake_llm_call(): return failure if provider_rejected else None - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) proxy_logging_obj.arelease_responses_websocket_slot = AsyncMock(return_value=None) user_api_key_dict = MagicMock() @@ -2768,7 +3333,7 @@ class TestResponsesWSFirstFrameModelAuth: return relay() row_written = asyncio.Event() - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock(side_effect=pre_call_hook) proxy_logging_obj.post_call_failure_hook = AsyncMock(side_effect=lambda **kwargs: row_written.set()) proxy_logging_obj.arelease_responses_websocket_slot = AsyncMock(return_value=None) @@ -2858,7 +3423,7 @@ class TestResponsesWSFirstFrameModelAuth: return relay() - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock( side_effect=lambda user_api_key_dict, data, call_type, guardrails_only: data ) @@ -2948,7 +3513,7 @@ class TestResponsesWSFirstFrameModelAuth: return relay() - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock( side_effect=lambda user_api_key_dict, data, call_type, guardrails_only: data ) @@ -3022,7 +3587,7 @@ class TestResponsesWSFirstFrameModelAuth: return relay() - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.pre_call_hook = AsyncMock( side_effect=lambda user_api_key_dict, data, call_type, guardrails_only: { **data, @@ -3083,7 +3648,7 @@ class TestResponsesWSFirstFrameModelAuth: rejection = litellm.RateLimitError( message="origin deployment is cooling down", model="gpt-4o-mini", llm_provider="openai" ) - proxy_logging_obj = MagicMock() + proxy_logging_obj = MagicMock(during_call_hook=AsyncMock()) proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) proxy_logging_obj.arelease_responses_websocket_slot = AsyncMock(return_value=None) user_api_key_dict = MagicMock() diff --git a/tests/unit/proxy/test_litellm_pre_call_utils.py b/tests/unit/proxy/test_litellm_pre_call_utils.py index c5136ff6e4a..3194124ad1c 100644 --- a/tests/unit/proxy/test_litellm_pre_call_utils.py +++ b/tests/unit/proxy/test_litellm_pre_call_utils.py @@ -5360,6 +5360,88 @@ async def test_add_guardrails_from_policy_engine_keeps_a_policy_added_guardrail_ assert [pipeline.mode for _policy_name, pipeline in data["metadata"]["_guardrail_pipelines"]] == ["post_call"] +@pytest.mark.asyncio +async def test_add_guardrails_from_policy_engine_keeps_configured_guardrails_and_skips_ones_already_selected(): + from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.types.proxy.policy_engine import Policy, PolicyAttachment, PolicyGuardrails + + configured = {"pii_blocker": {"extra_body": {"threshold": 0.9}}} + data = { + "model": "gpt-4", + "messages": [{"role": "user", "content": "Hello"}], + "metadata": {"guardrails": [configured, "audit"]}, + } + policy_registry = get_policy_registry() + policy_registry._policies = { + "baseline": Policy(guardrails=PolicyGuardrails(add=["pii_blocker", "audit", "hipaa_audit"])), + } + policy_registry._initialized = True + attachment_registry = get_attachment_registry() + attachment_registry._attachments = [PolicyAttachment(policy="baseline", scope="*")] + attachment_registry._initialized = True + + try: + await add_guardrails_from_policy_engine( + data=data, + metadata_variable_name="metadata", + user_api_key_dict=UserAPIKeyAuth(api_key="test-key"), + ) + finally: + policy_registry._policies = {} + policy_registry._initialized = False + attachment_registry._attachments = [] + attachment_registry._initialized = False + + assert data["metadata"]["guardrails"] == [configured, "audit", "hipaa_audit"] + + +def test_add_guardrails_from_auth_metadata_lists_each_key_and_team_guardrail_once_keeping_the_key_s_settings(): + from litellm.proxy.litellm_pre_call_utils import add_guardrails_from_auth_metadata + + configured = {"pii_blocker": {"extra_body": {"threshold": 0.9}}} + data = {"model": "gpt-4", "metadata": {}} + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", + metadata={"guardrails": [configured, "audit"]}, + team_metadata={"guardrails": ["audit", "pii_blocker", "toxic"]}, + ) + + with patch("litellm.proxy.utils.premium_user_check"): + add_guardrails_from_auth_metadata( + user_api_key_dict=user_api_key_dict, data=data, metadata_variable_name="metadata" + ) + + assert data["metadata"]["guardrails"] == [configured, "audit", "toxic"] + + +def test_add_guardrails_from_auth_metadata_keeps_a_selected_guardrail_s_settings_when_a_key_policy_adds_it(): + from litellm.proxy.litellm_pre_call_utils import add_guardrails_from_auth_metadata + from litellm.proxy.policy_engine.policy_registry import get_policy_registry + from litellm.types.proxy.policy_engine import Policy, PolicyGuardrails + + configured = {"pii_blocker": {"extra_body": {"threshold": 0.9}}} + data = {"model": "gpt-4", "metadata": {"guardrails": [configured, "audit"]}} + policy_registry = get_policy_registry() + policy_registry._policies = { + "baseline": Policy(guardrails=PolicyGuardrails(add=["pii_blocker", "audit", "hipaa_audit"])), + } + policy_registry._initialized = True + + try: + with patch("litellm.proxy.utils.premium_user_check"): + add_guardrails_from_auth_metadata( + user_api_key_dict=UserAPIKeyAuth(api_key="test-key", metadata={"policies": ["baseline"]}), + data=data, + metadata_variable_name="metadata", + ) + finally: + policy_registry._policies = {} + policy_registry._initialized = False + + assert data["metadata"]["guardrails"] == [configured, "audit", "hipaa_audit"] + + @pytest.mark.asyncio async def test_add_guardrails_from_policy_engine_applies_inherited_parent_guardrail_when_child_condition_misses(): from litellm.proxy.policy_engine.attachment_registry import get_attachment_registry diff --git a/tests/unit/proxy/utils/helpers/test_guardrail_merge.py b/tests/unit/proxy/utils/helpers/test_guardrail_merge.py index e2c1a7c2c09..dd933c508cb 100644 --- a/tests/unit/proxy/utils/helpers/test_guardrail_merge.py +++ b/tests/unit/proxy/utils/helpers/test_guardrail_merge.py @@ -362,7 +362,17 @@ def test_merge_guardrails_with_existing_creates_metadata_when_missing(): } -def test_merge_guardrails_with_existing_raises_on_unhashable_guardrail(): - data = {"metadata": {"guardrails": [{"unhashable": True}]}} - with pytest.raises(TypeError): - _merge_guardrails_with_existing(data, ["g1"]) +def test_merge_guardrails_with_existing_keeps_a_selected_guardrail_s_settings_over_the_model_s_plain_entry(): + configured = {"pii-redact": {"extra_body": {"threshold": 0.9}}} + + result = _merge_guardrails_with_existing({"metadata": {"guardrails": [configured]}}, ["pii-redact", "toxic-filter"]) + + assert result["metadata"]["guardrails"] == [configured, "toxic-filter"] + + +def test_merge_guardrails_with_existing_keeps_a_selected_guardrail_s_settings_listed_after_the_key_s_plain_entry(): + configured = {"pii-redact": {"extra_body": {"threshold": 0.9}}} + + result = _merge_guardrails_with_existing({"metadata": {"guardrails": ["pii-redact", configured]}}, ["toxic-filter"]) + + assert result["metadata"]["guardrails"] == [configured, "toxic-filter"] From 35d3a1784b90cb2dd6e20afcb74ebf1f4bedbac4 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 30 Oct 2026 14:27:23 -0700 Subject: [PATCH 24/46] fix(proxy): preserve access windows when editing models (#49316) --- .../model_management_endpoints.py | 6 +-- .../test_model_management_endpoints.py | 51 ++++++++++++++++++- 2 files changed, 50 insertions(+), 7 deletions(-) diff --git a/litellm/proxy/management_endpoints/model_management_endpoints.py b/litellm/proxy/management_endpoints/model_management_endpoints.py index dbd77e45d57..2c7453ebb1c 100644 --- a/litellm/proxy/management_endpoints/model_management_endpoints.py +++ b/litellm/proxy/management_endpoints/model_management_endpoints.py @@ -1515,14 +1515,10 @@ def update_db_model( # convert to prisma compatible format - for key, value in merged_model_info.items(): - if isinstance(value, datetime.datetime): - merged_model_info[key] = value.isoformat() - prisma_compatible_model_dict: Final = PrismaCompatibleUpdateDBModel( model_name=merged_model_name, litellm_params=json.dumps(merged_litellm_params), - model_info=json.dumps(merged_model_info), + model_info=TypeAdapter(dict[str, object]).dump_json(merged_model_info).decode(), ) if updated_patch.blocked is not None: diff --git a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py index abaa802de2c..ab6cfcfebfe 100644 --- a/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py +++ b/tests/unit/proxy/management_endpoints/test_model_management_endpoints.py @@ -4,8 +4,9 @@ import inspect import json import socket from collections.abc import AsyncIterator, Iterator, Mapping +from datetime import datetime, time, timezone from types import MappingProxyType, SimpleNamespace -from typing import Dict, Final, Optional +from typing import Dict, Final, Literal, Optional from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -46,7 +47,14 @@ from litellm.proxy.management_helpers.model_quota_policy import REQUIRE_BOTH_LIM from litellm.proxy.utils import PrismaClient from litellm.router import Router from litellm.router_api import ApplyResult, ChangeSet -from litellm.types.router import Deployment, LiteLLM_Params, ModelInfo, updateDeployment, updateLiteLLMParams +from litellm.types.router import ( + Deployment, + LiteLLM_Params, + ModelAccessWindow, + ModelInfo, + updateDeployment, + updateLiteLLMParams, +) async def _passthrough_row(update_data): @@ -4048,6 +4056,45 @@ class TestUpdateDBModelBlocked: assert "blocked" not in result +@pytest.mark.parametrize("window_patch", ["omit", "echo", "replace", "clear"]) +def test_model_edit_preserves_or_explicitly_changes_access_windows( + window_patch: Literal["omit", "echo", "replace", "clear"], +) -> None: + from litellm.proxy.management_endpoints.model_management_endpoints import update_db_model + + saved_window: Final = ModelAccessWindow( + start=time(13), end=time(17), timezone="Asia/Almaty", team_ids=("team-window", "team-shared") + ) + replacement_window: Final = ModelAccessWindow( + start=time(22, 0, 0, 123456), end=time(6, 0, 0, 654321), + timezone="America/New_York", team_ids=("team-other",), + ) + created_at: Final = datetime(2026, 1, 2, 3, 4, 5, 123456, tzinfo=timezone.utc) + requested_windows: Final = { + "omit": None, + "echo": (saved_window,), + "replace": (replacement_window,), + "clear": (), + }[window_patch] + expected_windows: Final = (saved_window,) if requested_windows is None else requested_windows + result: Final = update_db_model( + db_model=Deployment( + model_name="window-model", + litellm_params=LiteLLM_Params(model="openai/test-model"), + model_info=ModelInfo(id="window-deployment", access_windows=(saved_window,), created_at=created_at), + ), + updated_patch=updateDeployment( + litellm_params=updateLiteLLMParams(max_parallel_requests=7), + model_info=None if requested_windows is None else ModelInfo(access_windows=requested_windows), + ), + ) + + assert json.loads(result["litellm_params"])["max_parallel_requests"] == 7 + saved_info: Final = ModelInfo.model_validate_json(result["model_info"]) + assert saved_info.access_windows == expected_windows + assert saved_info.created_at == created_at + + class TestUpdateDBModelKeepsLegacyDropParams: def test_partial_patch_keeps_encrypted_string_drop_params(self, monkeypatch): from litellm.proxy.common_utils.encrypt_decrypt_utils import decrypt_value_helper From ec4f66bb94ff706bb7f2fbb55261228baed29f77 Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 30 Oct 2026 14:38:14 -0700 Subject: [PATCH 25/46] fix(router): budget classifier input by provider limits (#49306) --- litellm/decisions/call.py | 8 +- litellm/integrations/otel/routing.py | 5 + .../llms/azure_ai/decisions/transformation.py | 12 +- .../llms/base_llm/decisions/transformation.py | 21 +- .../llms/typesafe/decisions/transformation.py | 13 +- litellm/router_api.py | 3 + .../complexity_router/complexity_router.py | 119 ++++++---- .../complexity_router/jev_classifier.py | 157 ++++++++++++- litellm/types/utils.py | 10 + .../integrations/otel/test_otel_v2_logger.py | 5 + tests/unit/integrations/test_opentelemetry.py | 5 + .../test_base_decisions_transformation.py | 60 +++++ .../complexity_router/test_jev_classifier.py | 216 +++++++++++++++++- ui/litellm-dashboard/src/lib/http/schema.d.ts | 10 + 14 files changed, 589 insertions(+), 55 deletions(-) diff --git a/litellm/decisions/call.py b/litellm/decisions/call.py index 3b73f65c917..ae01aa8d0f7 100644 --- a/litellm/decisions/call.py +++ b/litellm/decisions/call.py @@ -218,7 +218,7 @@ def supported_decisions_providers() -> tuple[str, ...]: ) -def _provider_config(model: str, custom_llm_provider: str) -> BaseDecisionsConfig: +def get_decisions_provider_config(model: str, custom_llm_provider: str) -> BaseDecisionsConfig: provider: Final = next((member for member in LlmProviders if member.value == custom_llm_provider), None) provider_config: Final = ( None if provider is None else ProviderConfigManager.get_provider_decisions_config(model, provider) @@ -408,7 +408,7 @@ def _provider_ir_request( return supported -def _outbound_headers( +def decisions_outbound_headers( extra_headers: Mapping[str, str] | None, custom_llm_provider: str, kwargs: Mapping[str, object] ) -> Mapping[str, str]: forwarded: Final = ProviderSpecificHeaderUtils.get_provider_specific_headers( @@ -466,7 +466,7 @@ def prepare_call( api_base=call_api_base, api_key=api_key, ) - provider_config: Final = _provider_config(upstream_model, provider) + provider_config: Final = get_decisions_provider_config(upstream_model, provider) if not upstream_model: raise litellm.BadRequestError( message="A model name is required for the Decisions API", @@ -478,7 +478,7 @@ def prepare_call( except ValueError as error: raise litellm.BadRequestError(message=str(error), model=model, llm_provider=provider) from error try: - headers: Final = _outbound_headers(extra_headers, provider, kwargs) + headers: Final = decisions_outbound_headers(extra_headers, provider, kwargs) request_safety_identifier: Final = _request_safety_identifier(safety_identifier, provider_config, kwargs) request: Final = _fields_without_unsupported_media( fields, provider_config, canonical_model, provider, kwargs diff --git a/litellm/integrations/otel/routing.py b/litellm/integrations/otel/routing.py index 82dd6237edc..339775855be 100644 --- a/litellm/integrations/otel/routing.py +++ b/litellm/integrations/otel/routing.py @@ -21,6 +21,11 @@ _ROUTING_SCALAR_FIELDS: Final = frozenset( "classifier_cost", "classifier_failure_reason", "classifier_error_type", + "classifier_input_budget_tokens", + "classifier_input_tokens", + "classifier_input_count_method", + "classifier_input_truncated", + "classifier_input_budget_status", "classifier_confidence", "classifier_primary_rule", "classifier_capability_boundary", diff --git a/litellm/llms/azure_ai/decisions/transformation.py b/litellm/llms/azure_ai/decisions/transformation.py index fb5330ed715..22c4b59f5e8 100644 --- a/litellm/llms/azure_ai/decisions/transformation.py +++ b/litellm/llms/azure_ai/decisions/transformation.py @@ -12,7 +12,7 @@ from litellm.llms.azure_ai.common_utils import ( AZURE_ENTRA_LITELLM_PARAM_KEYS, get_azure_ai_auth_headers, ) -from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig +from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig, DecisionsInputLimits _FOUNDRY_ROUTE_SUFFIX: Final = re.compile(r"(/api/projects/[^/]+)?(/openai/v1|/openai|/models|/v1)?$") _AZURE_AI_CREDENTIAL_PARAM_KEYS: Final = AZURE_ENTRA_LITELLM_PARAM_KEYS | frozenset({"azure_ad_token_provider"}) @@ -26,6 +26,16 @@ class AzureAIDecisionsConfig(BaseDecisionsConfig): api_base_env = ("AZURE_AI_API_BASE",) api_key_required = False + def get_input_limits( + self, model: str, custom_llm_provider: str, headers: Mapping[str, str] + ) -> DecisionsInputLimits: + limits: Final = super().get_input_limits(model, custom_llm_provider, headers) + return ( + limits + if limits.max_request_tokens is not None + else super().get_input_limits("Microsoft-Decision-1", custom_llm_provider, headers) + ) + def missing_api_base_message(self, custom_llm_provider: str) -> str: return "Missing AZURE_AI_API_BASE - set AZURE_AI_API_BASE or pass api_base, e.g. https://.services.ai.azure.com" diff --git a/litellm/llms/base_llm/decisions/transformation.py b/litellm/llms/base_llm/decisions/transformation.py index e8e4b77e757..efd3e3fe9b5 100644 --- a/litellm/llms/base_llm/decisions/transformation.py +++ b/litellm/llms/base_llm/decisions/transformation.py @@ -1,7 +1,7 @@ import json from abc import ABC from collections.abc import Callable, Mapping, Sequence -from dataclasses import replace +from dataclasses import dataclass, replace from itertools import chain from types import MappingProxyType from typing import Final @@ -511,6 +511,12 @@ def _is_decisions_model(model: str, custom_llm_provider: str) -> bool: return is_decisions_model_mode(model_mode) +@dataclass(frozen=True, slots=True) +class DecisionsInputLimits: + max_request_tokens: int | None = None + max_question_tokens: int | None = None + + class BaseDecisionsConfig(ABC): path: str = "/v1/systemone" api_key_env: tuple[str, ...] = () @@ -533,6 +539,19 @@ class BaseDecisionsConfig(ABC): def provider_reported_cost(self, response: DecisionsIRResponse) -> float | None: return None + def get_input_limits( + self, model: str, custom_llm_provider: str, headers: Mapping[str, str] + ) -> DecisionsInputLimits: + try: + limit: Final = litellm.get_model_info(model=model, custom_llm_provider=custom_llm_provider).get( + "max_input_tokens" + ) + except Exception: # noqa: BLE001 # get_model_info raises bare Exception for an unmapped model + return DecisionsInputLimits() + return DecisionsInputLimits( + max_request_tokens=limit if isinstance(limit, int) and not isinstance(limit, bool) and limit > 0 else None + ) + def missing_api_base_message(self, custom_llm_provider: str) -> str: return f"api_base is required for Decisions provider '{custom_llm_provider}'" diff --git a/litellm/llms/typesafe/decisions/transformation.py b/litellm/llms/typesafe/decisions/transformation.py index 952c874bc8e..871596b7f00 100644 --- a/litellm/llms/typesafe/decisions/transformation.py +++ b/litellm/llms/typesafe/decisions/transformation.py @@ -6,7 +6,7 @@ from types import MappingProxyType from typing import Final, Literal from litellm.constants import TYPESAFE_DANGER_ZONE_HEADER_PREFIX -from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig, image_media_type +from litellm.llms.base_llm.decisions.transformation import BaseDecisionsConfig, DecisionsInputLimits, image_media_type from litellm.types.decisions import ( DecisionsAsset, DecisionsIRMessages, @@ -19,6 +19,7 @@ from litellm.types.decisions import ( ) _IMAGES_HEADER: Final = f"{TYPESAFE_DANGER_ZONE_HEADER_PREFIX}images" +_LONG_CONTEXT_HEADER: Final = f"{TYPESAFE_DANGER_ZONE_HEADER_PREFIX}long-context" _ASSET_FORMATS: Final[Mapping[str, Literal["png", "jpg"]]] = MappingProxyType( {"image/png": "png", "image/jpeg": "jpg", "image/jpg": "jpg"} ) @@ -84,6 +85,16 @@ class TypeSafeDecisionsConfig(BaseDecisionsConfig): def get_default_api_base(self) -> str | None: return "https://api.typesafe.ai" + def get_input_limits( + self, model: str, custom_llm_provider: str, headers: Mapping[str, str] + ) -> DecisionsInputLimits: + long_context_values: Final = tuple( + value.strip().lower() for name, value in headers.items() if name.lower() == _LONG_CONTEXT_HEADER + ) + if long_context_values and all(value == "true" for value in long_context_values): + return DecisionsInputLimits(max_request_tokens=262144, max_question_tokens=262144) + return DecisionsInputLimits(max_request_tokens=65536, max_question_tokens=32768) + def supports_images_in( self, decision_input: DecisionsIRState | DecisionsIRMessages, headers: Mapping[str, str] ) -> bool: diff --git a/litellm/router_api.py b/litellm/router_api.py index d62999cf261..19f12cec40d 100644 --- a/litellm/router_api.py +++ b/litellm/router_api.py @@ -1216,6 +1216,9 @@ class RoutingContext: def registry(self) -> RouterRegistry: return self._router.registry() + def deployment_params(self, request: Mapping[str, object], deployment: "Deployment | None") -> Mapping[str, object]: + return self._router._params_sent_to_deployment(request, deployment) + def settings(self) -> RouterSettings: return self._router.settings() diff --git a/litellm/router_strategy/complexity_router/complexity_router.py b/litellm/router_strategy/complexity_router/complexity_router.py index 30f1021251f..97ad5ffa05b 100644 --- a/litellm/router_strategy/complexity_router/complexity_router.py +++ b/litellm/router_strategy/complexity_router/complexity_router.py @@ -126,10 +126,13 @@ from .jev_classifier import ( DEFAULT_JEV_INSTRUCTIONS, HttpJevClassifierClient, JevClassifierClient, + JevSystemOneRequest, JevVerdict, RouterJevClassifierClient, build_jev_request, + fit_jev_request, jev_classifier_cost, + jev_input_targets, ) from .llm_v2 import LLM_V2_PROMPT_VERSION, LLMV2Decision, LLMV2TaskContext, LLMV2Verdict, llm_v2_response_format from .stall_detector import detect_stalled_task @@ -1166,6 +1169,7 @@ class ClassificationOutcome(NamedTuple): heuristic_v2_forecast: StandardLoggingHeuristicV2Forecast | None = None classifier_failure_reason: ClassifierFailureReason | None = None classifier_error_type: str | None = None + classifier_input_telemetry: StandardLoggingRoutingDecision | None = None def _with_classifier_failure( @@ -1207,23 +1211,24 @@ def _with_classifier_forecast( decision: StandardLoggingRoutingDecision, outcome: ClassificationOutcome ) -> StandardLoggingRoutingDecision: """Attach validated forecasts and their applied policy to the routing decision.""" + with_input: Final[StandardLoggingRoutingDecision] = {**decision, **(outcome.classifier_input_telemetry or {})} if outcome.heuristic_v2_forecast is not None: - return {**decision, "heuristic_v2_forecast": outcome.heuristic_v2_forecast} + return {**with_input, "heuristic_v2_forecast": outcome.heuristic_v2_forecast} if outcome.jev_verdict is not None: forecasted_decision: Final[StandardLoggingRoutingDecision] = { - **decision, + **with_input, "classifier_probabilities": outcome.jev_verdict.probabilities, "classifier_confidence": outcome.jev_verdict.confidence, } return forecasted_decision if outcome.llm_v2_forecast is not None: - return _with_llm_v2_forecast(decision, outcome.llm_v2_forecast) + return _with_llm_v2_forecast(with_input, outcome.llm_v2_forecast) forecast: Final = outcome.capability_forecast if forecast is None: - return decision + return with_input verdict: Final = forecast.verdict enriched: Final[StandardLoggingRoutingDecision] = { - **decision, + **with_input, "classifier_crux": verdict.crux, "classifier_primary_rule": verdict.primary_rule, "classifier_capability_boundary": verdict.capability_boundary, @@ -1974,19 +1979,21 @@ class ComplexityRouter(CustomLogger): `score` because the score is only interpretable against the boundaries that mapped it to a tier. """ - decision: Final = StandardLoggingRoutingDecision( - router_model_name=self.model_name, - router_type="complexity", - routed_model=routed_model, - cause=cause, - conversation_continuing=conversation_continuing, - **( - {"classifier_failure_reason": classifier_failure_reason} - if classifier_failure_reason is not None - else {} - ), - **({"classifier_error_type": classifier_error_type} if classifier_error_type is not None else {}), + failure: Final[StandardLoggingRoutingDecision] = ( + {"classifier_failure_reason": classifier_failure_reason} if classifier_failure_reason is not None else {} ) + error: Final[StandardLoggingRoutingDecision] = ( + {"classifier_error_type": classifier_error_type} if classifier_error_type is not None else {} + ) + decision: Final[StandardLoggingRoutingDecision] = { + "router_model_name": self.model_name, + "router_type": "complexity", + "routed_model": routed_model, + "cause": cause, + "conversation_continuing": conversation_continuing, + **failure, + **error, + } if (baseline := self.savings_baseline) is not None: decision["savings_baseline_model"] = baseline.model if baseline.deployment_id is not None: @@ -2371,16 +2378,6 @@ class ComplexityRouter(CustomLogger): system_prompt, failure_reason="unsupported_input", ) - breaker: Final = self._classifier_circuit_breaker - permit: Final = breaker.acquire_permit() if breaker is not None else None - if breaker is not None and permit is None: - return self._classifier_failure_outcome( - "jev classifier circuit is open", - prompt, - system_prompt, - failure_reason="circuit_open", - signal=_CLASSIFIER_CIRCUIT_OPEN_SIGNAL, - ) criteria: Final[Mapping[str, str]] = ( MappingProxyType( { @@ -2395,13 +2392,40 @@ class ComplexityRouter(CustomLogger): ) ) timeout_s: Final = config.timeout_ms / 1000 - request: Final = build_jev_request( - prompt=self._classifier_context_payload(prompt, system_prompt, request_kwargs, messages), - system_prompt=None, - model=config.model, - instructions=config.instructions or DEFAULT_JEV_INSTRUCTIONS, - criteria=criteria, + from litellm.litellm_core_utils.token_counter import offload_token_count + + def build_request(current_prompt: str, caller_constraints: str | None) -> JevSystemOneRequest: + return build_jev_request( + prompt=self._classifier_context_payload( + current_prompt, caller_constraints, request_kwargs, messages, current_ask_for_history=prompt + ), + system_prompt=None, + model=config.model, + instructions=config.instructions or DEFAULT_JEV_INSTRUCTIONS, + criteria=criteria, + ) + + budgeted: Final = await offload_token_count(fit_jev_request)( + build_request, prompt, system_prompt, jev_input_targets(self.litellm_router_instance, config), _truncate ) + request: Final = budgeted.request + if request is None: + return self._classifier_failure_outcome( + "classifier input cannot fit the provider budget", + prompt, + system_prompt, + failure_reason="unsupported_input", + )._replace(classifier_input_telemetry=budgeted.telemetry) + breaker: Final = self._classifier_circuit_breaker + permit: Final = breaker.acquire_permit() if breaker is not None else None + if breaker is not None and permit is None: + return self._classifier_failure_outcome( + "jev classifier circuit is open", + prompt, + system_prompt, + failure_reason="circuit_open", + signal=_CLASSIFIER_CIRCUIT_OPEN_SIGNAL, + )._replace(classifier_input_telemetry=budgeted.telemetry) try: response: Final = await asyncio.wait_for(client.evaluate(request, timeout_s, request_kwargs), timeout_s) answer: Final = response.answers.get("tier") @@ -2449,6 +2473,7 @@ class ComplexityRouter(CustomLogger): cause="jev_classifier", classifier_cost=verdict.cost, jev_verdict=verdict, + classifier_input_telemetry=budgeted.telemetry, ) except asyncio.CancelledError: if breaker is not None and permit is not None: @@ -2463,7 +2488,7 @@ class ComplexityRouter(CustomLogger): system_prompt, failure_reason=_classifier_failure_reason(e), error_type=type(e).__name__, - ) + )._replace(classifier_input_telemetry=budgeted.telemetry) def _classifier_failure_outcome( self, @@ -2654,6 +2679,7 @@ class ComplexityRouter(CustomLogger): messages: Sequence[Mapping[str, object]] | None, *, encrypted_task: bool = False, + current_ask_for_history: str | None = None, ) -> str: include_assistant: Final = self.config.classifier_context_include_assistant_turns marker_pairs: Final = self._reminder_markers_for_request(request_kwargs or EMPTY_MAPPING) @@ -2661,7 +2687,7 @@ class ComplexityRouter(CustomLogger): prior_turns: Final = ( _extract_prior_turns( messages, - current_ask=prompt, + current_ask=prompt if current_ask_for_history is None else current_ask_for_history, window_size=self.config.classifier_context_window_size, budget_chars=self.config.classifier_context_budget_chars, per_turn_chars=self.config.classifier_context_per_turn_chars, @@ -4983,16 +5009,19 @@ class ComplexityRouter(CustomLogger): return PreRoutingHookResponse( model=fallback_model, messages=messages if has_original_messages else None, - routing_decision=self._build_routing_decision( - routed_model=fallback_model, - conversation_continuing=conversation_continuing, - cause=outcome.cause, - signals=outcome.signals, - classifier_failure_reason=outcome.classifier_failure_reason, - classifier_error_type=outcome.classifier_error_type, - escalation_keyword=escalation_keyword, - escalated=False, - tier_litellm_params=fallback_tier_params, + routing_decision=_with_classifier_forecast( + self._build_routing_decision( + routed_model=fallback_model, + conversation_continuing=conversation_continuing, + cause=outcome.cause, + signals=outcome.signals, + classifier_failure_reason=outcome.classifier_failure_reason, + classifier_error_type=outcome.classifier_error_type, + escalation_keyword=escalation_keyword, + escalated=False, + tier_litellm_params=fallback_tier_params, + ), + outcome, ), litellm_params=fallback_tier_params, ) diff --git a/litellm/router_strategy/complexity_router/jev_classifier.py b/litellm/router_strategy/complexity_router/jev_classifier.py index 87bd031490a..25ea813229c 100644 --- a/litellm/router_strategy/complexity_router/jev_classifier.py +++ b/litellm/router_strategy/complexity_router/jev_classifier.py @@ -1,4 +1,5 @@ -from collections.abc import Mapping +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass from datetime import datetime, timezone from types import MappingProxyType from typing import TYPE_CHECKING, Annotated, Final, Literal, NamedTuple, Protocol, TypeAlias @@ -25,10 +26,13 @@ from litellm.llms.pass_through.typesafe_logging_handler import TypeSafePassthrou from litellm.router_strategy.complexity_router.config import DEFAULT_JEV_INSTRUCTIONS as _DEFAULT_JEV_INSTRUCTIONS from litellm.types.decisions import DecisionsResponse from litellm.types.llms.base import LiteLLMBaseModel -from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN +from litellm.types.utils import AUTOROUTER_CLASSIFIER_CALL_ORIGIN, StandardLoggingRoutingDecision if TYPE_CHECKING: from litellm.router import Router + from litellm.router_api import RoutingContext + from litellm.router_strategy.complexity_router.config import OpenSourceClassifierConfig + from litellm.types.router import Deployment JevProbability: TypeAlias = Annotated[float, Field(ge=0.0, le=1.0)] ClassifierProvider: TypeAlias = Literal["typesafe", "laya", "bespoke", "databricks"] @@ -52,6 +56,155 @@ class JevSystemOneRequest(LiteLLMBaseModel): questions: Mapping[str, JevChoiceQuestion] +@dataclass(frozen=True, slots=True) +class JevInputTarget: + model: str + max_tokens: int | None + + +@dataclass(frozen=True, slots=True) +class JevBudgetedRequest: + request: JevSystemOneRequest | None + telemetry: StandardLoggingRoutingDecision + + +def jev_input_targets(router: "Router", config: "OpenSourceClassifierConfig") -> tuple[JevInputTarget, ...]: + from litellm.router_api import RoutingContext + from litellm.types.router import Deployment + + if config.deployment_name is None: + provider: Final = "typesafe" if config.provider == "jev" else config.provider + return (_jev_input_target({"model": f"{provider}/{config.model}"}, None),) + context: Final = RoutingContext(router) + deployments: Final = context.registry().deployments(model_name=config.deployment_name) or () + return tuple( + _jev_deployment_input_target(context, Deployment.model_validate(deployment)) for deployment in deployments + ) + + +def _jev_deployment_input_target(context: "RoutingContext", deployment: "Deployment") -> JevInputTarget: + params: Final = context.deployment_params({}, deployment) + explicit_limit: Final = deployment.model_info.model_dump().get("max_input_tokens") + return _jev_input_target(params, explicit_limit) + + +def _jev_input_target(params: Mapping[str, object], explicit_limit: object) -> JevInputTarget: + from litellm.decisions.call import decisions_outbound_headers, get_decisions_provider_config + from litellm.litellm_core_utils.get_llm_provider_logic import resolve_model_and_provider_for_metadata + + qualified_model: Final = str(params.get("model") or "") + override: Final = params.get("custom_llm_provider") + try: + model, provider = resolve_model_and_provider_for_metadata( + qualified_model, override if isinstance(override, str) else None + ) + headers: Final = decisions_outbound_headers( + TypeAdapter[Mapping[str, str] | None](Mapping[str, str] | None).validate_python( + params.get("extra_headers") + ), + provider, + params, + ) + limits: Final = get_decisions_provider_config(model, provider).get_input_limits(model, provider, headers) + candidates: Final = tuple( + value + for value in (limits.max_request_tokens, limits.max_question_tokens, explicit_limit) + if isinstance(value, int) and not isinstance(value, bool) and value > 0 + ) + return JevInputTarget(f"{provider}/{model}", min(candidates) if candidates else None) + except Exception: # noqa: BLE001 # unavailable metadata must not turn an existing classifier into a request failure + return JevInputTarget(qualified_model, None) + + +def _jev_request_count(request: JevSystemOneRequest, models: Sequence[str], budget: int) -> tuple[int, str]: + from litellm._lazy_imports import get_token_counter_new + + payload: Final = request.model_dump_json(exclude={"model"}) + byte_count: Final = len(payload.encode("utf-8")) + if byte_count <= budget or litellm.disable_token_counter: + return byte_count, "utf8_bytes" + try: + counts: Final = tuple(get_token_counter_new()(model=model, text=payload) for model in models) + if counts and all(count > 0 for count in counts): + return max(counts), "estimated_tokens" + except Exception: # noqa: BLE001 # tokenization is best-effort; byte counting remains available + return byte_count, "utf8_bytes" + return byte_count, "utf8_bytes" + + +def fit_jev_request( + build_request: Callable[[str, str | None], JevSystemOneRequest], + prompt: str, + system_prompt: str | None, + targets: Sequence[JevInputTarget], + truncate: Callable[[str, int], str], +) -> JevBudgetedRequest: + known: Final = tuple(target for target in targets if target.max_tokens is not None) + if not known: + return JevBudgetedRequest( + build_request(prompt, system_prompt), + {"classifier_input_budget_status": "unknown_limits", "classifier_input_truncated": False}, + ) + budget: Final = max(1, int(min(target.max_tokens for target in known if target.max_tokens is not None) * 0.75)) + models: Final = tuple(dict.fromkeys(target.model for target in known)) + initial_cap: Final = max(120, budget * 8) + bounded_prompt: Final = truncate(prompt, initial_cap) + bounded_system: Final = truncate(system_prompt, initial_cap) if system_prompt is not None else None + return _fit_jev_components( + build_request, + bounded_prompt, + bounded_system, + models, + budget, + truncate, + len(known) != len(targets), + bounded_prompt != prompt or bounded_system != system_prompt, + ) + + +def _fit_jev_components( + build_request: Callable[[str, str | None], JevSystemOneRequest], + prompt: str, + system_prompt: str | None, + models: Sequence[str], + budget: int, + truncate: Callable[[str, int], str], + unknown_limits: bool, + truncated: bool, +) -> JevBudgetedRequest: + attempts: Final = len(prompt).bit_length() + len(system_prompt or "").bit_length() + 1 + components: tuple[str, str | None, bool, int] = ( # rebind-ok: bounded iterative fitting state + prompt, + system_prompt, + truncated, + attempts, + ) + while True: + current_prompt, current_system, was_truncated, remaining = components + request = build_request(current_prompt, current_system) + count, method = _jev_request_count(request, models, budget) + telemetry: StandardLoggingRoutingDecision = { + "classifier_input_budget_tokens": budget, + "classifier_input_tokens": count, + "classifier_input_count_method": method, + "classifier_input_truncated": was_truncated, + "classifier_input_budget_status": "unknown_limits" + if unknown_limits + else "truncated" + if was_truncated + else "within_budget", + } + if count <= budget: + return JevBudgetedRequest(request, telemetry) + if remaining == 0 or (len(current_system or "") <= 123 and len(current_prompt) <= 123): + return JevBudgetedRequest(None, {**telemetry, "classifier_input_budget_status": "untrimmable_input"}) + components = ( + (current_prompt, truncate(current_system, max(120, len(current_system) // 2)), True, remaining - 1) + if current_system is not None and len(current_system) > 123 + else (truncate(current_prompt, max(120, len(current_prompt) // 2)), current_system, True, remaining - 1) + ) + + class JevChoiceAnswer(LiteLLMBaseModel): model_config = ConfigDict(frozen=True, allow_inf_nan=False) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index db865c7488a..1a8f5adf154 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3323,6 +3323,11 @@ class StandardLoggingRoutingDecision(TypedDict, total=False): classifier_cost: float classifier_failure_reason: ReadOnly[ClassifierFailureReason] classifier_error_type: ReadOnly[str] + classifier_input_budget_tokens: ReadOnly[int] + classifier_input_tokens: ReadOnly[int] + classifier_input_count_method: ReadOnly[str] + classifier_input_truncated: ReadOnly[bool] + classifier_input_budget_status: ReadOnly[str] classifier_probabilities: ReadOnly[Mapping[str, float]] classifier_confidence: ReadOnly[float] heuristic_v2_forecast: ReadOnly[StandardLoggingHeuristicV2Forecast] @@ -3381,6 +3386,11 @@ DERIVED_ROUTING_DECISION_FIELDS: Final[frozenset[str]] = frozenset( "classifier_cost", "classifier_failure_reason", "classifier_error_type", + "classifier_input_budget_tokens", + "classifier_input_tokens", + "classifier_input_count_method", + "classifier_input_truncated", + "classifier_input_budget_status", "classifier_probabilities", "classifier_confidence", "heuristic_v2_forecast", diff --git a/tests/unit/integrations/otel/test_otel_v2_logger.py b/tests/unit/integrations/otel/test_otel_v2_logger.py index 1f9d3af71d9..ac1f5557bf6 100644 --- a/tests/unit/integrations/otel/test_otel_v2_logger.py +++ b/tests/unit/integrations/otel/test_otel_v2_logger.py @@ -207,6 +207,11 @@ def test_routing_diagnostics_reach_llm_spans_without_promoting_prompt_text(fail: "classifier_cost": 0.001, "classifier_failure_reason": "timeout", "classifier_error_type": "TimeoutError", + "classifier_input_budget_tokens": 24576, + "classifier_input_tokens": 23000, + "classifier_input_count_method": "estimated_tokens", + "classifier_input_truncated": True, + "classifier_input_budget_status": "truncated", "escalated": False, } payload: Final = _payload( diff --git a/tests/unit/integrations/test_opentelemetry.py b/tests/unit/integrations/test_opentelemetry.py index f0fca94c370..5f749e951b1 100644 --- a/tests/unit/integrations/test_opentelemetry.py +++ b/tests/unit/integrations/test_opentelemetry.py @@ -344,6 +344,11 @@ def test_v1_routing_diagnostics_add_scalar_attributes_and_preserve_legacy_metada "cause": "heuristic", "classifier_failure_reason": "timeout", "classifier_error_type": "TimeoutError", + "classifier_input_budget_tokens": 24576, + "classifier_input_tokens": 23000, + "classifier_input_count_method": "estimated_tokens", + "classifier_input_truncated": True, + "classifier_input_budget_status": "truncated", "classifier_cost": 0.001, } with provider.get_tracer(__name__).start_as_current_span("inference") as span: diff --git a/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py b/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py index 03459ba16c8..ec6f118837f 100644 --- a/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py +++ b/tests/unit/llms/base_llm/decisions/test_base_decisions_transformation.py @@ -6,14 +6,19 @@ import httpx import pytest from pydantic import TypeAdapter, ValidationError +import litellm +from litellm.llms.azure_ai.decisions.transformation import AzureAIDecisionsConfig from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.decisions.transformation import ( + BaseDecisionsConfig, + DecisionsInputLimits, ir_to_systemone_request, ir_to_systemone_response, parse_systemone_response, systemone_request_to_ir, ) from litellm.llms.openai.decisions.transformation import openai_request_to_ir +from litellm.llms.typesafe.decisions.transformation import TypeSafeDecisionsConfig from litellm.types.decisions import ( MAX_DECISION_QUESTIONS, DecisionsIRRefusal, @@ -49,6 +54,61 @@ _SYSTEMONE_REQUEST: Final[Mapping[str, object]] = { } +@pytest.mark.parametrize("declared", (4096, 0, -1, True, None)) +def test_decision_input_limits_use_only_positive_catalog_limits( + local_model_cost_map: None, monkeypatch: pytest.MonkeyPatch, declared: int | None +) -> None: + monkeypatch.setitem( + litellm.model_cost, + "bespoke/test-input-limit", + {"litellm_provider": "bespoke", "mode": "evaluation", "max_input_tokens": declared}, + ) + + limits: Final = BaseDecisionsConfig().get_input_limits("test-input-limit", "bespoke", {}) + + assert limits == DecisionsInputLimits(max_request_tokens=4096 if declared == 4096 else None) + + +def test_unmapped_decision_model_limits_remain_unknown(local_model_cost_map: None) -> None: + assert BaseDecisionsConfig().get_input_limits("unmapped-input-limit", "databricks", {}) == DecisionsInputLimits() + + +@pytest.mark.parametrize( + ("headers", "long_context"), + ( + ({}, False), + ({"TypeSafe-Danger-Zone-Long-Context": "false"}, False), + ({"typesafe-danger-zone-long-context": " TRUE "}, True), + ({"TypeSafe-Danger-Zone-Long-Context": "true", "typesafe-danger-zone-long-context": "true"}, True), + ({"TypeSafe-Danger-Zone-Long-Context": "true", "typesafe-danger-zone-long-context": "false"}, False), + ({"typesafe-danger-zone-long-context": "true,false"}, False), + ), +) +def test_typesafe_limits_require_unambiguous_long_context_opt_in( + headers: Mapping[str, str], long_context: bool +) -> None: + """Provider contract from https://docs.typesafe.ai/models, verified 2026-10-30""" + limits: Final = TypeSafeDecisionsConfig().get_input_limits("jev-latest", "typesafe", headers) + + assert limits == ( + DecisionsInputLimits(max_request_tokens=262144, max_question_tokens=262144) + if long_context + else DecisionsInputLimits(max_request_tokens=65536, max_question_tokens=32768) + ) + + +@pytest.mark.parametrize("model", ("Microsoft-Decision-1", "unmapped-decision-deployment")) +def test_azure_decision_deployments_use_catalog_limits_without_typesafe_modes( + local_model_cost_map: None, model: str +) -> None: + catalog_limit: Final = litellm.get_model_info("Microsoft-Decision-1", "azure_ai")["max_input_tokens"] + limits: Final = AzureAIDecisionsConfig().get_input_limits( + model, "azure_ai", {"TypeSafe-Danger-Zone-Long-Context": "true"} + ) + + assert limits == DecisionsInputLimits(max_request_tokens=catalog_limit) + + def _openai_ir(raw: Mapping[str, object]) -> DecisionsIRRequest: return openai_request_to_ir(_OPENAI_BODY.validate_python(raw)) diff --git a/tests/unit/router_strategy/complexity_router/test_jev_classifier.py b/tests/unit/router_strategy/complexity_router/test_jev_classifier.py index 77d465526a4..a71f572254e 100644 --- a/tests/unit/router_strategy/complexity_router/test_jev_classifier.py +++ b/tests/unit/router_strategy/complexity_router/test_jev_classifier.py @@ -9,6 +9,7 @@ from unittest.mock import create_autospec import httpx import pytest import respx +import tiktoken import litellm from litellm._logging import verbose_router_logger @@ -16,12 +17,14 @@ from litellm.constants import INTERNAL_CALL_ORIGIN_METADATA_KEY from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter +from litellm.llms.typesafe.decisions.transformation import TypeSafeDecisionsConfig +from litellm.router_strategy.complexity_router.complexity_router import ComplexityRouter, _ClassifierCircuitBreaker from litellm.router_strategy.complexity_router.config import ComplexityRouterConfig, JevClassifierConfig from litellm.router_strategy.complexity_router.jev_classifier import ( DEFAULT_JEV_INSTRUCTIONS, HttpJevClassifierClient, JevChoiceAnswer, + JevSystemOneRequest, JevSystemOneResponse, JevUsage, build_jev_request, @@ -1055,3 +1058,214 @@ def test_switching_to_deployed_classifier_does_not_copy_legacy_credentials() -> ) assert write.effective is not None assert write.effective["opensource_classifier_config"] == {"deployment_name": "registered"} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("large_field", ["user", "system"]) +@pytest.mark.parametrize("disable_counter", [False, True]) +async def test_jev_bounds_complete_input_without_changing_forwarded_messages( + large_field: str, disable_counter: bool, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_token_counter", disable_counter) + limit: Final = TypeSafeDecisionsConfig().get_input_limits("jev-latest", "typesafe", {}).max_question_tokens + assert limit is not None + rubric: Final = "Operator routing rubric " + "document the task category " * 500 + criterion: Final = ("Simple tasks " + "the classification criteria " * 15).rstrip() + large_text: Final = "BEGIN_SOURCE " + "budget 预算 function(x) return x; " * 18000 + " END_SOURCE" + ask: Final = large_text if large_field == "user" else "CURRENT_ASK Explain the example" + constraints: Final = large_text if large_field == "system" else "Use concise language" + messages: Final = [ + {"role": "system", "content": constraints}, + {"role": "user", "content": "PRIOR_TURN Explain the budget"}, + {"role": "user", "content": ask}, + ] + original: Final = deepcopy(messages) + handler: Final = AsyncHTTPHandler() + handler.client = httpx.AsyncClient() + classifier: Final = ComplexityRouter( + "bounded-classifier", litellm.Router(model_list=[]), + { + "classifier_type": "oss_classifier", + "opensource_classifier_config": {"instructions": rubric}, + "tier_definitions": [{"name": "SIMPLE", "description": criterion}, {"name": "REASONING"}], + "tiers": {"SIMPLE": "cheap", "REASONING": "deep"}, + "fallback_tier": "SIMPLE", + "session_affinity": False, + "deployment_affinity": False, + }, + jev_client=HttpJevClassifierClient("test", "https://budget.test", handler), + derive_savings_baseline=False, + ) + try: + with respx.mock(assert_all_called=True) as upstream: + route: Final = upstream.post("https://budget.test/v1/systemone").respond( + json={"answers": {"tier": _answer().model_dump()}} + ) + result: Final = await classifier.async_pre_routing_hook( + model="bounded-classifier", request_kwargs={"messages": messages}, messages=messages + ) + sent: Final = JevSystemOneRequest.model_validate_json(route.calls[0].request.content) + complete_input: Final = sent.model_dump_json(exclude={"model"}) + measured: Final = len(tiktoken.get_encoding("o200k_base").encode(complete_input)) + assert measured < limit + assert sent.questions["tier"].instructions == rubric + assert sent.questions["tier"].criteria["SIMPLE"] == criterion + assert sent.state.count("BEGIN_SOURCE") == sent.state.count("END_SOURCE") == 1 + assert sent.state.count("PRIOR_TURN") == 1 + assert result is not None and result.routing_decision is not None + decision: Final = result.routing_decision + assert decision["cause"] == "jev_classifier" + assert decision["classifier_input_budget_status"] == "truncated" + assert decision["classifier_input_truncated"] is True + assert measured <= decision["classifier_input_tokens"] <= decision["classifier_input_budget_tokens"] < limit + assert messages == original + assert result.messages is None or result.messages == original + if disable_counter: + assert len(complete_input.encode()) < limit + assert decision["classifier_input_count_method"] == "utf8_bytes" + finally: + await GLOBAL_LOGGING_WORKER.flush() + await handler.client.aclose() + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("mode", "default_header", "explicit_cap", "other_member", "expected_status"), + [ + ("false", None, None, None, "truncated"), + ("true", None, None, None, "within_budget"), + ("true", "false", None, None, "truncated"), + ("true", None, 8192, None, "truncated"), + ("true", None, 8192, "default", "truncated"), + ("true", None, None, "default", "truncated"), + ("false", None, None, "unknown", "unknown_limits"), + ], +) +async def test_jev_budget_uses_effective_headers_and_every_deployment( + mode: str, default_header: str | None, explicit_cap: int | None, + other_member: str | None, expected_status: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setitem(litellm.model_cost, "typesafe/jev-latest", deepcopy(litellm.model_cost["typesafe/jev-latest"])) + header: Final = "TypeSafe-Danger-Zone-Long-Context" + deployment: Final = { + "model_name": "budget-model", + "litellm_params": { + "model": "typesafe/jev-latest", "api_key": "test", "api_base": "https://budget.test", + "extra_headers": {header: mode}, + }, + "model_info": {"max_input_tokens": explicit_cap} if explicit_cap is not None else {}, + } + sibling: Final = { + "model_name": "budget-model", + "litellm_params": { + "model": "databricks/unknown-budget-endpoint" if other_member == "unknown" else "typesafe/jev-latest", + "api_key": "test", "api_base": "https://budget.test", + }, + } + router: Final = litellm.Router( + model_list=[deployment, sibling] if other_member else [deployment], + default_litellm_params={"extra_headers": {header: default_header}} if default_header else {}, + num_retries=0, + ) + classifier: Final = ComplexityRouter( + "bounded-classifier", router, + {"classifier_type": "oss_classifier", "opensource_classifier_config": {"deployment_name": "budget-model"}, + "tiers": {"SIMPLE": "cheap"}}, + derive_savings_baseline=False, + ) + ask: Final = "Explain this code " + "word001 word002 word003 " * 12000 + with respx.mock(assert_all_called=True) as upstream: + route: Final = upstream.route(method="POST").respond( + json={"answers": {"tier": _answer().model_dump()}} + ) + outcome: Final = await classifier.aclassify(ask) + assert outcome.cause == "jev_classifier" + telemetry: Final = outcome.classifier_input_telemetry + assert telemetry is not None and telemetry["classifier_input_budget_status"] == expected_status + assert telemetry["classifier_input_tokens"] <= telemetry["classifier_input_budget_tokens"] + sent: Final = JevSystemOneRequest.model_validate_json(route.calls[0].request.content) + if other_member is None: + assert route.calls[0].request.headers[header] == (default_header or mode) + if explicit_cap is not None: + assert telemetry["classifier_input_budget_tokens"] < explicit_cap + assert (ask in sent.state) is (expected_status == "within_budget") + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("oversized_field", "fallback"), + [("instructions", "heuristic"), ("instructions", "default_model"), ("criteria", "custom")], +) +async def test_jev_untrimmable_rubric_uses_configured_fallback_without_a_provider_call( + oversized_field: str, fallback: str, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setitem(litellm.model_cost, "typesafe/jev-latest", deepcopy(litellm.model_cost["typesafe/jev-latest"])) + router: Final = litellm.Router(model_list=[{ + "model_name": "budget-model", + "litellm_params": {"model": "typesafe/jev-latest", "api_key": "test", "api_base": "https://budget.test"}, + "model_info": {"max_input_tokens": 128 if oversized_field == "criteria" else 1024}, + }]) + rubric: Final = "word001 word002 word003 " * 2000 + classifier: Final = ComplexityRouter( + "bounded-classifier", router, + {"classifier_type": "oss_classifier", "opensource_classifier_config": { + "deployment_name": "budget-model", "instructions": rubric if oversized_field == "instructions" else "Choose", + }, "tiers": {"SIMPLE": "cheap", "REASONING": "deep"}, + **({"tier_definitions": [{"name": name, "description": "word001 " * 50} + for name in ("SIMPLE", "REASONING")], "fallback_tier": "SIMPLE"} + if fallback == "custom" else {"classifier_fallback": fallback, "default_model": "default"})}, + derive_savings_baseline=False, + ) + with respx.mock(assert_all_called=False) as upstream: + route: Final = upstream.route(method="POST").respond(json={"answers": {"tier": _answer().model_dump()}}) + outcome: Final = await classifier.aclassify("What is two plus two?") + assert route.call_count == 0 + assert outcome.cause == { + "heuristic": "heuristic_scorer", "default_model": "default_model_fallback", "custom": "classifier_fallback", + }[fallback] + assert outcome.classifier_failure_reason == "unsupported_input" + assert outcome.classifier_input_telemetry is not None + assert outcome.classifier_input_telemetry["classifier_input_budget_status"] == "untrimmable_input" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fallback", ["heuristic", "default_model"]) +async def test_jev_open_breaker_preserves_input_budget_in_fallback_decision(fallback: str) -> None: + def clock() -> float: + return 100.0 + + transport: Final = create_autospec(httpx.AsyncBaseTransport, instance=True) + transport.handle_async_request.side_effect = TimeoutError("controlled classifier timeout") + handler: Final = AsyncHTTPHandler() + handler.client = httpx.AsyncClient(transport=transport) + classifier: Final = ComplexityRouter( + "breaker-budget", litellm.Router(model_list=[]), + {"classifier_type": "oss_classifier", "opensource_classifier_config": {}, + "tiers": {"SIMPLE": "cheap", "REASONING": "deep"}, "classifier_fallback": fallback, + "default_model": "default", "session_affinity": False, "deployment_affinity": False}, + jev_client=HttpJevClassifierClient("test", "https://budget.test", handler), + derive_savings_baseline=False, + ) + classifier._classifier_circuit_breaker = _ClassifierCircuitBreaker(30.0, clock=clock) + try: + first: Final = await classifier.async_pre_routing_hook( + model="breaker-budget", request_kwargs={}, messages=[{"role": "user", "content": "Explain budgets"}] + ) + second: Final = await classifier.async_pre_routing_hook( + model="breaker-budget", request_kwargs={}, + messages=[{"role": "user", "content": "Explain this code " + "word001 word002 word003 " * 12000}], + ) + assert first is not None and first.routing_decision is not None + assert first.routing_decision["classifier_failure_reason"] == "timeout" + assert second is not None and second.routing_decision is not None + decision: Final = second.routing_decision + assert decision["classifier_failure_reason"] == "circuit_open" + assert decision["cause"] == ("heuristic_scorer" if fallback == "heuristic" else "default_model_fallback") + assert decision.get("classifier_input_budget_status") == "truncated" + assert decision["classifier_input_truncated"] is True + assert 0 < decision["classifier_input_tokens"] <= decision["classifier_input_budget_tokens"] + transport.handle_async_request.assert_awaited_once() + finally: + await handler.client.aclose() diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 96c7c34621b..24736cdb34e 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -46205,6 +46205,16 @@ export interface components { * @enum {string} */ classifier_failure_reason?: "timeout" | "circuit_open" | "not_configured" | "unsupported_input" | "invalid_response" | "declined" | "classifier_error"; + /** Classifier Input Budget Status */ + classifier_input_budget_status?: string; + /** Classifier Input Budget Tokens */ + classifier_input_budget_tokens?: number; + /** Classifier Input Count Method */ + classifier_input_count_method?: string; + /** Classifier Input Tokens */ + classifier_input_tokens?: number; + /** Classifier Input Truncated */ + classifier_input_truncated?: boolean; /** Classifier Max Quality Gap */ classifier_max_quality_gap?: number; /** Classifier Model */ From e8f572b39bff31e8b59cc01e00d5adbb8e278405 Mon Sep 17 00:00:00 2001 From: jesus-berri Date: Fri, 30 Oct 2026 14:39:25 -0700 Subject: [PATCH 26/46] feat(sso): create the team_id_jwt_field team on UI SSO login when team_id_upsert is set (#49256) * feat(sso): create the team_id_jwt_field team on UI SSO login when team_id_upsert is set * fix(sso): keep the upserted team_id_jwt_field team out of result.team_ids so restricted_sso_group and the access-token group fallback are unchanged * fix(sso): only upsert the team_id_jwt_field team on UI SSO login when enable_jwt_auth is on The API JWT path only reaches team_id_upsert inside the enable_jwt_auth branch, so a deployment with a leftover litellm_jwtauth block and JWT auth off never created teams. Gate the UI SSO upsert the same way * fix(sso): read enable_jwt_auth through the proxy_server module so the basedpyright gate stays at its ceiling --- litellm/proxy/management_endpoints/types.py | 1 + litellm/proxy/management_endpoints/ui_sso.py | 67 +++- .../management/test_sso_sign_in.py | 110 ++++++- .../proxy/management_endpoints/test_ui_sso.py | 309 ++++++++++++++++-- 4 files changed, 448 insertions(+), 39 deletions(-) diff --git a/litellm/proxy/management_endpoints/types.py b/litellm/proxy/management_endpoints/types.py index 221a6221fab..d4fd7a31d13 100644 --- a/litellm/proxy/management_endpoints/types.py +++ b/litellm/proxy/management_endpoints/types.py @@ -93,5 +93,6 @@ def get_litellm_user_role(role_str: object) -> LitellmUserRoles | None: class CustomOpenID(OpenID): team_ids: list[str] + team_id_to_upsert: str | None = None user_role: LitellmUserRoles | None = None extra_fields: dict[str, object] | None = None diff --git a/litellm/proxy/management_endpoints/ui_sso.py b/litellm/proxy/management_endpoints/ui_sso.py index 59c2df91d77..53079d84fdd 100644 --- a/litellm/proxy/management_endpoints/ui_sso.py +++ b/litellm/proxy/management_endpoints/ui_sso.py @@ -87,7 +87,7 @@ from litellm.proxy._types import ( TeamMemberAddRequest, UserAPIKeyAuth, ) -from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_user_object +from litellm.proxy.auth.auth_checks import ExperimentalUIJWTToken, get_team_object, get_user_object from litellm.proxy.auth.auth_utils import ( # noqa: F401 # legacy module exports _get_request_ip_address, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export get_request_ip_address, @@ -1266,6 +1266,7 @@ def generic_response_convertor( sso_jwt_handler: JWTHandler | None = None, role_mappings: Optional["RoleMappings"] = None, team_mappings: Optional["TeamMappings"] = None, + jwt_auth_enabled: bool = False, ) -> CustomOpenID: generic_user_id_attribute_name: Final = os.getenv("GENERIC_USER_ID_ATTRIBUTE", "preferred_username") generic_user_display_name_attribute_name: Final = os.getenv("GENERIC_USER_DISPLAY_NAME_ATTRIBUTE", "sub") @@ -1286,14 +1287,15 @@ def generic_response_convertor( generic_user_email_attribute_name, ) - all_teams: Final = [] + token: Final = cast(dict[str, object], response) # cast-ok: the IdP userinfo JSON object fastapi-sso hands back + all_teams: Final = list[str]() if sso_jwt_handler is not None: - team_ids = sso_jwt_handler.get_all_jwt_team_ids(cast(dict, response)) + team_ids = sso_jwt_handler.get_all_jwt_team_ids(token) all_teams.extend(team_ids) if team_mappings is not None and team_mappings.team_ids_jwt_field is not None: team_ids_from_db_mapping: Final[list[str] | None] = get_nested_value( - data=cast(dict, response), + data=token, key_path=team_mappings.team_ids_jwt_field, default=[], ) @@ -1305,8 +1307,17 @@ def generic_response_convertor( team_ids_from_db_mapping, ) else: - team_ids = jwt_handler.get_all_jwt_team_ids(cast(dict, response)) + team_ids = jwt_handler.get_all_jwt_team_ids(token) all_teams.extend(team_ids) + team_id_jwt_field: Final = jwt_handler.litellm_jwtauth.team_id_jwt_field + upsert_candidate: Final[object] = ( + get_nested_value(data=token, key_path=team_id_jwt_field) + if jwt_auth_enabled and jwt_handler.litellm_jwtauth.team_id_upsert and team_id_jwt_field is not None + else None + ) + team_id_to_upsert: Final = ( + upsert_candidate if isinstance(upsert_candidate, str) and upsert_candidate.strip() else None + ) # Determine user role based on role_mappings if available # Only apply role_mappings for GENERIC SSO provider @@ -1375,6 +1386,7 @@ def generic_response_convertor( last_name=get_nested_value(response, generic_user_last_name_attribute_name), provider=get_nested_value(response, generic_provider_attribute_name), team_ids=all_teams, + team_id_to_upsert=team_id_to_upsert, user_role=user_role, extra_fields=extra_fields, ) @@ -1614,6 +1626,8 @@ async def get_generic_sso_response( # make generic sso provider from fastapi_sso.sso.base import DiscoveryDocument + from litellm.proxy import proxy_server + received_response: dict | None = None sso_assertion: SSOIdentityAssertion | None = None @@ -1665,6 +1679,7 @@ async def get_generic_sso_response( sso_jwt_handler=sso_jwt_handler, role_mappings=role_mappings, team_mappings=team_mappings, + jwt_auth_enabled=proxy_server.general_settings.get("enable_jwt_auth", False) is True, ) SSOProvider: Final = create_generic_sso_provider( @@ -1884,6 +1899,29 @@ async def find_existing_team_ids(team_ids: Sequence[str]) -> frozenset[str] | No return frozenset(row.team_id for row in rows) +async def create_sso_team(team_id: str) -> None: + """Create the ``team_id_jwt_field`` team on UI SSO login the way the API JWT path does under ``team_id_upsert``.""" + from litellm.proxy.proxy_server import prisma_client, user_api_key_cache + + if prisma_client is None: + return + await get_team_object( + team_id=team_id, + prisma_client=prisma_client, + user_api_key_cache=user_api_key_cache, + check_db_only=True, + team_id_upsert=True, + ) + + +async def _create_sso_team_or_warn(team_id: str, create_team: Callable[[str], Awaitable[object]]) -> None: + """A concurrent login may create the team first, so a failed create is not fatal: the join still runs.""" + try: + await create_team(team_id) + except Exception as e: # noqa: BLE001 # fail-safe: a failed team create must not fail the login or skip the join + verbose_proxy_logger.warning("SSO login could not create team %s, trying to join it anyway: %s", team_id, e) + + async def add_missing_team_member( user_info: NewUserResponse | LiteLLM_UserTable, sso_teams: Sequence[str], @@ -1892,21 +1930,30 @@ async def add_missing_team_member( [str, NewUserResponse | LiteLLM_UserTable], Awaitable[object] ] = create_team_member_add_task, existing_team_ids: Callable[[Sequence[str]], Awaitable[frozenset[str] | None]] = find_existing_team_ids, + team_id_to_upsert: str | None = None, + create_team: Callable[[str], Awaitable[object]] = create_sso_team, ) -> None: """ - Get missing teams (diff b/w user_info.team_ids and sso_teams) - Drop the IdP groups that are not teams with one lookup, since group claims often carry many of them + - Join the ``team_id_to_upsert`` team too, creating it when it is missing (``team_id_upsert``), the way the API JWT path does; + it stays out of ``result.team_ids`` so restricted_sso_group and the access-token group fallback see only the claimed groups - Add missing user to missing teams, one at a time in sso_teams order, so user.teams keeps the claim's order """ held_teams: Final = frozenset(user_info.teams or ()) - missing_teams: Final = tuple(team_id for team_id in dict.fromkeys(sso_teams) if team_id not in held_teams) + claimed_teams: Final = sso_teams if team_id_to_upsert is None else (*sso_teams, team_id_to_upsert) + missing_teams: Final = tuple(team_id for team_id in dict.fromkeys(claimed_teams) if team_id not in held_teams) if not missing_teams: return existing: Final = await existing_team_ids(missing_teams) teams_to_add: Final = ( - missing_teams if existing is None else tuple(team_id for team_id in missing_teams if team_id in existing) + missing_teams + if existing is None + else tuple(team_id for team_id in missing_teams if team_id in existing or team_id == team_id_to_upsert) ) for team_id in teams_to_add: + if team_id == team_id_to_upsert and (existing is None or team_id not in existing): + await _create_sso_team_or_warn(team_id, create_team) await add_team_member(team_id, user_info) @@ -3663,7 +3710,11 @@ class SSOAuthenticationHandler: verbose_proxy_logger.debug("User not found in LiteLLM DB, skipping team member addition") return sso_teams: Final = result.team_ids if isinstance(result, CustomOpenID) else [] - await add_missing_team_member(user_info=user_info, sso_teams=sso_teams) + await add_missing_team_member( + user_info=user_info, + sso_teams=sso_teams, + team_id_to_upsert=result.team_id_to_upsert if isinstance(result, CustomOpenID) else None, + ) @staticmethod def verify_user_in_restricted_sso_group( diff --git a/tests/integration/management/test_sso_sign_in.py b/tests/integration/management/test_sso_sign_in.py index d0b6e2579c8..93cfa00ea13 100644 --- a/tests/integration/management/test_sso_sign_in.py +++ b/tests/integration/management/test_sso_sign_in.py @@ -16,7 +16,7 @@ import yaml from pydantic import JsonValue from tests.integration._support.client import Gateway, gateway_from_environment -from tests.integration._support.database import read_rows +from tests.integration._support.database import read_rows, write_rows from tests.integration._support.process import OWNED_STATE_VARIABLE, OwnedState, owned_proxy from tests.integration._support.wire import Reply, Request, wire_server @@ -52,14 +52,29 @@ def _idp_reply(request: Request) -> Reply: "preferred_username": signed_in, "email": f"{signed_in}@example.com", "provider": IDP_PROVIDER, + "dept": _dept(signed_in), + "groups": [_group(signed_in)], } ).encode() ) return Reply(status=404, body=b'{"error": "not_found"}') +def _dept(subject: str) -> str: + return f"dept-{subject}" + + +def _group(subject: str) -> str: + return f"group-{subject}" + + @contextmanager -def _sso_gateway(directory: Path, litellm_settings: Mapping[str, JsonValue]) -> Iterator[Gateway]: +def _sso_gateway( + directory: Path, + litellm_settings: Mapping[str, JsonValue], + general_settings: Mapping[str, JsonValue] | None = None, + db_sso_settings: Mapping[str, JsonValue] | None = None, +) -> Iterator[Gateway]: config: Final = directory / "sso_proxy_config.yaml" config.write_text( yaml.safe_dump( @@ -68,6 +83,7 @@ def _sso_gateway(directory: Path, litellm_settings: Mapping[str, JsonValue]) -> "general_settings": { "master_key": "os.environ/LITELLM_MASTER_KEY", "database_url": "os.environ/DATABASE_URL", + **(general_settings or {}), }, "litellm_settings": dict(litellm_settings), } @@ -83,7 +99,33 @@ def _sso_gateway(directory: Path, litellm_settings: Mapping[str, JsonValue]) -> "OAUTHLIB_INSECURE_TRANSPORT": "1", } with owned_proxy(rig, directory, environment, config=config, remove_environment=("PROXY_BASE_URL",)) as gw: - yield gw + if db_sso_settings is None: + yield gw + return + with contextlib.ExitStack() as cleanups: + cleanups.callback( + write_rows, + 'DELETE FROM "LiteLLM_SSOConfig" WHERE id = %s', + ("sso_config",), + database_url=gw.database_url, + ) + stored: Final = _store_sso_settings(gw, {**_generic_sso_settings(environment), **db_sso_settings}) + assert stored.status_code == 200, stored.text + yield gw + + +def _generic_sso_settings(environment: Mapping[str, str]) -> Mapping[str, JsonValue]: + return { + "generic_client_id": environment["GENERIC_CLIENT_ID"], + "generic_client_secret": environment["GENERIC_CLIENT_SECRET"], + "generic_authorization_endpoint": environment["GENERIC_AUTHORIZATION_ENDPOINT"], + "generic_token_endpoint": environment["GENERIC_TOKEN_ENDPOINT"], + "generic_userinfo_endpoint": environment["GENERIC_USERINFO_ENDPOINT"], + } + + +def _store_sso_settings(gateway: Gateway, settings: Mapping[str, JsonValue]) -> httpx.Response: + return gateway.request("PATCH", "/update/sso_settings", settings) def _sign_in(gateway: Gateway, subject: str) -> httpx.Response: @@ -165,3 +207,65 @@ def test_free_sso_admission_preserves_existing_users_and_serializes_the_last_sea refused: Final = _sign_in(gateway, "free-sso-extra") assert (above_capacity.status_code, refused.status_code) == (303, 403) assert gateway.rows('SELECT COUNT(*) AS count FROM "LiteLLM_UserTable"', ()) == [{"count": 6}] + + +@pytest.mark.parametrize( + ("enable_jwt_auth", "team_id_upsert", "team_exists_before", "expect_member"), + [(True, True, False, True), (False, True, False, False), (True, False, False, False), (True, False, True, True)], + ids=[ + "upsert-creates-missing-team", + "upsert-without-jwt-auth-skips-missing-team", + "no-upsert-skips-missing-team", + "no-upsert-joins-existing-team", + ], +) +def test_sso_sign_in_upserts_the_team_id_jwt_field_team_only_when_team_id_upsert_is_on( + tmp_path: Path, enable_jwt_auth: bool, team_id_upsert: bool, team_exists_before: bool, expect_member: bool +) -> None: + subject: Final = f"sso-dept-user-{uuid.uuid4().hex[:12]}" + team_id: Final = _dept(subject) + jwtauth: Final = { + "enable_jwt_auth": enable_jwt_auth, + "litellm_jwtauth": {"team_id_jwt_field": "dept", "team_id_upsert": team_id_upsert}, + } + with _sso_gateway(tmp_path, {}, jwtauth) as gateway, contextlib.ExitStack() as cleanups: + cleanups.callback(gateway.request, "POST", "/team/delete", {"team_ids": [team_id]}) + cleanups.callback(_delete_user, gateway, subject) + if team_exists_before: + created: Final = gateway.request("POST", "/team/new", {"team_id": team_id}) + assert created.status_code == 200, created.text + + callback: Final = _sign_in(gateway, subject) + + assert callback.status_code == 303, f"{callback.status_code} {callback.text}" + teams: Final = gateway.rows('SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id = %s', (team_id,)) + user: Final = gateway.rows('SELECT teams FROM "LiteLLM_UserTable" WHERE user_id = %s', (subject,)) + assert (len(teams), user) == (int(expect_member), [{"teams": [team_id] if expect_member else []}]) + + +def test_sso_sign_in_with_db_team_mappings_creates_the_team_id_jwt_field_team_and_only_joins_groups( + tmp_path: Path, +) -> None: + subject: Final = f"sso-mapped-user-{uuid.uuid4().hex[:12]}" + team_id: Final = _dept(subject) + group_team_id: Final = _group(subject) + jwtauth: Final = { + "enable_jwt_auth": True, + "litellm_jwtauth": {"team_id_jwt_field": "dept", "team_id_upsert": True}, + } + team_mappings: Final = {"team_mappings": {"team_ids_jwt_field": "groups"}} + with ( + _sso_gateway(tmp_path, {}, jwtauth, db_sso_settings=team_mappings) as gateway, + contextlib.ExitStack() as cleanups, + ): + cleanups.callback(gateway.request, "POST", "/team/delete", {"team_ids": [team_id, group_team_id]}) + cleanups.callback(_delete_user, gateway, subject) + + callback: Final = _sign_in(gateway, subject) + + assert callback.status_code == 303, f"{callback.status_code} {callback.text}" + teams: Final = gateway.rows( + 'SELECT team_id FROM "LiteLLM_TeamTable" WHERE team_id IN (%s, %s)', (team_id, group_team_id) + ) + user: Final = gateway.rows('SELECT teams FROM "LiteLLM_UserTable" WHERE user_id = %s', (subject,)) + assert (teams, user) == ([{"team_id": team_id}], [{"teams": [team_id]}]) diff --git a/tests/unit/proxy/management_endpoints/test_ui_sso.py b/tests/unit/proxy/management_endpoints/test_ui_sso.py index 136cc86e67a..ba78fefd653 100644 --- a/tests/unit/proxy/management_endpoints/test_ui_sso.py +++ b/tests/unit/proxy/management_endpoints/test_ui_sso.py @@ -15,9 +15,15 @@ from fastapi import HTTPException, Request from pydantic import SecretStr import litellm +from litellm._logging import verbose_proxy_logger from litellm._uuid import uuid from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler -from litellm.proxy._types import LiteLLM_UserTable, NewUserResponse, ProxyException +from litellm.proxy._types import ( + LiteLLM_JWTAuth, + LiteLLM_UserTable, + NewUserResponse, + ProxyException, +) from litellm.proxy.auth.handle_jwt import JWTHandler from litellm.proxy.management_endpoints.sso import CustomMicrosoftSSO from litellm.proxy.management_endpoints.types import CustomOpenID @@ -287,6 +293,13 @@ def test_get_microsoft_callback_response(stubbed_graph_api): assert sso_assertion.issuer == "https://idp.example.com" + +def _jwt_handler_double() -> MagicMock: + handler: Final = MagicMock(spec=JWTHandler) + handler.litellm_jwtauth = LiteLLM_JWTAuth() + return handler + + GRAPH_ERROR_BODIES: Final = ( {"error": {"code": "Authentication_MissingOrMalformed", "message": "Access Token missing or malformed."}}, {"error": {"code": "InvalidAuthenticationToken", "message": "IDX14100: JWT is not well formed."}}, @@ -1158,7 +1171,7 @@ def test_generic_response_convertor_normalizes_email(): } # Mock JWT handler - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] result = generic_response_convertor( @@ -2226,7 +2239,7 @@ async def test_get_generic_sso_response_with_additional_headers(): # Arrange mock_request = MagicMock(spec=Request) - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] generic_client_id = "test_client_id" @@ -2299,7 +2312,7 @@ async def test_get_generic_sso_response_with_empty_headers(): # Arrange mock_request = MagicMock(spec=Request) - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] generic_client_id = "test_client_id" @@ -2359,8 +2372,8 @@ async def test_get_generic_sso_response_includes_token_claims_when_enabled(monke from litellm.proxy.management_endpoints.ui_sso import get_generic_sso_response mock_request = MagicMock(spec=Request) - mock_jwt_handler = MagicMock(spec=JWTHandler) - mock_sso_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() + mock_sso_jwt_handler = _jwt_handler_double() mock_sso_jwt_handler.get_all_jwt_team_ids.return_value = ["team-from-userinfo"] mock_sso_jwt_handler.get_team_ids_from_jwt.return_value = [] @@ -2435,8 +2448,8 @@ async def test_get_generic_sso_response_does_not_include_token_claims_when_disab from litellm.proxy.management_endpoints.ui_sso import get_generic_sso_response mock_request = MagicMock(spec=Request) - mock_jwt_handler = MagicMock(spec=JWTHandler) - mock_sso_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() + mock_sso_jwt_handler = _jwt_handler_double() mock_sso_jwt_handler.get_all_jwt_team_ids.return_value = ["team-from-userinfo"] mock_sso_jwt_handler.get_team_ids_from_jwt.return_value = [] access_token = pyjwt.encode({"upn": "token-user@example.com"}, "test-secret", algorithm="HS256") @@ -2578,7 +2591,7 @@ async def test_get_generic_sso_response_pkce_merges_token_claims_and_excludes_cr ): result, received_response, _, _ = await get_generic_sso_response( request=request, - jwt_handler=MagicMock(spec=JWTHandler), + jwt_handler=_jwt_handler_double(), generic_client_id="test-client", redirect_url="http://test.com/callback", sso_jwt_handler=None, @@ -2600,7 +2613,7 @@ async def test_get_generic_sso_response_ignores_opaque_and_empty_token_claims(mo from litellm.proxy.management_endpoints.ui_sso import get_generic_sso_response mock_request = MagicMock(spec=Request) - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() userinfo = { "preferred_username": "userinfo-user", "email": "userinfo@example.com", @@ -4961,7 +4974,7 @@ class TestProcessSSOJWTAccessToken: @pytest.fixture def mock_jwt_handler(self): """Create a mock JWT handler for testing""" - mock_handler = MagicMock(spec=JWTHandler) + mock_handler = _jwt_handler_double() mock_handler.get_team_ids_from_jwt.return_value = ["team1", "team2", "team3"] return mock_handler @@ -5348,7 +5361,7 @@ class TestGenericResponseConvertorNestedAttributes: from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor # Mock JWT handler - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] # Payload with nested attributes structure @@ -5419,7 +5432,7 @@ class TestGenericResponseConvertorUserRole: from litellm.proxy._types import LitellmUserRoles from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] sso_response = { @@ -5449,7 +5462,7 @@ class TestGenericResponseConvertorUserRole: """ from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] sso_response = { @@ -6813,7 +6826,9 @@ class TestAddMissingTeamMember: # Verify add_missing_team_member was called with correct args mock_add_member.assert_called_once_with( - user_info=new_user_info, sso_teams=["entra-group-1", "entra-group-2"] + user_info=new_user_info, + sso_teams=["entra-group-1", "entra-group-2"], + team_id_to_upsert=None, ) @pytest.mark.asyncio @@ -7897,7 +7912,7 @@ def test_generic_response_convertor_with_extra_attributes(monkeypatch): "custom_field3": {"nested": "data"}, } - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] result = generic_response_convertor( @@ -7931,7 +7946,7 @@ def test_generic_response_convertor_without_extra_attributes(monkeypatch): "custom_field2": "value2", } - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] result = generic_response_convertor( @@ -7959,7 +7974,7 @@ def test_generic_response_convertor_extra_attributes_with_nested_paths(monkeypat "org_info": {"department": "Engineering", "manager": "Jane Smith"}, } - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] result = generic_response_convertor( @@ -7986,7 +8001,7 @@ def test_generic_response_convertor_extra_attributes_missing_field(monkeypatch): "email": "user@example.com", } - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] result = generic_response_convertor( @@ -8865,7 +8880,7 @@ class TestPKCEStateCookieBinding: ): await get_generic_sso_response( request=mock_request, - jwt_handler=MagicMock(spec=JWTHandler), + jwt_handler=_jwt_handler_double(), generic_client_id="cid", redirect_url="https://proxy.example.com/sso/callback", sso_jwt_handler=None, @@ -8917,7 +8932,7 @@ class TestPKCEStateCookieBinding: ): await get_generic_sso_response( request=mock_request, - jwt_handler=MagicMock(spec=JWTHandler), + jwt_handler=_jwt_handler_double(), generic_client_id="cid", redirect_url="https://proxy.example.com/sso/callback", sso_jwt_handler=None, @@ -8980,7 +8995,7 @@ class TestPKCEStateCookieBinding: }, ), ): - jwt_handler = MagicMock(spec=JWTHandler) + jwt_handler = _jwt_handler_double() jwt_handler.get_team_ids_from_jwt.return_value = [] result, _, _, _ = await get_generic_sso_response( request=mock_request, @@ -9087,7 +9102,7 @@ async def test_debug_sso_callback_renders_full_jwt_claims(): patch("litellm.proxy.proxy_server.general_settings", {}), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), - patch("litellm.proxy.proxy_server.jwt_handler", MagicMock(spec=JWTHandler)), + patch("litellm.proxy.proxy_server.jwt_handler", _jwt_handler_double()), ): # Microsoft / Google envs may leak in from other tests — ensure only # the generic path runs. @@ -9173,7 +9188,7 @@ async def test_debug_sso_callback_handles_missing_raw_response(): patch("litellm.proxy.proxy_server.general_settings", {}), patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), - patch("litellm.proxy.proxy_server.jwt_handler", MagicMock(spec=JWTHandler)), + patch("litellm.proxy.proxy_server.jwt_handler", _jwt_handler_double()), ): for var in ("GENERIC_CLIENT_ID", "GOOGLE_CLIENT_ID"): os.environ.pop(var, None) @@ -9224,7 +9239,7 @@ async def _render_debug_page(provider_env, id_jag_registered, force_inert=False) patch("litellm.proxy.proxy_server.general_settings", {}), # test-quality-ok: debug endpoint reads proxy globals patch("litellm.proxy.proxy_server.prisma_client", MagicMock()), # test-quality-ok: debug endpoint reads proxy DB patch("litellm.proxy.proxy_server.user_api_key_cache", MagicMock()), # test-quality-ok: debug endpoint reads proxy globals - patch("litellm.proxy.proxy_server.jwt_handler", MagicMock(spec=JWTHandler)), # test-quality-ok: debug endpoint reads proxy globals + patch("litellm.proxy.proxy_server.jwt_handler", _jwt_handler_double()), # test-quality-ok: debug endpoint reads proxy globals ] if force_inert: stack.append( @@ -9589,7 +9604,7 @@ async def test_pkce_arm_captures_sso_assertion(): }, ), ): - jwt_handler = MagicMock(spec=JWTHandler) + jwt_handler = _jwt_handler_double() jwt_handler.get_team_ids_from_jwt.return_value = [] result, received_response, _, sso_assertion = await get_generic_sso_response( request=mock_request, @@ -9615,7 +9630,7 @@ async def test_verify_and_process_arm_captures_sso_assertion(): id_token = _ema_id_token() mock_request = MagicMock(spec=Request) - mock_jwt_handler = MagicMock(spec=JWTHandler) + mock_jwt_handler = _jwt_handler_double() mock_jwt_handler.get_team_ids_from_jwt.return_value = [] mock_sso_instance = MagicMock() @@ -9680,7 +9695,7 @@ def _generic_sso_callback_request() -> Request: def _generic_sso_jwt_handler() -> MagicMock: - jwt_handler = MagicMock(spec=JWTHandler) + jwt_handler = _jwt_handler_double() jwt_handler.get_team_ids_from_jwt.return_value = [] return jwt_handler @@ -11144,3 +11159,241 @@ async def test_cli_sso_login_never_serves_the_admin_login_form_on_a_ui_disabled_ body = response.body.decode() assert "Admin UI is Disabled" in body assert 'name="username"' not in body + + +def _upsert_jwt_handler(**jwtauth: object) -> JWTHandler: + from litellm.caching.caching import DualCache + from litellm.proxy._types import LiteLLM_JWTAuth + + handler: Final = JWTHandler() + handler.update_environment( + prisma_client=None, user_api_key_cache=DualCache(), litellm_jwtauth=LiteLLM_JWTAuth(**jwtauth) + ) + return handler + + +@pytest.mark.parametrize( + ("jwtauth", "userinfo", "expected_team_ids", "expected_upsert"), + [ + ({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": "finance"}, ["finance"], "finance"), + ({"team_id_jwt_field": "dept", "team_id_upsert": False}, {"dept": "finance"}, ["finance"], None), + ({"team_id_jwt_field": "dept", "team_id_upsert": True}, {}, [], None), + ({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": ""}, [], None), + ({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": " "}, [" "], None), + ({"team_id_jwt_field": "dept", "team_id_upsert": True}, {"dept": ["a", "b"]}, ["a", "b"], None), + ( + {"team_id_jwt_field": "dept", "team_id_upsert": True, "team_id_default": "everyone"}, + {}, + [], + None, + ), + ( + {"team_ids_jwt_field": "groups", "team_id_upsert": True}, + {"groups": ["g1", "g2"]}, + ["g1", "g2"], + None, + ), + ( + {"team_ids_jwt_field": "groups", "team_id_jwt_field": "dept", "team_id_upsert": True}, + {"groups": ["g1"], "dept": "finance"}, + ["g1", "finance"], + "finance", + ), + ], + ids=[ + "upsert-on", + "upsert-off", + "claim-missing", + "blank-claim-not-upserted", + "whitespace-claim-not-upserted", + "list-valued-claim-join-only", + "default-team-not-upserted", + "plural-claim-join-only", + "mixed", + ], +) +def test_generic_sso_login_upserts_only_the_team_id_jwt_field_claim( + jwtauth: dict[str, object], + userinfo: dict[str, object], + expected_team_ids: list[str], + expected_upsert: str | None, +) -> None: + from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor + + result: Final = generic_response_convertor( + response={"preferred_username": "alice", "email": "alice@example.com", **userinfo}, + jwt_handler=_upsert_jwt_handler(**jwtauth), + jwt_auth_enabled=True, + ) + + assert (result.team_ids, result.team_id_to_upsert) == (expected_team_ids, expected_upsert) + + +def test_generic_sso_login_does_not_upsert_the_team_id_jwt_field_claim_when_jwt_auth_is_disabled() -> None: + from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor + + result: Final = generic_response_convertor( + response={"preferred_username": "alice", "email": "alice@example.com", "dept": "finance"}, + jwt_handler=_upsert_jwt_handler(team_id_jwt_field="dept", team_id_upsert=True), + jwt_auth_enabled=False, + ) + + assert (result.team_ids, result.team_id_to_upsert) == (["finance"], None) + + +def test_generic_sso_login_marks_the_team_id_jwt_field_team_for_upsert_without_adding_it_to_team_ids() -> None: + from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor + + result: Final = generic_response_convertor( + response={"preferred_username": "alice", "email": "alice@example.com", "groups": ["g1"], "dept": "finance"}, + jwt_handler=_upsert_jwt_handler(team_id_jwt_field="dept", team_id_upsert=True), + team_mappings=TeamMappings(team_ids_jwt_field="groups"), + jwt_auth_enabled=True, + ) + + assert (result.team_ids, result.team_id_to_upsert) == (["g1"], "finance") + + +def test_generic_sso_login_keeps_team_ids_empty_for_a_dept_only_user_so_the_access_token_group_fallback_runs() -> None: + from litellm.proxy.management_endpoints.ui_sso import generic_response_convertor + + result: Final = generic_response_convertor( + response={"preferred_username": "alice", "email": "alice@example.com", "dept": "d1"}, + jwt_handler=_upsert_jwt_handler(team_id_jwt_field="dept", team_id_upsert=True), + team_mappings=TeamMappings(team_ids_jwt_field="groups"), + jwt_auth_enabled=True, + ) + + assert (result.team_ids, result.team_id_to_upsert) == ([], "d1") + + +def test_restricted_sso_group_still_rejects_a_user_whose_only_team_is_the_upserted_dept() -> None: + from litellm.proxy._types import ProxyException + from litellm.proxy.management_endpoints.ui_sso import SSOAuthenticationHandler, generic_response_convertor + + response: Final = {"preferred_username": "alice", "email": "alice@example.com", "dept": "X", "groups": []} + result: Final = generic_response_convertor( + response=response, + jwt_handler=_upsert_jwt_handler(team_id_jwt_field="dept", team_id_upsert=True), + team_mappings=TeamMappings(team_ids_jwt_field="groups"), + jwt_auth_enabled=True, + ) + + with pytest.raises(ProxyException) as denied: + SSOAuthenticationHandler.verify_user_in_restricted_sso_group( + general_settings={"ui_access_mode": {"type": "restricted_sso_group", "restricted_sso_group": "X"}}, + result=result, + received_response=response, + ) + + assert denied.value.code == "403" + + +def _recording_team_calls( + events: list[tuple[str, str]], +) -> tuple[Callable[[str, NewUserResponse | LiteLLM_UserTable], Awaitable[None]], Callable[[str], Awaitable[None]]]: + async def add(team_id: str, user_info: NewUserResponse | LiteLLM_UserTable) -> None: + events.append(("join", team_id)) + + async def create(team_id: str) -> None: + events.append(("create", team_id)) + + return add, create + + +async def _only_group_teams_exist(team_ids: Sequence[str]) -> frozenset[str] | None: + return frozenset(team_id for team_id in team_ids if team_id.startswith("g")) + + +@pytest.mark.asyncio +async def test_sso_login_joins_the_upsert_team_from_team_id_to_upsert_and_creates_it_before_the_join() -> None: + from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member + + events: Final[list[tuple[str, str]]] = [] # mutable-ok: records call order across awaited tasks + add, create = _recording_team_calls(events) + + await add_missing_team_member( + NewUserResponse(user_id="alice", key="sk-x", teams=None), + ["g1", "g2"], + add_team_member=add, + existing_team_ids=_only_group_teams_exist, + team_id_to_upsert="finance", + create_team=create, + ) + + assert events == [("join", "g1"), ("join", "g2"), ("create", "finance"), ("join", "finance")] + + +@pytest.mark.asyncio +async def test_sso_login_joins_only_the_claimed_groups_when_team_id_upsert_is_off() -> None: + from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member + + events: Final[list[tuple[str, str]]] = [] # mutable-ok: records call order across awaited tasks + add, create = _recording_team_calls(events) + + await add_missing_team_member( + NewUserResponse(user_id="alice", key="sk-x", teams=None), + ["g1", "g2"], + add_team_member=add, + existing_team_ids=_only_group_teams_exist, + create_team=create, + ) + + assert events == [("join", "g1"), ("join", "g2")] + + +@pytest.mark.asyncio +async def test_sso_login_does_not_create_the_team_id_jwt_field_team_when_it_already_exists() -> None: + from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member + + events: Final[list[tuple[str, str]]] = [] # mutable-ok: records call order across awaited tasks + add, create = _recording_team_calls(events) + + async def every_team_exists(team_ids: Sequence[str]) -> frozenset[str] | None: + return frozenset(team_ids) + + await add_missing_team_member( + NewUserResponse(user_id="alice", key="sk-x", teams=None), + ["finance"], + add_team_member=add, + existing_team_ids=every_team_exists, + team_id_to_upsert="finance", + create_team=create, + ) + + assert events == [("join", "finance")] + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "create_error", + [ + HTTPException(status_code=400, detail={"error": "Team id = finance already exists"}), + HTTPException(status_code=403, detail={"error": "team limit reached for this license"}), + ], + ids=["lost-the-create-race", "create-rejected"], +) +async def test_sso_login_still_joins_the_team_and_warns_when_creating_it_fails( + caplog: pytest.LogCaptureFixture, create_error: HTTPException +) -> None: + from litellm.proxy.management_endpoints.ui_sso import add_missing_team_member + + events: Final[list[tuple[str, str]]] = [] # mutable-ok: records call order across awaited tasks + add, _ = _recording_team_calls(events) + + async def failed_create(team_id: str) -> None: + raise create_error + + with caplog.at_level(logging.WARNING, logger=verbose_proxy_logger.name): + await add_missing_team_member( + NewUserResponse(user_id="alice", key="sk-x", teams=None), + ["finance"], + add_team_member=add, + existing_team_ids=_only_group_teams_exist, + team_id_to_upsert="finance", + create_team=failed_create, + ) + + warnings: Final = [r.getMessage() for r in caplog.records if r.levelno == logging.WARNING] + assert events == [("join", "finance")] + assert len(warnings) == 1 and "finance" in warnings[0] and str(create_error.detail) in warnings[0] From 7e5a653c0129c5ba99e828d15e5e2ff1f5690e90 Mon Sep 17 00:00:00 2001 From: shrey Date: Fri, 30 Oct 2026 15:03:49 -0700 Subject: [PATCH 27/46] chore(cost-map): add Deepgram Aura-1 and Aura-2 text-to-speech rates (#49280) --- ...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 1d18903bb07..843ab574da4 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -23078,6 +23078,24 @@ "output_cost_per_token_batches": 1e-06, "source": "https://developers.openai.com/api/docs/pricing" }, + "deepgram/aura": { + "input_cost_per_character": 1.5e-05, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.0150/1000 characters = $0.000015 per character", + "note": "Aura-1 text-to-speech voices (aura--), pay as you go" + }, + "source": "https://deepgram.com/pricing" + }, + "deepgram/aura-2": { + "input_cost_per_character": 3e-05, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.030/1000 characters = $0.00003 per character", + "note": "Aura-2 text-to-speech voices (aura-2--), pay as you go" + }, + "source": "https://deepgram.com/pricing" + }, "deepgram/base": { "input_cost_per_second": 0.00020833, "litellm_provider": "deepgram", diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 1d18903bb07..843ab574da4 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -23078,6 +23078,24 @@ "output_cost_per_token_batches": 1e-06, "source": "https://developers.openai.com/api/docs/pricing" }, + "deepgram/aura": { + "input_cost_per_character": 1.5e-05, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.0150/1000 characters = $0.000015 per character", + "note": "Aura-1 text-to-speech voices (aura--), pay as you go" + }, + "source": "https://deepgram.com/pricing" + }, + "deepgram/aura-2": { + "input_cost_per_character": 3e-05, + "litellm_provider": "deepgram", + "metadata": { + "calculation": "$0.030/1000 characters = $0.00003 per character", + "note": "Aura-2 text-to-speech voices (aura-2--), pay as you go" + }, + "source": "https://deepgram.com/pricing" + }, "deepgram/base": { "input_cost_per_second": 0.00020833, "litellm_provider": "deepgram", From faf9ee18e962c0dde19de46fc9190ab70f34a70f Mon Sep 17 00:00:00 2001 From: tin-berri Date: Fri, 30 Oct 2026 15:05:32 -0700 Subject: [PATCH 28/46] fix(cost): preserve zero prices across batch, speech and search (#49294) * fix(cost): preserve zero cache-write prices in batch fallback * fix(cost): retain batch cache charges when ordinary input is free * fix(cost): preserve zero speech and Perplexity search rates * fix(cost): reject unset speech prices before zero defaults * test(cost): consolidate zero-price regression coverage --- CONTRIBUTING.md | 2 + litellm/cost_calculator.py | 13 +++- .../litellm_core_utils/llm_cost_calc/utils.py | 7 +- litellm/llms/perplexity/cost_calculator.py | 5 +- litellm/types/utils.py | 2 + litellm/utils.py | 9 ++- .../test_perplexity_cost_calculator.py | 51 ++++++++++++- tests/unit/test_cost_calculator.py | 73 ++++++++++++++----- ui/litellm-dashboard/src/lib/http/schema.d.ts | 15 ++++ 9 files changed, 147 insertions(+), 30 deletions(-) diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md index 531a78d8982..e956b27d5da 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -276,6 +276,8 @@ When the provider reports its own usage or spend numbers, compare the proxy's nu Provider work (a new provider, a new parameter, a new usage or cost field) is the most common kind of contribution, and it is where PRs get reworked most often. Follow these rules +**Preserve zero prices.** In [`model_prices_and_context_window.json`](model_prices_and_context_window.json), an explicit `0.0` means that the priced operation is free. A missing field or `null` means no rate is specified, so the relevant calculator may use a fallback. For example, when no batch input rate is configured, batch cache writes use `cache_creation_input_token_cost` if set, otherwise the regular input rate, then apply the batch discount. A zero ordinary input rate does not make separately priced cache reads or writes free. Test `rate is None` before substituting a fallback, since `row.get("cache_creation_input_token_cost") or fallback` would discard a valid zero. Speech keeps character pricing ahead of token pricing even when the character rate is zero. If neither rate is configured, speech cost calculation raises an error. Its model lookup preserves unset rates instead of treating compatibility defaults as explicit zero prices. Perplexity keeps a configured legacy search rate ahead of the current search rate, including zero. When registering custom prices, `None` leaves an existing rate unchanged and `0.0` overrides it + **Find the existing machinery before you add plumbing.** LiteLLM already has generic paths that move data between the endpoints, the provider transformations, and cost tracking. A second, hand-rolled path for the same data is a common reason a provider PR gets rewritten before merge. [`ARCHITECTURE.md`](ARCHITECTURE.md) shows where the layers live, and the decision records in [`docs/adrs/`](docs/adrs/AGENTS.md) explain why they are shaped that way. For example, [the ADR on provider usage fields](docs/adrs/adr_001_provider_usage_fields_ride_as_usage_extras.md) covers how provider-specific usage fields reach cost tracking. If you cannot find a mechanism for what you need, say so in the PR or ask in [#pr-review on Slack](https://join.slack.com/t/litellmossslack/shared_invite/zt-3o7nkuyfr-p_kbNJj8taRfXGgQI1~YyA) before you build one **Bill a new provider usage field from the usage object.** Leave the field on the usage object as an extra and read it in that provider's `cost_calculator.py`. A provider without one gets a new `cost_calculator.py` and an arm for it in `cost_per_token` (`litellm/cost_calculator.py`), which picks the calculator by provider. If the provider's transform builds `Usage` from named fields, as the Anthropic, Gemini, and Bedrock Converse ones do, pass the new field there too. Streamed requests rebuild usage in `ChunkProcessor.calculate_usage` (`litellm/litellm_core_utils/streaming_chunk_builder_utils.py`), which does not copy extras yet the way the Responses-to-chat bridge does (minus a short exclusion list), so until it does, bill the field on streams by naming it there, as Anthropic's `inference_geo` is. A provider that reports its own charge for the request, as OpenRouter does in `usage.cost`, is billed at that charge from `_hidden_params` before any calculator runs, unless the deployment sets its own prices. When the field counts a built-in tool such as web search, add the provider to `get_cost_for_web_search_request` in `litellm/llms/__init__.py`, and add a gate to `litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py` only when no existing gate fires. Then show the field's cost in your proof on every endpoint it reaches, streaming included diff --git a/litellm/cost_calculator.py b/litellm/cost_calculator.py index ee102e82ca4..f5b9014d87c 100644 --- a/litellm/cost_calculator.py +++ b/litellm/cost_calculator.py @@ -623,7 +623,11 @@ def cost_per_token( ) if lyria_generation_cost is not None: return 0.0, lyria_generation_cost - speech_model_info = litellm.get_model_info(model=model_without_prefix, custom_llm_provider=custom_llm_provider) + speech_model_info: Final = cached_get_model_info_helper( + model=model_without_prefix, + custom_llm_provider=custom_llm_provider, + preserve_unset_costs=True, + ) cost_metric: Final = select_cost_metric_for_model(speech_model_info) prompt_cost: float = 0.0 completion_cost: float = 0.0 @@ -3084,7 +3088,7 @@ def batch_cost_calculator( batch_rates.cache_creation, batch_rates.cache_creation_rates_by_ttl, ) - elif input_cost_per_token: + elif input_cost_per_token is not None: details: Final = parse_prompt_tokens_details(usage) cache_read_tokens: Final = details["cache_hit_tokens"] cache_creation_tokens: Final = details["cache_creation_tokens"] @@ -3100,7 +3104,10 @@ def batch_cost_calculator( cache_read_cost_key: Final = get_service_tier_cost_key("cache_read_input_token_cost", None) total_prompt_cost += calculate_cost_component(model_info, cache_read_cost_key, cache_read_tokens) / 2 - cache_creation_cost: Final = model_info.get("cache_creation_input_token_cost") or input_cost_per_token + configured_cache_creation_cost: Final = model_info.get("cache_creation_input_token_cost") + cache_creation_cost: Final = ( + input_cost_per_token if configured_cache_creation_cost is None else configured_cache_creation_cost + ) total_prompt_cost += cache_creation_tokens * cache_creation_cost / 2 text_rate: Final = _batch_or_half(batch_rates.output, output_cost_per_token) image_rate: Final = _batch_or_half( diff --git a/litellm/litellm_core_utils/llm_cost_calc/utils.py b/litellm/litellm_core_utils/llm_cost_calc/utils.py index 404dfab0a6d..50c26f4d4ca 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/utils.py +++ b/litellm/litellm_core_utils/llm_cost_calc/utils.py @@ -28,6 +28,7 @@ from litellm.types.utils import ( DataResidency, ImageResponse, ModelInfo, + ModelInfoBase, PassthroughCallTypes, PromptTokensDetailsWrapper, ServiceTier, @@ -197,15 +198,15 @@ def get_billable_input_tokens(usage: Usage) -> int: def select_cost_metric_for_model( - model_info: ModelInfo, + model_info: ModelInfoBase, ) -> Literal["cost_per_character", "cost_per_token"]: """ Select 'cost_per_character' if model_info has 'input_cost_per_character' Select 'cost_per_token' if model_info has 'input_cost_per_token' """ - if model_info.get("input_cost_per_character"): + if model_info.get("input_cost_per_character") is not None: return "cost_per_character" - elif model_info.get("input_cost_per_token"): + elif model_info.get("input_cost_per_token") is not None: return "cost_per_token" else: raise ValueError( diff --git a/litellm/llms/perplexity/cost_calculator.py b/litellm/llms/perplexity/cost_calculator.py index c3d178771a5..c40a86b221b 100644 --- a/litellm/llms/perplexity/cost_calculator.py +++ b/litellm/llms/perplexity/cost_calculator.py @@ -100,8 +100,9 @@ def cost_per_token(model: str, usage: Usage, current_time: datetime | None = Non num_search_queries = getattr(usage.prompt_tokens_details, "web_search_requests", 0) or 0 # Check both possible keys for search cost (legacy and current) - search_cost_value: Final = model_info.get("search_queries_cost_per_query") or model_info.get( - "search_context_cost_per_query" + legacy_search_cost: Final = model_info.get("search_queries_cost_per_query") + search_cost_value: Final = ( + model_info.get("search_context_cost_per_query") if legacy_search_cost is None else legacy_search_cost ) if num_search_queries > 0 and search_cost_value is not None: # Handle both dict and float formats diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 1a8f5adf154..0c3440e88f7 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -464,6 +464,7 @@ class ModelInfoBase(ProviderSpecificModelInfo, total=False): annotation_cost_per_page: float | None # for OCR models annotation_cost_per_page_batches: ReadOnly[float | None] search_context_cost_per_query: SearchContextCostPerQuery | None # Cost for using web search tool + search_queries_cost_per_query: ReadOnly[float | SearchContextCostPerQuery | None] web_search_billing_unit: ( Literal["per_query", "per_prompt"] | None ) # "per_query" (Gemini 3.x) or "per_prompt" (Gemini 2.x) @@ -4168,6 +4169,7 @@ class CustomPricingLiteLLMParams(MirroredPricingParams): output_cost_per_video_per_second: float | None = None output_cost_per_audio_per_second: float | None = None search_context_cost_per_query: dict[str, Any] | None = None + search_queries_cost_per_query: float | SearchContextCostPerQuery | None = None google_maps_grounding_cost_per_query: float | None = None citation_cost_per_token: float | None = None cache_read_input_token_cost_above_272k_tokens: float | None = None diff --git a/litellm/utils.py b/litellm/utils.py index 560d0c31e4a..c2633d7eb81 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -6889,6 +6889,8 @@ def cached_get_model_info_helper( model: str, custom_llm_provider: str | None, api_base: str | None = None, + *, + preserve_unset_costs: bool = False, ) -> ModelInfoBase: """ _get_model_info_helper wrapped with lru_cache @@ -6899,6 +6901,7 @@ def cached_get_model_info_helper( model=model, custom_llm_provider=custom_llm_provider, api_base=api_base, + preserve_unset_costs=preserve_unset_costs, ) @@ -6962,6 +6965,7 @@ def get_model_info_helper( api_key: str | None = None, *, catalog_only: bool = False, + preserve_unset_costs: bool = False, ) -> ModelInfoBase: """ Helper for 'get_model_info'. Separated out to avoid infinite loop caused by returning 'supported_openai_param's @@ -7079,7 +7083,7 @@ def get_model_info_helper( verbose_logger.debug("model not found in litellm.model_cost, tried: %s", potential_model_names) raise ModelNotMappedError(_model_not_mapped_message(model, custom_llm_provider)) _input_cost_per_token: float | None = _model_info.get("input_cost_per_token") - if _input_cost_per_token is None: + if _input_cost_per_token is None and not preserve_unset_costs: # default value to 0, be noisy about this verbose_logger.debug( "model=%s, custom_llm_provider=%s has no input_cost_per_token in model_cost_map. Defaulting to 0.", @@ -7089,7 +7093,7 @@ def get_model_info_helper( _input_cost_per_token = 0 _output_cost_per_token: float | None = _model_info.get("output_cost_per_token") - if _output_cost_per_token is None: + if _output_cost_per_token is None and not preserve_unset_costs: # default value to 0, be noisy about this verbose_logger.debug( "model=%s, custom_llm_provider=%s has no output_cost_per_token in model_cost_map. Defaulting to 0.", @@ -7371,6 +7375,7 @@ def get_model_info_helper( supports_regex_lookaround=_model_info.get("supports_regex_lookaround", None), supports_computer_use=_model_info.get("supports_computer_use", None), search_context_cost_per_query=_model_info.get("search_context_cost_per_query", None), + search_queries_cost_per_query=_model_info.get("search_queries_cost_per_query"), web_search_billing_unit=_model_info.get("web_search_billing_unit", None), search_free_requests=_model_info.get("search_free_requests", None), google_maps_grounding_cost_per_query=_model_info.get("google_maps_grounding_cost_per_query", None), diff --git a/tests/unit/llms/perplexity/test_perplexity_cost_calculator.py b/tests/unit/llms/perplexity/test_perplexity_cost_calculator.py index 380f85445e5..5e774b05207 100644 --- a/tests/unit/llms/perplexity/test_perplexity_cost_calculator.py +++ b/tests/unit/llms/perplexity/test_perplexity_cost_calculator.py @@ -9,6 +9,7 @@ import json import logging import math import os +from collections.abc import Mapping from datetime import datetime, timezone from typing import Final from unittest.mock import patch @@ -119,7 +120,8 @@ class TestPerplexityCostCalculator: assert completion_cost == 0.008 assert prompt_cost + completion_cost == 0.008 - def test_uses_perplexity_provided_cost_when_normalized_to_float(self): + @pytest.mark.parametrize("reported_cost", (0.0, 0.008)) + def test_uses_perplexity_provided_cost_when_normalized_to_float(self, reported_cost: float): """ Regression: for Responses API / Agent API models, `ResponseAPIUsage.parse_cost` (litellm/types/llms/openai.py) already flattens Perplexity's @@ -135,12 +137,55 @@ class TestPerplexityCostCalculator: perplexity/openai/gpt-5.2 before rates existed). """ usage = Usage(prompt_tokens=100, completion_tokens=50, total_tokens=150) - usage.cost = 0.008 + usage.cost = reported_cost prompt_cost, completion_cost = perplexity_cost_per_token(model="sonar-pro", usage=usage) assert prompt_cost == 0.0 - assert completion_cost == 0.008 + assert completion_cost == reported_cost + + @pytest.mark.parametrize( + ("legacy_pricing", "expected_search_rate"), + ( + pytest.param({}, 0.01, id="missing"), + pytest.param({"search_queries_cost_per_query": None}, 0.01, id="null"), + pytest.param({"search_queries_cost_per_query": 0.0}, 0.0, id="free"), + pytest.param({"search_queries_cost_per_query": 0.004}, 0.004, id="paid"), + ), + ) + def test_search_cost_preserves_zero_and_defaults_only_when_unset( + self, + monkeypatch: pytest.MonkeyPatch, + legacy_pricing: Mapping[str, float | None], + expected_search_rate: float, + ) -> None: + model: Final = "perplexity/search-pricing-fixture" + monkeypatch.setitem( + litellm.model_cost, + model, + { + "litellm_provider": "perplexity", + "mode": "chat", + "input_cost_per_token": 2e-6, + "output_cost_per_token": 8e-6, + "output_cost_per_reasoning_token": 3e-6, + "citation_cost_per_token": 2e-6, + "search_context_cost_per_query": {"search_context_size_low": 0.01}, + **legacy_pricing, + }, + ) + usage: Final = Usage( + prompt_tokens=100, + completion_tokens=50, + prompt_tokens_details=PromptTokensDetailsWrapper(web_search_requests=2), + completion_tokens_details=CompletionTokensDetailsWrapper(reasoning_tokens=10), + citation_tokens=25, + ) + + prompt_cost, completion_cost = litellm.cost_per_token(model=model, usage_object=usage) + + assert prompt_cost == pytest.approx((100 + 25) * 2e-6) + assert completion_cost == pytest.approx(40 * 8e-6 + 10 * 3e-6 + 2 * expected_search_rate) @pytest.mark.parametrize( "reported_cost", diff --git a/tests/unit/test_cost_calculator.py b/tests/unit/test_cost_calculator.py index f10baccd1cb..c536522e736 100644 --- a/tests/unit/test_cost_calculator.py +++ b/tests/unit/test_cost_calculator.py @@ -37,7 +37,7 @@ from litellm.types.llms.xai import RealtimeMeteredUsage from litellm.types.rerank import RerankResponse from litellm.types.utils import CacheCreationTokenDetails, CallTypes, Choices, CompletionTokensDetailsWrapper, EmbeddingResponse, ImageObject, ImageResponse, ImageUsage, ImageUsageInputTokensDetails, LiteLLMRealtimeStreamLoggingObject, Message, ModelInfo, ModelResponse, PromptTokensDetails, PromptTokensDetailsWrapper, Usage from litellm.types.videos.main import VideoObject -from litellm.utils import _invalidate_model_cost_lowercase_map, supports_prompt_caching +from litellm.utils import _invalidate_model_cost_lowercase_map, cached_get_model_info_helper, supports_prompt_caching import os import threading from unittest.mock import MagicMock, patch @@ -4598,23 +4598,40 @@ def test_batch_cost_calculator_prices_cache_creation_tokens_at_cache_write_rate( assert completion_cost_value == pytest.approx(200 * 15e-6 / 2) -def test_batch_cost_calculator_cache_creation_falls_back_to_input_rate(): +@pytest.mark.parametrize("input_rate", (3e-6, 0.0)) +@pytest.mark.parametrize( + ("cache_pricing", "expected_write_rate"), + [ + pytest.param({}, None, id="missing"), + pytest.param({"cache_creation_input_token_cost": None}, None, id="null"), + pytest.param({"cache_creation_input_token_cost": 0.0}, 0.0, id="free"), + pytest.param({"cache_creation_input_token_cost": 3.75e-6}, 3.75e-6, id="paid"), + ], +) +def test_batch_cost_calculator_preserves_cache_write_prices_and_defaults_only_when_unset( + cache_pricing: ModelInfo, + expected_write_rate: float | None, + input_rate: float, +) -> None: from litellm.cost_calculator import batch_cost_calculator - model_info: ModelInfo = { + model_info: Final[ModelInfo] = { "supported_openai_params": [], - "input_cost_per_token": 3e-6, + "input_cost_per_token": input_rate, "output_cost_per_token": 15e-6, "cache_read_input_token_cost": 3e-7, + **cache_pricing, } - prompt_cost, _ = batch_cost_calculator( + prompt_cost, completion_cost_value = batch_cost_calculator( usage=_batch_cache_usage(), - model="claude-sonnet-4-5-20250929", + model="batch-pricing-fixture", custom_llm_provider="anthropic", model_info=model_info, ) - assert prompt_cost == pytest.approx((1000 * 3e-6 + 8000 * 3e-7 + 2000 * 3e-6) / 2) + write_rate: Final = input_rate if expected_write_rate is None else expected_write_rate + assert prompt_cost == pytest.approx((1000 * input_rate + 8000 * 3e-7 + 2000 * write_rate) / 2) + assert completion_cost_value == pytest.approx(200 * 15e-6 / 2) def test_completion_cost_bills_interactions_api_response(): @@ -7817,20 +7834,42 @@ def test_vertex_ai_mistral_predict_cost(usage): assert predictive_cost > 0 +@pytest.mark.usefixtures("_local_model_cost_map") +@pytest.mark.parametrize("call_type", ("speech", "aspeech")) +@pytest.mark.parametrize("model", ("speech-fixture", "openai/speech-fixture", "azure/speech-fixture")) @pytest.mark.parametrize( - "model", ["openai/tts-1", "azure/tts-1", "openai/gpt-4o-mini-tts"] + ("pricing", "expected_cost"), + ( + pytest.param({}, None, id="missing"), + pytest.param({"input_cost_per_character": None, "input_cost_per_token": None}, None, id="null"), + pytest.param({"input_cost_per_token": 0.0}, 0.0, id="free-tokens"), + pytest.param({"input_cost_per_character": 0.0, "input_cost_per_token": 2e-6}, 0.0, id="free-characters"), + pytest.param({"input_cost_per_token": 2e-6}, 100 * 2e-6, id="missing-character-rate"), + pytest.param({"input_cost_per_character": None, "input_cost_per_token": 2e-6}, 100 * 2e-6, id="null-character-rate"), + pytest.param({"input_cost_per_character": 3e-6, "input_cost_per_token": 2e-6}, 1000 * 3e-6, id="paid-characters"), + ), ) -def test_completion_cost_tts(model): - os.environ["LITELLM_LOCAL_MODEL_COST_MAP"] = "True" - litellm.model_cost = litellm.get_model_cost_map(url="") - - cost = completion_cost( - model=model, - prompt="the quick brown fox jumped over the lazy dogs", - call_type="speech", +def test_registered_speech_prices_distinguish_unset_from_zero( + call_type: str, model: str, pricing: ModelInfo, expected_cost: float | None +) -> None: + provider: Final = "azure" if model.startswith("azure/") else "openai" + litellm.register_model( + {model: {"litellm_provider": provider, "mode": "audio_speech", **pricing}}, + persist_across_reloads=False, ) + compatibility_info: Final = cached_get_model_info_helper(model, provider) + assert compatibility_info["input_cost_per_token"] == (pricing.get("input_cost_per_token") or 0.0) + litellm.get_model_info(model) + response: Final = ModelResponse(usage=Usage(prompt_tokens=100, completion_tokens=0, total_tokens=100)) - assert cost > 0 + if expected_cost is None: + with pytest.raises(ValueError, match="does not have 'input_cost_per_character' or 'input_cost_per_token'"): + completion_cost(completion_response=response, model=model, prompt="a" * 1000, call_type=call_type) + litellm.register_model({model: {"input_cost_per_token": 0.0}}, persist_across_reloads=False) + + assert completion_cost( + completion_response=response, model=model, prompt="a" * 1000, call_type=call_type + ) == pytest.approx(expected_cost or 0.0) def test_completion_cost_anthropic(): diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 24736cdb34e..30d61f181cb 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -35894,6 +35894,8 @@ export interface components { search_context_cost_per_query?: { [key: string]: unknown; } | null; + /** Search Queries Cost Per Query */ + search_queries_cost_per_query?: number | components["schemas"]["SearchContextCostPerQuery"] | null; /** Stream Timeout */ stream_timeout?: number | string | null; /** Stream Timeout By Service Tier */ @@ -45445,6 +45447,17 @@ export interface components { } & { [key: string]: unknown; }; + /** SearchContextCostPerQuery */ + SearchContextCostPerQuery: { + /** Search Context Size High */ + search_context_size_high?: number; + /** Search Context Size Low */ + search_context_size_low?: number; + /** Search Context Size Medium */ + search_context_size_medium?: number; + } & { + [key: string]: unknown; + }; /** * SearchTool * @description Search tool configuration. @@ -51409,6 +51422,8 @@ export interface components { search_context_cost_per_query?: { [key: string]: unknown; } | null; + /** Search Queries Cost Per Query */ + search_queries_cost_per_query?: number | components["schemas"]["SearchContextCostPerQuery"] | null; /** Stream Timeout */ stream_timeout?: number | string | null; /** Stream Timeout By Service Tier */ From 298911d61bded3b81aa7d0face6e8538cb7114ed Mon Sep 17 00:00:00 2001 From: Kerry Lu Date: Fri, 30 Oct 2026 15:25:47 -0700 Subject: [PATCH 29/46] fix(batches): cost an unmanaged gs:// batch through the deployment that created it (#49301) * fix(batches): cost an unmanaged gs:// batch through the deployment that created it The create now stores the router deployment id as model_id inside the managed object row's file_object JSON and CheckBatchCost reads it before falling back to the first deployment matching the base model in the gs:// path. Rows without the field keep the old lookup * fix(batches): tolerate a non-string or malformed file_object when reading the stored deployment * fix(batches): read the stored creating deployment from the writer, not a lagging read replica * fix(batches): resolve the writer database once so the type gate count does not grow * fix(batches): decode a dict-shaped file_object when reading the stored deployment, accumulate wire samples, annotate test locals --------- Co-authored-by: devin-ai-integration[bot] <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../proxy/common_utils/check_batch_cost.py | 31 ++- .../proxy/hooks/managed_files.py | 50 ++++- ...ed_vertex_batch_deployment_routing_wire.py | 181 ++++++++++++++++++ .../proxy/test_managed_files_hook.py | 107 +++++++++++ .../common_utils/test_check_batch_cost.py | 47 +++++ 5 files changed, 408 insertions(+), 8 deletions(-) create mode 100644 tests/integration/spend/test_unmanaged_vertex_batch_deployment_routing_wire.py diff --git a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py index a1b4804ae29..a31e926256a 100644 --- a/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py +++ b/enterprise/litellm_enterprise/proxy/common_utils/check_batch_cost.py @@ -10,6 +10,7 @@ from dataclasses import replace as dataclasses_replace from datetime import datetime, timedelta from types import MappingProxyType from typing import TYPE_CHECKING, Final, List, Literal, Optional, Protocol, Tuple, cast +from pydantic import TypeAdapter, ValidationError from litellm_enterprise.proxy.common_utils.managed_object_attribution import ( ManagedObjectCreatorAttribution, @@ -50,6 +51,7 @@ def _error_names_resource(error: Exception, resource_id: str) -> bool: CHECK_BATCH_COST_USER_AGENT = "LiteLLM Proxy/CheckBatchCost" RECENT_BATCH_POLL_WINDOW: Final = timedelta(hours=48) GIVE_UP_POLLING_AFTER: Final = timedelta(days=2 * MANAGED_OBJECT_STALENESS_CUTOFF_DAYS) +_FILE_OBJECT_ADAPTER: Final[TypeAdapter[dict[str, object]]] = TypeAdapter(dict[str, object]) PROVIDER_TERMINAL_BATCH_STATUSES: Final[Tuple[str, ...]] = ( @@ -550,6 +552,10 @@ class CheckBatchCost: llm_provider: str, bare_model_name: str, ) -> Optional[Tuple[str, str]]: + stored_model_id: Final = self._get_stored_model_id(job) + if stored_model_id is not None and self._batch_deployment_exists(stored_model_id): + return stored_model_id, job.unified_object_id + deployment_id = self._get_deployment_id_for_bare_model(bare_model_name, llm_provider) if deployment_id is None: verbose_proxy_logger.info( @@ -662,23 +668,36 @@ class CheckBatchCost: ) @staticmethod - def _get_input_file_id(job: "ManagedObjectRow") -> Optional[str]: - from litellm.types.utils import LiteLLMBatch - + def _parse_file_object(job: "ManagedObjectRow") -> Optional[dict[str, object]]: file_object: Final = job.file_object try: - parsed_file_object: Final[object] = ( + return _FILE_OBJECT_ADAPTER.validate_python( json.loads(file_object) if isinstance(file_object, str) else file_object ) - except (json.JSONDecodeError, ValueError): + except (json.JSONDecodeError, ValueError, ValidationError): return None - if not isinstance(parsed_file_object, dict): + + @classmethod + def _get_input_file_id(cls, job: "ManagedObjectRow") -> Optional[str]: + from litellm.types.utils import LiteLLMBatch + + parsed_file_object: Final = cls._parse_file_object(job) + if parsed_file_object is None: return None try: return LiteLLMBatch.model_validate(parsed_file_object).input_file_id except Exception: return None + @classmethod + def _get_stored_model_id(cls, job: "ManagedObjectRow") -> Optional[str]: + """The deployment that created the batch, written into file_object by the managed + files hook at create time. Older rows have no such key and fall back to matching + the bare model name against the router.""" + parsed_file_object: Final = cls._parse_file_object(job) + model_id: Final = parsed_file_object.get("model_id") if parsed_file_object else None + return model_id if isinstance(model_id, str) and model_id else None + async def _track_completed_batch_cost( self, job: "ManagedObjectRow", diff --git a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py index 5b54e7d72a0..0d46efbbb08 100644 --- a/enterprise/litellm_enterprise/proxy/hooks/managed_files.py +++ b/enterprise/litellm_enterprise/proxy/hooks/managed_files.py @@ -21,6 +21,7 @@ from typing import ( TypedDict, Union, cast, + runtime_checkable, ) from uuid import NAMESPACE_URL, uuid5 @@ -374,6 +375,12 @@ class _ManagedResourcePrismaClient(Protocol): def db(self) -> _ManagedResourceDatabase: ... +@runtime_checkable +class _WriterRoutedDatabase(Protocol): + @property + def writer(self) -> _ManagedResourceDatabase: ... + + class _SchedulerWithJobLookup(Protocol): def get_job(self, job_id: str) -> object: ... @@ -479,6 +486,11 @@ def _managed_object_table(prisma_client: _ManagedResourcePrismaClient) -> _Manag return prisma_client.db.litellm_managedobjecttable +def _writer_database(db: _ManagedResourceDatabase) -> _ManagedResourceDatabase: + """Reads that feed a write must see the create that just happened; a lagging read replica may not""" + return db.writer if isinstance(db, _WriterRoutedDatabase) else db + + _JSON_VALUE_ADAPTER: Final[TypeAdapter[JsonValue]] = TypeAdapter(JsonValue) _POSTGRES_UNSTORABLE_CHARACTER: Final = "\x00" @@ -512,6 +524,25 @@ def _without_model_output( ) +def _with_model_id( + file_object: LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse, + model_id: str | None, +) -> LiteLLMBatch | LiteLLMFineTuningJob | ResponsesAPIResponse: + return file_object if model_id is None else file_object.model_copy(update={"model_id": model_id}) + + +async def _stored_model_id(managed_object_table: _ManagedObjectTableActions, unified_object_id: str) -> str | None: + row: Final = await managed_object_table.find_first(where={"unified_object_id": unified_object_id}) + if row is None: + return None + try: + parsed: Final = _JSON_VALUE_ADAPTER.validate_python(_decode_json_blob(row.file_object)) + except (json.JSONDecodeError, ValidationError): + return None + model_id: Final = parsed.get("model_id") if isinstance(parsed, dict) else None + return model_id if isinstance(model_id, str) and model_id else None + + def _storage_metadata_of(file_object: OpenAIFileObject | None) -> Mapping[str, str]: hidden_params: Final = cast( # cast-ok: _hidden_params is an untyped attribute the upload path sets "Mapping[str, object]", getattr(file_object, "_hidden_params", None) or {} @@ -697,9 +728,17 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): persist_attribution: bool = False, create_if_missing: bool = True, batch_processed: bool = False, + creating_model_id: str | None = None, ) -> None: """Persist a managed object row, caching it and upserting it in the DB. + creating_model_id is the router deployment that created the object. It is kept + inside the stored file_object JSON as `model_id`, because the raw provider id an + unmanaged batch stores as unified_object_id says nothing about which deployment + made it, and the cost poller needs that deployment to read the result back. A + later observe of the row (create_if_missing=False) keeps the stored value rather + than overwriting it with whichever deployment served the retrieve. + persist_attribution is set only by the batch and background response creates, the callers that can speak for the creator; it gates the api_key, request_tags, and, for a response, end_user_id columns that the cost poll jobs bill against, so a later poll or @@ -720,7 +759,14 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): It is written only in the upsert create branch. """ verbose_logger.info(f"Storing LiteLLM Managed {file_purpose} object with id={unified_object_id} in cache") - stored_object: Final = _without_model_output(file_object) + managed_database: Final = _managed_database(self.prisma_client) + managed_object_table: Final = managed_database.litellm_managedobjecttable + preserved_model_id: Final = ( + creating_model_id + if creating_model_id is not None or create_if_missing + else await _stored_model_id(_writer_database(managed_database).litellm_managedobjecttable, unified_object_id) + ) + stored_object: Final = _with_model_id(_without_model_output(file_object), preserved_model_id) litellm_managed_object = LiteLLM_ManagedObjectTable( unified_object_id=unified_object_id, model_object_id=model_object_id, @@ -752,7 +798,6 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): "status": file_object.status, "updated_by": user_api_key_dict.user_id, } - managed_object_table: Final = _managed_object_table(self.prisma_client) if not create_if_missing: await managed_object_table.update_many( where={"unified_object_id": unified_object_id}, @@ -2170,6 +2215,7 @@ class _PROXY_LiteLLMManagedFiles(CustomLogger, BaseFileEndpoints): request_tags=request_tags_from_metadata(request_metadata if isinstance(request_metadata, dict) else {}), persist_attribution=is_batch_create, create_if_missing=is_batch_create, + creating_model_id=model_id if is_batch_create else None, ) if not is_batch_create: await self.store_batch_output_file_ownership( diff --git a/tests/integration/spend/test_unmanaged_vertex_batch_deployment_routing_wire.py b/tests/integration/spend/test_unmanaged_vertex_batch_deployment_routing_wire.py new file mode 100644 index 00000000000..91269e6557f --- /dev/null +++ b/tests/integration/spend/test_unmanaged_vertex_batch_deployment_routing_wire.py @@ -0,0 +1,181 @@ +"""Two Vertex deployments share a base model but live in different projects with different +buckets. An unmanaged gs:// batch created on the second one must be polled through that +deployment (LIT-10326). Before the fix the cost poller matched the base model name in the +gs:// path to the first deployment in the list and asked the wrong project for the job.""" + +import json +import re +from collections.abc import Callable +from contextlib import ExitStack +from pathlib import Path +from typing import Final +from uuid import uuid4 + +import pytest +import yaml +from integration._support.client import JSON_OBJECT, Gateway, eventually, object_value, string_value +from integration._support.process import owned_proxy_process +from integration._support.vertex import service_account_json +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +LOCATION: Final = "us-central1" +BASE_MODEL: Final = "gemini-2.5-flash" +TEAMS: Final = ("a", "b") +_JOB_PATH: Final = re.compile( + r"^/v1/projects/(?P[^/]+)/locations/(?P[^/]+)/batchPredictionJobs(?:/(?P\d+))?$" +) + + +def _project(team: str) -> str: + return f"scripted-project-{team}" + + +def _deployment(team: str) -> str: + return f"{BASE_MODEL}-team-{team}" + + +def _input_file(team: str) -> str: + return f"gs://scripted-bucket-{team}/litellm-vertex-files/publishers/google/models/{BASE_MODEL}/{uuid4().hex}/input.jsonl" + + +def _job_number(team: str) -> str: + return f"{TEAMS.index(team) + 1}{uuid4().int % 10**17:017d}" + + +def _vertex_peer(jobs: dict[str, tuple[str, str]]) -> Callable[[Request], Reply]: + """A Vertex double that remembers which project each job was created in. A retrieve + from any other project answers the 404 real Vertex gives (the job does not exist + there), a retrieve from the right one reports the job still running.""" + + def respond(request: Request) -> Reply: + path: Final = _JOB_PATH.match(request.target.split("?", 1)[0]) + assert path is not None, request.target + assert request.headers.get("authorization") == "Bearer scripted-token", request.headers + project, location, job = path.group("project"), path.group("location"), path.group("job") + if request.method == "POST": + team: Final = project.rsplit("-", 1)[-1] + uri: Final = re.search(r'"(gs://[^"]+)"', request.body.decode()) + assert uri is not None and uri.group(1).startswith(f"gs://scripted-bucket-{team}/"), (project, request.body) + created, created_uri = jobs.setdefault(project, (_job_number(team), uri.group(1))) + return Reply(body=json.dumps(_job(project, location, created, created_uri, "JOB_STATE_PENDING")).encode()) + assert request.method == "GET" and job is not None, (request.method, request.target) + known: Final = jobs.get(project) + if known is None or known[0] != job: + return Reply( + status=404, + body=json.dumps( + { + "error": { + "code": 404, + "message": f"The BatchPredictionJob does not exist: {job}", + "status": "NOT_FOUND", + } + } + ).encode(), + ) + return Reply(body=json.dumps(_job(project, location, job, known[1], "JOB_STATE_RUNNING")).encode()) + + return respond + + +def _job(project: str, location: str, job: str, input_uri: str, state: str) -> dict[str, JsonValue]: + return { + "name": f"projects/{project}/locations/{location}/batchPredictionJobs/{job}", + "displayName": "litellm-vertex-batch-scripted", + "model": f"publishers/google/models/{BASE_MODEL}", + "inputConfig": {"gcsSource": {"uris": [input_uri]}, "instancesFormat": "jsonl"}, + "outputConfig": {"predictionsFormat": "jsonl", "gcsDestination": {"outputUriPrefix": input_uri.rsplit("/", 1)[0]}}, + "state": state, + "createTime": "2026-10-30T20:00:00.000000Z", + "updateTime": "2026-10-30T20:00:00.000000Z", + } + + +def _owned_config(path: Path, gateway: Gateway, wire_url: str) -> Path: + config: Final = JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + path.write_text( + yaml.safe_dump( + { + **config, + "model_list": [ + { + "model_name": _deployment(team), + "litellm_params": { + "model": f"vertex_ai/{BASE_MODEL}", + "api_base": wire_url, + "vertex_project": _project(team), + "vertex_location": LOCATION, + "vertex_credentials": service_account_json(_project(team), gateway.upstream_url.rstrip("/")), + "gcs_bucket_name": f"scripted-bucket-{team}", + }, + } + for team in TEAMS + ], + "general_settings": { + **object_value(config.get("general_settings") or {}), + "track_unmanaged_batch_cost": True, + }, + "litellm_settings": { + **object_value(config.get("litellm_settings") or {}), + "enable_loadbalancing_on_batch_endpoints": True, + "user_url_allowed_hosts": ["127.0.0.1"], + }, + } + ) + ) + return path + + +def _create_batch(gateway: Gateway, key: str, team: str) -> str: + response: Final = gateway.request( + "POST", + "/v1/batches", + { + "input_file_id": _input_file(team), + "endpoint": "/v1/chat/completions", + "completion_window": "24h", + "model": _deployment(team), + }, + key=key, + ) + assert response.status_code == 200, response.text + return string_value(JSON_OBJECT.validate_json(response.content)["id"]) + + +def _retrieve_targets(requests: tuple[Request, ...], job: str) -> tuple[str, ...]: + return tuple(request.target for request in requests if request.method == "GET" and request.target.endswith(f"/{job}")) + + +@pytest.mark.timeout(240) +def test_unmanaged_batch_on_the_second_deployment_sharing_a_base_model_is_polled_through_its_own_project( + gateway: Gateway, tmp_path: Path +) -> None: + jobs: Final[dict[str, tuple[str, str]]] = {} # mutable-ok: the wire double records which project created each job + with ExitStack() as stack: + wire: Final = stack.enter_context(wire_server(_vertex_peer(jobs))) + config: Final = _owned_config(tmp_path / "two-vertex-teams.yaml", gateway, wire.url) + owned: Final = stack.enter_context(owned_proxy_process(gateway, tmp_path, {}, config=config)) + scenario: Final = stack.enter_context(owned.gateway.scenario()) + key: Final = scenario.key(models=[_deployment(team) for team in TEAMS]) + + batch_a: Final = _create_batch(owned.gateway, key, "a") + batch_b: Final = _create_batch(owned.gateway, key, "b") + assert (batch_a, batch_b) == (jobs[_project("a")][0], jobs[_project("b")][0]), (batch_a, batch_b, jobs) + + seen: Final[list[Request]] = [] # mutable-ok: drain empties the wire queue, so samples accumulate across retries + + def polled_so_far() -> tuple[Request, ...]: + seen.extend(wire.drain()) + return tuple(seen) + + polled: Final = eventually( + polled_so_far, + lambda requests: bool(_retrieve_targets(requests, batch_b)) and bool(_retrieve_targets(requests, batch_a)), + seconds=90, + ) + for team, batch in (("a", batch_a), ("b", batch_b)): + assert all( + target.startswith(f"/v1/projects/{_project(team)}/") for target in _retrieve_targets(polled, batch) + ), (team, _retrieve_targets(polled, batch)) + diff --git a/tests/unit/enterprise/proxy/test_managed_files_hook.py b/tests/unit/enterprise/proxy/test_managed_files_hook.py index f2df14cc8a3..89fff365c3f 100644 --- a/tests/unit/enterprise/proxy/test_managed_files_hook.py +++ b/tests/unit/enterprise/proxy/test_managed_files_hook.py @@ -1437,6 +1437,8 @@ def _make_object_store_instance(): mock_prisma = MagicMock() mock_prisma.db.litellm_managedobjecttable.upsert = AsyncMock() mock_prisma.db.litellm_managedobjecttable.update_many = AsyncMock() + mock_prisma.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None) + mock_prisma.db.writer = mock_prisma.db return ( PROXY_LiteLLMManagedFiles( @@ -2161,3 +2163,108 @@ async def test_store_unified_file_id_caches_the_storage_location_the_db_row_gets assert (cached.storage_backend, cached.storage_url) == ("litellm_db", "litellm_db://content-row-1") create_data = file_table.upsert.await_args.kwargs["data"]["create"] assert (create_data["storage_backend"], create_data["storage_url"]) == ("litellm_db", "litellm_db://content-row-1") + + +@pytest.mark.asyncio +async def test_create_stores_the_creating_deployment_in_file_object(): + """Regression (LIT-10326): an unmanaged gs:// batch stores the raw Vertex job id, which + says nothing about which of two deployments sharing a base model created it. The create + must keep the router's deployment id inside the stored file_object so the cost poller + reads the result through that deployment instead of the first base-model match.""" + managed_files, mock_prisma = _make_object_store_instance() + creator: Final = UserAPIKeyAuth(api_key="sk-team-b", user_id="bob", team_id="team-b", parent_otel_span=None) + + await managed_files.store_unified_object_id( + unified_object_id="projects/123/locations/us/batchPredictionJobs/456", + file_object=_make_batch_response(status="validating"), + litellm_parent_otel_span=None, + model_object_id="projects/123/locations/us/batchPredictionJobs/456", + file_purpose="batch", + user_api_key_dict=creator, + persist_attribution=True, + create_if_missing=True, + creating_model_id="deploy-team-b", + ) + + created: Final = mock_prisma.db.litellm_managedobjecttable.upsert.await_args.kwargs["data"]["create"] + assert json.loads(created["file_object"])["model_id"] == "deploy-team-b" + + +@pytest.mark.asyncio +async def test_poll_refresh_keeps_the_stored_creating_deployment(): + """A later observe of the row (a GET through whichever deployment the client named) + rewrites file_object; it must carry the create-time model_id forward, not drop it or + replace it with the deployment that served the retrieve.""" + managed_files, mock_prisma = _make_object_store_instance() + stored_row: Final = MagicMock() + stored_row.file_object = json.dumps({**json.loads(_make_batch_response(status="validating").model_dump_json()), "model_id": "deploy-team-b"}) + mock_prisma.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=stored_row) + observer: Final = UserAPIKeyAuth(api_key="sk-team-b", user_id="bob", team_id="team-b", parent_otel_span=None) + + await managed_files.store_unified_object_id( + unified_object_id="projects/123/locations/us/batchPredictionJobs/456", + file_object=_make_batch_response(status="completed"), + litellm_parent_otel_span=None, + model_object_id="projects/123/locations/us/batchPredictionJobs/456", + file_purpose="batch", + user_api_key_dict=observer, + persist_attribution=False, + create_if_missing=False, + ) + + written: Final = json.loads(mock_prisma.db.litellm_managedobjecttable.update_many.await_args.kwargs["data"]["file_object"]) + assert written["model_id"] == "deploy-team-b" + assert written["status"] == "completed" + + +@pytest.mark.asyncio +async def test_poll_refresh_reads_the_stored_creating_deployment_from_the_writer(): + """With a read replica configured, a retrieve that lands before the replica has the create + must still find the create-time model_id, so the lookup goes to the writer. Prisma may hand + the Json column back already decoded, so the row here is a dict rather than JSON text.""" + managed_files, mock_prisma = _make_object_store_instance() + stored_row: Final = MagicMock() + stored_row.file_object = {**json.loads(_make_batch_response(status="validating").model_dump_json()), "model_id": "deploy-team-b"} + mock_prisma.db.litellm_managedobjecttable.find_first = AsyncMock(return_value=None) + mock_prisma.db.writer = MagicMock() + mock_prisma.db.writer.litellm_managedobjecttable.find_first = AsyncMock(return_value=stored_row) + observer: Final = UserAPIKeyAuth(api_key="sk-team-b", user_id="bob", team_id="team-b", parent_otel_span=None) + + await managed_files.store_unified_object_id( + unified_object_id="projects/123/locations/us/batchPredictionJobs/456", + file_object=_make_batch_response(status="completed"), + litellm_parent_otel_span=None, + model_object_id="projects/123/locations/us/batchPredictionJobs/456", + file_purpose="batch", + user_api_key_dict=observer, + persist_attribution=False, + create_if_missing=False, + ) + + written: Final = json.loads(mock_prisma.db.litellm_managedobjecttable.update_many.await_args.kwargs["data"]["file_object"]) + assert written["model_id"] == "deploy-team-b" + mock_prisma.db.litellm_managedobjecttable.find_first.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_batch_create_hook_passes_the_router_deployment_as_creating_model_id(): + """The post-call hook on a batch create hands the router's model_id to the object store; + a retrieve of the same batch hands None so it cannot overwrite the creator.""" + from litellm_enterprise.proxy.hooks.managed_files import BATCH_CREATE_HIDDEN_PARAM + + managed_files: Final = _make_managed_files_instance() + created: Final = _make_batch_response(model_id="deploy-team-b", output_file_id=None, status="validating") + created._hidden_params[BATCH_CREATE_HIDDEN_PARAM] = True + retrieved: Final = _make_batch_response(model_id="deploy-team-a", output_file_id=None, status="in_progress") + + with patch("litellm.proxy.proxy_server.llm_router", MagicMock()): + await managed_files.async_post_call_success_hook( + data={}, user_api_key_dict=_make_user_api_key_dict(), response=created + ) + await managed_files.async_post_call_success_hook( + data={}, user_api_key_dict=_make_user_api_key_dict(), response=retrieved + ) + + first, second = managed_files.store_unified_object_id.await_args_list + assert first.kwargs["creating_model_id"] == "deploy-team-b" + assert second.kwargs["creating_model_id"] is None diff --git a/tests/unit/proxy/common_utils/test_check_batch_cost.py b/tests/unit/proxy/common_utils/test_check_batch_cost.py index 4861b2687de..5b37cee88c5 100644 --- a/tests/unit/proxy/common_utils/test_check_batch_cost.py +++ b/tests/unit/proxy/common_utils/test_check_batch_cost.py @@ -1771,6 +1771,53 @@ class TestUnmanagedVertexRouting: assert result == ("deploy-ft", "8823717160934178816") + def _two_team_deployments(self): + """Two Vertex deployments sharing the base model, team A listed first.""" + return [ + { + "model_name": f"gemini-2.5-flash-team-{team}", + "litellm_params": {"model": "vertex_ai/gemini-2.5-flash", "custom_llm_provider": "vertex_ai"}, + "model_info": {"id": f"deploy-team-{team}"}, + } + for team in ("a", "b") + ] + + def test_flag_on_prefers_the_deployment_stored_on_the_row_over_the_first_base_model_match(self): + """Regression (LIT-10326): two deployments share gemini-2.5-flash with different + buckets. A batch created on team B's deployment stores deploy-team-b in file_object, + so the poller must read it back through team B, not the first match (team A).""" + router: Final = MagicMock() + router.registry.return_value.resolve.return_value = None + router.registry.return_value.deployments.return_value = self._two_team_deployments() + router.registry.return_value.deployment = MagicMock(return_value=self._vertex_deployment()) + instance: Final = self._instance(track_unmanaged=True, router=router) + job: Final = self._job( + file_object=json.dumps({**json.loads(_unmanaged_vertex_file_object()), "model_id": "deploy-team-b"}) + ) + + with patch(_IS_B64, return_value=False): + result: Final = instance._resolve_job_routing(job, MagicMock()) + + assert result == ("deploy-team-b", "8823717160934178816") + router.registry.return_value.deployment.assert_called_once_with("deploy-team-b") + + def test_flag_on_falls_back_to_base_model_match_when_the_stored_deployment_is_gone(self): + """The stored deployment was removed from the config after the batch was created: + the row still routes through today's base-model lookup instead of being skipped.""" + router: Final = MagicMock() + router.registry.return_value.resolve.return_value = None + router.registry.return_value.deployments.return_value = self._two_team_deployments() + router.registry.return_value.deployment = MagicMock(return_value=None) + instance: Final = self._instance(track_unmanaged=True, router=router) + job: Final = self._job( + file_object=json.dumps({**json.loads(_unmanaged_vertex_file_object()), "model_id": "deploy-removed"}) + ) + + with patch(_IS_B64, return_value=False): + result: Final = instance._resolve_job_routing(job, MagicMock()) + + assert result == ("deploy-team-a", "8823717160934178816") + def test_flag_on_skips_non_vertex_deployment_sharing_model_group(self): """Flag on, but the only deployment for the model group is a non-vertex_ai provider: must not be selected, even though the model group name matches.""" From 991ea7f7e8419feb9be44c1aa67fdc95289c1762 Mon Sep 17 00:00:00 2001 From: yuneng-jiang Date: Fri, 30 Oct 2026 15:47:06 -0700 Subject: [PATCH 30/46] refactor(ui): move policies, tool policies and guardrails monitor text onto i18n (#49343) Moves the page-authored English on /policies, /tool-policies and /guardrails-monitor into react-i18next t() calls backed by en.json, and adds a small data test for en.json. The rendered English is unchanged. --- .../_components/EvaluationSettingsModal.tsx | 36 +- .../_components/GuardrailDetail.tsx | 46 +- .../_components/GuardrailUsageBreakdown.tsx | 111 ++- .../_components/GuardrailsOverview.tsx | 91 +- .../_components/ScoreChart.tsx | 6 +- .../(dashboard)/guardrails-monitor/page.tsx | 4 +- .../policies/_components/AttachmentTable.tsx | 15 +- .../_components/AttachmentTableColumns.tsx | 76 +- .../policies/_components/PolicyTable.tsx | 15 +- .../_components/PolicyTableColumns.tsx | 65 +- .../policies/_components/TokenSelect.tsx | 6 +- .../_components/add_attachment_form.tsx | 240 ++--- .../policies/_components/add_policy_form.tsx | 215 +++-- .../_components/ai_suggestion_modal.tsx | 154 +-- .../_components/guardrail_selection_modal.tsx | 78 +- .../policies/_components/impact_popover.tsx | 28 +- .../_components/impact_preview_alert.tsx | 58 +- .../policies/_components/index.tsx | 157 +-- .../_components/pipeline_flow_builder.tsx | 530 +++++----- .../policies/_components/policy_info.tsx | 56 +- .../policies/_components/policy_templates.tsx | 31 +- .../_components/policy_test_panel.tsx | 99 +- .../_components/template_parameter_modal.tsx | 61 +- .../GuardrailsMonitor/CalcPopover.tsx | 7 +- .../GuardrailsMonitor/LogViewer.tsx | 38 +- .../GuardrailsMonitor/UnpricedNote.tsx | 26 +- .../GuardrailsMonitor/usageUnits.ts | 21 +- .../src/components/ToolDetail.tsx | 86 +- .../ToolPolicies/ToolPoliciesPanel.tsx | 51 +- .../ToolPolicies/ToolPoliciesTable.tsx | 82 +- .../ToolPolicies/ToolPoliciesTableColumns.tsx | 41 +- .../src/components/ToolPoliciesView.tsx | 6 +- .../policies/useDeletePolicyAttachment.ts | 6 +- ui/litellm-dashboard/src/i18n/locales.test.ts | 54 ++ ui/litellm-dashboard/src/i18n/locales/en.json | 911 ++++++++++++++++++ 35 files changed, 2455 insertions(+), 1052 deletions(-) create mode 100644 ui/litellm-dashboard/src/i18n/locales.test.ts diff --git a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/EvaluationSettingsModal.tsx b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/EvaluationSettingsModal.tsx index 78091584cde..12ff10bc0c5 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/EvaluationSettingsModal.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/guardrails-monitor/_components/EvaluationSettingsModal.tsx @@ -1,6 +1,7 @@ import { useTerminology } from "@/contexts/TerminologyContext"; import { Play } from "lucide-react"; import React, { useEffect, useMemo, useState } from "react"; +import { useTranslation } from "react-i18next"; import { fetchAvailableModels, type ModelGroup } from "@/components/llm_calls/fetch_models"; import { SearchSelect } from "@/components/shared/SearchSelect"; import { Button } from "@/components/ui/button"; @@ -48,6 +49,7 @@ export function EvaluationSettingsModal({ accessToken, onRunEvaluation, }: EvaluationSettingsModalProps) { + const { t } = useTranslation(); const term = useTerminology(); const [prompt, setPrompt] = useState(DEFAULT_PROMPT); const [schema, setSchema] = useState(DEFAULT_SCHEMA); @@ -94,11 +96,11 @@ export function EvaluationSettingsModal({ !nextOpen && onClose()}> - Evaluation Settings + {t("guardrailsMonitor.evaluationSettings.title")} {guardrailName - ? `${term("Configure AI evaluation for")} ${guardrailName}` - : term("Configure AI evaluation for re-running on logs")} + ? `${term(t("guardrailsMonitor.evaluationSettings.descriptionPrefix"))} ${guardrailName}` + : term(t("guardrailsMonitor.evaluationSettings.descriptionGeneric"))} @@ -106,10 +108,10 @@ export function EvaluationSettingsModal({