From 79accda0900ce2dbe3de9da5f6e7290eda360e03 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 11 Aug 2026 02:35:49 +0000 Subject: [PATCH 01/24] fix(bedrock): add text block to converse user messages carrying documents --- .../prompt_templates/factory.py | 34 ++++- .../pdf_input/test_bedrock_converse.py | 4 - ...llm_core_utils_prompt_templates_factory.py | 138 ++++++++++++++++++ 3 files changed, 170 insertions(+), 6 deletions(-) diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 3a1a426eaa9..bf86ff3aea8 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -3967,6 +3967,36 @@ def _rename_duplicate_bedrock_document_names( return contents +BEDROCK_DOCUMENT_PLACEHOLDER_TEXT: Final = "." + + +def _with_text_when_document_only(message: BedrockMessageBlock) -> BedrockMessageBlock: + blocks: Final = message["content"] + needs_text: Final = ( + message["role"] == "user" + and any("document" in block for block in blocks) + and all("text" not in block for block in blocks) + ) + if not needs_text: + return message + placeholder: Final = BedrockContentBlock(text=BEDROCK_DOCUMENT_PLACEHOLDER_TEXT) + cut: Final = len(blocks) - 1 if "cachePoint" in blocks[-1] else len(blocks) + return BedrockMessageBlock(role="user", content=[*blocks[:cut], placeholder, *blocks[cut:]]) + + +def _ensure_document_messages_have_text( + contents: list[BedrockMessageBlock], +) -> list[BedrockMessageBlock]: + """ + Bedrock Converse rejects any user message that carries a document block + without a sibling text block ("A text block must be included when using + documents"), e.g. Claude Code sends the PDF as a document-only user turn. + Inject a placeholder text block, kept ahead of a trailing cachePoint so + the caller's cache boundary stays the final block. + """ + return [_with_text_when_document_only(message) for message in contents] + + def _sort_bedrock_assistant_content_blocks( blocks: list[BedrockContentBlock], ) -> list[BedrockContentBlock]: @@ -4535,7 +4565,7 @@ class BedrockConverseMessagesProcessor: llm_provider=llm_provider, ) - return _rename_duplicate_bedrock_document_names(contents) + return _ensure_document_messages_have_text(_rename_duplicate_bedrock_document_names(contents)) @staticmethod def translate_thinking_blocks_to_reasoning_content_blocks( @@ -4911,7 +4941,7 @@ def _bedrock_converse_messages_pt( llm_provider=llm_provider, ) - return _rename_duplicate_bedrock_document_names(contents) + return _ensure_document_messages_have_text(_rename_duplicate_bedrock_document_names(contents)) def make_valid_bedrock_tool_name(input_tool_name: str) -> str: diff --git a/tests/e2e/claude_code/pdf_input/test_bedrock_converse.py b/tests/e2e/claude_code/pdf_input/test_bedrock_converse.py index 5725255ed8b..76aa84f0f47 100644 --- a/tests/e2e/claude_code/pdf_input/test_bedrock_converse.py +++ b/tests/e2e/claude_code/pdf_input/test_bedrock_converse.py @@ -88,10 +88,6 @@ def _build_minimal_pdf(marker: str) -> bytes: return bytes(out) -@pytest.mark.skip( - reason="product bug LIT-4523: Bedrock Converse requires a text block with document; " - "re-enable when document-only content is handled" -) @pytest.mark.covers("llm.messages.bedrock_converse.pdf_input.nonstream.works") def test_pdf_input_bedrock_converse(compat_result, tmp_path): base_url, api_key = require_proxy(compat_result) diff --git a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index dc745abb9e7..8edc6a91cbf 100644 --- a/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/test_litellm/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -8,6 +8,7 @@ import pytest import litellm from litellm.litellm_core_utils.prompt_templates.factory import ( BAD_MESSAGE_ERROR_STR, + BEDROCK_DOCUMENT_PLACEHOLDER_TEXT, BedrockConverseMessagesProcessor, BedrockImageProcessor, _bedrock_converse_messages_pt, @@ -3269,3 +3270,140 @@ def test_group_tool_exchanges_is_linear_in_message_count(): assert len(groups) == 100_000 assert elapsed < 3.0, f"grouping 100k messages took {elapsed:.2f}s; suspect superlinear accumulation" + + +_PDF_DATA_URI = "data:application/pdf;base64," + base64.b64encode(b"%PDF-1.4 regression fixture").decode() +_PNG_DATA_URI = ( + "data:image/png;base64," + "iVBORw0KGgoAAAANSUhEUgAAAAEAAAABCAYAAAAfFcSJAAAADUlEQVR42mNkYPhfDwAChwGA60e6kgAAAABJRU5ErkJggg==" +) + + +def _text_blocks(message): + return [block["text"] for block in message["content"] if "text" in block] + + +def test_bedrock_converse_pdf_only_user_message_gets_text_block(): + """ + Regression for LIT-4523: Claude Code sends a PDF as a user turn whose only + content is the document (an image_url part with a pdf data URI after the + /v1/messages -> completion bridge). Bedrock Converse rejects any user + message carrying a document without a sibling text block, so the builder + must inject a placeholder text block. + """ + messages = [ + { + "role": "user", + "content": [{"type": "image_url", "image_url": {"url": _PDF_DATA_URI}}], + } + ] + + result = _bedrock_converse_messages_pt( + messages, "anthropic.claude-haiku-4-5", "bedrock" + ) + + assert len(result) == 1 + assert any("document" in block for block in result[0]["content"]) + assert _text_blocks(result[0]) == [BEDROCK_DOCUMENT_PLACEHOLDER_TEXT] + + +def test_bedrock_converse_document_with_text_gets_no_extra_text_block(): + messages = [ + { + "role": "user", + "content": [ + {"type": "image_url", "image_url": {"url": _PDF_DATA_URI}}, + {"type": "text", "text": "summarize this"}, + ], + } + ] + + result = _bedrock_converse_messages_pt( + messages, "anthropic.claude-haiku-4-5", "bedrock" + ) + + assert _text_blocks(result[0]) == ["summarize this"] + + +def test_bedrock_converse_image_only_user_message_gets_no_text_block(): + messages = [ + { + "role": "user", + "content": [{"type": "image_url", "image_url": {"url": _PNG_DATA_URI}}], + } + ] + + result = _bedrock_converse_messages_pt( + messages, "anthropic.claude-haiku-4-5", "bedrock" + ) + + assert any("image" in block for block in result[0]["content"]) + assert _text_blocks(result[0]) == [] + + +def test_bedrock_converse_tool_round_trip_document_injects_text_before_cache_point(): + """ + Claude Code shape: after a Read tool round trip, the document-only user + turn (with cache_control) merges into the toolResult message. The injected + text block must land before the trailing cachePoint so the cache boundary + stays the final block, and earlier turns must stay untouched. + """ + messages = [ + {"role": "user", "content": "read the pdf"}, + { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "tooluse_pdf1", + "type": "function", + "function": {"name": "Read", "arguments": "{}"}, + } + ], + }, + {"role": "tool", "tool_call_id": "tooluse_pdf1", "content": "read ok"}, + { + "role": "user", + "content": [ + { + "type": "document", + "source": { + "type": "base64", + "media_type": "application/pdf", + "data": "dGVzdA==", + }, + "cache_control": {"type": "ephemeral"}, + } + ], + }, + ] + + result = _bedrock_converse_messages_pt( + messages, "anthropic.claude-haiku-4-5", "bedrock" + ) + + assert _text_blocks(result[0]) == ["read the pdf"] + document_message = result[-1] + block_keys = [next(iter(block)) for block in document_message["content"]] + assert block_keys == ["toolResult", "document", "text", "cachePoint"] + assert _text_blocks(document_message) == [BEDROCK_DOCUMENT_PLACEHOLDER_TEXT] + + +@pytest.mark.asyncio +async def test_bedrock_converse_pdf_only_user_message_gets_text_block_async(): + messages = [ + { + "role": "user", + "content": [{"type": "image_url", "image_url": {"url": _PDF_DATA_URI}}], + } + ] + + result = await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( + messages=messages, + model="anthropic.claude-haiku-4-5", + llm_provider="bedrock", + ) + + assert len(result) == 1 + assert any("document" in block for block in result[0]["content"]) + assert _text_blocks(result[0]) == [BEDROCK_DOCUMENT_PLACEHOLDER_TEXT] From bc98c67028235aafd893e607bbe3d09a02ae663b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 11 Aug 2026 02:46:27 +0000 Subject: [PATCH 02/24] fix(bedrock): send tool-search beta header for Haiku 4.5 on Invoke /v1/messages Bedrock InvokeModel rejects tool_search_tool_* tool types unless the request body carries the tool-search-tool-2025-10-19 beta. The model allowlist gating that beta omitted Haiku 4.5 (and Opus 4.7, supported since launch per live verification), so every tool-search request on those models got a Bedrock 400. Add both to the allowlist and re-enable the e2e compat cell that caught it. --- .../anthropic_claude3_transformation.py | 29 ++++++----- .../tool_search/test_bedrock_invoke.py | 4 -- .../test_anthropic_claude3_transformation.py | 49 +++++++++++++++++++ 3 files changed, 67 insertions(+), 15 deletions(-) diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index fad7e7558c2..9a6fdf78090 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -372,8 +372,9 @@ class AmazonAnthropicClaudeMessagesConfig( """ Check if the model supports tool search on Bedrock. - On Amazon Bedrock, server-side tool search is supported on Claude Opus 4.5 - and Claude Sonnet 4.5 with the tool-search-tool-2025-10-19 beta header. + On Amazon Bedrock, server-side tool search is supported on Claude + Opus 4.5/4.6/4.7, Sonnet 4.5/4.6, and Haiku 4.5 with the + tool-search-tool-2025-10-19 beta header. Ref: https://platform.claude.com/docs/en/agents-and-tools/tool-use/tool-search-tool @@ -407,10 +408,17 @@ class AmazonAnthropicClaudeMessagesConfig( "sonnet_4.6", "sonnet-4-6", "sonnet_4_6", - # NOTE: Opus 4.7 on Bedrock does not support server-side tool search - # as of launch (2026-04-16). Bedrock rejects the tool type with: - # "tool type 'tool_search_tool_..._20251119' is not supported for this model". - # Re-add the opus-4.7 patterns here once AWS announces support. + # Opus 4.7 (unsupported at its 2026-04-16 launch; verified live + # 2026-08-11 that Bedrock now accepts the beta on it) + "opus-4.7", + "opus_4.7", + "opus-4-7", + "opus_4_7", + # Haiku 4.5 + "haiku-4.5", + "haiku_4.5", + "haiku-4-5", + "haiku_4_5", ] return any(pattern in model_lower for pattern in supported_patterns) @@ -426,11 +434,10 @@ class AmazonAnthropicClaudeMessagesConfig( """ Adjust tool search beta header for Bedrock. - Bedrock requires a different beta header for tool search on Opus 4 models - when tool search is used without programmatic tool calling or input examples. - - Note: On Amazon Bedrock, server-side tool search is only supported on Claude Opus 4 - with the `tool-search-tool-2025-10-19` beta header. + Bedrock requires a different beta header for tool search than the + Anthropic API when tool search is used without programmatic tool + calling or input examples: `tool-search-tool-2025-10-19`, and only on + the models listed in `_supports_tool_search_on_bedrock`. Ref: https://platform.claude.com/docs/en/agents-and-tools/tool-use/tool-search-tool diff --git a/tests/e2e/claude_code/tool_search/test_bedrock_invoke.py b/tests/e2e/claude_code/tool_search/test_bedrock_invoke.py index 12f8909e3e8..c4735c78f0c 100644 --- a/tests/e2e/claude_code/tool_search/test_bedrock_invoke.py +++ b/tests/e2e/claude_code/tool_search/test_bedrock_invoke.py @@ -59,10 +59,6 @@ BEDROCK_INVOKE_MODELS = [ ] -@pytest.mark.skip( - reason="product bug LIT-4522: Bedrock Invoke /v1/messages does not normalize " - "tool_search_tool_regex_20251119; re-enable when messages path matches chat path" -) @pytest.mark.covers("llm.messages.bedrock_invoke.tool_search.nonstream.works") def test_tool_search_bedrock_invoke(compat_result): """Probe `/v1/messages` with a `tool_search_tool_regex_20251119` diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 2d1a5bc4c2a..67409cc6cd0 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -2474,6 +2474,55 @@ def test_filter_and_transform_beta_headers_passes_context_management_for_bedrock assert out_converse == [] +@pytest.mark.parametrize( + "model", + [ + "us.anthropic.claude-haiku-4-5-20251001-v1:0", + "us.anthropic.claude-sonnet-4-5-20250929-v1:0", + "us.anthropic.claude-opus-4-7", + ], +) +def test_bedrock_messages_tool_search_adds_beta_header(local_beta_headers_config, model): + """ + LIT-4522: Bedrock InvokeModel only admits ``tool_search_tool_*`` tool types + when the request body carries the ``tool-search-tool-2025-10-19`` beta; + without it Bedrock 400s with "Input tag 'tool_search_tool_regex_20251119' + ... does not match any of the expected tags". The allowlist in + ``_supports_tool_search_on_bedrock`` previously omitted Haiku 4.5 and + Opus 4.7, so the beta was silently dropped for those models and every + tool-search request failed. Verified live 2026-08-11: Bedrock returns 200 + with ``server_tool_use`` for all three models once the beta is sent. + """ + from litellm.types.router import GenericLiteLLMParams + + cfg = AmazonAnthropicClaudeMessagesConfig() + messages = [{"role": "user", "content": [{"type": "text", "text": "Hi"}]}] + optional_params = { + "max_tokens": 64, + "tools": [ + {"type": "tool_search_tool_regex_20251119", "name": "tool_search_tool_regex"}, + { + "name": "add_numbers", + "description": "Add two integers", + "input_schema": { + "type": "object", + "properties": {"a": {"type": "integer"}, "b": {"type": "integer"}}, + "required": ["a", "b"], + }, + }, + ], + } + + result = cfg.transform_anthropic_messages_request( + model=model, + messages=messages, + anthropic_messages_optional_request_params=optional_params, + litellm_params=GenericLiteLLMParams(), + headers={}, + ) + + assert "tool-search-tool-2025-10-19" in (result.get("anthropic_beta") or []) + def test_bedrock_messages_thinking_shape_follows_exact_bedrock_entry_flag( local_model_cost_map, monkeypatch From 929ee52b873d1f70b6395b607b7ae6d169270dc8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 11 Aug 2026 03:23:16 +0000 Subject: [PATCH 03/24] fix(bedrock): preserve adaptive thinking effort through the /v1/messages bridge Claude Code drives Opus 4.7 with thinking {"type": "adaptive"} plus output_config {"effort": "max"}. The anthropic-to-openai adapter forwarded thinking verbatim for Claude models but dropped output_config, and Bedrock Converse streams zero reasoningContent blocks for adaptive thinking without an effort tier. Forward the effort subset of output_config for Bedrock targets, accept it in the converse supported params, and map it with the model's effort ceiling applied. Re-enable the skipped e2e compat cell that catches this --- .../adapters/transformation.py | 11 +++ .../bedrock/chat/converse_transformation.py | 6 ++ litellm/types/llms/openai.py | 1 + .../thinking/test_bedrock_converse.py | 4 - ...al_pass_through_adapters_transformation.py | 82 +++++++++++++++++++ .../chat/test_converse_transformation.py | 40 +++++++++ 6 files changed, 140 insertions(+), 4 deletions(-) diff --git a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py index 22f9bfd30ea..3640f9fd2b8 100644 --- a/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/experimental_pass_through/adapters/transformation.py @@ -976,6 +976,17 @@ class LiteLLMAnthropicMessagesAdapter: model: Final = new_kwargs.get("model", "") if self.is_anthropic_claude_model(model) or self.is_bedrock_arn_model(model): new_kwargs["thinking"] = thinking + # Adaptive thinking without its effort tier makes Bedrock Converse + # return zero reasoning blocks, so forward output_config (minus + # `format`, already translated to response_format) for Bedrock + # targets only: other bridged providers reject the raw param, and + # get_llm_provider strips the `bedrock/` prefix before this runs. + if model.startswith(("bedrock/", "converse/", "invoke/")) or self.is_bedrock_arn_model(model): + claude_output_config: Final = anthropic_message_request.get("output_config") + if isinstance(claude_output_config, dict): + effort_config: Final = {k: v for k, v in claude_output_config.items() if k != "format"} + if effort_config: + new_kwargs["output_config"] = effort_config # rebind-ok: out-param store like thinking above return reasoning_effort = self.translate_anthropic_thinking_to_reasoning_effort(cast(dict[str, Any], thinking)) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index 193987a3543..f7f240af54f 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -514,6 +514,7 @@ class AmazonConverseConfig(BaseConfig): supported_params.append("tool_choice") supported_params.append("thinking") supported_params.append("reasoning_effort") + supported_params.append("output_config") # For nova imported models, also add web_search_options if "nova" in model.lower(): supported_params.append("web_search_options") @@ -564,6 +565,7 @@ class AmazonConverseConfig(BaseConfig): ): supported_params.append("thinking") supported_params.append("reasoning_effort") + supported_params.append("output_config") if base_model.startswith("anthropic"): supported_params.append("context_management") @@ -919,6 +921,10 @@ class AmazonConverseConfig(BaseConfig): self._handle_reasoning_effort_parameter( model=model, reasoning_effort=value, optional_params=optional_params ) + elif param == "output_config" and isinstance(value, dict): + mapped_output_config = dict(value) + normalize_bedrock_opus_output_config_effort(model=model, output_config=mapped_output_config) + optional_params["output_config"] = mapped_output_config # rebind-ok: out-param store like siblings elif param == "context_management" and isinstance(value, (dict, list)): self._map_context_management_param(value, optional_params) if param == "requestMetadata": diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index da0592e6bb2..4eec48c9c89 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -929,6 +929,7 @@ class ChatCompletionRequest(TypedDict, total=False): user: str metadata: dict # litellm specific param reasoning_effort: str # OpenAI o1/o3 reasoning parameter + output_config: Mapping[str, object] # Anthropic adaptive-thinking effort, bridged for Bedrock Claude class ChatCompletionDeltaChunk(TypedDict, total=False): diff --git a/tests/e2e/claude_code/thinking/test_bedrock_converse.py b/tests/e2e/claude_code/thinking/test_bedrock_converse.py index 3b1449d8cb7..0b409f18ea7 100644 --- a/tests/e2e/claude_code/thinking/test_bedrock_converse.py +++ b/tests/e2e/claude_code/thinking/test_bedrock_converse.py @@ -54,10 +54,6 @@ def _has_thinking_block(events: Sequence[Mapping[str, Any]]) -> bool: return False -@pytest.mark.skip( - reason="product bug LIT-4524: Bedrock Converse streaming Content block is not a text block; " - "re-enable when empty/mismatched content_block_delta is fixed" -) @pytest.mark.covers("llm.messages.bedrock_converse.thinking.nonstream.works") def test_thinking_bedrock_converse(compat_result): """Drive the `claude` CLI against the LiteLLM proxy with thinking diff --git a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index c0c6e315b5b..1ade9a19f35 100644 --- a/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/test_litellm/llms/anthropic/experimental_pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -1610,6 +1610,88 @@ def test_thinking_disabled_stays_plain_string_when_auto_summary_enabled(): assert new_kwargs["reasoning_effort"] == "none" +@pytest.mark.parametrize( + "model", + [ + # SDK-style model with the provider prefix intact + "bedrock/converse/us.anthropic.claude-opus-4-7", + # what the bridge actually sees in the proxy: get_llm_provider has + # already stripped the `bedrock/` prefix by the time it translates + "converse/us.anthropic.claude-opus-4-7", + ], +) +def test_adaptive_thinking_output_config_effort_preserved_for_claude_model(model): + """ + Regression: Claude Code drives adaptive thinking as `thinking: {"type": "adaptive"}` + plus `output_config: {"effort": "max"}`. The Claude branch of the thinking translator + forwarded `thinking` verbatim but returned early without reading `output_config`, and + the handler strips the raw key from extra_kwargs, so the effort tier never reached the + backend. On Bedrock Converse, adaptive thinking without effort streams zero reasoning + blocks. The `format` subkey must still be excluded (it is translated to + `response_format` separately). + """ + from litellm.types.llms.anthropic import AnthropicMessagesRequest + + anthropic_request = AnthropicMessagesRequest( + model=model, + max_tokens=1024, + messages=[{"role": "user", "content": "hi"}], + thinking={"type": "adaptive"}, + output_config={ + "effort": "max", + "format": {"type": "json_schema", "schema": {"type": "object", "properties": {}}}, + }, + ) + + adapter = LiteLLMAnthropicMessagesAdapter() + openai_request, _ = adapter.translate_anthropic_to_openai(anthropic_message_request=anthropic_request) + + assert openai_request["thinking"] == {"type": "adaptive"} + assert openai_request["output_config"] == {"effort": "max"} + assert "response_format" in openai_request + + +def test_adaptive_thinking_format_only_output_config_not_forwarded_for_claude_model(): + """When `output_config` carries only `format`, nothing effort-bearing remains, so the + translator must not forward an empty `output_config` dict.""" + from litellm.types.llms.anthropic import AnthropicMessagesRequest + + anthropic_request = AnthropicMessagesRequest( + model="bedrock/converse/us.anthropic.claude-opus-4-7", + max_tokens=1024, + messages=[{"role": "user", "content": "hi"}], + thinking={"type": "adaptive"}, + output_config={"format": {"type": "json_schema", "schema": {"type": "object", "properties": {}}}}, + ) + + adapter = LiteLLMAnthropicMessagesAdapter() + openai_request, _ = adapter.translate_anthropic_to_openai(anthropic_message_request=anthropic_request) + + assert openai_request["thinking"] == {"type": "adaptive"} + assert "output_config" not in openai_request + + +def test_adaptive_thinking_output_config_not_forwarded_for_non_bedrock_claude_model(): + """`output_config` is forwarded only for Bedrock-destined Claude models. Other + Claude-through-bridge providers (e.g. openrouter) accept `thinking` but reject a raw + `output_config` param with UnsupportedParamsError when drop_params is off.""" + from litellm.types.llms.anthropic import AnthropicMessagesRequest + + anthropic_request = AnthropicMessagesRequest( + model="openrouter/anthropic/claude-opus-4-7", + max_tokens=1024, + messages=[{"role": "user", "content": "hi"}], + thinking={"type": "adaptive"}, + output_config={"effort": "max"}, + ) + + adapter = LiteLLMAnthropicMessagesAdapter() + openai_request, _ = adapter.translate_anthropic_to_openai(anthropic_message_request=anthropic_request) + + assert openai_request["thinking"] == {"type": "adaptive"} + assert "output_config" not in openai_request + + def test_stop_sequences_translated_to_stop_for_non_claude_model(): from litellm.types.llms.anthropic import AnthropicMessagesRequest diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index 7b7796fb01b..cba87427bee 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -370,6 +370,46 @@ def test_output_config_effort_forwarded_into_additional_request_fields(model): assert additional.get("output_config") == {"effort": "high"} +@pytest.mark.parametrize( + "model,effort,expected_effort", + [ + ("bedrock/converse/us.anthropic.claude-opus-4-7", "max", "max"), + ("bedrock/converse/us.anthropic.claude-opus-4-6-v1", "xhigh", "max"), + ], +) +def test_explicit_output_config_effort_mapped_for_adaptive_thinking_converse(model, effort, expected_effort): + """Regression: Claude Code drives adaptive thinking as ``thinking: {"type": + "adaptive"}`` plus ``output_config: {"effort": ...}``. ``output_config`` must + be a supported openai param and survive ``map_openai_params`` (clamped to the + model's Bedrock effort ceiling), otherwise the Converse request carries + adaptive thinking without an effort tier and Bedrock streams zero + ``reasoningContent`` blocks.""" + config = AmazonConverseConfig() + + assert "output_config" in config.get_supported_openai_params(model) + + optional_params = config.map_openai_params( + non_default_params={ + "thinking": {"type": "adaptive"}, + "output_config": {"effort": effort}, + }, + optional_params={}, + model=model, + drop_params=False, + ) + + assert optional_params["thinking"] == {"type": "adaptive"} + assert optional_params["output_config"] == {"effort": expected_effort} + + +def test_output_config_supported_param_for_arn_models_converse(): + """ARN model ids hide the underlying Claude model, so ``output_config`` must + be in the blanket ARN supported-params list too.""" + config = AmazonConverseConfig() + arn_model = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abcdef123456" + assert "output_config" in config.get_supported_openai_params(arn_model) + + def test_output_config_format_translated_to_native_output_config_converse(): """``output_config.format`` becomes Bedrock ``outputConfig`` and is not forwarded raw.""" config = AmazonConverseConfig() From 0f41365c3473115bb5e2e946af855dfbc023178b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 11 Aug 2026 09:11:22 +0000 Subject: [PATCH 04/24] fix(bedrock): forward output_config effort for application inference profile ARNs --- .../bedrock/chat/converse_transformation.py | 8 +++++- .../chat/test_converse_transformation.py | 27 +++++++++++++++++++ 2 files changed, 34 insertions(+), 1 deletion(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index f7f240af54f..85918d40e12 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -80,6 +80,7 @@ from ..common_utils import ( bedrock_converse_supports_parallel_tool_use_config, get_anthropic_beta_from_headers, get_bedrock_tool_name, + is_bedrock_application_inference_profile_arn, is_claude_4_5_on_bedrock, normalize_bedrock_opus_output_config_effort, ) @@ -1318,7 +1319,12 @@ class AmazonConverseConfig(BaseConfig): additional_request_params = filter_exceptions_from_params(additional_request_params) if anthropic_output_config is not None and isinstance(anthropic_output_config, dict): - if base_model.startswith("anthropic"): + # Application inference profile ARNs hide the underlying model, so the + # effort ceiling and capability gates below cannot run; forward + # verbatim (like ``thinking``) and let Bedrock enforce. + if is_bedrock_application_inference_profile_arn(model): + additional_request_params["output_config"] = anthropic_output_config + elif base_model.startswith("anthropic"): if litellm.drop_params is True and not AnthropicConfig._model_supports_effort_param(model, "bedrock"): litellm.verbose_logger.warning( DROP_UNSUPPORTED_OUTPUT_CONFIG_WARNING, diff --git a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py index cba87427bee..d1d1f9ab489 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -410,6 +410,33 @@ def test_output_config_supported_param_for_arn_models_converse(): assert "output_config" in config.get_supported_openai_params(arn_model) +def test_output_config_effort_forwarded_for_application_inference_profile_arn(): + """Regression: opaque application inference profile ARNs cannot resolve a + base model, so the anthropic-only serialization gate dropped ``output_config`` + while still sending ``thinking``: adaptive thinking with no effort tier, and + Bedrock streams zero ``reasoningContent`` blocks. The effort must be forwarded + verbatim (ceilings and capability gates are unknowable behind the alias) for + Bedrock to enforce.""" + config = AmazonConverseConfig() + arn_model = "arn:aws:bedrock:us-east-1:123456789012:application-inference-profile/abcdef123456" + + result = config._transform_request( + model=arn_model, + messages=[{"role": "user", "content": "hi"}], + optional_params={ + "maxTokens": 256, + "thinking": {"type": "adaptive"}, + "output_config": {"effort": "max"}, + }, + litellm_params={}, + headers={}, + ) + + additional = result.get("additionalModelRequestFields", {}) + assert additional.get("thinking") == {"type": "adaptive"} + assert additional.get("output_config") == {"effort": "max"} + + def test_output_config_format_translated_to_native_output_config_converse(): """``output_config.format`` becomes Bedrock ``outputConfig`` and is not forwarded raw.""" config = AmazonConverseConfig() From 7d4488d2e8a9b9c497be95dc5af7cfe547fb72cc Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 11 Aug 2026 10:06:04 +0000 Subject: [PATCH 05/24] refactor(bedrock): read tool search support from the model map Record supports_tool_search on the Bedrock Claude entries in both cost map files and have _supports_tool_search_on_bedrock read it first via the provider-resolved capability lookup, keeping the name patterns as a fallback for ARNs and ids the map cannot resolve. Threads the flag through ModelInfoBase and drops a dated remark from the pattern list --- .../anthropic_claude3_transformation.py | 14 +++--- ...odel_prices_and_context_window_backup.json | 43 +++++++++++++++++++ litellm/types/utils.py | 1 + litellm/utils.py | 1 + model_prices_and_context_window.json | 43 +++++++++++++++++++ .../test_anthropic_claude3_transformation.py | 36 ++++++++++++++++ 6 files changed, 132 insertions(+), 6 deletions(-) diff --git a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py index 9a6fdf78090..8d039d95bb1 100644 --- a/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py +++ b/litellm/llms/bedrock/messages/invoke_transformations/anthropic_claude3_transformation.py @@ -372,9 +372,9 @@ class AmazonAnthropicClaudeMessagesConfig( """ Check if the model supports tool search on Bedrock. - On Amazon Bedrock, server-side tool search is supported on Claude - Opus 4.5/4.6/4.7, Sonnet 4.5/4.6, and Haiku 4.5 with the - tool-search-tool-2025-10-19 beta header. + The model map's ``supports_tool_search`` flag is authoritative when + ``model`` resolves to an entry that sets it; the name patterns below + cover ids the map cannot resolve (ARNs, unlisted regional variants). Ref: https://platform.claude.com/docs/en/agents-and-tools/tool-use/tool-search-tool @@ -384,9 +384,12 @@ class AmazonAnthropicClaudeMessagesConfig( Returns: True if the model supports tool search on Bedrock """ + catalog: Final = AnthropicModelInfo._get_provider_resolved_capability(model, "supports_tool_search", "bedrock") + if catalog is not None: + return catalog + model_lower: Final = model.lower() - # Supported models for tool search on Bedrock supported_patterns: Final = [ # Opus 4.5 "opus-4.5", @@ -408,8 +411,7 @@ class AmazonAnthropicClaudeMessagesConfig( "sonnet_4.6", "sonnet-4-6", "sonnet_4_6", - # Opus 4.7 (unsupported at its 2026-04-16 launch; verified live - # 2026-08-11 that Bedrock now accepts the beta on it) + # Opus 4.7 "opus-4.7", "opus_4.7", "opus-4-7", diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 951c114b0a9..c8b34927e7d 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -730,6 +730,7 @@ "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 1e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -755,6 +756,7 @@ "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 1e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -973,6 +975,7 @@ "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -1005,6 +1008,7 @@ "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, @@ -1038,6 +1042,7 @@ "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, @@ -1071,6 +1076,7 @@ "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 5.5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, @@ -1104,6 +1110,7 @@ "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 5.5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, @@ -1137,6 +1144,7 @@ "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 5.5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, @@ -1171,6 +1179,7 @@ "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, @@ -1222,6 +1231,7 @@ "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, @@ -1258,6 +1268,7 @@ "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 5.5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, @@ -1294,6 +1305,7 @@ "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 5.5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, @@ -1330,6 +1342,7 @@ "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 5.5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, @@ -1946,6 +1959,7 @@ "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 5.5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, @@ -2203,6 +2217,7 @@ "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -2235,6 +2250,7 @@ "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -2267,6 +2283,7 @@ "cache_read_input_token_cost": 3.3e-07, "input_cost_per_token": 3.3e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -2299,6 +2316,7 @@ "cache_read_input_token_cost": 3.3e-07, "input_cost_per_token": 3.3e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -2331,6 +2349,7 @@ "cache_read_input_token_cost": 3.3e-07, "input_cost_per_token": 3.3e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -2363,6 +2382,7 @@ "cache_read_input_token_cost": 3.3e-07, "input_cost_per_token": 3.3e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -2430,6 +2450,7 @@ "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -2686,6 +2707,7 @@ "cache_read_input_token_cost": 1.1e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -2775,6 +2797,7 @@ "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.32e-05, "cache_read_input_token_cost_above_200k_tokens": 6.6e-07, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -10826,6 +10849,7 @@ "cache_read_input_token_cost": 3.6e-07, "input_cost_per_token": 3.6e-06, "litellm_provider": "bedrock", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 8192, "max_tokens": 8192, @@ -10850,6 +10874,7 @@ "cache_read_input_token_cost": 3.6e-07, "input_cost_per_token": 3.6e-06, "litellm_provider": "bedrock", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 8192, "max_tokens": 8192, @@ -11005,6 +11030,7 @@ "cache_read_input_token_cost": 3.6e-07, "input_cost_per_token": 3.6e-06, "litellm_provider": "bedrock", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 8192, "max_tokens": 8192, @@ -11029,6 +11055,7 @@ "cache_read_input_token_cost": 3.6e-07, "input_cost_per_token": 3.6e-06, "litellm_provider": "bedrock", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 8192, "max_tokens": 8192, @@ -11811,6 +11838,7 @@ "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "bedrock", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -15998,6 +16026,7 @@ "input_cost_per_token": 1.1e-06, "deprecation_date": "2026-10-15", "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -16213,6 +16242,7 @@ "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.32e-05, "cache_read_input_token_cost_above_200k_tokens": 6.6e-07, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -22031,6 +22061,7 @@ "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -22091,6 +22122,7 @@ "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 1e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -26534,6 +26566,7 @@ "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.32e-05, "cache_read_input_token_cost_above_200k_tokens": 6.6e-07, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -26563,6 +26596,7 @@ "cache_read_input_token_cost": 1.1e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -35259,6 +35293,7 @@ "cache_read_input_token_cost": 1.1e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -35417,6 +35452,7 @@ "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.32e-05, "cache_read_input_token_cost_above_200k_tokens": 6.6e-07, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -35451,6 +35487,7 @@ "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.44e-05, "cache_read_input_token_cost_above_200k_tokens": 7.2e-07, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -35475,6 +35512,7 @@ "cache_read_input_token_cost": 1.1e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -35525,6 +35563,7 @@ "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 5.5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -35556,6 +35595,7 @@ "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -35586,6 +35626,7 @@ "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -45975,6 +46016,7 @@ "cache_read_input_token_cost": 1.2e-07, "input_cost_per_token": 1.2e-06, "litellm_provider": "bedrock", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -46000,6 +46042,7 @@ "cache_read_input_token_cost": 1.2e-07, "input_cost_per_token": 1.2e-06, "litellm_provider": "bedrock", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 74311d59d8e..d0d65c1a614 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -152,6 +152,7 @@ class ProviderSpecificModelInfo(TypedDict, total=False): supports_web_search: bool | None supports_reasoning: bool | None supports_adaptive_thinking: bool | None + supports_tool_search: bool | None supports_mid_conversation_system: bool | None supports_url_context: bool | None supports_none_reasoning_effort: bool | None diff --git a/litellm/utils.py b/litellm/utils.py index 87937c99a0c..d534a2e505b 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -5679,6 +5679,7 @@ def _get_model_info_helper( supports_url_context=_model_info.get("supports_url_context", None), supports_reasoning=_model_info.get("supports_reasoning", None), supports_adaptive_thinking=_model_info.get("supports_adaptive_thinking", None), + supports_tool_search=_model_info.get("supports_tool_search", None), supports_mid_conversation_system=_model_info.get("supports_mid_conversation_system", None), supports_none_reasoning_effort=_model_info.get("supports_none_reasoning_effort", None), supports_minimal_reasoning_effort=_model_info.get("supports_minimal_reasoning_effort", None), diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 951c114b0a9..c8b34927e7d 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -730,6 +730,7 @@ "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 1e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -755,6 +756,7 @@ "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 1e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -973,6 +975,7 @@ "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -1005,6 +1008,7 @@ "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, @@ -1038,6 +1042,7 @@ "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, @@ -1071,6 +1076,7 @@ "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 5.5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, @@ -1104,6 +1110,7 @@ "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 5.5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, @@ -1137,6 +1144,7 @@ "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 5.5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, @@ -1171,6 +1179,7 @@ "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, @@ -1222,6 +1231,7 @@ "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, @@ -1258,6 +1268,7 @@ "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 5.5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, @@ -1294,6 +1305,7 @@ "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 5.5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, @@ -1330,6 +1342,7 @@ "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 5.5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, @@ -1946,6 +1959,7 @@ "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 5.5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 128000, "max_tokens": 128000, @@ -2203,6 +2217,7 @@ "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -2235,6 +2250,7 @@ "cache_read_input_token_cost": 3e-07, "input_cost_per_token": 3e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -2267,6 +2283,7 @@ "cache_read_input_token_cost": 3.3e-07, "input_cost_per_token": 3.3e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -2299,6 +2316,7 @@ "cache_read_input_token_cost": 3.3e-07, "input_cost_per_token": 3.3e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -2331,6 +2349,7 @@ "cache_read_input_token_cost": 3.3e-07, "input_cost_per_token": 3.3e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -2363,6 +2382,7 @@ "cache_read_input_token_cost": 3.3e-07, "input_cost_per_token": 3.3e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 1000000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -2430,6 +2450,7 @@ "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -2686,6 +2707,7 @@ "cache_read_input_token_cost": 1.1e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -2775,6 +2797,7 @@ "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.32e-05, "cache_read_input_token_cost_above_200k_tokens": 6.6e-07, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -10826,6 +10849,7 @@ "cache_read_input_token_cost": 3.6e-07, "input_cost_per_token": 3.6e-06, "litellm_provider": "bedrock", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 8192, "max_tokens": 8192, @@ -10850,6 +10874,7 @@ "cache_read_input_token_cost": 3.6e-07, "input_cost_per_token": 3.6e-06, "litellm_provider": "bedrock", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 8192, "max_tokens": 8192, @@ -11005,6 +11030,7 @@ "cache_read_input_token_cost": 3.6e-07, "input_cost_per_token": 3.6e-06, "litellm_provider": "bedrock", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 8192, "max_tokens": 8192, @@ -11029,6 +11055,7 @@ "cache_read_input_token_cost": 3.6e-07, "input_cost_per_token": 3.6e-06, "litellm_provider": "bedrock", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 8192, "max_tokens": 8192, @@ -11811,6 +11838,7 @@ "cache_creation_input_token_cost_above_200k_tokens": 7.5e-06, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "bedrock", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -15998,6 +16026,7 @@ "input_cost_per_token": 1.1e-06, "deprecation_date": "2026-10-15", "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -16213,6 +16242,7 @@ "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.32e-05, "cache_read_input_token_cost_above_200k_tokens": 6.6e-07, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -22031,6 +22061,7 @@ "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.2e-05, "cache_read_input_token_cost_above_200k_tokens": 6e-07, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -22091,6 +22122,7 @@ "cache_read_input_token_cost": 1e-07, "input_cost_per_token": 1e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -26534,6 +26566,7 @@ "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.32e-05, "cache_read_input_token_cost_above_200k_tokens": 6.6e-07, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -26563,6 +26596,7 @@ "cache_read_input_token_cost": 1.1e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -35259,6 +35293,7 @@ "cache_read_input_token_cost": 1.1e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -35417,6 +35452,7 @@ "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.32e-05, "cache_read_input_token_cost_above_200k_tokens": 6.6e-07, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -35451,6 +35487,7 @@ "cache_creation_input_token_cost_above_1hr_above_200k_tokens": 1.44e-05, "cache_read_input_token_cost_above_200k_tokens": 7.2e-07, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -35475,6 +35512,7 @@ "cache_read_input_token_cost": 1.1e-07, "input_cost_per_token": 1.1e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -35525,6 +35563,7 @@ "cache_read_input_token_cost": 5.5e-07, "input_cost_per_token": 5.5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -35556,6 +35595,7 @@ "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -35586,6 +35626,7 @@ "cache_read_input_token_cost": 5e-07, "input_cost_per_token": 5e-06, "litellm_provider": "bedrock_converse", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -45975,6 +46016,7 @@ "cache_read_input_token_cost": 1.2e-07, "input_cost_per_token": 1.2e-06, "litellm_provider": "bedrock", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, @@ -46000,6 +46042,7 @@ "cache_read_input_token_cost": 1.2e-07, "input_cost_per_token": 1.2e-06, "litellm_provider": "bedrock", + "supports_tool_search": true, "max_input_tokens": 200000, "max_output_tokens": 64000, "max_tokens": 64000, diff --git a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py index 67409cc6cd0..02bd3535c0a 100644 --- a/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py +++ b/tests/test_litellm/llms/bedrock/messages/invoke_transformations/test_anthropic_claude3_transformation.py @@ -2524,6 +2524,42 @@ def test_bedrock_messages_tool_search_adds_beta_header(local_beta_headers_config assert "tool-search-tool-2025-10-19" in (result.get("anthropic_beta") or []) +def test_bedrock_messages_tool_search_model_map_flag_is_authoritative(local_model_cost_map, monkeypatch): + """``supports_tool_search`` lives in the model map; the name patterns in + ``_supports_tool_search_on_bedrock`` are only a fallback for ids the map + cannot resolve. Flipping the mapped entry's flag to ``False`` must win even + though the model name still matches the ``haiku-4-5`` pattern.""" + import litellm + from litellm.llms.anthropic.common_utils import AnthropicModelInfo + + model = "us.anthropic.claude-haiku-4-5-20251001-v1:0" + cfg = AmazonAnthropicClaudeMessagesConfig() + + assert AnthropicModelInfo._get_provider_resolved_capability(model, "supports_tool_search", "bedrock") is True + assert cfg._supports_tool_search_on_bedrock(model) is True + + monkeypatch.setitem(litellm.model_cost[model], "supports_tool_search", False) + litellm.get_model_info.cache_clear() + + assert cfg._supports_tool_search_on_bedrock(model) is False + + +@pytest.mark.parametrize( + "model, expected", + [ + pytest.param("us.anthropic.claude-opus-4-6-v99:9", True, id="unmapped_id_falls_back_to_patterns"), + pytest.param("anthropic.claude-3-5-sonnet-20240620-v1:0", False, id="mapped_entry_without_flag_no_pattern"), + ], +) +def test_bedrock_messages_tool_search_pattern_fallback(local_model_cost_map, model, expected): + """Ids the model map cannot resolve (or resolves without a + ``supports_tool_search`` opinion) fall through to the name patterns, so + ARNs and unlisted regional variants of supported families keep working.""" + cfg = AmazonAnthropicClaudeMessagesConfig() + + assert cfg._supports_tool_search_on_bedrock(model) is expected + + def test_bedrock_messages_thinking_shape_follows_exact_bedrock_entry_flag( local_model_cost_map, monkeypatch ): From e5ac4e0068e37fc2a1f2c897c892eae581edec8e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 11 Aug 2026 06:47:39 -0700 Subject: [PATCH 06/24] chore(typing): clear 1.6k basedpyright Any errors across 56 files reportAny 16720 -> 15482 and reportExplicitAny 5689 -> 5316 with real types only: no casts, no ignores, no new Any. Whole-tree basedpyright drops 2173 diagnostics with zero per-rule or per-file regressions. Budgets ratcheted: basedpyright -2173, ruff-strict -188, type-discipline -55 --- basedpyright-code-budget.json | 30 +- .../handler.py | 28 +- .../transformation.py | 101 ++++--- litellm/images/main.py | 7 +- .../integrations/braintrust_mock_client.py | 3 +- .../gcs_bucket/gcs_bucket_mock_client.py | 3 +- .../generic_api/generic_api_callback.py | 19 +- litellm/integrations/mock_client_factory.py | 5 +- litellm/integrations/opentelemetry.py | 64 ++-- .../websearch_interception/handler.py | 13 +- litellm/litellm_core_utils/litellm_logging.py | 57 ++-- .../prompt_templates/factory.py | 50 ++-- .../streaming_chunk_builder_utils.py | 10 +- .../litellm_core_utils/streaming_handler.py | 108 +++++-- litellm/llms/anthropic/chat/handler.py | 5 +- litellm/llms/azure/azure.py | 6 +- litellm/llms/azure/chat/o_series_handler.py | 5 +- litellm/llms/azure/completion/handler.py | 9 +- litellm/llms/azure_ai/anthropic/handler.py | 3 +- litellm/llms/azure_ai/chat/transformation.py | 2 +- litellm/llms/bedrock/chat/converse_handler.py | 2 +- litellm/llms/codestral/completion/handler.py | 4 +- litellm/llms/custom_httpx/aiohttp_handler.py | 2 +- litellm/llms/custom_httpx/http_handler.py | 61 ++-- litellm/llms/custom_httpx/llm_http_handler.py | 64 ++-- .../github_copilot/chat/transformation.py | 2 +- litellm/llms/oci/chat/transformation.py | 23 +- litellm/llms/openai/completion/handler.py | 4 +- litellm/llms/openai/openai.py | 72 +++-- litellm/llms/openai_like/chat/handler.py | 5 +- litellm/llms/predibase/chat/handler.py | 5 +- litellm/llms/replicate/chat/handler.py | 5 +- litellm/llms/sagemaker/completion/handler.py | 8 +- .../vertex_and_google_ai_studio_gemini.py | 41 +-- .../llms/vertex_ai/vertex_ai_non_gemini.py | 2 +- litellm/main.py | 276 +++++++++++------- .../mcp_server/oauth2_token_cache.py | 3 +- litellm/proxy/auth/auth_checks.py | 234 ++++++++++++--- litellm/proxy/common_request_processing.py | 113 +++++-- litellm/proxy/db/spend_log_tool_index.py | 4 +- .../proxy/hooks/parallel_request_limiter.py | 43 +-- .../key_management_endpoints.py | 18 +- .../gemini_passthrough_logging_handler.py | 2 +- .../openai_passthrough_logging_handler.py | 2 +- .../vertex_passthrough_logging_handler.py | 2 +- .../pass_through_endpoints.py | 96 +++--- .../streaming_handler.py | 8 +- .../pass_through_endpoints/success_handler.py | 16 +- litellm/proxy/proxy_server.py | 220 ++++++++++---- .../proxy/response_api_endpoints/endpoints.py | 9 +- litellm/proxy/utils.py | 101 +++---- litellm/responses/main.py | 121 ++++---- .../responses/mcp/chat_completions_handler.py | 6 +- litellm/responses/streaming_iterator.py | 32 +- litellm/router.py | 71 +++-- litellm/types/completion.py | 2 +- litellm/utils.py | 123 +++++--- ruff-strict-budget.json | 20 +- type-discipline-budget.json | 10 +- 59 files changed, 1513 insertions(+), 847 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 96b689aed74..14fdd3ed0bc 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,15 +1,15 @@ { "reportAny": { - "limit": 26391 + "limit": 25153 }, "reportArgumentType": { - "limit": 2614 + "limit": 2596 }, "reportAssignmentType": { "limit": 327 }, "reportAttributeAccessIssue": { - "limit": 514 + "limit": 510 }, "reportCallIssue": { "limit": 114 @@ -18,13 +18,13 @@ "limit": 40 }, "reportDeprecated": { - "limit": 215 + "limit": 214 }, "reportDuplicateImport": { "limit": 19 }, "reportExplicitAny": { - "limit": 8319 + "limit": 7946 }, "reportFunctionMemberAccess": { "limit": 7 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5825 + "limit": 5770 }, "reportMissingTypeArgument": { - "limit": 15695 + "limit": 15676 }, "reportMissingTypeStubs": { "limit": 40 @@ -72,7 +72,7 @@ "limit": 0 }, "reportOptionalMemberAccess": { - "limit": 1077 + "limit": 1073 }, "reportOptionalOperand": { "limit": 0 @@ -99,31 +99,31 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44996 + "limit": 44911 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 39643 + "limit": 39464 }, "reportUnknownParameterType": { - "limit": 20132 + "limit": 20058 }, "reportUnknownVariableType": { - "limit": 31153 + "limit": 31038 }, "reportUnnecessaryCast": { "limit": 118 }, "reportUnnecessaryComparison": { - "limit": 701 + "limit": 700 }, "reportUnnecessaryContains": { "limit": 5 }, "reportUnnecessaryIsInstance": { - "limit": 857 + "limit": 855 }, "reportUntypedBaseClass": { "limit": 0 @@ -138,7 +138,7 @@ "limit": 139 }, "reportUnusedImport": { - "limit": 555 + "limit": 550 }, "reportUnusedVariable": { "limit": 146 diff --git a/litellm/completion_extras/litellm_responses_transformation/handler.py b/litellm/completion_extras/litellm_responses_transformation/handler.py index f290bc631b4..33206629b41 100644 --- a/litellm/completion_extras/litellm_responses_transformation/handler.py +++ b/litellm/completion_extras/litellm_responses_transformation/handler.py @@ -22,6 +22,7 @@ class ResponsesToCompletionBridgeHandlerInputKwargs(TypedDict): model_response: "ModelResponse" logging_obj: "LiteLLMLoggingObj" custom_llm_provider: str + encoding: object class ResponsesToCompletionBridgeHandler: @@ -102,35 +103,37 @@ class ResponsesToCompletionBridgeHandler: from litellm import LiteLLMLoggingObj from litellm.types.utils import ModelResponse - model: Final = kwargs.get("model") + typed_kwargs: Final[dict[str, object]] = kwargs + + model: Final = typed_kwargs.get("model") if model is None or not isinstance(model, str): raise ValueError("model is required") - custom_llm_provider: Final = kwargs.get("custom_llm_provider") + custom_llm_provider: Final = typed_kwargs.get("custom_llm_provider") if custom_llm_provider is None or not isinstance(custom_llm_provider, str): raise ValueError("custom_llm_provider is required") - messages: Final = kwargs.get("messages") + messages: Final = typed_kwargs.get("messages") if messages is None or not isinstance(messages, list): raise ValueError("messages is required") - optional_params: Final = kwargs.get("optional_params") + optional_params: Final = typed_kwargs.get("optional_params") if optional_params is None or not isinstance(optional_params, dict): raise ValueError("optional_params is required") - litellm_params: Final = kwargs.get("litellm_params") + litellm_params: Final = typed_kwargs.get("litellm_params") if litellm_params is None or not isinstance(litellm_params, dict): raise ValueError("litellm_params is required") - headers: Final = kwargs.get("headers") + headers: Final = typed_kwargs.get("headers") if headers is None or not isinstance(headers, dict): raise ValueError("headers is required") - model_response: Final = kwargs.get("model_response") + model_response: Final = typed_kwargs.get("model_response") if model_response is None or not isinstance(model_response, ModelResponse): raise ValueError("model_response is required") - logging_obj: Final = kwargs.get("logging_obj") + logging_obj: Final = typed_kwargs.get("logging_obj") if logging_obj is None or not isinstance(logging_obj, LiteLLMLoggingObj): raise ValueError("logging_obj is required") @@ -143,6 +146,7 @@ class ResponsesToCompletionBridgeHandler: model_response=model_response, logging_obj=logging_obj, custom_llm_provider=custom_llm_provider, + encoding=typed_kwargs.get("encoding"), ) def completion( @@ -205,7 +209,7 @@ class ResponsesToCompletionBridgeHandler: messages=messages, optional_params=optional_params, litellm_params=litellm_params, - encoding=kwargs.get("encoding"), + encoding=validated_kwargs["encoding"], api_key=kwargs.get("api_key"), json_mode=kwargs.get("json_mode"), ) @@ -230,7 +234,7 @@ class ResponsesToCompletionBridgeHandler: messages=messages, optional_params=optional_params, litellm_params=litellm_params, - encoding=kwargs.get("encoding"), + encoding=validated_kwargs["encoding"], api_key=kwargs.get("api_key"), json_mode=kwargs.get("json_mode"), ) @@ -303,7 +307,7 @@ class ResponsesToCompletionBridgeHandler: messages=messages, optional_params=optional_params, litellm_params=litellm_params, - encoding=kwargs.get("encoding"), + encoding=validated_kwargs["encoding"], api_key=kwargs.get("api_key"), json_mode=kwargs.get("json_mode"), ) @@ -328,7 +332,7 @@ class ResponsesToCompletionBridgeHandler: messages=messages, optional_params=optional_params, litellm_params=litellm_params, - encoding=kwargs.get("encoding"), + encoding=validated_kwargs["encoding"], api_key=kwargs.get("api_key"), json_mode=kwargs.get("json_mode"), ) diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index f31e228e456..579cf83bffa 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -4,8 +4,8 @@ Handler for transforming /chat/completions api requests to litellm.responses req import json import os -from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping -from typing import TYPE_CHECKING, Any, Final, Literal, Union, cast +from collections.abc import AsyncIterator, Callable, Iterable, Iterator, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict, Union, cast from openai.types.responses.custom_tool_param import CustomToolParam from openai.types.responses.response_input_param import ( @@ -45,6 +45,9 @@ from litellm.types.utils import GenericStreamingChunk, ModelResponseStream if TYPE_CHECKING: from openai.types.responses import ResponseInputImageParam + from openai.types.responses.response_text_config_param import ( + ResponseTextConfigParam as ResponseText, + ) from pydantic import BaseModel from litellm import LiteLLMLoggingObj, ModelResponse @@ -57,6 +60,19 @@ if TYPE_CHECKING: ChatCompletionThinkingBlock, OpenAIMessageContentListBlock, ) + from litellm.types.utils import Choices + + +class _ReasoningSummaryText(TypedDict): + type: str + text: str + + +class _BuiltReasoningItem(TypedDict): + type: Literal["reasoning"] + id: str + encrypted_content: str | None + summary: Sequence[_ReasoningSummaryText] def _get_reasoning_items( @@ -72,13 +88,13 @@ def _get_reasoning_items( def _build_reasoning_item( item_id: str, encrypted_content: str | None, - summary_raw: Any, -) -> dict[str, Any]: + summary_raw: Iterable[object] | None, +) -> _BuiltReasoningItem: """Build a ChatCompletionReasoningItem-shaped dict from raw response data. Handles both pydantic objects (attribute access) and plain dicts. """ - summary: Final[list[dict[str, Any]]] = [] + summary: Final[list[_ReasoningSummaryText]] = [] for s in summary_raw or []: if isinstance(s, dict): summary.append({"type": s.get("type", "summary_text"), "text": s.get("text", "")}) @@ -98,7 +114,7 @@ def _build_reasoning_item( class _ChatToolCallDict(ChatCompletionToolCallChunk, total=False): - provider_specific_fields: Mapping[str, Any] + provider_specific_fields: Mapping[str, object] def _tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _ChatToolCallDict: @@ -142,10 +158,10 @@ def _flat_responses_tool_choice(choice_type: str, name: str) -> ToolChoiceFuncti def _reasoning_item_to_response_input( - r_item: ChatCompletionReasoningItem | dict[str, Any], -) -> dict[str, Any]: + r_item: ChatCompletionReasoningItem, +) -> dict[str, object]: """Convert a stored ChatCompletionReasoningItem back to a Responses API input item.""" - r_input: Final[dict[str, Any]] = { + r_input: Final[dict[str, object]] = { "type": "reasoning", "id": r_item.get("id") or f"rs_{id(r_item)}", # summary is always required by the Responses API, even when empty @@ -181,7 +197,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): return _flat_responses_tool_choice(choice_type, nested_name) return tool_choice - def _handle_raw_dict_response_item(self, item: dict[str, Any], index: int) -> tuple[Any | None, int]: + def _handle_raw_dict_response_item(self, item: dict[str, Any], index: int) -> tuple["Choices | None", int]: """ Handle raw dict response items from Responses API (e.g., GPT-5 Codex format). @@ -228,8 +244,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): def convert_chat_completion_messages_to_responses_api( self, messages: list["AllMessageValues"] - ) -> tuple[list[Any], str | None]: - input_items: Final[list[Any]] = [] + ) -> tuple[list[object], str | None]: + input_items: Final[list[object]] = [] instructions: str | None = None custom_tool_call_ids: Final = frozenset( tool_call["id"] @@ -270,7 +286,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): # Convert tool message to function call output format # The Responses API expects 'output' to be a list with input_text/input_image types # Using list format for consistency across text and multimodal content - tool_output: list[dict[str, Any]] + tool_output: list[dict[str, object]] if content is None: tool_output = [] elif isinstance(content, str): @@ -308,7 +324,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): function = tool_call.get("function") custom = tool_call.get("custom") if function: - input_tool_call: dict[str, Any] = { + input_tool_call: dict[str, object] = { "type": "function_call", "call_id": tool_call["id"], } @@ -376,15 +392,15 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): elif key == "web_search_options": self._add_web_search_tool(responses_api_request, value) - def _build_sanitized_litellm_params(self, litellm_params: dict) -> dict[str, Any]: + def _build_sanitized_litellm_params(self, litellm_params: dict) -> dict[str, object]: """Build sanitized litellm_params with merged metadata.""" responses_optional_param_keys: Final = set(ResponsesAPIOptionalRequestParams.__annotations__.keys()) - sanitized: Final[dict[str, Any]] = { + sanitized: Final[dict[str, object]] = { key: value for key, value in litellm_params.items() if key not in responses_optional_param_keys } legacy_metadata: Final = litellm_params.get("metadata") existing_litellm_metadata: Final = litellm_params.get("litellm_metadata") - merged_litellm_metadata: Final[dict[str, Any]] = {} + merged_litellm_metadata: Final[dict[str, object]] = {} if isinstance(legacy_metadata, dict): merged_litellm_metadata.update(legacy_metadata) if isinstance(existing_litellm_metadata, dict): @@ -424,7 +440,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): litellm_params: dict, headers: dict, litellm_logging_obj: "LiteLLMLoggingObj", - client: Any | None = None, + client: object | None = None, ) -> dict: ( input_items, @@ -498,9 +514,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): @staticmethod def _convert_response_output_to_choices( - output_items: list[Any], - handle_raw_dict_callback: Callable | None = None, - ) -> list[Any]: + output_items: Sequence[object], + handle_raw_dict_callback: Callable[..., tuple["Choices | None", int]] | None = None, + ) -> list["Choices"]: """ Convert Responses API output items to chat completion choices. @@ -529,11 +545,11 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): choices: Final[list[Choices]] = [] index = 0 reasoning_content: str | None = None - pending_reasoning_item: dict[str, Any] | None = None + pending_reasoning_item: _BuiltReasoningItem | None = None # Collect all tool calls to put them in a single choice # (Chat Completions API expects all tool calls in one message) - accumulated_tool_calls: Final[list[dict[str, Any]]] = [] + accumulated_tool_calls: Final[list[Mapping[str, object]]] = [] tool_call_index = 0 for item in output_items: @@ -640,7 +656,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): return choices @classmethod - def _extract_output_from_completed_event(cls, parsed_chunk: dict[str, Any]) -> list[dict[str, Any]] | None: + def _extract_output_from_completed_event(cls, parsed_chunk: Mapping[str, object]) -> list[dict[str, object]] | None: response_payload: Final = parsed_chunk.get("response") if not isinstance(response_payload, dict): return None @@ -650,12 +666,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): return cast(list[dict[str, Any]], response_output) @classmethod - def _recover_output_items_from_raw_sse(cls, raw_sse: str | None) -> list[dict[str, Any]]: + def _recover_output_items_from_raw_sse(cls, raw_sse: str | None) -> list[dict[str, object]]: if not raw_sse or not isinstance(raw_sse, str): return [] - recovered_output_items: Final[dict[int, dict[str, Any]]] = {} - recovered_text_only_items: Final[dict[int, dict[str, Any]]] = {} + recovered_output_items: Final[dict[int, dict[str, object]]] = {} + recovered_text_only_items: Final[dict[int, dict[str, object]]] = {} for chunk in raw_sse.splitlines(): parsed_chunk = parse_sse_json_chunk(chunk) @@ -690,7 +706,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): # but text-only items at indices without a matching OUTPUT_ITEM_DONE # must still be preserved (e.g. multi-output responses where some # indices only emitted OUTPUT_TEXT_DONE). - merged_items: Final[dict[int, dict[str, Any]]] = {**recovered_text_only_items} + merged_items: Final[dict[int, dict[str, object]]] = {**recovered_text_only_items} merged_items.update(recovered_output_items) if merged_items: @@ -699,7 +715,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): return [] @classmethod - def _recover_output_items_from_logging(cls, logging_obj: "LiteLLMLoggingObj") -> list[dict[str, Any]]: + def _recover_output_items_from_logging(cls, logging_obj: "LiteLLMLoggingObj") -> list[dict[str, object]]: model_call_details: Final = getattr(logging_obj, "model_call_details", {}) or {} original_response: Final = model_call_details.get("original_response") return cls._recover_output_items_from_raw_sse(original_response) @@ -714,7 +730,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): messages: list["AllMessageValues"], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: object, api_key: str | None = None, json_mode: bool | None = None, ) -> "ModelResponse": @@ -788,7 +804,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ) -> BaseModelResponseIterator: return OpenAiResponsesToChatCompletionStreamIterator(streaming_response, sync_stream, json_mode) - def _convert_content_str_to_input_text(self, content: str, role: str) -> dict[str, Any]: + def _convert_content_str_to_input_text(self, content: str, role: str) -> dict[str, object]: if role == "user" or role == "system" or role == "tool": return {"type": "input_text", "text": content} else: @@ -825,13 +841,13 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): def _convert_content_to_responses_format( self, content: str - | list[Any] + | list[object] | Iterable[ Union["OpenAIMessageContentListBlock", "ChatCompletionThinkingBlock", "ChatCompletionRedactedThinkingBlock"] ] | None, role: str, - ) -> list[dict[str, Any]]: + ) -> list[dict[str, object]]: """Convert chat completion content to responses API format""" from litellm.types.llms.openai import ChatCompletionImageObject @@ -973,7 +989,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): return optional_params - def _map_reasoning_effort(self, reasoning_effort: str | dict[str, Any]) -> Reasoning | None: + def _map_reasoning_effort(self, reasoning_effort: str | Reasoning) -> Reasoning | None: # If dict is passed, convert it directly to Reasoning object if isinstance(reasoning_effort, dict): return Reasoning(**reasoning_effort) @@ -1006,7 +1022,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): def _add_web_search_tool( self, responses_api_request: ResponsesAPIOptionalRequestParams, - web_search_options: Any, + web_search_options: object, ) -> None: """ Add web search tool to responses API request. @@ -1024,14 +1040,14 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): tools = [] responses_api_request["tools"] = tools - web_search_tool: Final[dict[str, Any]] = {"type": "web_search"} + web_search_tool: Final[dict[str, object]] = {"type": "web_search"} if isinstance(web_search_options, dict): web_search_tool.update(web_search_options) # Cast to Any to match the expected union type for tools list items tools.append(cast(Any, web_search_tool)) - def _transform_response_format_to_text_format(self, response_format: dict[str, Any] | Any) -> dict[str, Any] | None: + def _transform_response_format_to_text_format(self, response_format: object) -> "ResponseText | None": """ Transform Chat Completion response_format parameter to Responses API text.format parameter. @@ -1130,7 +1146,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): - def __init__(self, streaming_response, sync_stream: bool, json_mode: bool | None = False): + def __init__( + self, + streaming_response: Union[Iterator[str], AsyncIterator[str], "ModelResponse", "BaseModel"], + sync_stream: bool, + json_mode: bool | None = False, + ): super().__init__(streaming_response, sync_stream, json_mode) self._chat_completion_id: str | None = None self._tool_call_index_map: dict[int, int] = {} # mutable-ok: per-stream accumulator state @@ -1387,7 +1408,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): finish_reason: Final = "tool_calls" if has_function_calls else "stop" # Extract reasoning items with encrypted_content for round-tripping - completed_reasoning_items: list[dict[str, Any]] | None = None + completed_reasoning_items: list[_BuiltReasoningItem] | None = None for item in output_items: if not isinstance(item, dict) or item.get("type") != "reasoning": continue @@ -1439,7 +1460,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): ] ) - def chunk_parser(self, chunk: dict) -> "ModelResponseStream": + def chunk_parser(self, chunk: dict[str, object]) -> "ModelResponseStream": """ Parse a Responses API streaming chunk and convert to OpenAI format. diff --git a/litellm/images/main.py b/litellm/images/main.py index f04e0e21ecd..ae4818b1967 100644 --- a/litellm/images/main.py +++ b/litellm/images/main.py @@ -315,7 +315,12 @@ def image_generation( or get_secret_str("AZURE_API_KEY") ) - azure_ad_token: Final = optional_params.pop("azure_ad_token", None) or get_secret_str("AZURE_AD_TOKEN") + azure_ad_token_param: Final = optional_params.pop("azure_ad_token", None) + azure_ad_token: Final = ( + azure_ad_token_param + if isinstance(azure_ad_token_param, str) and azure_ad_token_param + else get_secret_str("AZURE_AD_TOKEN") + ) # Create azure_ad_token_provider from tenant_id, client_id, client_secret if not already provided if azure_ad_token_provider is None: diff --git a/litellm/integrations/braintrust_mock_client.py b/litellm/integrations/braintrust_mock_client.py index 795bcff5b56..3840eabdd20 100644 --- a/litellm/integrations/braintrust_mock_client.py +++ b/litellm/integrations/braintrust_mock_client.py @@ -10,6 +10,7 @@ Usage: import os import time +from collections.abc import AsyncIterable, Iterable from typing import Final from urllib.parse import urlparse @@ -84,7 +85,7 @@ def _mock_http_handler_post( timeout=None, stream=False, files=None, - content=None, + content: str | bytes | Iterable[bytes] | AsyncIterable[bytes] | None = None, logging_obj=None, ): """Monkey-patched HTTPHandler.post that intercepts Braintrust calls with endpoint-specific responses.""" diff --git a/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py b/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py index 24bdd535576..9dfd75e5559 100644 --- a/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py +++ b/litellm/integrations/gcs_bucket/gcs_bucket_mock_client.py @@ -9,6 +9,7 @@ Usage: """ import asyncio +from collections.abc import AsyncIterable, Iterable from typing import Final from litellm._logging import verbose_logger @@ -113,7 +114,7 @@ async def _mock_async_handler_delete( headers=None, timeout=None, stream=False, - content=None, + content: str | bytes | Iterable[bytes] | AsyncIterable[bytes] | None = None, ): """Monkey-patched AsyncHTTPHandler.delete that intercepts GCS calls.""" # Only mock GCS API calls diff --git a/litellm/integrations/generic_api/generic_api_callback.py b/litellm/integrations/generic_api/generic_api_callback.py index 268fa7f4374..dfedc3a3cc9 100644 --- a/litellm/integrations/generic_api/generic_api_callback.py +++ b/litellm/integrations/generic_api/generic_api_callback.py @@ -11,7 +11,7 @@ import json import os import re import traceback -from typing import Any, Final, Literal +from typing import Final, Literal import httpx @@ -158,7 +158,7 @@ class GenericAPILogger(CustomBatchLogger): "endpoint not set for GenericAPILogger, GENERIC_LOGGER_ENDPOINT not found in environment variables" ) - self.headers: dict = self._get_headers(headers) + self.headers: dict[str, str] = self._get_headers(headers) self.endpoint: str = endpoint self.event_types: list[API_EVENT_TYPES] | None = event_types self.callback_name: str | None = callback_name @@ -248,18 +248,15 @@ class GenericAPILogger(CustomBatchLogger): await asyncio.sleep(delay) async def _post_with_retries(self, data: str) -> httpx.Response: - post_kwargs: Final[dict[str, Any]] = { - "url": self.endpoint, - "headers": self.headers, - "data": data, - } - if self.timeout is not None: - post_kwargs["timeout"] = self.timeout - total_attempts: Final = self.max_retries + 1 for attempt in range(total_attempts): try: - return await self.async_httpx_client.post(**post_kwargs) + return await self.async_httpx_client.post( + url=self.endpoint, + headers=self.headers, + data=data, + timeout=self.timeout, + ) except Exception as e: is_last_attempt = attempt == self.max_retries should_retry = self._should_retry_exception(e) diff --git a/litellm/integrations/mock_client_factory.py b/litellm/integrations/mock_client_factory.py index 9377bc18475..59f0279dc7c 100644 --- a/litellm/integrations/mock_client_factory.py +++ b/litellm/integrations/mock_client_factory.py @@ -8,6 +8,7 @@ making actual network calls. import asyncio import json +from collections.abc import AsyncIterable, Iterable from dataclasses import dataclass from datetime import timedelta from typing import Final, cast @@ -140,7 +141,7 @@ def create_mock_client_factory(config: MockClientConfig): stream=False, logging_obj=None, files=None, - content=None, + content: str | bytes | Iterable[bytes] | AsyncIterable[bytes] | None = None, ): """Monkey-patched AsyncHTTPHandler.post that intercepts API calls.""" if isinstance(url, str) and _is_mock_url(url): @@ -193,7 +194,7 @@ def create_mock_client_factory(config: MockClientConfig): timeout=None, stream=False, files=None, - content=None, + content: str | bytes | Iterable[bytes] | AsyncIterable[bytes] | None = None, logging_obj=None, ): """Monkey-patched HTTPHandler.post that intercepts API calls.""" diff --git a/litellm/integrations/opentelemetry.py b/litellm/integrations/opentelemetry.py index 39dbf8ed487..c3461c849dc 100644 --- a/litellm/integrations/opentelemetry.py +++ b/litellm/integrations/opentelemetry.py @@ -1,7 +1,8 @@ import os +from collections.abc import Mapping from dataclasses import dataclass, field from datetime import datetime -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, TypedDict, cast import litellm from litellm._logging import verbose_logger @@ -37,9 +38,11 @@ from litellm.types.utils import ( # OpenTelemetry imports moved to individual functions to avoid import errors when not installed if TYPE_CHECKING: + from opentelemetry.sdk.trace import TracerProvider as _SDKTracerProvider from opentelemetry.sdk.trace.export import SpanExporter as _SpanExporter from opentelemetry.trace import Context as _Context from opentelemetry.trace import Span as _Span + from opentelemetry.trace import SpanKind as _SpanKind from opentelemetry.trace import Tracer as _Tracer from litellm.proxy._types import ( @@ -61,6 +64,25 @@ else: ManagementEndpointLoggingPayload = Any Context = Any + +class _StartSpanRequiredKwargs(TypedDict): + name: str + start_time: int + context: "Context | None" + + +class _StartSpanKwargs(_StartSpanRequiredKwargs, total=False): + kind: "_SpanKind" + + +class _UsageCompletionTokensView(TypedDict, total=False): + completion_tokens: int + + +class _ResponseWithUsageView(TypedDict, total=False): + usage: "_UsageCompletionTokensView | None" + + LITELLM_TRACER_NAME: Final = os.getenv("OTEL_TRACER_NAME", "litellm") LITELLM_METER_NAME: Final = os.getenv("LITELLM_METER_NAME", "litellm") LITELLM_LOGGER_NAME: Final = os.getenv("LITELLM_LOGGER_NAME", "litellm") @@ -297,9 +319,9 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): config: OpenTelemetryConfig | None = None, callback_name: str | None = None, # injection points for testing - tracer_provider: Any | None = None, - logger_provider: Any | None = None, - meter_provider: Any | None = None, + tracer_provider: object | None = None, + logger_provider: object | None = None, + meter_provider: object | None = None, **kwargs, ): team_metadata_keys_override: Final = kwargs.pop("baggage_team_metadata_keys", None) @@ -325,7 +347,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): self.OTEL_EXPORTER = self.config.exporter self.OTEL_ENDPOINT = self.config.endpoint self.OTEL_HEADERS = self.config.headers - self._tracer_provider_cache: dict[str, Any] = {} + self._tracer_provider_cache: dict[str, _SDKTracerProvider] = {} self._init_tracing(tracer_provider) _debug_otel: Final = str(os.getenv("DEBUG_OTEL", "False")).lower() @@ -870,7 +892,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): def _emit_guardrail_spans_from_request_data( self, request_data: dict, - parent_span: Any | None, + parent_span: "Span | None", ) -> None: """Emit ``guardrail`` spans from the request's proxy-internal metadata bucket (``standard_logging_guardrail_information``). @@ -896,7 +918,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): # kwargs["litellm_params"]["metadata"]["_otel_internal"]. Pass the # SAME metadata dict the proxy populated so _handle_failure and # this hook see the same dedupe markers. - kwargs: Final[dict[str, Any]] = { + kwargs: Final[dict[str, object]] = { "litellm_params": {"metadata": metadata}, "standard_logging_object": { "guardrail_information": guardrail_information, @@ -1257,13 +1279,13 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): response_obj, start_time, end_time, - context, + context: "Context | None", ): from opentelemetry.trace import Status, StatusCode otel_tracer: Final[Tracer] = self.get_tracer_to_use_for_request(kwargs) - span_kwargs: Final[dict[str, Any]] = { + span_kwargs: Final[_StartSpanKwargs] = { "name": self._get_span_name(kwargs), "start_time": self._to_ns(start_time), "context": context, @@ -1454,7 +1476,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): ) = _resolve_metric_attribute_filter(attributes) self._metric_attr_filter_resolved = True - def _filter_metric_attributes(self, attrs: dict[str, Any]) -> dict[str, Any]: + def _filter_metric_attributes(self, attrs: dict[str, str]) -> dict[str, str]: if not self._metric_attr_filter_resolved: self._ensure_metric_attribute_filter() if self._metric_attr_include is not None: @@ -1559,7 +1581,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): def _record_time_per_output_token_metric( self, kwargs: dict, - response_obj: Any | None, + response_obj: "_ResponseWithUsageView | None", end_time: datetime, duration_s: float, common_attrs: dict, @@ -1775,10 +1797,10 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): @staticmethod def _resolve_guardrail_context( - span: Any | None, - parent_span: Any | None, - fallback_ctx: Any | None, - ) -> Any | None: + span: "Span | None", + parent_span: "Span | None", + fallback_ctx: "Context | None", + ) -> "Context | None": """ Return a valid OTEL context for guardrail child spans so they are never orphaned (Issue #5). Priority: @@ -1945,7 +1967,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): if should_create_primary_span: # Span 1: Request sent to litellm SDK otel_tracer: Final[Tracer] = self.get_tracer_to_use_for_request(kwargs) - span_kwargs: Final[dict[str, Any]] = { + span_kwargs: Final[_StartSpanKwargs] = { "name": self._get_span_name(kwargs), "start_time": self._to_ns(start_time), "context": _parent_context, @@ -2131,10 +2153,10 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): @staticmethod def _tool_calls_kv_pair( tool_calls: list[ChatCompletionMessageToolCall], - ) -> dict[str, Any]: + ) -> dict[str, object]: from litellm.proxy._types import SpanAttributes - kv_pairs: Final[dict[str, Any]] = {} + kv_pairs: Final[dict[str, object]] = {} for idx, tool_call in enumerate(tool_calls): _function = tool_call.get("function") if not _function: @@ -2691,8 +2713,8 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): import json try: - _raw_response = json.loads(_raw_response) - for param, val in _raw_response.items(): + _parsed: Final[Mapping[str, object]] = json.loads(_raw_response) + for param, val in _parsed.items(): self.safe_set_attribute( span=span, key=f"llm.{custom_llm_provider}.{param}", @@ -2722,7 +2744,7 @@ class OpenTelemetry(OTELGenAISemconvMixin, CustomLogger): return int(dt * 1e9) return int(dt.timestamp() * 1e9) - def _get_span_name(self, kwargs): + def _get_span_name(self, kwargs) -> str: litellm_params: Final = kwargs.get("litellm_params", {}) metadata: Final = litellm_params.get("metadata") or {} generation_name: Final = metadata.get("generation_name") diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index f7f27459768..972ae1d9856 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -85,6 +85,11 @@ class _WebSearchSettingsView(TypedDict): websearch_interception_params: WebSearchInterceptionConfig +class _SearchToolConfig(TypedDict, total=False): + search_tool_name: str + litellm_params: Mapping[str, object] | None + + class WebSearchInterceptionLogger(CustomLogger): """ CustomLogger that intercepts WebSearch tool calls for models that don't @@ -1487,7 +1492,7 @@ class WebSearchInterceptionLogger(CustomLogger): return None - def _select_search_tool_from_router(self, llm_router: object) -> dict[str, Any] | None: + def _select_search_tool_from_router(self, llm_router: object) -> "_SearchToolConfig | None": if llm_router is None or not hasattr(llm_router, "search_tools"): return None search_tools: Final = list(getattr(llm_router, "search_tools") or []) @@ -1495,9 +1500,9 @@ class WebSearchInterceptionLogger(CustomLogger): def _select_search_tool_from_list( self, - search_tools: list[dict[str, Any]], + search_tools: list[_SearchToolConfig], source: str, - ) -> dict[str, Any] | None: + ) -> "_SearchToolConfig | None": if self.search_tool_name: matching_tools = [tool for tool in search_tools if tool.get("search_tool_name") == self.search_tool_name] if matching_tools: @@ -1692,7 +1697,7 @@ class WebSearchInterceptionLogger(CustomLogger): @staticmethod def initialize_from_proxy_config( - litellm_settings: dict[str, Any], + litellm_settings: Mapping[str, WebSearchInterceptionConfig], callback_specific_params: Mapping[str, object], ) -> "WebSearchInterceptionLogger": """ diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index a3ff048e92a..05d278094ea 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -10,7 +10,7 @@ import subprocess import sys import time import traceback -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence from datetime import datetime as dt_object from functools import lru_cache from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast @@ -176,6 +176,9 @@ from .initialize_dynamic_callback_params import ( from .specialty_caches.dynamic_logging_cache import DynamicLoggingCache if TYPE_CHECKING: + from mcp.types import EmbeddedResource, ImageContent, TextContent + + from litellm.integrations.otel.logger import OpenTelemetryV2 from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig try: from litellm_enterprise.enterprise_callbacks.callback_controls import ( @@ -211,14 +214,30 @@ except Exception as e: PagerDutyAlerting = CustomLogger EnterpriseCallbackControls = None EnterpriseStandardLoggingPayloadSetupVAR = None -_in_memory_loggers: Final[list[Any]] = [] +if TYPE_CHECKING: + from litellm.integrations.generic_api.generic_api_callback import ( + GenericAPILogger as _GenericAPILoggerCls, + ) -_STANDARD_LOGGING_METADATA_KEYS: Final[frozenset] = frozenset(StandardLoggingMetadata.__annotations__.keys()) + _GENERIC_API_LOGGER_CLS: Final = _GenericAPILoggerCls + _RESEND_EMAIL_LOGGER_FACTORY: Final = CustomLogger + _SENDGRID_EMAIL_LOGGER_FACTORY: Final = CustomLogger + _SMTP_EMAIL_LOGGER_FACTORY: Final = CustomLogger + _PAGERDUTY_ALERTING_FACTORY: Final = CustomLogger +else: + _GENERIC_API_LOGGER_CLS: Final = GenericAPILogger + _RESEND_EMAIL_LOGGER_FACTORY: Final = ResendEmailLogger + _SENDGRID_EMAIL_LOGGER_FACTORY: Final = SendGridEmailLogger + _SMTP_EMAIL_LOGGER_FACTORY: Final = SMTPEmailLogger + _PAGERDUTY_ALERTING_FACTORY: Final = PagerDutyAlerting +_in_memory_loggers: Final[list[CustomLogger]] = [] + +_STANDARD_LOGGING_METADATA_KEYS: Final[frozenset[str]] = frozenset(StandardLoggingMetadata.__annotations__.keys()) ### GLOBAL VARIABLES ### # Cache custom pricing keys as frozenset for O(1) lookups instead of looping through 49 keys -_CUSTOM_PRICING_KEYS: Final[frozenset] = frozenset(CustomPricingLiteLLMParams.model_fields.keys()) +_CUSTOM_PRICING_KEYS: Final[frozenset[str]] = frozenset(CustomPricingLiteLLMParams.model_fields.keys()) sentry_sdk_instance = None capture_exception = None @@ -1285,7 +1304,9 @@ class Logging(LiteLLMLoggingBaseClass): verbose_logger.exception("LiteLLM.LoggingError: [Non-Blocking] Exception occurred while logging %s", e) return response_obj - def _parse_post_mcp_call_hook_response(self, response: MCPPostCallResponseObject | None) -> Any: + def _parse_post_mcp_call_hook_response( + self, response: MCPPostCallResponseObject | None + ) -> "Sequence[TextContent | ImageContent | EmbeddedResource] | None": """ Parse the response from the post_mcp_tool_call_hook @@ -1729,7 +1750,7 @@ class Logging(LiteLLMLoggingBaseClass): self.completion_start_time = completion_start_time self.model_call_details["completion_start_time"] = self.completion_start_time - def normalize_logging_result(self, result: Any) -> Any: + def normalize_logging_result(self, result: Any) -> object: """ Some endpoints return a different type of result than what is expected by the logging system. This function is used to normalize the result to the expected type. @@ -1765,7 +1786,7 @@ class Logging(LiteLLMLoggingBaseClass): ) return logging_result - def _merge_hidden_params_from_response_into_metadata(self, logging_result: Any) -> None: + def _merge_hidden_params_from_response_into_metadata(self, logging_result: object) -> None: """ Copy response._hidden_params into litellm_params.metadata['hidden_params']. @@ -1826,7 +1847,9 @@ class Logging(LiteLLMLoggingBaseClass): if (standard_logging_payload := self.model_call_details.get("standard_logging_object")) is not None: emit_standard_logging_payload(standard_logging_payload) - def _build_standard_logging_payload(self, init_response_obj: Any, start_time: Any, end_time: Any) -> Any: + def _build_standard_logging_payload( + self, init_response_obj: object, start_time: Any, end_time: Any + ) -> StandardLoggingPayload | None: """Build StandardLoggingPayload and accumulate its construction time.""" _start: Final = time.time() payload: Final = get_standard_logging_object_payload( @@ -1947,7 +1970,7 @@ class Logging(LiteLLMLoggingBaseClass): def _is_recognized_call_type_for_logging( self, - logging_result: Any, + logging_result: object, ): """ Returns True if the call type is recognized for logging (eg. ModelResponse, ModelResponseStream, etc.) @@ -4216,7 +4239,7 @@ def _init_custom_logger_compatible_class( for callback in _in_memory_loggers: if isinstance(callback, PagerDutyAlerting): return callback - pagerduty_logger: Final = PagerDutyAlerting(**custom_logger_init_args) + pagerduty_logger: Final = _PAGERDUTY_ALERTING_FACTORY(**custom_logger_init_args) _in_memory_loggers.append(pagerduty_logger) return pagerduty_logger elif logging_integration == "anthropic_cache_control_hook": @@ -4246,7 +4269,7 @@ def _init_custom_logger_compatible_class( return _gcs_pubsub_logger elif logging_integration == "generic_api": for callback in _in_memory_loggers: - if isinstance(callback, GenericAPILogger): + if isinstance(callback, _GENERIC_API_LOGGER_CLS): return callback generic_api_logger: Final = GenericAPILogger() _in_memory_loggers.append(generic_api_logger) @@ -4255,21 +4278,21 @@ def _init_custom_logger_compatible_class( for callback in _in_memory_loggers: if isinstance(callback, ResendEmailLogger): return callback - resend_email_logger: Final = ResendEmailLogger() + resend_email_logger: Final = _RESEND_EMAIL_LOGGER_FACTORY() _in_memory_loggers.append(resend_email_logger) return resend_email_logger elif logging_integration == "sendgrid_email": for callback in _in_memory_loggers: if isinstance(callback, SendGridEmailLogger): return callback - sendgrid_email_logger: Final = SendGridEmailLogger() + sendgrid_email_logger: Final = _SENDGRID_EMAIL_LOGGER_FACTORY() _in_memory_loggers.append(sendgrid_email_logger) return sendgrid_email_logger elif logging_integration == "smtp_email": for callback in _in_memory_loggers: if isinstance(callback, SMTPEmailLogger): return callback - smtp_email_logger: Final = SMTPEmailLogger() + smtp_email_logger: Final = _SMTP_EMAIL_LOGGER_FACTORY() _in_memory_loggers.append(smtp_email_logger) return smtp_email_logger elif logging_integration == "humanloop": @@ -4336,7 +4359,7 @@ def _init_custom_logger_compatible_class( return None -def _maybe_construct_otel_v2(callback_name: str, _in_memory_loggers: list) -> Any | None: +def _maybe_construct_otel_v2(callback_name: str, _in_memory_loggers: list[CustomLogger]) -> "OpenTelemetryV2 | None": """If ``LITELLM_OTEL_V2`` is on, build (or reuse) a single ``OpenTelemetryV2`` instance configured via the preset for ``callback_name``. @@ -4367,7 +4390,7 @@ def _maybe_construct_otel_v2(callback_name: str, _in_memory_loggers: list) -> An return v2_logger -def _maybe_auto_initialize_arize_phoenix(_in_memory_loggers: list) -> None: +def _maybe_auto_initialize_arize_phoenix(_in_memory_loggers: list[CustomLogger]) -> None: """ Auto-initialize ArizePhoenixLogger when Phoenix env vars are detected. @@ -4594,7 +4617,7 @@ def get_custom_logger_compatible_class( return callback elif logging_integration == "generic_api": for callback in _in_memory_loggers: - if isinstance(callback, GenericAPILogger): + if isinstance(callback, _GENERIC_API_LOGGER_CLS): return callback elif logging_integration == "resend_email": for callback in _in_memory_loggers: diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 3a1a426eaa9..b24e694c6fd 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -549,7 +549,7 @@ def _fetch_and_extract_template( return chat_template, bos_token, eos_token -async def ahf_chat_template(model: str, messages: list, chat_template: Any | None = None): +async def ahf_chat_template(model: str, messages: list, chat_template: str | None = None): """HuggingFace chat template (async version)""" from litellm.litellm_core_utils.prompt_templates.huggingface_template_handler import ( _aget_chat_template_file, @@ -576,7 +576,7 @@ async def ahf_chat_template(model: str, messages: list, chat_template: Any | Non ) -def hf_chat_template(model: str, messages: list, chat_template: Any | None = None): +def hf_chat_template(model: str, messages: list, chat_template: str | None = None): """HuggingFace chat template (sync version)""" from litellm.litellm_core_utils.prompt_templates.huggingface_template_handler import ( _get_chat_template_file, @@ -1130,7 +1130,7 @@ def convert_to_azure_openai_messages( def infer_protocol_value( - value: Any, + value: object, ) -> Literal[ "string_value", "number_value", @@ -1702,7 +1702,9 @@ def convert_function_to_anthropic_tool_invoke( _name: Final = get_attribute_or_key(function_call, "name") or "" _arguments: Final = get_attribute_or_key(function_call, "arguments") - tool_input = parse_tool_call_arguments(_arguments, tool_name=_name, context="Anthropic function to tool invoke") + tool_input: Final = parse_tool_call_arguments( + _arguments, tool_name=_name, context="Anthropic function to tool invoke" + ) anthropic_tool_invoke: Final = [ AnthropicMessagesToolUseParam( @@ -1764,7 +1766,7 @@ def convert_to_anthropic_tool_invoke( Fixes: https://github.com/BerriAI/litellm/issues/17737 """ - anthropic_tool_invoke: Final[list[AnthropicMessagesToolUseParam | dict[str, Any]]] = [] + anthropic_tool_invoke: Final[list[AnthropicMessagesToolUseParam | dict[str, object]]] = [] for tool in tool_calls: if not get_attribute_or_key(tool, "type") == "function": @@ -1785,7 +1787,7 @@ def convert_to_anthropic_tool_invoke( # Server tool IDs start with "srvtoolu_" if tool_id.startswith("srvtoolu_"): # Create server_tool_use block instead of tool_use - _anthropic_server_tool_use: dict[str, Any] = { + _anthropic_server_tool_use: dict[str, object] = { "type": "server_tool_use", "id": tool_id, "name": tool_name, @@ -2177,7 +2179,7 @@ def _is_orphaned_tool_result( return False -def _declared_tool_call_ids(message: Mapping[str, Any]) -> frozenset[str]: +def _declared_tool_call_ids(message: Mapping[str, object]) -> frozenset[str]: tool_calls: Final = message.get("tool_calls") if not isinstance(tool_calls, list): return frozenset() @@ -2186,7 +2188,7 @@ def _declared_tool_call_ids(message: Mapping[str, Any]) -> frozenset[str]: ) -def group_tool_exchanges(messages: Sequence[Mapping[str, Any]]) -> tuple[tuple[int, ...], ...]: +def group_tool_exchanges(messages: Sequence[Mapping[str, object]]) -> tuple[tuple[int, ...], ...]: """Group message indices into tool exchanges: an assistant row that made tool calls, together with the tool rows answering the ids it declared. @@ -2204,7 +2206,7 @@ def group_tool_exchanges(messages: Sequence[Mapping[str, Any]]) -> tuple[tuple[i return tuple(_iter_tool_exchange_groups(messages)) -def _iter_tool_exchange_groups(messages: Sequence[Mapping[str, Any]]) -> Iterator[tuple[int, ...]]: +def _iter_tool_exchange_groups(messages: Sequence[Mapping[str, object]]) -> Iterator[tuple[int, ...]]: index = 0 while index < len(messages): declared = _declared_tool_call_ids(messages[index]) @@ -2409,7 +2411,7 @@ def anthropic_messages_pt( # Convert ChatCompletionImageUrlObject to dict if needed image_url_value = m["image_url"] if isinstance(image_url_value, str): - image_url_input: str | dict[str, Any] = image_url_value + image_url_input: str | dict[str, object] = image_url_value else: # ChatCompletionImageUrlObject or dict case - convert to dict image_url_input = { @@ -3179,7 +3181,7 @@ def _load_image_from_url(image_url): try: # Send a GET request to the image URL client: Final = HTTPHandler(concurrent_limit=1) - response: Final = safe_get(client, image_url) + response: Final[httpx.Response] = safe_get(client, image_url) response.raise_for_status() # Raise an exception for HTTP errors # Check the response's content type to ensure it is an image @@ -3382,7 +3384,7 @@ class BedrockImageProcessor: @staticmethod def _post_call_image_processing(response: httpx.Response, image_url: str = "") -> tuple[str, str]: # Check the response's content type to ensure it is an image - content_type = response.headers.get("content-type") + content_type: str | None = response.headers.get("content-type") # Use helper function to infer content type with fallback logic content_type = infer_content_type_from_url_and_content( @@ -3406,7 +3408,7 @@ class BedrockImageProcessor: params={"concurrent_limit": 1}, ) # Send a GET request to the image URL - response: Final = await async_safe_get(client, image_url) + response: Final[httpx.Response] = await async_safe_get(client, image_url) response.raise_for_status() # Raise an exception for HTTP errors return BedrockImageProcessor._post_call_image_processing(response, image_url) @@ -3419,7 +3421,7 @@ class BedrockImageProcessor: try: client: Final = HTTPHandler(concurrent_limit=1) # Send a GET request to the image URL - response: Final = safe_get(client, image_url) + response: Final[httpx.Response] = safe_get(client, image_url) response.raise_for_status() # Raise an exception for HTTP errors return BedrockImageProcessor._post_call_image_processing(response, image_url) @@ -5328,10 +5330,10 @@ def get_attribute_or_key(tool_or_function, attribute, default=None): class NormalizedToolCall(TypedDict): id: str | None name: str | None - arguments: dict[str, Any] + arguments: dict[str, object] -def _parse_tool_call_arguments(raw: Any, tool_name: str | None, context: str) -> dict[str, Any]: +def _parse_tool_call_arguments(raw: Any, tool_name: str | None, context: str) -> dict[str, object]: # Anthropic's tool_use blocks already carry a parsed dict in "input"; # chat completions and the Responses API carry a JSON string that may be # truncated by the model, so route those through the repair-aware parser. @@ -5352,12 +5354,12 @@ def _parse_tool_call_arguments(raw: Any, tool_name: str | None, context: str) -> def _tool_calls_from_chat_completion_response( - response: Any, include_all_choices: bool = False + response: object, include_all_choices: bool = False ) -> list[NormalizedToolCall]: choices: Final = get_attribute_or_key(response, "choices", None) if not (isinstance(choices, list) and choices): return [] - tool_calls: Final[list[Any]] = [] + tool_calls: Final[list[object]] = [] for choice in choices if include_all_choices else choices[:1]: message = get_attribute_or_key(choice, "message", None) choice_tool_calls = get_attribute_or_key(message, "tool_calls", None) if message else None @@ -5383,7 +5385,7 @@ def _tool_calls_from_chat_completion_response( return result -def _tool_calls_from_responses_api_response(response: Any) -> list[NormalizedToolCall]: +def _tool_calls_from_responses_api_response(response: object) -> list[NormalizedToolCall]: output: Final = get_attribute_or_key(response, "output", None) if not isinstance(output, list): return [] @@ -5406,7 +5408,7 @@ def _tool_calls_from_responses_api_response(response: Any) -> list[NormalizedToo return result -def _tool_calls_from_anthropic_messages_response(response: Any) -> list[NormalizedToolCall]: +def _tool_calls_from_anthropic_messages_response(response: object) -> list[NormalizedToolCall]: content: Final = get_attribute_or_key(response, "content", None) if not isinstance(content, list): return [] @@ -5425,7 +5427,7 @@ def _tool_calls_from_anthropic_messages_response(response: Any) -> list[Normaliz return result -def get_tool_calls_from_response(response: Any, include_all_choices: bool = False) -> list[NormalizedToolCall]: +def get_tool_calls_from_response(response: object, include_all_choices: bool = False) -> list[NormalizedToolCall]: """ Extract tool/function calls from a response object into a normalized ``{"id", "name", "arguments"}`` shape, regardless of which API surface @@ -5456,7 +5458,7 @@ def get_tool_calls_from_response(response: Any, include_all_choices: bool = Fals return [] -def has_tool_with_name(tools: Any, tool_name: str) -> bool: +def has_tool_with_name(tools: object, tool_name: str) -> bool: """ Check whether a tools list (as sent to an LLM) includes a tool with the given name, regardless of shape: OpenAI-style function tools @@ -5482,9 +5484,9 @@ def has_tool_with_name(tools: Any, tool_name: str) -> bool: def resolve_structured_messages( - messages: list[dict[str, Any]] | None, + messages: list[dict[str, object]] | None, request_kwargs: dict[str, Any], -) -> list[dict[str, Any]] | None: +) -> list[dict[str, object]] | None: """ Normalize a request's messages to OpenAI-spec chat-completions shape, regardless of which API surface produced them (chat completions, diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 886ba6a3a18..e1da36ac8cd 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -145,7 +145,7 @@ class ChunkProcessor: if first_hidden_params.get("created_at"): - def _created_at(chunk: Any) -> int | float: + def _created_at(chunk: object) -> int | float: if isinstance(chunk, dict): params = chunk.get("_hidden_params", {}) else: @@ -158,7 +158,7 @@ class ChunkProcessor: return chunks def update_model_response_with_hidden_params( - self, model_response: ModelResponse, chunk: dict[str, Any] | None = None + self, model_response: ModelResponse, chunk: Mapping[str, dict[str, object]] | None = None ) -> ModelResponse: if chunk is None: return model_response @@ -176,7 +176,7 @@ class ChunkProcessor: if not chunks: return - model: Final = getattr(response, "model", None) + model: Final[str | None] = getattr(response, "model", None) if not model: return @@ -214,7 +214,7 @@ class ChunkProcessor: ) @staticmethod - def _get_chunk_id(chunks: list[dict[str, Any]]) -> str: + def _get_chunk_id(chunks: Sequence[Mapping[str, str]]) -> str: """ Chunks: [{"id": ""}, {"id": "1"}, {"id": "1"}] @@ -225,7 +225,7 @@ class ChunkProcessor: return "" @staticmethod - def _get_model_from_chunks(chunks: list[dict[str, Any]], first_chunk_model: str) -> str: + def _get_model_from_chunks(chunks: Sequence[Mapping[str, str]], first_chunk_model: str) -> str: """ Get the actual model from chunks, preferring a model that differs from the first chunk. diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 2dc71abee3e..570c34169d0 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -6,13 +6,14 @@ import logging import threading import time import traceback -from collections.abc import AsyncIterator, Callable, Iterator +from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence from dataclasses import dataclass -from typing import Any, Final, NoReturn, TypeVar, cast +from typing import Any, Final, NoReturn, Protocol, TypeVar, cast import anyio import httpx from pydantic import BaseModel +from typing_extensions import NotRequired, TypedDict import litellm from litellm import verbose_logger @@ -54,7 +55,7 @@ _SYNC_ITER_EXHAUSTED: Final = object() _GCHUNK_FIELDS: Final[frozenset] = frozenset(GChunk.__annotations__) -def _next_sync_or_exhausted(it: Any) -> Any: +def _next_sync_or_exhausted(it: Any) -> object: """ Call next(it) from a thread and return _SYNC_ITER_EXHAUSTED on StopIteration. @@ -68,7 +69,7 @@ def _next_sync_or_exhausted(it: Any) -> Any: return _SYNC_ITER_EXHAUSTED -def is_async_iterable(obj: Any) -> bool: +def is_async_iterable(obj: object) -> bool: """ Check if an object is an async iterable (can be used with 'async for'). @@ -81,7 +82,7 @@ def is_async_iterable(obj: Any) -> bool: return isinstance(obj, collections.abc.AsyncIterable) -def print_verbose(print_statement): +def print_verbose(print_statement: object): try: if litellm.set_verbose: print(print_statement) # noqa: T201 @@ -96,18 +97,78 @@ class _ProviderChunkParsed: @dataclass(frozen=True, slots=True) class _ProviderChunkEarlyReturn: - value: Any + value: "ModelResponseStream | None" _ProviderChunkResult = _ProviderChunkParsed | _ProviderChunkEarlyReturn +class _PredibaseStreamData(TypedDict): + token: NotRequired[Mapping[str, str]] + details: Mapping[str, str] + generated_text: str | None + error: str | None + + +class _Ai21StreamData(TypedDict): + completions: Sequence[Mapping[str, Mapping[str, str]]] + + +class _MaritalkStreamData(TypedDict): + answer: str + + +class _NlpCloudStreamData(TypedDict): + generated_text: str + + +class _AlephAlphaStreamData(TypedDict): + completions: Sequence[Mapping[str, str]] + + +class _AzureStreamChoice(TypedDict): + delta: Mapping[str, str] | None + finish_reason: str | None + + +class _AzureStreamData(TypedDict): + choices: Sequence[_AzureStreamChoice] + + +class _BasetenModelOutput(TypedDict): + data: NotRequired[Sequence[str]] + + +class _BasetenStreamData(TypedDict): + token: NotRequired[Mapping[str, str]] + model_output: NotRequired["_BasetenModelOutput | str"] + completion: NotRequired[object] + + +class _DeltaDumpDict(TypedDict): + role: NotRequired[str | None] + tool_calls: NotRequired[Sequence[Mapping[str, object]]] + + +class _TextCompletionChoiceLike(Protocol): + text: str + finish_reason: str | None + + +def validated_stream_logging_obj(candidate: object) -> LiteLLMLoggingObject: + from litellm.litellm_core_utils.litellm_logging import Logging + + if isinstance(candidate, Logging): + return candidate + raise TypeError("CustomStreamWrapper requires a LiteLLMLoggingObject") + + class CustomStreamWrapper: def __init__( self, completion_stream, model, - logging_obj: Any, + logging_obj: LiteLLMLoggingObject, custom_llm_provider: str | None = None, stream_options=None, make_call: Callable | None = None, @@ -186,7 +247,7 @@ class CustomStreamWrapper: # Snapshot assumes self._hidden_params is populated from litellm_params # at init and never mutated during the stream. If that ever changes, # this cache must be removed. - self._base_hidden_params: dict[str, Any] = { + self._base_hidden_params: dict[str, object] = { **self._hidden_params, "response_cost": None, } @@ -416,7 +477,7 @@ class CustomStreamWrapper: finish_reason = "" print_verbose(f"chunk: {chunk}") if chunk.startswith("data:"): - data_json: Final = json.loads(chunk[5:]) + data_json: Final[_PredibaseStreamData] = json.loads(chunk[5:]) print_verbose(f"data json: {data_json}") if "token" in data_json and "text" in data_json["token"]: text = data_json["token"]["text"] @@ -446,7 +507,7 @@ class CustomStreamWrapper: def handle_ai21_chunk(self, chunk): # fake streaming chunk = chunk.decode("utf-8") - data_json: Final = json.loads(chunk) + data_json: Final[_Ai21StreamData] = json.loads(chunk) try: text: Final = data_json["completions"][0]["data"]["text"] is_finished: Final = True @@ -461,7 +522,7 @@ class CustomStreamWrapper: def handle_maritalk_chunk(self, chunk): # fake streaming chunk = chunk.decode("utf-8") - data_json: Final = json.loads(chunk) + data_json: Final[_MaritalkStreamData] = json.loads(chunk) try: text: Final = data_json["answer"] is_finished: Final = True @@ -482,7 +543,7 @@ class CustomStreamWrapper: if self.model and "dolphin" in self.model: chunk = self.process_chunk(chunk=chunk) else: - data_json: Final = json.loads(chunk) + data_json: Final[_NlpCloudStreamData] = json.loads(chunk) chunk = data_json["generated_text"] text = chunk if "[DONE]" in text: @@ -499,7 +560,7 @@ class CustomStreamWrapper: def handle_aleph_alpha_chunk(self, chunk): chunk = chunk.decode("utf-8") - data_json: Final = json.loads(chunk) + data_json: Final[_AlephAlphaStreamData] = json.loads(chunk) try: text: Final = data_json["completions"][0]["completion"] is_finished: Final = True @@ -527,7 +588,7 @@ class CustomStreamWrapper: "finish_reason": finish_reason, } elif chunk.startswith("data:"): - data_json: Final = json.loads(chunk[5:]) # chunk.startswith("data:"): + data_json: Final[_AzureStreamData] = json.loads(chunk[5:]) # chunk.startswith("data:"): try: if len(data_json["choices"]) > 0: delta: Final = data_json["choices"][0]["delta"] @@ -616,7 +677,7 @@ class CustomStreamWrapper: text = "" is_finished = False finish_reason = None - choices: Final = getattr(chunk, "choices", []) + choices: Final[Sequence[_TextCompletionChoiceLike]] = getattr(chunk, "choices", []) if len(choices) > 0: text = choices[0].text if choices[0].finish_reason is not None: @@ -637,7 +698,7 @@ class CustomStreamWrapper: is_finished = False finish_reason = None usage = None - choices: Final = getattr(chunk, "choices", []) + choices: Final[Sequence[_TextCompletionChoiceLike]] = getattr(chunk, "choices", []) if len(choices) > 0: text = choices[0].text if choices[0].finish_reason is not None: @@ -654,12 +715,12 @@ class CustomStreamWrapper: except Exception as e: raise e - def handle_baseten_chunk(self, chunk): + def handle_baseten_chunk(self, chunk) -> str: try: chunk = chunk.decode("utf-8") if len(chunk) > 0: if chunk.startswith("data:"): - data_json = json.loads(chunk[5:]) + data_json: _BasetenStreamData = json.loads(chunk[5:]) if "token" in data_json and "text" in data_json["token"]: return data_json["token"]["text"] else: @@ -1325,13 +1386,14 @@ class CustomStreamWrapper: if response_obj["is_finished"]: self.received_finish_reason = response_obj["finish_reason"] if "usage" in response_obj is not None: + _codestral_usage: Final[Usage] = response_obj["usage"] setattr( model_response, "usage", litellm.Usage( - prompt_tokens=response_obj["usage"].prompt_tokens, - completion_tokens=response_obj["usage"].completion_tokens, - total_tokens=response_obj["usage"].total_tokens, + prompt_tokens=_codestral_usage.prompt_tokens, + completion_tokens=_codestral_usage.completion_tokens, + total_tokens=_codestral_usage.total_tokens, ), ) elif self.custom_llm_provider == "azure_text": @@ -1474,7 +1536,7 @@ class CustomStreamWrapper: is None ): t.function.arguments = "" - _json_delta: Final = delta.model_dump() + _json_delta: Final[_DeltaDumpDict] = delta.model_dump() if "role" not in _json_delta or _json_delta["role"] is None: _json_delta["role"] = "assistant" # mistral's api returns role as None if "tool_calls" in _json_delta and isinstance(_json_delta["tool_calls"], list): @@ -1744,7 +1806,7 @@ class CustomStreamWrapper: usage.cost, copy it into _hidden_params so litellm's cost calculator uses it instead of a token-based estimate. """ - _usage: Final = getattr(response, "usage", None) + _usage: Final[Usage | None] = getattr(response, "usage", None) if _usage is not None and hasattr(_usage, "cost") and _usage.cost is not None: if "additional_headers" not in response._hidden_params: response._hidden_params["additional_headers"] = {} diff --git a/litellm/llms/anthropic/chat/handler.py b/litellm/llms/anthropic/chat/handler.py index 8c4facc1ba2..39d3947c07c 100644 --- a/litellm/llms/anthropic/chat/handler.py +++ b/litellm/llms/anthropic/chat/handler.py @@ -58,6 +58,7 @@ from ..common_utils import AnthropicError, process_anthropic_headers from .transformation import ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY, AnthropicConfig if TYPE_CHECKING: + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.base_llm.chat.transformation import BaseConfig @@ -206,7 +207,7 @@ class AnthropicChatCompletion(BaseLLM): client: AsyncHTTPHandler | None, encoding, api_key, - logging_obj, + logging_obj: "LiteLLMLoggingObj", stream, _is_function_call, data: dict, @@ -324,7 +325,7 @@ class AnthropicChatCompletion(BaseLLM): print_verbose: Callable, encoding, api_key, - logging_obj, + logging_obj: "LiteLLMLoggingObj", optional_params: dict, timeout: float | httpx.Timeout, litellm_params: dict, diff --git a/litellm/llms/azure/azure.py b/litellm/llms/azure/azure.py index 91cd683d5a9..3438e835faf 100644 --- a/litellm/llms/azure/azure.py +++ b/litellm/llms/azure/azure.py @@ -228,7 +228,7 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): litellm_params=litellm_params, ) - data = {"model": None, "messages": messages, **optional_params} + data: dict[str, object] = {"model": None, "messages": messages, **optional_params} elif litellm.AzureOpenAIGPT5Config.is_model_gpt_5_model(model=litellm_params.get("base_model") or model): data = litellm.AzureOpenAIGPT5Config().transform_request( model=model, @@ -482,12 +482,12 @@ class AzureChatCompletion(BaseAzureLLM, BaseLLM): def streaming( self, - logging_obj, + logging_obj: LiteLLMLoggingObj, api_base: str, api_key: str | None, api_version: str, dynamic_params: bool, - data: dict, + data: dict[str, object], model: str, timeout: Any, max_retries: int, diff --git a/litellm/llms/azure/chat/o_series_handler.py b/litellm/llms/azure/chat/o_series_handler.py index 64b6025f6ea..30de68e40ef 100644 --- a/litellm/llms/azure/chat/o_series_handler.py +++ b/litellm/llms/azure/chat/o_series_handler.py @@ -5,10 +5,11 @@ Written separately to handle faking streaming for o1 and o3 models. """ from collections.abc import Callable -from typing import TYPE_CHECKING, Any, Optional +from typing import TYPE_CHECKING, Optional import httpx +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.types.utils import ModelResponse from ...openai.openai import OpenAIChatCompletion @@ -25,7 +26,7 @@ class AzureOpenAIO1ChatCompletion(BaseAzureLLM, OpenAIChatCompletion): timeout: float | httpx.Timeout, optional_params: dict, litellm_params: dict, - logging_obj: Any, + logging_obj: LiteLLMLoggingObj, model: str | None = None, messages: list | None = None, print_verbose: Callable | None = None, diff --git a/litellm/llms/azure/completion/handler.py b/litellm/llms/azure/completion/handler.py index 79fbd0a5f86..728968e12e7 100644 --- a/litellm/llms/azure/completion/handler.py +++ b/litellm/llms/azure/completion/handler.py @@ -3,6 +3,7 @@ from typing import Any, Final from openai import AsyncAzureOpenAI, AzureOpenAI +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.prompt_templates.factory import prompt_factory from litellm.utils import CustomStreamWrapper, ModelResponse, TextCompletionResponse @@ -39,9 +40,9 @@ class AzureTextCompletion(BaseAzureLLM): azure_ad_token_provider: Callable | None, print_verbose: Callable, timeout, - logging_obj, + logging_obj: LiteLLMLoggingObj, optional_params, - litellm_params, + litellm_params: dict[str, object], logger_fn, acompletion: bool = False, headers: dict | None = None, @@ -246,7 +247,7 @@ class AzureTextCompletion(BaseAzureLLM): def streaming( self, - logging_obj, + logging_obj: LiteLLMLoggingObj, api_base: str, api_key: str | None, api_version: str, @@ -299,7 +300,7 @@ class AzureTextCompletion(BaseAzureLLM): async def async_streaming( self, - logging_obj, + logging_obj: LiteLLMLoggingObj, api_base: str, api_key: str | None, api_version: str, diff --git a/litellm/llms/azure_ai/anthropic/handler.py b/litellm/llms/azure_ai/anthropic/handler.py index 80471c9060a..24ee76b31d0 100644 --- a/litellm/llms/azure_ai/anthropic/handler.py +++ b/litellm/llms/azure_ai/anthropic/handler.py @@ -9,6 +9,7 @@ from typing import Final import httpx +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.anthropic.chat.handler import AnthropicChatCompletion from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, @@ -40,7 +41,7 @@ class AzureAnthropicChatCompletion(AnthropicChatCompletion): print_verbose: Callable, encoding, api_key, - logging_obj, + logging_obj: LiteLLMLoggingObj, optional_params: dict, timeout: float | httpx.Timeout, litellm_params: dict, diff --git a/litellm/llms/azure_ai/chat/transformation.py b/litellm/llms/azure_ai/chat/transformation.py index 5540d79f667..8545d646035 100644 --- a/litellm/llms/azure_ai/chat/transformation.py +++ b/litellm/llms/azure_ai/chat/transformation.py @@ -248,7 +248,7 @@ class AzureAIStudioConfig(OpenAIConfig): messages=messages, optional_params=optional_params, litellm_params=litellm_params, - encoding=encoding, + encoding=encoding if encoding is not None else None, api_key=api_key, json_mode=json_mode, ) diff --git a/litellm/llms/bedrock/chat/converse_handler.py b/litellm/llms/bedrock/chat/converse_handler.py index 6970e324db7..25e544f4521 100644 --- a/litellm/llms/bedrock/chat/converse_handler.py +++ b/litellm/llms/bedrock/chat/converse_handler.py @@ -89,7 +89,7 @@ class BedrockConverseLLM(BaseAWSLLM): model_response: ModelResponse, timeout: float | httpx.Timeout | None, encoding, - logging_obj, + logging_obj: LiteLLMLoggingObject, stream, optional_params: dict, litellm_params: dict, diff --git a/litellm/llms/codestral/completion/handler.py b/litellm/llms/codestral/completion/handler.py index 25a51927e22..8c08b2bc33c 100644 --- a/litellm/llms/codestral/completion/handler.py +++ b/litellm/llms/codestral/completion/handler.py @@ -195,7 +195,7 @@ class CodestralTextCompletion: print_verbose: Callable, encoding, api_key: str, - logging_obj, + logging_obj: LiteLLMLogging, optional_params: dict, timeout: float | httpx.Timeout, acompletion=None, @@ -383,7 +383,7 @@ class CodestralTextCompletion: print_verbose: Callable, encoding, api_key, - logging_obj, + logging_obj: LiteLLMLogging, data: dict, timeout: float | httpx.Timeout, optional_params=None, diff --git a/litellm/llms/custom_httpx/aiohttp_handler.py b/litellm/llms/custom_httpx/aiohttp_handler.py index 3cc43cb6072..9f579fd6f55 100644 --- a/litellm/llms/custom_httpx/aiohttp_handler.py +++ b/litellm/llms/custom_httpx/aiohttp_handler.py @@ -221,7 +221,7 @@ class BaseLLMAIOHTTPHandler: timeout=timeout, stream=stream, files=files, - content=content, + content=content if content is not None else None, params=params, ) except httpx.HTTPStatusError as e: diff --git a/litellm/llms/custom_httpx/http_handler.py b/litellm/llms/custom_httpx/http_handler.py index 9ada3674d33..52f30e31641 100644 --- a/litellm/llms/custom_httpx/http_handler.py +++ b/litellm/llms/custom_httpx/http_handler.py @@ -7,9 +7,9 @@ import ssl import sys import threading import time -from collections.abc import Callable, Mapping +from collections.abc import AsyncIterable, Callable, Iterable, Mapping from http.cookiejar import CookieJar, DefaultCookiePolicy -from typing import TYPE_CHECKING, Any, Final, Optional +from typing import TYPE_CHECKING, Any, Final, Optional, TypeAlias, TypedDict import certifi import httpx @@ -62,8 +62,23 @@ except Exception: # https://docs.aiohttp.org/en/stable/client_reference.html#aiohttp.TCPConnector _AIOHTTP_SUPPORTS_SOCKET_FACTORY: Final = "socket_factory" in inspect.signature(TCPConnector.__init__).parameters +_AddrInfo: TypeAlias = tuple[int | socket.AddressFamily, int | socket.SocketKind, int, str, tuple[object, ...]] -def _build_aiohttp_keepalive_socket_factory() -> Callable[[tuple[Any, ...]], socket.socket] | None: +_RequestContent: TypeAlias = str | bytes | Iterable[bytes] | AsyncIterable[bytes] + + +class _TCPConnectorKwargs(TypedDict, total=False): + local_addr: tuple[str, int] | None + ssl: "ssl.SSLContext | bool" + keepalive_timeout: float + ttl_dns_cache: int + enable_cleanup_closed: bool + limit: int + limit_per_host: int + socket_factory: Callable[[_AddrInfo], socket.socket] + + +def _build_aiohttp_keepalive_socket_factory() -> Callable[[_AddrInfo], socket.socket] | None: """ Build a socket_factory that enables SO_KEEPALIVE on aiohttp TCP sockets. @@ -78,7 +93,7 @@ def _build_aiohttp_keepalive_socket_factory() -> Callable[[tuple[Any, ...]], soc if not AIOHTTP_SO_KEEPALIVE or not _AIOHTTP_SUPPORTS_SOCKET_FACTORY: return None - def factory(addr_info: tuple[Any, ...]) -> socket.socket: + def factory(addr_info: _AddrInfo) -> socket.socket: family, type_, proto = addr_info[0], addr_info[1], addr_info[2] sock: Final = socket.socket(family=family, type=type_, proto=proto) sock.setblocking(False) @@ -163,8 +178,8 @@ _STREAMING_ERROR_BODY_READ_EXECUTOR: Final = concurrent.futures.ThreadPoolExecut def _prepare_request_data_and_content( data: dict | str | bytes | None = None, - content: Any = None, -) -> tuple[dict | Mapping | None, Any]: + content: _RequestContent | None = None, +) -> tuple[dict | Mapping | None, _RequestContent | None]: """ Helper function to route data/content parameters correctly for httpx requests @@ -528,7 +543,7 @@ class AsyncHTTPHandler: def __init__( self, timeout: float | httpx.Timeout | None = None, - event_hooks: Mapping[str, list[Callable[..., Any]]] | None = None, + event_hooks: Mapping[str, list[Callable[..., object]]] | None = None, concurrent_limit=None, # Kept for backward compatibility, but ignored (no limits) client_alias: str | None = None, # name for client in logs ssl_verify: VerifyTypes | None = None, @@ -566,7 +581,7 @@ class AsyncHTTPHandler: def create_client( self, timeout: float | httpx.Timeout | None, - event_hooks: Mapping[str, list[Callable[..., Any]]] | None, + event_hooks: Mapping[str, list[Callable[..., object]]] | None, ssl_verify: VerifyTypes | None = None, shared_session: Optional["ClientSession"] = None, ) -> httpx.AsyncClient: @@ -648,7 +663,7 @@ class AsyncHTTPHandler: stream: bool = False, logging_obj: LiteLLMLoggingObject | None = None, files: RequestFiles | None = None, - content: Any = None, + content: _RequestContent | None = None, ): start_time: Final = time.time() try: @@ -691,7 +706,7 @@ class AsyncHTTPHandler: end_time: Final = time.time() time_delta: Final = round(end_time - start_time, 3) headers = {} - error_response: Final = getattr(e, "response", None) + error_response: Final[httpx.Response | None] = getattr(e, "response", None) if error_response is not None: for key, value in error_response.headers.items(): headers[f"response_headers-{key}"] = value @@ -716,7 +731,7 @@ class AsyncHTTPHandler: headers: dict | None = None, timeout: float | httpx.Timeout | None = None, stream: bool = False, - content: Any = None, + content: _RequestContent | None = None, ): try: if timeout is None: @@ -755,7 +770,7 @@ class AsyncHTTPHandler: await new_client.aclose() except httpx.TimeoutException as e: headers = {} - error_response: Final = getattr(e, "response", None) + error_response: Final[httpx.Response | None] = getattr(e, "response", None) if error_response is not None: for key, value in error_response.headers.items(): headers[f"response_headers-{key}"] = value @@ -780,7 +795,7 @@ class AsyncHTTPHandler: headers: dict | None = None, timeout: float | httpx.Timeout | None = None, stream: bool = False, - content: Any = None, + content: _RequestContent | None = None, ): try: if timeout is None: @@ -819,7 +834,7 @@ class AsyncHTTPHandler: await new_client.aclose() except httpx.TimeoutException as e: headers = {} - error_response: Final = getattr(e, "response", None) + error_response: Final[httpx.Response | None] = getattr(e, "response", None) if error_response is not None: for key, value in error_response.headers.items(): headers[f"response_headers-{key}"] = value @@ -844,7 +859,7 @@ class AsyncHTTPHandler: headers: dict | None = None, timeout: float | httpx.Timeout | None = None, stream: bool = False, - content: Any = None, + content: _RequestContent | None = None, ): try: if timeout is None: @@ -895,7 +910,7 @@ class AsyncHTTPHandler: params: dict | None = None, headers: dict | None = None, stream: bool = False, - content: Any = None, + content: _RequestContent | None = None, ): """ Making POST request for a single connection client. @@ -993,7 +1008,7 @@ class AsyncHTTPHandler: def _get_ssl_connector_kwargs( ssl_verify: bool | None = None, ssl_context: ssl.SSLContext | None = None, - ) -> dict[str, Any]: + ) -> _TCPConnectorKwargs: """ Helper method to get SSL connector initialization arguments for aiohttp TCPConnector. @@ -1004,7 +1019,7 @@ class AsyncHTTPHandler: Returns: Dict with appropriate SSL configuration for TCPConnector """ - connector_kwargs: Final[dict[str, Any]] = { + connector_kwargs: Final[_TCPConnectorKwargs] = { "local_addr": ("0.0.0.0", 0) if litellm.force_ipv4 else None, } @@ -1054,7 +1069,7 @@ class AsyncHTTPHandler: verbose_logger.debug("Creating AiohttpTransport...") - transport_connector_kwargs: Final = { + transport_connector_kwargs: Final[_TCPConnectorKwargs] = { "keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT, "ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE, **connector_kwargs, @@ -1212,7 +1227,7 @@ class HTTPHandler: stream: bool = False, timeout: float | httpx.Timeout | None = None, files: dict | RequestFiles | None = None, - content: Any = None, + content: _RequestContent | None = None, logging_obj: LiteLLMLoggingObject | None = None, ): try: @@ -1265,7 +1280,7 @@ class HTTPHandler: headers: dict | None = None, stream: bool = False, timeout: float | httpx.Timeout | None = None, - content: Any = None, + content: _RequestContent | None = None, ): try: # Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix) @@ -1315,7 +1330,7 @@ class HTTPHandler: headers: dict | None = None, stream: bool = False, timeout: float | httpx.Timeout | None = None, - content: Any = None, + content: _RequestContent | None = None, ): try: # Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix) @@ -1364,7 +1379,7 @@ class HTTPHandler: headers: dict | None = None, timeout: float | httpx.Timeout | None = None, stream: bool = False, - content: Any = None, + content: _RequestContent | None = None, ): try: # Prepare data/content parameters to prevent httpx DeprecationWarning (memory leak fix) diff --git a/litellm/llms/custom_httpx/llm_http_handler.py b/litellm/llms/custom_httpx/llm_http_handler.py index a58397c9184..721b9545ac1 100644 --- a/litellm/llms/custom_httpx/llm_http_handler.py +++ b/litellm/llms/custom_httpx/llm_http_handler.py @@ -5,7 +5,8 @@ import ssl from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping from contextlib import asynccontextmanager from functools import lru_cache -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeVar, Union, cast, get_type_hints +from types import ModuleType +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypedDict, TypeVar, Union, cast, get_type_hints from urllib.parse import parse_qs, urlencode, urlparse, urlunparse import httpx @@ -148,6 +149,7 @@ from .http_handler import get_shared_realtime_ssl_context if TYPE_CHECKING: from aiohttp import ClientSession + from websockets.asyncio.client import ClientConnection from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging as _LiteLLMLoggingObj @@ -176,6 +178,19 @@ else: _ResponseT = TypeVar("_ResponseT") +class _DeleteRequestKwargs(TypedDict, total=False): + url: str + headers: dict[str, str] + timeout: float | httpx.Timeout | None + json: dict[str, object] + + +class _MediaUploadKwargs(TypedDict, total=False): + headers: dict[str, str] + content: Iterator[bytes] | AsyncIterator[bytes] + timeout: float | httpx.Timeout + + def _google_genai_streaming_hidden_params( *, api_base: str, @@ -1413,7 +1428,7 @@ class BaseLLMHTTPHandler: headers: dict[str, object] | None, provider_config: BaseOCRConfig, litellm_params: dict, - ) -> tuple[dict[str, Any], str, dict[str, Any], None]: + ) -> tuple[dict[str, object], str, dict[str, object], None]: """ Shared logic for preparing OCR requests. Returns: (headers, complete_url, data, files) @@ -1479,7 +1494,7 @@ class BaseLLMHTTPHandler: headers: dict[str, object] | None, provider_config: BaseOCRConfig, litellm_params: dict, - ) -> tuple[dict[str, Any], str, dict[str, Any], None]: + ) -> tuple[dict[str, object], str, dict[str, object], None]: """ Async version of _prepare_ocr_request for providers that need async transforms. Returns: (headers, complete_url, data, files) @@ -2361,14 +2376,14 @@ class BaseLLMHTTPHandler: model: str, input: str | ResponseInputParam, custom_llm_provider: str, - response_api_optional_request_params: dict[str, Any], + response_api_optional_request_params: dict[str, object], litellm_params: GenericLiteLLMParams, logging_obj: LiteLLMLoggingObj, ) -> tuple[ str, str | ResponseInputParam, str, - dict[str, Any], + dict[str, object], GenericLiteLLMParams, ]: if not _has_pre_call_deployment_hook(logging_obj): @@ -2894,7 +2909,7 @@ class BaseLLMHTTPHandler: }, ) - delete_kwargs: Final[dict[str, Any]] = { + delete_kwargs: Final[_DeleteRequestKwargs] = { "url": url, "headers": headers, "timeout": timeout, @@ -2984,7 +2999,7 @@ class BaseLLMHTTPHandler: }, ) - delete_kwargs: Final[dict[str, Any]] = { + delete_kwargs: Final[_DeleteRequestKwargs] = { "url": url, "headers": headers, "timeout": timeout, @@ -3725,7 +3740,7 @@ class BaseLLMHTTPHandler: timeout: float | httpx.Timeout | None, ) -> httpx.Response: headers: Final = {**base_headers, "Content-Type": content_type} - kwargs: Final[dict[str, Any]] = { + kwargs: Final[_MediaUploadKwargs] = { "headers": headers, "content": self._iter_in_blocks(body_stream.iter_bytes(), self._MEDIA_UPLOAD_BLOCK_SIZE), } @@ -3762,7 +3777,7 @@ class BaseLLMHTTPHandler: break yield cast(bytes, block) - kwargs: Final[dict[str, Any]] = {"headers": headers, "content": _abody()} + kwargs: Final[_MediaUploadKwargs] = {"headers": headers, "content": _abody()} if timeout is not None: kwargs["timeout"] = timeout resp: Final = await client.client.post(url, **kwargs) @@ -5242,7 +5257,7 @@ class BaseLLMHTTPHandler: def _wrap_responses_response_as_fake_stream( self, - result: Any, + result: ResponsesAPIResponse, model: str, responses_api_provider_config: BaseResponsesAPIConfig, logging_obj: "LiteLLMLoggingObj", @@ -5365,7 +5380,7 @@ class BaseLLMHTTPHandler: async def _call_agentic_completion_hooks( self, - response: Any, + response: object, model: str, messages: list[dict], anthropic_messages_provider_config: "BaseAnthropicMessagesConfig", @@ -5536,7 +5551,7 @@ class BaseLLMHTTPHandler: async def _call_agentic_chat_completion_hooks( self, - response: Any, + response: ModelResponse, model: str, messages: list[dict], optional_params: dict, @@ -5760,14 +5775,14 @@ class BaseLLMHTTPHandler: @staticmethod async def _open_realtime_backend_ws( - websockets_module: Any, + websockets_module: ModuleType, url: str, headers: dict, - ssl_context: Any, + ssl_context: bool | str | ssl.SSLContext, *, open_timeout: float = 8.0, max_attempts: int = 3, - ) -> Any: + ) -> "ClientConnection": """Open the backend realtime websocket, retrying a hung open handshake. The upstream Live handshake (e.g. Gemini Live) intermittently hangs on @@ -5826,7 +5841,6 @@ class BaseLLMHTTPHandler: query_params: RealtimeQueryParams | None = None, ): import websockets - from websockets.asyncio.client import ClientConnection url: Final = provider_config.get_complete_url(api_base, model, api_key) headers = provider_config.validate_environment( @@ -5844,12 +5858,12 @@ class BaseLLMHTTPHandler: ssl_context.verify_mode = ssl.CERT_NONE backend_ws: Final = await self._open_realtime_backend_ws(websockets, url, headers, ssl_context) async with backend_ws: - _request_data: Final[dict[str, Any]] = {} + _request_data: Final[dict[str, object]] = {} if litellm_metadata: _request_data["litellm_metadata"] = litellm_metadata realtime_streaming: Final = RealTimeStreaming( websocket, - cast(ClientConnection, backend_ws), + backend_ws, logging_obj, provider_config, model, @@ -6008,7 +6022,7 @@ class BaseLLMHTTPHandler: ) else: url = provider_config.get_complete_url(api_base=api_base, model=model or "", api_version=api_version) - headers: dict[str, Any] = provider_config.validate_environment( + headers: dict[str, object] = provider_config.validate_environment( headers={}, model=model or "", api_key=api_key ) else: @@ -6079,7 +6093,7 @@ class BaseLLMHTTPHandler: if provider_config is not None: url = provider_config.get_realtime_calls_url(api_base=api_base, model=model or "", api_version=api_version) - headers: dict[str, Any] = provider_config.get_realtime_calls_headers(ephemeral_key=openai_ephemeral_key) + headers: dict[str, object] = provider_config.get_realtime_calls_headers(ephemeral_key=openai_ephemeral_key) else: url = f"{api_base.rstrip('/')}/v1/realtime/calls" headers = { @@ -6247,7 +6261,7 @@ class BaseLLMHTTPHandler: yield backend async with _backend_connection() as backend_ws: - _request_data: Final[dict[str, Any]] = {} + _request_data: Final[dict[str, object]] = {} if litellm_metadata: _request_data["litellm_metadata"] = litellm_metadata @@ -9444,7 +9458,7 @@ class BaseLLMHTTPHandler: litellm_params=dict(litellm_params), extra_body=extra_body, ) - all_optional_params: Final[dict[str, Any]] = dict(litellm_params) + all_optional_params: Final[dict[str, object]] = dict(litellm_params) all_optional_params.update(vector_store_search_optional_params or {}) headers, signed_json_body = vector_store_provider_config.sign_request( headers=headers, @@ -9540,7 +9554,7 @@ class BaseLLMHTTPHandler: extra_body=extra_body, ) - all_optional_params: Final[dict[str, Any]] = dict(litellm_params) + all_optional_params: Final[dict[str, object]] = dict(litellm_params) all_optional_params.update(vector_store_search_optional_params or {}) headers, signed_json_body = vector_store_provider_config.sign_request( @@ -9860,7 +9874,7 @@ class BaseLLMHTTPHandler: url: Final = api_base - params: Final[dict[str, Any]] = {} + params: Final[dict[str, object]] = {} if after is not None: params["after"] = after if before is not None: @@ -9938,7 +9952,7 @@ class BaseLLMHTTPHandler: url: Final = api_base - params: Final[dict[str, Any]] = {} + params: Final[dict[str, object]] = {} if after is not None: params["after"] = after if before is not None: diff --git a/litellm/llms/github_copilot/chat/transformation.py b/litellm/llms/github_copilot/chat/transformation.py index b3df6f14d84..27a0028ce4a 100644 --- a/litellm/llms/github_copilot/chat/transformation.py +++ b/litellm/llms/github_copilot/chat/transformation.py @@ -277,7 +277,7 @@ class GithubCopilotConfig(OpenAIConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: object, api_key: str | None = None, json_mode: bool | None = None, ) -> "ModelResponse": diff --git a/litellm/llms/oci/chat/transformation.py b/litellm/llms/oci/chat/transformation.py index 6615ad46944..94494a87bba 100644 --- a/litellm/llms/oci/chat/transformation.py +++ b/litellm/llms/oci/chat/transformation.py @@ -10,7 +10,7 @@ implement the LiteLLM BaseConfig interface. Heavy-lifting lives in: """ import json -from collections.abc import AsyncIterator, Iterator +from collections.abc import AsyncIterator, Callable, Iterator from typing import TYPE_CHECKING, Any, Final import httpx @@ -713,8 +713,25 @@ class OCIChatConfig(BaseConfig): class OCIStreamWrapper(CustomStreamWrapper): """Custom stream wrapper that dispatches OCI SSE chunks to the correct handler.""" - def __init__(self, **kwargs: Any): - super().__init__(**kwargs) + def __init__( + self, + completion_stream: object, + model: str, + logging_obj: LiteLLMLoggingObj, + custom_llm_provider: str | None = None, + stream_options: object = None, + make_call: Callable[..., object] | None = None, + _response_headers: dict[str, object] | None = None, + ) -> None: + super().__init__( + completion_stream=completion_stream, + model=model, + logging_obj=logging_obj, + custom_llm_provider=custom_llm_provider, + stream_options=stream_options, + make_call=make_call, + _response_headers=_response_headers, + ) # Tracks whether any prior Cohere chunk in this stream has emitted # tool calls. The Cohere handler uses this to decide whether the # terminal consolidation chunk's tool calls are duplicates (suppress) diff --git a/litellm/llms/openai/completion/handler.py b/litellm/llms/openai/completion/handler.py index 7f29e3f4114..c7b59509eb0 100644 --- a/litellm/llms/openai/completion/handler.py +++ b/litellm/llms/openai/completion/handler.py @@ -217,7 +217,7 @@ class OpenAITextCompletion(BaseLLM): def streaming( self, - logging_obj, + logging_obj: LiteLLMLoggingObj, api_key: str, data: dict, headers: dict, @@ -274,7 +274,7 @@ class OpenAITextCompletion(BaseLLM): async def async_streaming( self, - logging_obj, + logging_obj: LiteLLMLoggingObj, api_key: str, data: dict, headers: dict, diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index e8a6e5a7450..e96b61d8204 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -1,6 +1,6 @@ import time import types -from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Iterator +from collections.abc import AsyncIterator, Callable, Coroutine, Iterable, Iterator, Mapping from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast from urllib.parse import urlparse @@ -61,16 +61,17 @@ class MistralEmbeddingConfig: def __init__( self, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = locals().copy() for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) @classmethod def get_config(cls): + config_attrs: Final[Mapping[str, object]] = cls.__dict__ return { k: v - for k, v in cls.__dict__.items() + for k, v in config_attrs.items() if not k.startswith("__") and not isinstance( v, @@ -153,7 +154,7 @@ class OpenAIConfig(BaseConfig): top_p: int | None = None, response_format: dict | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = locals().copy() for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) @@ -261,7 +262,7 @@ class OpenAIConfig(BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: object, api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: @@ -299,7 +300,7 @@ class OpenAIConfig(BaseConfig): streaming_response: Iterator[str] | AsyncIterator[str] | ModelResponse, sync_stream: bool, json_mode: bool | None = False, - ) -> Any: + ) -> "OpenAIChatCompletionResponseIterator": return OpenAIChatCompletionResponseIterator( streaming_response=streaming_response, sync_stream=sync_stream, @@ -478,14 +479,14 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): async def _call_agentic_completion_hooks_openai( self, - response: Any, + response: object, model: str, messages: list[dict], optional_params: dict, logging_obj: LiteLLMLoggingObj, stream: bool, litellm_params: dict, - ) -> Any | None: + ) -> object | None: """ Call agentic completion hooks for all custom loggers (OpenAI Chat Completions API). @@ -536,7 +537,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): kwargs_with_provider["custom_llm_provider"] = custom_llm_provider # For OpenAI Chat Completions, use the chat completion agentic loop method - agentic_response = await callback.async_run_chat_completion_agentic_loop( + agentic_response: object = await callback.async_run_chat_completion_agentic_loop( tools=tool_calls, model=model, messages=messages, @@ -580,7 +581,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): timeout: float | httpx.Timeout, optional_params: dict, litellm_params: dict, - logging_obj: Any, + logging_obj: LiteLLMLoggingObj, model: str | None = None, messages: list | None = None, print_verbose: Callable | None = None, @@ -1590,7 +1591,7 @@ class OpenAIFilesAPI(BaseLLM): client: OpenAI | AsyncOpenAI | None = None, _is_async: bool = False, ) -> OpenAI | AsyncOpenAI | None: - received_args: Final = locals() + received_args: Final[Mapping[str, object]] = locals() openai_client: OpenAI | AsyncOpenAI | None = None if client is None: data: Final = {} @@ -1628,7 +1629,7 @@ class OpenAIFilesAPI(BaseLLM): max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | None = None, - ) -> OpenAIFileObject | Coroutine[Any, Any, OpenAIFileObject]: + ) -> OpenAIFileObject | Coroutine[None, None, OpenAIFileObject]: openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client( api_key=api_key, api_base=api_base, @@ -1670,7 +1671,7 @@ class OpenAIFilesAPI(BaseLLM): max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | None = None, - ) -> HttpxBinaryResponseContent | Coroutine[Any, Any, HttpxBinaryResponseContent]: + ) -> HttpxBinaryResponseContent | Coroutine[None, None, HttpxBinaryResponseContent]: openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client( api_key=api_key, api_base=api_base, @@ -1948,7 +1949,7 @@ class OpenAIBatchesAPI(BaseLLM): client: OpenAI | AsyncOpenAI | None = None, _is_async: bool = False, ) -> OpenAI | AsyncOpenAI | None: - received_args: Final = locals() + received_args: Final[Mapping[str, object]] = locals() openai_client: OpenAI | AsyncOpenAI | None = None if client is None: data: Final = {} @@ -1986,7 +1987,7 @@ class OpenAIBatchesAPI(BaseLLM): max_retries: int | None, organization: str | None, client: OpenAI | AsyncOpenAI | None = None, - ) -> LiteLLMBatch | Coroutine[Any, Any, LiteLLMBatch]: + ) -> LiteLLMBatch | Coroutine[None, None, LiteLLMBatch]: openai_client: Final[OpenAI | AsyncOpenAI | None] = self.get_openai_client( api_key=api_key, api_base=api_base, @@ -2160,7 +2161,7 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: OpenAI | None = None, ) -> OpenAI: - received_args: Final = locals() + received_args: Final[Mapping[str, object]] = locals() if client is None: data: Final = {} for k, v in received_args.items(): @@ -2185,7 +2186,7 @@ class OpenAIAssistantsAPI(BaseLLM): organization: str | None, client: AsyncOpenAI | None = None, ) -> AsyncOpenAI: - received_args: Final = locals() + received_args: Final[Mapping[str, object]] = locals() if client is None: data: Final = {} for k, v in received_args.items(): @@ -2848,7 +2849,7 @@ class OpenAIAssistantsAPI(BaseLLM): assistant_id: str, additional_instructions: str | None, instructions: str | None, - metadata: dict | None, + metadata: dict[str, str] | None, model: str | None, stream: bool | None, tools: Iterable[AssistantToolParam] | None, @@ -2912,23 +2913,32 @@ class OpenAIAssistantsAPI(BaseLLM): assistant_id: str, additional_instructions: str | None, instructions: str | None, - metadata: dict | None, + metadata: dict[str, str] | None, model: str | None, tools: Iterable[AssistantToolParam] | None, event_handler: AssistantEventHandler | None, ) -> AssistantStreamManager[AssistantEventHandler]: - data: Final[dict[str, Any]] = { - "thread_id": thread_id, - "assistant_id": assistant_id, - "additional_instructions": additional_instructions, - "instructions": instructions, - "metadata": metadata, - "model": model, - "tools": tools, - } + runs_stream: Final = client.beta.threads.runs.stream if event_handler is not None: - data["event_handler"] = event_handler - return client.beta.threads.runs.stream(**data) + return runs_stream( + thread_id=thread_id, + assistant_id=assistant_id, + additional_instructions=additional_instructions, + instructions=instructions, + metadata=metadata, + model=model, + tools=tools, + event_handler=event_handler, + ) + return runs_stream( + thread_id=thread_id, + assistant_id=assistant_id, + additional_instructions=additional_instructions, + instructions=instructions, + metadata=metadata, + model=model, + tools=tools, + ) # fmt: off @@ -2984,7 +2994,7 @@ class OpenAIAssistantsAPI(BaseLLM): assistant_id: str, additional_instructions: str | None, instructions: str | None, - metadata: dict | None, + metadata: dict[str, str] | None, model: str | None, stream: bool | None, tools: Iterable[AssistantToolParam] | None, diff --git a/litellm/llms/openai_like/chat/handler.py b/litellm/llms/openai_like/chat/handler.py index 3ce7a63c532..8c548b6b0d6 100644 --- a/litellm/llms/openai_like/chat/handler.py +++ b/litellm/llms/openai_like/chat/handler.py @@ -12,6 +12,7 @@ import httpx import litellm from litellm import LlmProviders +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.bedrock.chat.invoke_handler import MockResponseIterator from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.databricks.streaming_utils import ModelResponseIterator @@ -112,7 +113,7 @@ class OpenAILikeChatHandler(OpenAILikeBase): print_verbose: Callable, encoding, api_key, - logging_obj, + logging_obj: LiteLLMLoggingObj, stream, data: dict, optional_params=None, @@ -214,7 +215,7 @@ class OpenAILikeChatHandler(OpenAILikeBase): print_verbose: Callable, encoding, api_key: str | None, - logging_obj, + logging_obj: LiteLLMLoggingObj, optional_params: dict, acompletion=None, litellm_params: dict = {}, diff --git a/litellm/llms/predibase/chat/handler.py b/litellm/llms/predibase/chat/handler.py index b4cbf1e2e05..d0a61ea6e00 100644 --- a/litellm/llms/predibase/chat/handler.py +++ b/litellm/llms/predibase/chat/handler.py @@ -9,6 +9,7 @@ from typing import Final import httpx import litellm +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, get_async_httpx_client, @@ -59,7 +60,7 @@ class PredibaseChatCompletion: print_verbose: Callable, encoding, api_key: str, - logging_obj, + logging_obj: LiteLLMLoggingObj, optional_params: dict, litellm_params: dict, tenant_id: str, @@ -250,7 +251,7 @@ class PredibaseChatCompletion: print_verbose: Callable, encoding, api_key, - logging_obj, + logging_obj: LiteLLMLoggingObj, data: dict, timeout: float | httpx.Timeout, optional_params=None, diff --git a/litellm/llms/replicate/chat/handler.py b/litellm/llms/replicate/chat/handler.py index 8d6ba6c8a65..fc114104d32 100644 --- a/litellm/llms/replicate/chat/handler.py +++ b/litellm/llms/replicate/chat/handler.py @@ -6,6 +6,7 @@ from typing import Final import litellm from litellm.constants import REPLICATE_POLLING_DELAY_SECONDS +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.custom_httpx.http_handler import ( AsyncHTTPHandler, HTTPHandler, @@ -128,7 +129,7 @@ def completion( print_verbose: Callable, optional_params: dict, litellm_params: dict, - logging_obj, + logging_obj: LiteLLMLoggingObj, api_key, encoding, custom_prompt_dict={}, @@ -246,7 +247,7 @@ async def async_completion( input_data, api_key, api_base, - logging_obj, + logging_obj: LiteLLMLoggingObj, print_verbose, headers: dict, ) -> ModelResponse | CustomStreamWrapper: diff --git a/litellm/llms/sagemaker/completion/handler.py b/litellm/llms/sagemaker/completion/handler.py index 8d81d16d5eb..84cad56f0d4 100644 --- a/litellm/llms/sagemaker/completion/handler.py +++ b/litellm/llms/sagemaker/completion/handler.py @@ -8,6 +8,7 @@ import httpx import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.asyncify import asyncify +from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.bedrock.base_aws_llm import BaseAWSLLM from litellm.llms.custom_httpx.http_handler import ( _get_httpx_client, @@ -138,7 +139,7 @@ class SagemakerLLM(BaseAWSLLM): model_response: ModelResponse, print_verbose: Callable, encoding, - logging_obj, + logging_obj: LiteLLMLoggingObj, optional_params: dict, litellm_params: dict, timeout: float | httpx.Timeout | None = None, @@ -431,17 +432,18 @@ class SagemakerLLM(BaseAWSLLM): if not prepared_request.body: raise ValueError("Prepared request body is empty") + stream_logging_obj: Final[LiteLLMLoggingObj] = logging_obj completion_stream: Final = await self.make_async_call( api_base=prepared_request.url, headers=prepared_request.headers, data=cast(str, prepared_request.body), - logging_obj=logging_obj, + logging_obj=stream_logging_obj, ) streaming_response: Final = CustomStreamWrapper( completion_stream=completion_stream, model=model, custom_llm_provider="sagemaker", - logging_obj=logging_obj, + logging_obj=stream_logging_obj, ) # LOGGING diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index ff51f1a013e..d298670aa7a 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -3,7 +3,7 @@ ## Initial implementation - covers gemini + image gen calls import json import time -from collections.abc import Callable, Mapping +from collections.abc import Callable, Mapping, Sequence from copy import deepcopy from functools import partial from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast @@ -208,7 +208,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): presence_penalty: float | None = None, seed: int | None = None, ) -> None: - locals_: Final = locals().copy() + locals_: Final[Mapping[str, object]] = locals().copy() for key, value in locals_.items(): if key != "self" and value is not None: setattr(self.__class__, key, value) @@ -1427,7 +1427,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): @staticmethod def _extract_server_side_tool_invocations( parts: list[HttpxPartType], - ) -> list[dict[str, Any]] | None: + ) -> list[dict[str, object]] | None: """Extract server-side tool invocations (toolCall/toolResponse) from parts. These are returned by Gemini when context circulation is enabled @@ -1438,15 +1438,15 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): Returns: List of server-side invocation dicts if any found, None otherwise. """ - invocations: Final[list[dict[str, Any]]] = [] + invocations: Final[list[dict[str, object]]] = [] # Index toolCalls by id so we can pair them with responses - tool_calls_by_id: Final[dict[str, dict[str, Any]]] = {} - tool_responses_by_id: Final[dict[str, dict[str, Any]]] = {} + tool_calls_by_id: Final[dict[str, dict[str, object]]] = {} + tool_responses_by_id: Final[dict[str, dict[str, object]]] = {} for part in parts: if "toolCall" in part: tc = part["toolCall"] - entry: dict[str, Any] = { + entry: dict[str, object] = { "tool_type": tc.get("toolType"), "id": tc.get("id"), "args": tc.get("args"), @@ -1753,7 +1753,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): response_tokens_details: CompletionTokensDetailsWrapper | None = None usage_metadata: Final = completion_response["usageMetadata"] - def _get_token_count(detail: Mapping[str, Any]) -> int: + def _get_token_count(detail: Mapping[str, object]) -> int: raw_token_count: Final = detail.get("tokenCount", detail.get("token_count", 0)) return raw_token_count if isinstance(raw_token_count, int) else 0 @@ -2068,7 +2068,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ) @staticmethod - def _get_stream_chunk_attr(chunk: Any, field_name: str) -> Any: + def _get_stream_chunk_attr(chunk: object, field_name: str) -> object: if isinstance(chunk, dict): value = chunk.get(field_name) if value is not None: @@ -2110,10 +2110,10 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): def apply_assembled_streaming_response_metadata( self, response: ModelResponse, - chunks: list[Any], + chunks: list[object], ) -> None: for field_name in VERTEX_AI_PROVIDER_METADATA_FIELDS: - merged: list[Any] = [] + merged: list[object] = [] for chunk in chunks: value = VertexGeminiConfig._get_stream_chunk_attr(chunk, field_name) if not value: @@ -2214,8 +2214,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): functions: ChatCompletionToolCallFunctionChunk | None = None thinking_blocks: list[ChatCompletionThinkingBlock] | None = None reasoning_content: str | None = None - thought_signatures: Any | None = None - server_side_tool_invocations: list[dict[str, Any]] | None = None + thought_signatures: Sequence[str] | None = None + server_side_tool_invocations: list[dict[str, object]] | None = None for idx, candidate in enumerate(_candidates): if "content" not in candidate: @@ -2370,7 +2370,7 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): messages: list[AllMessageValues], optional_params: dict, litellm_params: dict, - encoding: Any, + encoding: object, api_key: str | None = None, json_mode: bool | None = None, ) -> ModelResponse: @@ -2486,7 +2486,8 @@ class VertexGeminiConfig(VertexAIBaseConfig, BaseConfig): ## ADD SERVICE TIER ## if getattr(raw_response, "headers", None): - if service_tier := raw_response.headers.get("x-gemini-service-tier"): + service_tier: Final[str | None] = raw_response.headers.get("x-gemini-service-tier") + if service_tier: if service_tier.lower() == "standard": setattr(model_response, "service_tier", "default") else: @@ -2660,7 +2661,7 @@ class VertexLLM(VertexBase): print_verbose: Callable, data: dict, timeout: float | httpx.Timeout | None, - encoding, + encoding: object, logging_obj, stream, optional_params: dict, @@ -2756,7 +2757,7 @@ class VertexLLM(VertexBase): "vertex_ai", "vertex_ai_beta", "gemini" ], # if it's vertex_ai or gemini (google ai studio) timeout: float | httpx.Timeout | None, - encoding, + encoding: object, logging_obj, stream, optional_params: dict, @@ -2873,7 +2874,7 @@ class VertexLLM(VertexBase): custom_llm_provider: Literal[ "vertex_ai", "vertex_ai_beta", "gemini" ], # if it's vertex_ai or gemini (google ai studio) - encoding, + encoding: object, logging_obj, optional_params: dict, acompletion: bool, @@ -3122,7 +3123,7 @@ class ModelResponseIterator: def _apply_stream_candidates( self, _candidates: list[Candidates], - model_response: Any, + model_response: "ModelResponseStream", ) -> tuple[list[dict], list[dict], list[dict], list[dict]]: ( grounding_metadata, @@ -3200,7 +3201,7 @@ class ModelResponseIterator: def _apply_stream_usage_metadata( self, - processed_chunk: Any, + processed_chunk: GenerateContentResponseBody, model_response: Any, grounding_metadata: list[dict], ) -> Usage | None: diff --git a/litellm/llms/vertex_ai/vertex_ai_non_gemini.py b/litellm/llms/vertex_ai/vertex_ai_non_gemini.py index 8916c0b8740..1c582c7c376 100644 --- a/litellm/llms/vertex_ai/vertex_ai_non_gemini.py +++ b/litellm/llms/vertex_ai/vertex_ai_non_gemini.py @@ -28,7 +28,7 @@ class TextStreamer: Fake streaming iterator for Vertex AI Model Garden calls """ - def __init__(self, text): + def __init__(self, text: str): self.text = text.split() # let's assume words as a streaming unit self.index = 0 diff --git a/litellm/main.py b/litellm/main.py index c70a41c891a..5d9144674de 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -19,12 +19,12 @@ import random import sys import time import traceback -from collections.abc import AsyncIterator, Coroutine, Iterable, Mapping +from collections.abc import AsyncIterator, Coroutine, Iterable, Mapping, Sequence from concurrent import futures from concurrent.futures import FIRST_COMPLETED, ThreadPoolExecutor, wait from copy import deepcopy from functools import partial -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast, get_args +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, Union, cast, get_args from litellm._logging import _redact_string from litellm._uuid import uuid @@ -96,6 +96,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.litellm_core_utils.request_timeout_resolver import ( get_configured_request_timeout, ) +from litellm.litellm_core_utils.streaming_handler import validated_stream_logging_obj from litellm.llms.base_llm import BaseConfig, BaseImageGenerationConfig from litellm.llms.base_llm.base_model_iterator import ( convert_model_response_to_streaming, @@ -633,10 +634,10 @@ async def acompletion( init_response: Final = await loop.run_in_executor(None, func_with_context) if isinstance(init_response, dict) or isinstance(init_response, ModelResponse): ## CACHING SCENARIO if isinstance(init_response, dict): - response = ModelResponse(**init_response) + response = _model_response_from_cached_dict(init_response) response = init_response elif asyncio.iscoroutine(init_response): - response = await init_response + response = await _resolve_dispatched_chat_response(init_response) else: response = init_response @@ -698,6 +699,20 @@ async def acompletion( ) +async def _resolve_dispatched_chat_response( + pending: Coroutine[object, object, ModelResponse | CustomStreamWrapper], +) -> ModelResponse | CustomStreamWrapper: + return await pending + + +def _model_response_from_cached_dict(cached_response_dict: Mapping[str, object]) -> ModelResponse: + return ModelResponse(**cached_response_dict) + + +def _transcription_response_from_cached_dict(cached_response_dict: Mapping[str, object]) -> TranscriptionResponse: + return TranscriptionResponse(**cached_response_dict) + + async def _async_streaming(response, model, custom_llm_provider, args): try: print_verbose(f"received response in _async_streaming: {response}") @@ -812,7 +827,7 @@ def mock_completion( mock_response: MOCK_RESPONSE_TYPE | None = "This is a mock request", mock_tool_calls: list | None = None, mock_timeout: bool | None = False, - logging=None, + logging: LiteLLMLoggingObj | None = None, custom_llm_provider=None, timeout: float | str | httpx.Timeout | None = None, **kwargs, @@ -896,7 +911,7 @@ def mock_completion( ), model=model, custom_llm_provider="openai", - logging_obj=logging, + logging_obj=validated_stream_logging_obj(logging), ) return CustomStreamWrapper( completion_stream=mock_completion_streaming_obj( @@ -904,7 +919,7 @@ def mock_completion( ), model=model, custom_llm_provider="openai", - logging_obj=logging, + logging_obj=validated_stream_logging_obj(logging), ) if isinstance(mock_response, litellm.MockException): raise mock_response @@ -983,12 +998,12 @@ def responses_api_bridge_check( model: str, custom_llm_provider: str, web_search_options: OpenAIWebSearchOptions | None = None, - tools: list[Any] | None = None, - reasoning_effort: Any | None = None, - reasoning_summary: Any | None = None, + tools: Sequence[Mapping[str, object]] | None = None, + reasoning_effort: str | Mapping[str, object] | None = None, + reasoning_summary: object | None = None, api_base: str | None = None, ) -> tuple[dict, str]: - model_info: dict[str, Any] = {} + model_info: dict[str, object] = {} # Global flag: route ALL OpenAI chat completions through Responses API. # Returns early with minimal model_info; callers only inspect the "mode" key. @@ -1110,6 +1125,22 @@ def _drop_input_examples_from_tools( return cleaned_tools +class _ProxyAuthHeadersProvider(Protocol): + def get_auth_headers(self) -> Mapping[str, str]: ... + + +def _proxy_auth_headers(proxy_auth: _ProxyAuthHeadersProvider) -> Mapping[str, str]: + return proxy_auth.get_auth_headers() + + +def _provider_config_items(config: Mapping[str, object]) -> Iterable[tuple[str, object]]: + return config.items() + + +def _locals_snapshot(values: Mapping[str, object]) -> Mapping[str, object]: + return values + + def _build_custom_pricing_entry( custom_llm_provider: str, kwargs: dict, @@ -1185,13 +1216,31 @@ def _register_custom_pricing_for_request( ) +def _dispatch_metadata(ctx: _CompletionDispatchContext) -> Mapping[str, object] | None: + return ctx.metadata + + +def _dispatch_client_http(ctx: _CompletionDispatchContext) -> HTTPHandler | AsyncHTTPHandler | None: + return ctx.client + + +def _dispatch_client_azure( + ctx: _CompletionDispatchContext, +) -> openai.AzureOpenAI | openai.AsyncAzureOpenAI | HTTPHandler | AsyncHTTPHandler | None: + return ctx.client + + +def _dispatch_client_openai(ctx: _CompletionDispatchContext) -> openai.OpenAI | openai.AsyncOpenAI | None: + return ctx.client + + def _complete_azure(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: _azure_detection_model: Final = ctx._azure_detection_model acompletion: Final = ctx.acompletion api_base = ctx.api_base api_key = ctx.api_key api_version = ctx.api_version - client: Final = ctx.client + client: Final = _dispatch_client_azure(ctx) custom_llm_provider: Final = ctx.custom_llm_provider extra_headers: Final = ctx.extra_headers headers = ctx.headers @@ -1232,7 +1281,8 @@ def _complete_azure(ctx: _CompletionDispatchContext) -> _CompletionDispatchResul "AZURE_AD_TOKEN" ) - azure_ad_token_provider: Final = litellm_params.get("azure_ad_token_provider", None) + azure_ad_token_provider_value: Final = litellm_params.get("azure_ad_token_provider", None) + azure_ad_token_provider: Final = azure_ad_token_provider_value if callable(azure_ad_token_provider_value) else None headers = headers or litellm.headers @@ -1244,7 +1294,7 @@ def _complete_azure(ctx: _CompletionDispatchContext) -> _CompletionDispatchResul if litellm.AzureOpenAIO1Config().is_o_series_model(model=_azure_detection_model): ## LOAD CONFIG - if set config = litellm.AzureOpenAIO1Config.get_config() - for k, v in config.items(): + for k, v in _provider_config_items(config): if ( k not in optional_params ): # completion(top_k=3) > azure_config(top_k=3) <- allows for dynamic variables to be passed in @@ -1273,7 +1323,7 @@ def _complete_azure(ctx: _CompletionDispatchContext) -> _CompletionDispatchResul else: ## LOAD CONFIG - if set config = litellm.AzureOpenAIConfig.get_config() - for k, v in config.items(): + for k, v in _provider_config_items(config): if ( k not in optional_params ): # completion(top_k=3) > azure_config(top_k=3) <- allows for dynamic variables to be passed in @@ -1323,7 +1373,7 @@ def _complete_azure_text(ctx: _CompletionDispatchContext) -> _CompletionDispatch api_base = ctx.api_base api_key = ctx.api_key api_version = ctx.api_version - client: Final = ctx.client + client: Final = _dispatch_client_azure(ctx) extra_headers: Final = ctx.extra_headers headers = ctx.headers litellm_params: Final = ctx.litellm_params @@ -1358,7 +1408,8 @@ def _complete_azure_text(ctx: _CompletionDispatchContext) -> _CompletionDispatch "AZURE_AD_TOKEN" ) - azure_ad_token_provider: Final = litellm_params.get("azure_ad_token_provider", None) + azure_ad_token_provider_value: Final = litellm_params.get("azure_ad_token_provider", None) + azure_ad_token_provider: Final = azure_ad_token_provider_value if callable(azure_ad_token_provider_value) else None headers = headers or litellm.headers @@ -1367,7 +1418,7 @@ def _complete_azure_text(ctx: _CompletionDispatchContext) -> _CompletionDispatch ## LOAD CONFIG - if set config: Final = litellm.AzureOpenAIConfig.get_config() - for k, v in config.items(): + for k, v in _provider_config_items(config): if ( k not in optional_params ): # completion(top_k=3) > azure_config(top_k=3) <- allows for dynamic variables to be passed in @@ -1415,7 +1466,7 @@ def _complete_deepseek(ctx: _CompletionDispatchContext) -> _CompletionDispatchRe acompletion: Final = ctx.acompletion api_base: Final = ctx.api_base api_key: Final = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params @@ -1466,7 +1517,7 @@ def _complete_azure_ai(ctx: _CompletionDispatchContext) -> _CompletionDispatchRe acompletion: Final = ctx.acompletion api_base = ctx.api_base api_key = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider extra_headers: Final = ctx.extra_headers headers = ctx.headers @@ -1622,7 +1673,7 @@ def _complete_text_completion_openai( acompletion: Final = ctx.acompletion api_base = ctx.api_base api_key = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_openai(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers = ctx.headers litellm_params: Final = ctx.litellm_params @@ -1654,7 +1705,7 @@ def _complete_text_completion_openai( ## LOAD CONFIG - if set config: Final = litellm.OpenAITextCompletionConfig.get_config() - for k, v in config.items(): + for k, v in _provider_config_items(config): if ( k not in optional_params ): # completion(top_k=3) > openai_text_config(top_k=3) <- allows for dynamic variables to be passed in @@ -1704,7 +1755,7 @@ def _complete_fireworks_ai( acompletion: Final = ctx.acompletion api_base: Final = ctx.api_base api_key: Final = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params @@ -1755,7 +1806,7 @@ def _complete_heroku(ctx: _CompletionDispatchContext) -> _CompletionDispatchResu acompletion: Final = ctx.acompletion api_base: Final = ctx.api_base api_key: Final = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params @@ -1805,7 +1856,7 @@ def _complete_ragflow(ctx: _CompletionDispatchContext) -> _CompletionDispatchRes acompletion: Final = ctx.acompletion api_base: Final = ctx.api_base api_key: Final = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params @@ -1855,7 +1906,7 @@ def _complete_xai(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: acompletion: Final = ctx.acompletion api_base: Final = ctx.api_base api_key: Final = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params @@ -1906,7 +1957,7 @@ def _complete_groq(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult acompletion: Final = ctx.acompletion api_base = ctx.api_base api_key = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers = ctx.headers litellm_params: Final = ctx.litellm_params @@ -1938,7 +1989,7 @@ def _complete_groq(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult ## LOAD CONFIG - if set config: Final = litellm.GroqChatConfig.get_config() - for k, v in config.items(): + for k, v in _provider_config_items(config): if ( k not in optional_params ): # completion(top_k=3) > openai_config(top_k=3) <- allows for dynamic variables to be passed in @@ -1970,7 +2021,7 @@ def _complete_bedrock_mantle( acompletion: Final = ctx.acompletion api_base = ctx.api_base api_key = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers = ctx.headers litellm_params: Final = ctx.litellm_params @@ -1987,7 +2038,7 @@ def _complete_bedrock_mantle( api_key = api_key or litellm.api_key or get_secret("BEDROCK_MANTLE_API_KEY") headers = headers or litellm.headers config: Final = litellm.BedrockMantleChatConfig.get_config() - for k, v in config.items(): + for k, v in _provider_config_items(config): if k not in optional_params: optional_params[k] = v return base_llm_http_handler.completion( @@ -2014,7 +2065,7 @@ def _complete_a2a(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: acompletion: Final = ctx.acompletion api_base = ctx.api_base api_key = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers = ctx.headers litellm_params: Final = ctx.litellm_params @@ -2077,7 +2128,7 @@ def _complete_gigachat(ctx: _CompletionDispatchContext) -> _CompletionDispatchRe acompletion: Final = ctx.acompletion api_base: Final = ctx.api_base api_key = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers = ctx.headers litellm_params: Final = ctx.litellm_params @@ -2139,7 +2190,7 @@ def _complete_sap(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: acompletion: Final = ctx.acompletion api_base: Final = ctx.api_base api_key: Final = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers = ctx.headers litellm_params: Final = ctx.litellm_params @@ -2155,7 +2206,7 @@ def _complete_sap(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: headers = headers or litellm.headers ## LOAD CONFIG - if set config: Final = litellm.GenAIHubOrchestrationConfig.get_config() - for k, v in config.items(): + for k, v in _provider_config_items(config): if ( k not in optional_params ): # completion(top_k=3) > openai_config(top_k=3) <- allows for dynamic variables to be passed in @@ -2187,7 +2238,7 @@ def _complete_aiohttp_openai( acompletion: Final = ctx.acompletion api_base = ctx.api_base api_key = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider extra_headers: Final = ctx.extra_headers headers = ctx.headers @@ -2242,7 +2293,7 @@ def _complete_cometapi(ctx: _CompletionDispatchContext) -> _CompletionDispatchRe acompletion: Final = ctx.acompletion api_base = ctx.api_base api_key = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params @@ -2291,7 +2342,7 @@ def _complete_minimax(ctx: _CompletionDispatchContext) -> _CompletionDispatchRes acompletion: Final = ctx.acompletion api_base = ctx.api_base api_key = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params @@ -2337,7 +2388,7 @@ def _complete_hosted_vllm(ctx: _CompletionDispatchContext) -> _CompletionDispatc acompletion: Final = ctx.acompletion api_base = ctx.api_base api_key: Final = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params @@ -2383,7 +2434,7 @@ def _complete_custom_openai( acompletion: Final = ctx.acompletion api_base = ctx.api_base api_key = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider custom_prompt_dict: Final = ctx.custom_prompt_dict extra_headers = ctx.extra_headers @@ -2392,7 +2443,7 @@ def _complete_custom_openai( logger_fn: Final = ctx.logger_fn logging: Final = ctx.logging messages: Final = ctx.messages - metadata: Final = ctx.metadata + metadata: Final = _dispatch_metadata(ctx) model: Final = ctx.model model_response: Final = ctx.model_response optional_params: Final = ctx.optional_params @@ -2445,7 +2496,7 @@ def _complete_custom_openai( ## LOAD CONFIG - if set config: Final = litellm.OpenAIConfig.get_config() - for k, v in config.items(): + for k, v in _provider_config_items(config): if ( k not in optional_params ): # completion(top_k=3) > openai_config(top_k=3) <- allows for dynamic variables to be passed in @@ -2522,7 +2573,7 @@ def _complete_mistral(ctx: _CompletionDispatchContext) -> _CompletionDispatchRes acompletion: Final = ctx.acompletion api_base = ctx.api_base api_key = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params @@ -2673,7 +2724,7 @@ def _complete_anthropic(ctx: _CompletionDispatchContext) -> _CompletionDispatchR acompletion: Final = ctx.acompletion api_base = ctx.api_base api_key = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider custom_prompt_dict = ctx.custom_prompt_dict headers: Final = ctx.headers @@ -2972,7 +3023,7 @@ def _complete_huggingface(ctx: _CompletionDispatchContext) -> _CompletionDispatc acompletion: Final = ctx.acompletion api_base: Final = ctx.api_base api_key: Final = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params @@ -3015,7 +3066,7 @@ def _complete_oci(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: acompletion: Final = ctx.acompletion api_base: Final = ctx.api_base api_key: Final = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params @@ -3050,7 +3101,7 @@ def _complete_compactifai(ctx: _CompletionDispatchContext) -> _CompletionDispatc acompletion: Final = ctx.acompletion api_base = ctx.api_base api_key = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params @@ -3126,7 +3177,7 @@ def _complete_databricks(ctx: _CompletionDispatchContext) -> _CompletionDispatch acompletion: Final = ctx.acompletion api_base = ctx.api_base api_key = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) headers = ctx.headers litellm_params: Final = ctx.litellm_params logging: Final = ctx.logging @@ -3198,7 +3249,7 @@ def _complete_datarobot(ctx: _CompletionDispatchContext) -> _CompletionDispatchR acompletion: Final = ctx.acompletion api_base: Final = ctx.api_base api_key: Final = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params @@ -3235,7 +3286,7 @@ def _complete_openrouter(ctx: _CompletionDispatchContext) -> _CompletionDispatch acompletion: Final = ctx.acompletion api_base = ctx.api_base api_key = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) headers = ctx.headers litellm_params: Final = ctx.litellm_params logging: Final = ctx.logging @@ -3273,7 +3324,7 @@ def _complete_openrouter(ctx: _CompletionDispatchContext) -> _CompletionDispatch ## Load Config config: Final = litellm.OpenrouterConfig.get_config() - for k, v in config.items(): + for k, v in _provider_config_items(config): if k == "extra_body": # we use openai 'extra_body' to pass openrouter specific params - transforms, route, models if "extra_body" in optional_params: @@ -3314,7 +3365,7 @@ def _complete_vercel_ai_gateway( acompletion: Final = ctx.acompletion api_base = ctx.api_base api_key = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) headers = ctx.headers litellm_params: Final = ctx.litellm_params logging: Final = ctx.logging @@ -3351,7 +3402,7 @@ def _complete_vercel_ai_gateway( ## Load Config config: Final = litellm.VercelAIGatewayConfig.get_config() - for k, v in config.items(): + for k, v in _provider_config_items(config): if k == "extra_body": # we use openai 'extra_body' to pass vercel specific params - providerOptions if "extra_body" in optional_params: @@ -3392,7 +3443,7 @@ def _complete_vertex_ai_beta( acompletion: Final = ctx.acompletion api_base = ctx.api_base api_key: Final = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params @@ -3457,7 +3508,7 @@ def _complete_vertex_ai_beta( def _complete_vertex_ai(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: acompletion: Final = ctx.acompletion api_base = ctx.api_base - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider custom_prompt_dict: Final = ctx.custom_prompt_dict headers: Final = ctx.headers @@ -3754,7 +3805,7 @@ def _complete_text_completion_inception( acompletion: Final = ctx.acompletion api_base = ctx.api_base api_key = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_openai(ctx) headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params logger_fn: Final = ctx.logger_fn @@ -3818,7 +3869,7 @@ def _complete_sagemaker_chat( acompletion: Final = ctx.acompletion api_base: Final = ctx.api_base api_key: Final = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params @@ -3881,7 +3932,7 @@ def _complete_bedrock(ctx: _CompletionDispatchContext) -> _CompletionDispatchRes acompletion: Final = ctx.acompletion api_base: Final = ctx.api_base api_key: Final = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_prompt_dict = ctx.custom_prompt_dict headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params @@ -4005,7 +4056,7 @@ def _complete_watsonx(ctx: _CompletionDispatchContext) -> _CompletionDispatchRes acompletion: Final = ctx.acompletion api_base: Final = ctx.api_base api_key: Final = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_prompt_dict: Final = ctx.custom_prompt_dict headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params @@ -4044,7 +4095,7 @@ def _complete_watsonx_text( acompletion: Final = ctx.acompletion api_base = ctx.api_base api_key = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params logging: Final = ctx.logging @@ -4156,7 +4207,7 @@ def _complete_ollama(ctx: _CompletionDispatchContext) -> _CompletionDispatchResu acompletion: Final = ctx.acompletion api_base = ctx.api_base api_key: Final = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params logging: Final = ctx.logging @@ -4196,7 +4247,7 @@ def _complete_ollama_chat(ctx: _CompletionDispatchContext) -> _CompletionDispatc acompletion: Final = ctx.acompletion api_base = ctx.api_base api_key = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params logging: Final = ctx.logging @@ -4311,7 +4362,7 @@ def _complete_cloudflare(ctx: _CompletionDispatchContext) -> _CompletionDispatch def _complete_petals(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: api_base = ctx.api_base - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) litellm_params: Final = ctx.litellm_params logger_fn: Final = ctx.logger_fn logging: Final = ctx.logging @@ -4353,7 +4404,7 @@ def _complete_snowflake(ctx: _CompletionDispatchContext) -> _CompletionDispatchR acompletion: Final = ctx.acompletion api_base: Final = ctx.api_base api_key: Final = ctx.api_key - client = ctx.client + client = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params @@ -4441,7 +4492,7 @@ def _complete_gdc(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: acompletion: Final = ctx.acompletion api_base = ctx.api_base api_key = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params @@ -4480,7 +4531,7 @@ def _complete_bytez(ctx: _CompletionDispatchContext) -> _CompletionDispatchResul acompletion: Final = ctx.acompletion api_base: Final = ctx.api_base api_key = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params @@ -4520,7 +4571,7 @@ def _complete_lemonade(ctx: _CompletionDispatchContext) -> _CompletionDispatchRe acompletion: Final = ctx.acompletion api_base: Final = ctx.api_base api_key = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params @@ -4560,7 +4611,7 @@ def _complete_ovhcloud(ctx: _CompletionDispatchContext) -> _CompletionDispatchRe acompletion: Final = ctx.acompletion api_base = ctx.api_base api_key = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers: Final = ctx.headers litellm_params: Final = ctx.litellm_params @@ -4603,6 +4654,10 @@ def _complete_ovhcloud(ctx: _CompletionDispatchContext) -> _CompletionDispatchRe return response +def _custom_api_first_output(resp: httpx.Response | None) -> str: + return resp.json()["data"][0]["output"][0] + + def _complete_custom(ctx: _CompletionDispatchContext) -> _CompletionDispatchResult: api_base: Final = ctx.api_base headers: Final = ctx.headers @@ -4651,7 +4706,6 @@ def _complete_custom(ctx: _CompletionDispatchContext) -> _CompletionDispatchResu **kwargs.get("extra_body", {}), }, ) - response_json: Final = resp.json() """ assume all responses from custom api_bases of this format: { @@ -4665,7 +4719,7 @@ def _complete_custom(ctx: _CompletionDispatchContext) -> _CompletionDispatchResu ] } """ - string_response: Final = response_json["data"][0]["output"][0] + string_response: Final = _custom_api_first_output(resp) ## RESPONSE OBJECT model_response.choices[0].message.content = string_response model_response.created = int(time.time()) @@ -4740,7 +4794,7 @@ def _complete_langgraph(ctx: _CompletionDispatchContext) -> _CompletionDispatchR acompletion: Final = ctx.acompletion api_base = ctx.api_base api_key = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers = ctx.headers litellm_params: Final = ctx.litellm_params @@ -4789,7 +4843,7 @@ def _complete_langflow(ctx: _CompletionDispatchContext) -> _CompletionDispatchRe acompletion: Final = ctx.acompletion api_base = ctx.api_base api_key = ctx.api_key - client: Final = ctx.client + client: Final = _dispatch_client_http(ctx) custom_llm_provider: Final = ctx.custom_llm_provider headers = ctx.headers litellm_params: Final = ctx.litellm_params @@ -4947,7 +5001,7 @@ def completion( thinking = validate_and_fix_thinking_param(thinking=thinking) ######### unpacking kwargs ##################### - args: Final = locals() + args: Final = _locals_snapshot(locals()) # Set by the responses->completion fallback so completion() does not bridge # back to the Responses API: that round-trip mutually recurses forever for a @@ -5038,7 +5092,7 @@ def completion( # Inject proxy auth headers if configured if litellm.proxy_auth is not None: try: - proxy_headers: Final = litellm.proxy_auth.get_auth_headers() + proxy_headers: Final = _proxy_auth_headers(litellm.proxy_auth) headers.update(proxy_headers) except Exception as e: verbose_logger.warning("Failed to get proxy auth headers: %s", e) @@ -5091,7 +5145,7 @@ def completion( ) ######## end of unpacking kwargs ########### non_default_params: Final = get_non_default_completion_params(kwargs=kwargs) - litellm_params = {} # used to prevent unbound var errors + litellm_params: dict[str, object] = {} # used to prevent unbound var errors ## PROMPT MANAGEMENT HOOKS ## from litellm.integrations.anthropic_cache_control_hook import ( @@ -5913,7 +5967,7 @@ def embedding( *, aembedding: Literal[True], **kwargs, -) -> Coroutine[Any, Any, EmbeddingResponse]: +) -> Coroutine[object, object, EmbeddingResponse]: ... @@ -5964,7 +6018,7 @@ def embedding( litellm_call_id=None, logger_fn=None, **kwargs, -) -> EmbeddingResponse | Coroutine[Any, Any, EmbeddingResponse]: +) -> EmbeddingResponse | Coroutine[object, object, EmbeddingResponse]: """ Embedding function that calls an API to generate embeddings for the given input. @@ -6007,7 +6061,7 @@ def embedding( # Inject proxy auth headers if configured if litellm.proxy_auth is not None: try: - proxy_headers: Final = litellm.proxy_auth.get_auth_headers() + proxy_headers: Final = _proxy_auth_headers(litellm.proxy_auth) headers.update(proxy_headers) except Exception as e: verbose_logger.warning("Failed to get proxy auth headers: %s", e) @@ -6084,7 +6138,7 @@ def embedding( if mock_response is not None: return mock_embedding(model=model, mock_response=mock_response) try: - response: EmbeddingResponse | Coroutine[Any, Any, EmbeddingResponse] | None = None + response: EmbeddingResponse | Coroutine[object, object, EmbeddingResponse] | None = None if azure is True or custom_llm_provider == "azure": # azure configs @@ -6387,7 +6441,7 @@ def embedding( response = huggingface_embed.embedding( model=model, input=input, - encoding=_get_encoding(), + encoding=sys.modules[__name__].encoding, api_key=api_key, api_base=api_base, logging_obj=logging, @@ -6990,6 +7044,20 @@ def embedding( ###### Text Completion ################ +async def _resolve_dispatched_text_completion_response( + pending: Coroutine[ + object, + object, + TextCompletionResponse | ModelResponse | CustomStreamWrapper | TextCompletionStreamWrapper, + ], +) -> TextCompletionResponse | ModelResponse | CustomStreamWrapper | TextCompletionStreamWrapper: + return await pending + + +async def _resolve_pending_chat_response(pending: Coroutine[object, object, ModelResponse]) -> ModelResponse: + return await pending + + @client async def atext_completion(*args, **kwargs) -> TextCompletionResponse | TextCompletionStreamWrapper: """ @@ -7015,7 +7083,7 @@ async def atext_completion(*args, **kwargs) -> TextCompletionResponse | TextComp else: response = init_response elif asyncio.iscoroutine(init_response): - response = await init_response + response = await _resolve_dispatched_text_completion_response(init_response) else: response = init_response @@ -7040,7 +7108,7 @@ async def atext_completion(*args, **kwargs) -> TextCompletionResponse | TextComp if isinstance(response, TextCompletionResponse): return response elif asyncio.iscoroutine(response): - response = await response + response = await _resolve_pending_chat_response(response) text_completion_response = TextCompletionResponse() text_completion_response = litellm.utils.LiteLLMResponseObjectHandler.convert_chat_to_text_completion( @@ -7330,11 +7398,11 @@ async def aadapter_completion(*, adapter_id: str, **kwargs) -> BaseModel | Adapt async def aadapter_generate_content( **kwargs, -) -> dict[str, Any] | AsyncIterator[bytes]: +) -> dict[str, object] | AsyncIterator[bytes]: from litellm.google_genai.adapters.handler import GenerateContentToCompletionHandler coro: Final = cast( - Coroutine[Any, Any, dict[str, Any] | AsyncIterator[bytes]], + Coroutine[object, object, dict[str, object] | AsyncIterator[bytes]], GenerateContentToCompletionHandler.generate_content_handler(**kwargs, _is_async=True), ) return await coro @@ -7486,7 +7554,7 @@ async def atranscription(*args, **kwargs) -> TranscriptionResponse: # Await normally init_response: Final = await loop.run_in_executor(None, func_with_context) if isinstance(init_response, dict): - response = TranscriptionResponse(**init_response) + response = _transcription_response_from_cached_dict(init_response) elif isinstance(init_response, TranscriptionResponse): ## CACHING SCENARIO response = init_response elif asyncio.iscoroutine(init_response): @@ -7541,7 +7609,7 @@ def transcription( max_retries: int | None = None, custom_llm_provider=None, **kwargs, -) -> TranscriptionResponse | Coroutine[Any, Any, TranscriptionResponse]: +) -> TranscriptionResponse | Coroutine[object, object, TranscriptionResponse]: """ Calls openai + azure whisper endpoints. @@ -7608,7 +7676,7 @@ def transcription( custom_llm_provider=custom_llm_provider, ) - response: TranscriptionResponse | Coroutine[Any, Any, TranscriptionResponse] | None = None + response: TranscriptionResponse | Coroutine[object, object, TranscriptionResponse] | None = None provider_config: Final = ProviderConfigManager.get_provider_audio_transcription_config( model=model, @@ -7842,7 +7910,7 @@ def speech( custom_llm_provider: str | None = None, aspeech: bool | None = None, **kwargs, -) -> HttpxBinaryResponseContent | Coroutine[Any, Any, HttpxBinaryResponseContent]: +) -> HttpxBinaryResponseContent | Coroutine[object, object, HttpxBinaryResponseContent]: user: Final = kwargs.get("user", None) litellm_call_id: Final[str | None] = kwargs.get("litellm_call_id", None) proxy_server_request: Final = kwargs.get("proxy_server_request", None) @@ -7901,7 +7969,7 @@ def speech( }, custom_llm_provider=custom_llm_provider, ) - response: HttpxBinaryResponseContent | Coroutine[Any, Any, HttpxBinaryResponseContent] | None = None + response: HttpxBinaryResponseContent | Coroutine[object, object, HttpxBinaryResponseContent] | None = None if custom_llm_provider == "openai" or custom_llm_provider in litellm.openai_compatible_providers: if voice is None or not (isinstance(voice, str)): raise litellm.BadRequestError( @@ -8663,7 +8731,7 @@ def stream_chunk_builder( ] if len(provider_specific_chunks) > 0: - combined_provider_fields: Final[dict[str, Any]] = {} + combined_provider_fields: Final[dict[str, object]] = {} for chunk in provider_specific_chunks: fields = chunk["choices"][0]["delta"]["provider_specific_fields"] if isinstance(fields, dict): @@ -8728,7 +8796,7 @@ def stream_chunk_builder( async def acount_tokens( model: str, - messages: list[dict[str, Any]] | None = None, + messages: list[dict[str, object]] | None = None, tools: list[dict[str, Any]] | None = None, system: str | None = None, api_key: str | None = None, @@ -8774,7 +8842,7 @@ async def acount_tokens( api_base = dynamic_api_base # Build deployment dict for the token counter - deployment: Final[dict[str, Any]] = { + deployment: Final[dict[str, object]] = { "litellm_params": { "model": model, "api_key": api_key, @@ -8825,29 +8893,37 @@ async def acount_tokens( # Cache for encoding to avoid repeated __getattr__ calls -_encoding_cache: Any | None = None +_encoding_cache: tiktoken.Encoding | None = None -def _get_encoding(): +def _load_module_encoding() -> tiktoken.Encoding: + import sys + + return sys.modules[__name__].encoding + + +def _get_encoding() -> tiktoken.Encoding: """Get encoding, loading it lazily if needed.""" global _encoding_cache if _encoding_cache is None: - import sys - # Access via module to trigger __getattr__ if not cached - _encoding_cache = sys.modules[__name__].encoding + _encoding_cache = _load_module_encoding() return _encoding_cache -def __getattr__(name: str) -> Any: +def _load_default_encoding() -> tiktoken.Encoding: + from litellm._lazy_imports import _get_default_encoding + + return _get_default_encoding() + + +def __getattr__(name: str) -> tiktoken.Encoding: """Lazy import handler for main module""" if name == "encoding": # Use _get_default_encoding which properly sets TIKTOKEN_CACHE_DIR # before loading tiktoken, ensuring the local cache is used # instead of downloading from the internet - from litellm._lazy_imports import _get_default_encoding - - _encoding: Final = _get_default_encoding() + _encoding: Final = _load_default_encoding() # Cache it in the module's __dict__ for subsequent accesses import sys diff --git a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py index 53f378e89e1..c76c933c5b5 100644 --- a/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py +++ b/litellm/proxy/_experimental/mcp_server/oauth2_token_cache.py @@ -145,9 +145,8 @@ class MCPOAuth2TokenCache(InMemoryCache): server.server_id, ) - post_kwargs: Final = {"data": data, **({"headers": token_request.headers} if token_request.headers else {})} try: - response: Final = await client.post(server.token_url, **post_kwargs) + response: Final = await client.post(server.token_url, data=data, headers=token_request.headers or None) response.raise_for_status() except httpx.HTTPStatusError as exc: raise ValueError( diff --git a/litellm/proxy/auth/auth_checks.py b/litellm/proxy/auth/auth_checks.py index d07ac0c5586..b1e444e55d6 100644 --- a/litellm/proxy/auth/auth_checks.py +++ b/litellm/proxy/auth/auth_checks.py @@ -13,8 +13,8 @@ import asyncio import math import re import time -from collections.abc import Sequence -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast +from collections.abc import Iterator, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, cast from fastapi import HTTPException, Request, status from pydantic import BaseModel @@ -87,6 +87,7 @@ from litellm.proxy.utils import PrismaClient, ProxyLogging, log_db_metrics from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.object_permission_repository import ObjectPermissionRepository from litellm.repositories.organization_repository import OrganizationRepository +from litellm.repositories.prisma_protocols import RowT_co from litellm.repositories.project_repository import ProjectRepository from litellm.repositories.table_repositories import ( AccessGroupRepository, @@ -110,11 +111,144 @@ from .auth_utils import get_model_from_request, get_request_route_template if TYPE_CHECKING: from opentelemetry.trace import Span as _Span - Span = _Span | Any + Span = _Span else: Span = Any +class _PrismaDictableRow(Protocol): + def dict(self) -> Mapping[str, object]: ... + + +class _PrismaJWTKeyMappingRow(Protocol): + token: str + + +class _PrismaModelDumpRow(Protocol): + def model_dump(self) -> Mapping[str, object]: ... + + +class _PrismaTeamRow(Protocol): + def dict(self) -> Mapping[str, object]: ... + + def model_dump(self) -> Mapping[str, object]: ... + + +class _PrismaVectorStoreRow(Protocol): + def dict(self) -> Mapping[str, object]: ... + + def model_dump(self) -> Mapping[str, object]: ... + + def __iter__(self) -> Iterator[tuple[str, object]]: ... + + +class _PrismaUserRow(Protocol): + user_id: str + organization_memberships: Sequence[LiteLLM_OrganizationMembershipTable | None] | None + + def __iter__(self) -> Iterator[tuple[str, object]]: ... + + +class _PrismaAuthTable(Protocol[RowT_co]): + async def find_unique( + self, *, where: Mapping[str, object], include: Mapping[str, object] | None = None + ) -> RowT_co | None: ... + + async def find_first( + self, *, where: Mapping[str, object], include: Mapping[str, object] | None = None + ) -> RowT_co | None: ... + + async def find_many( + self, + *, + where: Mapping[str, object], + include: Mapping[str, object] | None = None, + take: int | None = None, + ) -> Sequence[RowT_co]: ... + + async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> RowT_co | None: ... + + async def create(self, *, data: Mapping[str, object], include: Mapping[str, object] | None = None) -> RowT_co: ... + + +class _PrismaTableHolder(Protocol[RowT_co]): + @property + def table(self) -> _PrismaAuthTable[RowT_co]: ... + + +def _dictable_table(repo: _PrismaTableHolder[_PrismaDictableRow]) -> _PrismaAuthTable[_PrismaDictableRow]: + return repo.table + + +def _jwt_key_mapping_table( + repo: _PrismaTableHolder[_PrismaJWTKeyMappingRow], +) -> _PrismaAuthTable[_PrismaJWTKeyMappingRow]: + return repo.table + + +def _model_dump_table(repo: _PrismaTableHolder[_PrismaModelDumpRow]) -> _PrismaAuthTable[_PrismaModelDumpRow]: + return repo.table + + +def _team_table(repo: _PrismaTableHolder[_PrismaTeamRow]) -> _PrismaAuthTable[_PrismaTeamRow]: + return repo.table + + +def _vector_store_table(repo: _PrismaTableHolder[_PrismaVectorStoreRow]) -> _PrismaAuthTable[_PrismaVectorStoreRow]: + return repo.table + + +def _user_table(repo: _PrismaTableHolder[_PrismaUserRow]) -> _PrismaAuthTable[_PrismaUserRow]: + return repo.table + + +def _object_permission_table( + repo: _PrismaTableHolder[LiteLLM_ObjectPermissionTable], +) -> _PrismaAuthTable[LiteLLM_ObjectPermissionTable]: + return repo.table + + +class _PrismaTagRow(Protocol): + tag_name: str + + def dict(self) -> Mapping[str, object]: ... + + +def _tag_table(repo: _PrismaTableHolder[_PrismaTagRow]) -> _PrismaAuthTable[_PrismaTagRow]: + return repo.table + + +class _RawCacheRead(Protocol): + async def async_get_cache(self, *, key: str) -> object: ... + + +def _raw_cache(cache: _RawCacheRead) -> _RawCacheRead: + return cache + + +class _BudgetCacheRead(Protocol): + async def async_get_cache(self, *, key: str) -> "LiteLLM_BudgetTable | Mapping[str, object] | None": ... + + +def _budget_cache(cache: _BudgetCacheRead) -> _BudgetCacheRead: + return cache + + +def _typed_request_body(request_body: dict) -> Mapping[str, object]: + return request_body + + +class _JsonLoadsObj(Protocol): + def __call__(self, data: str) -> object: ... + + +def _typed_json_loads(fn: _JsonLoadsObj) -> _JsonLoadsObj: + return fn + + +_safe_json_loads_obj: Final = _typed_json_loads(safe_json_loads) + + last_db_access_time: Final = LimitedSizeOrderedDict(max_size=100) db_cache_expiry: Final = DEFAULT_IN_MEMORY_TTL # refresh every 5s @@ -384,7 +518,7 @@ _GUARDRAIL_MODIFICATION_KEYS: Final[tuple] = ( ) -def _guardrail_modification_check(request_body: dict, team_object: LiteLLM_TeamTable | None) -> 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 @@ -399,7 +533,7 @@ def _guardrail_modification_check(request_body: dict, team_object: LiteLLM_TeamT """ from litellm.proxy.guardrails.guardrail_helpers import can_modify_guardrails - def _coerce_to_dict(container: Any) -> dict | None: + 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 @@ -411,11 +545,11 @@ def _guardrail_modification_check(request_body: dict, team_object: LiteLLM_TeamT if isinstance(container, dict): return container if isinstance(container, str): - parsed: Final = safe_json_loads(container) + parsed: Final = _safe_json_loads_obj(container) return parsed if isinstance(parsed, dict) else None return None - def _user_requested_modification(container: Any) -> bool: + def _user_requested_modification(container: object) -> bool: coerced: Final = _coerce_to_dict(container) if coerced is None: return False @@ -731,7 +865,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(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) @@ -955,7 +1089,7 @@ async def get_default_end_user_budget( # Fetch from database try: - budget_record: Final = await BudgetRepository(prisma_client).table.find_unique( + budget_record: Final = await _dictable_table(BudgetRepository(prisma_client)).find_unique( where={"budget_id": litellm.max_end_user_budget_id} ) @@ -1007,14 +1141,16 @@ async def get_team_member_default_budget( cache_key: Final = f"team_member_default_budget:{budget_id}" - cached_budget: Final = await user_api_key_cache.async_get_cache(key=cache_key) + cached_budget: Final = await _budget_cache(user_api_key_cache).async_get_cache(key=cache_key) if isinstance(cached_budget, LiteLLM_BudgetTable): return cached_budget if isinstance(cached_budget, dict): return LiteLLM_BudgetTable.model_validate(cached_budget) try: - budget_record: Final = await BudgetRepository(prisma_client).table.find_unique(where={"budget_id": budget_id}) + budget_record: Final = await _dictable_table(BudgetRepository(prisma_client)).find_unique( + where={"budget_id": budget_id} + ) if budget_record is None: verbose_proxy_logger.warning("Team-default member budget not found in database: %s", budget_id) @@ -1171,7 +1307,7 @@ async def get_end_user_object( # Fetch from database try: - response: Final = await EndUserRepository(prisma_client).table.find_unique( + response: Final = await _dictable_table(EndUserRepository(prisma_client)).find_unique( where={"user_id": end_user_id}, include={"litellm_budget_table": True, "object_permission": True}, ) @@ -1243,7 +1379,7 @@ async def resolve_and_validate_end_user_id( return raw_end_user_id cache_key: Final = f"end_user_validation:{raw_end_user_id}" - cached: Final = await user_api_key_cache.async_get_cache(key=cache_key) + cached: Final = await _raw_cache(user_api_key_cache).async_get_cache(key=cache_key) if cached == "valid": return raw_end_user_id if cached == "invalid": @@ -1345,8 +1481,8 @@ async def get_tag_objects_batch( if not tag_names: return {} - tag_objects: Final = {} - uncached_tags: Final = [] + tag_objects: Final = dict[str, LiteLLM_TagTable]() + uncached_tags: Final = list[str]() # Try to get all tags from cache first for tag_name in tag_names: @@ -1363,7 +1499,7 @@ async def get_tag_objects_batch( # Batch fetch uncached tags from DB in one query if uncached_tags: try: - db_tags: Final = await TagRepository(prisma_client).table.find_many( + db_tags: Final = await _tag_table(TagRepository(prisma_client)).find_many( where={"tag_name": {"in": uncached_tags}}, include={"litellm_budget_table": True}, ) @@ -1457,7 +1593,7 @@ async def get_team_membership( # else, check db try: - response: Final = await TeamMembershipRepository(prisma_client).table.find_unique( + response: Final = await _dictable_table(TeamMembershipRepository(prisma_client)).find_unique( where={"user_id_team_id": {"user_id": user_id, "team_id": team_id}}, include={"litellm_budget_table": True}, ) @@ -1524,7 +1660,7 @@ def _should_check_db(key: str, last_db_access_time: LimitedSizeOrderedDict, db_c return False -def _update_last_db_access_time(key: str, value: Any | None, last_db_access_time: LimitedSizeOrderedDict): +def _update_last_db_access_time(key: str, value: object | None, last_db_access_time: LimitedSizeOrderedDict): last_db_access_time[key] = (value, time.time()) @@ -1545,7 +1681,7 @@ def _get_role_based_permissions( for role_based_permission in role_based_permissions: if role_based_permission.role == rbac_role: - return getattr(role_based_permission, key) + return role_based_permission.models if key == "models" else role_based_permission.routes return None @@ -1586,7 +1722,7 @@ async def _get_fuzzy_user_object( prisma_client: PrismaClient, sso_user_id: str | None = None, user_email: str | None = None, -) -> LiteLLM_UserTable | None: +) -> "_PrismaUserRow | None": """ Checks if sso user is in db. @@ -1600,7 +1736,7 @@ async def _get_fuzzy_user_object( response = None if sso_user_id is not None: - response = await UserRepository(prisma_client).table.find_unique( + response = await _user_table(UserRepository(prisma_client)).find_unique( where={"sso_user_id": sso_user_id}, include={"organization_memberships": True}, ) @@ -1608,14 +1744,14 @@ async def _get_fuzzy_user_object( if response is None and user_email is not None: # Use case-insensitive query to handle emails with different casing # This matches the pattern used in _check_duplicate_user_email - response = await UserRepository(prisma_client).table.find_first( + response = await _user_table(UserRepository(prisma_client)).find_first( where={"user_email": {"equals": user_email, "mode": "insensitive"}}, include={"organization_memberships": True}, ) if response is not None and sso_user_id is not None: # update sso_user_id asyncio.create_task( # background task to update user with sso id - UserRepository(prisma_client).table.update( + _user_table(UserRepository(prisma_client)).update( where={"user_id": response.user_id}, data={"sso_user_id": sso_user_id}, ) @@ -1698,7 +1834,7 @@ async def get_user_object( ) if should_check_db: - response = await UserRepository(prisma_client).table.find_unique( + response = await _user_table(UserRepository(prisma_client)).find_unique( where={"user_id": user_id}, include={"organization_memberships": True} ) @@ -1736,7 +1872,7 @@ async def get_user_object( budget_duration=new_user_params["budget_duration"] ) - response = await UserRepository(prisma_client).table.create( + response = await _user_table(UserRepository(prisma_client)).create( data=new_user_params, include={"organization_memberships": True}, ) @@ -1802,7 +1938,7 @@ async def get_user_object( async def _cache_management_object( key: str, - value: BaseModel | dict[str, Any], + value: BaseModel | Mapping[str, object], user_api_key_cache: UserApiKeyCache, proxy_logging_obj: ProxyLogging | None, *, @@ -1916,8 +2052,10 @@ async def _delete_cache_key_object( @log_db_metrics -async def _get_team_db_check(team_id: str, prisma_client: PrismaClient, team_id_upsert: bool | None = None): - response = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) +async def _get_team_db_check( + team_id: str, prisma_client: PrismaClient, team_id_upsert: bool | None = None +) -> "_PrismaTeamRow | None": + response = await _team_table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id}) if response is None and team_id_upsert: from litellm.proxy.management_endpoints.team_endpoints import new_team @@ -1936,8 +2074,8 @@ async def _get_team_db_check(team_id: str, prisma_client: PrismaClient, team_id_ return response -async def _get_team_object_from_db(team_id: str, prisma_client: PrismaClient): - return await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) +async def _get_team_object_from_db(team_id: str, prisma_client: PrismaClient) -> "_PrismaTeamRow | None": + return await _team_table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id}) async def _get_team_object_from_user_api_key_cache( @@ -2148,7 +2286,7 @@ async def get_access_object( # Not in cache - fetch from DB try: - response: Final = await AccessGroupRepository(prisma_client).table.find_unique( + response: Final = await _dictable_table(AccessGroupRepository(prisma_client)).find_unique( where={"access_group_id": access_group_id} ) @@ -2224,7 +2362,7 @@ async def get_team_object_by_alias( # Query database by team_alias try: - teams: Final = await TeamRepository(prisma_client).table.find_many(where={"team_alias": team_alias}) + teams: Final = await _team_table(TeamRepository(prisma_client)).find_many(where={"team_alias": team_alias}) if not teams: raise HTTPException( @@ -2329,7 +2467,9 @@ async def get_org_object_by_alias( # Query database by organization_alias try: - orgs = await OrganizationRepository(prisma_client).table.find_many(where={"organization_alias": org_alias}) + orgs = await _model_dump_table(OrganizationRepository(prisma_client)).find_many( + where={"organization_alias": org_alias} + ) if not orgs: raise HTTPException( @@ -2546,7 +2686,7 @@ async def get_jwt_key_mapping_object( Returns the hashed token (str) if a matching active mapping is found, else None. """ - mapping: Final = await JWTKeyMappingRepository(prisma_client).table.find_first( + mapping: Final = await _jwt_key_mapping_table(JWTKeyMappingRepository(prisma_client)).find_first( where={ "jwt_claim_name": jwt_claim_name, "jwt_claim_value": jwt_claim_value, @@ -2674,7 +2814,7 @@ async def get_object_permission( # else, check db try: - response: Final = await ObjectPermissionRepository(prisma_client).table.find_unique( + response: Final = await _dictable_table(ObjectPermissionRepository(prisma_client)).find_unique( where={"object_permission_id": object_permission_id} ) @@ -2730,7 +2870,7 @@ async def get_managed_vector_store_rows_by_uuids( if not cache_misses: return result - rows: Final = await ManagedVectorStoresRepository(prisma_client).table.find_many( + rows: Final = await _vector_store_table(ManagedVectorStoresRepository(prisma_client)).find_many( where={"vector_store_id": {"in": cache_misses}}, take=len(cache_misses), ) @@ -2804,11 +2944,11 @@ async def get_org_object( return deserialized_org # else, check db try: - query_kwargs: Final[dict[str, Any]] = {"where": {"organization_id": org_id}} + query_kwargs: Final[dict[str, Mapping[str, object]]] = {"where": {"organization_id": org_id}} if include_budget_table: query_kwargs["include"] = {"litellm_budget_table": True} - response: Final = await OrganizationRepository(prisma_client).table.find_unique(**query_kwargs) + response: Final = await _model_dump_table(OrganizationRepository(prisma_client)).find_unique(**query_kwargs) except Exception: # An operational failure (DB down, timeout, cache fault) is NOT the same fact as a confirmed # missing row, and relabelling it as "doesn't exist" made every caller unable to tell them @@ -3763,7 +3903,7 @@ async def _virtual_key_soft_budget_check( ) -def _parse_email_list(raw: Any) -> list[str]: +def _parse_email_list(raw: str | Sequence[object] | None) -> list[str]: """Parse emails from a list or comma-separated string.""" if isinstance(raw, list): return [e.strip() for e in raw if isinstance(e, str) and e.strip()] @@ -3773,7 +3913,7 @@ def _parse_email_list(raw: Any) -> list[str]: def _normalize_alert_emails( - cfg: dict[str, Any] | None, + cfg: Mapping[str, str | Sequence[object] | None] | None, ) -> dict[str, list[str]]: """Coerce user-supplied threshold→recipients mapping to Dict[str, List[str]]. @@ -3786,8 +3926,8 @@ def _normalize_alert_emails( def _merge_budget_alert_email_configs( - global_cfg: dict[str, Any] | None, - per_key_cfg: dict[str, Any] | None, + global_cfg: Mapping[str, str | Sequence[object] | None] | None, + per_key_cfg: Mapping[str, str | Sequence[object] | None] | None, ) -> dict[str, list[str]] | None: """ Per-threshold additive merge: each threshold's recipient list is the union @@ -4294,7 +4434,7 @@ async def get_project_object( return deserialized_project # Fetch from DB - project_row: Final = await ProjectRepository(prisma_client).table.find_unique( + project_row: Final = await _model_dump_table(ProjectRepository(prisma_client)).find_unique( where={"project_id": project_id}, include={"litellm_budget_table": True}, ) @@ -4621,7 +4761,9 @@ async def vector_store_access_check( ######################################################### # Check if the key can access the vector store if valid_token is not None and valid_token.object_permission_id is not None: - key_object_permission: Final = await ObjectPermissionRepository(prisma_client).table.find_unique( + key_object_permission: Final = await _object_permission_table( + ObjectPermissionRepository(prisma_client) + ).find_unique( where={"object_permission_id": valid_token.object_permission_id}, ) if key_object_permission is not None: @@ -4633,7 +4775,9 @@ async def vector_store_access_check( # Check if the team can access the vector store if team_object is not None and team_object.object_permission_id is not None: - team_object_permission: Final = await ObjectPermissionRepository(prisma_client).table.find_unique( + team_object_permission: Final = await _object_permission_table( + ObjectPermissionRepository(prisma_client) + ).find_unique( where={"object_permission_id": team_object.object_permission_id}, ) if team_object_permission is not None: diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 159d7508f4e..8e9a5fe2f00 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -7,7 +7,7 @@ import traceback from collections.abc import AsyncGenerator, Callable, Mapping from datetime import datetime from functools import lru_cache -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload import anyio import httpx @@ -311,8 +311,49 @@ def _stream_usage_tracking_updates( } +def _getattr_object(value: object, name: str, default: object = None) -> object: + return getattr(value, name, default) + + +class _UpstreamHttpResponse(Protocol): + @property + def status_code(self) -> int: ... + + @property + def headers(self) -> httpx.Headers: ... + + async def aread(self) -> bytes: ... + + +def _as_upstream_response(response: _UpstreamHttpResponse) -> _UpstreamHttpResponse: + return response + + +class _ReadsHeaderValues(Protocol): + def get(self, key: str, default: str = "") -> str: ... + + +def _as_header_reader(headers: _ReadsHeaderValues) -> _ReadsHeaderValues: + return headers + + +class _DispatchesSuccessHandlers(Protocol): + async def dispatch_success_handlers( + self, + result: object = None, + start_time: object = None, + end_time: object = None, + cache_hit: object = None, + prefer_async_handlers: bool = False, + ) -> None: ... + + +def _as_success_dispatcher(logging_obj: _DispatchesSuccessHandlers) -> _DispatchesSuccessHandlers: + return logging_obj + + def _serialize_http_exception_detail( - detail: Any, + detail: object, ) -> tuple[str, dict | None]: """ Convert an HTTPException.detail value into (message, structured_fields) @@ -342,7 +383,7 @@ def _serialize_http_exception_detail( return str(detail), None -def _collect_response_file_search_vector_store_ids(data: dict[str, Any]) -> set[str]: +def _collect_response_file_search_vector_store_ids(data: Mapping[str, object]) -> set[str]: vector_store_ids: Final[set[str]] = set() tools: Final = data.get("tools") if not isinstance(tools, list): @@ -369,7 +410,7 @@ def _collect_response_file_search_vector_store_ids(data: dict[str, Any]) -> set[ async def _authorize_response_file_search_vector_stores( - data: dict[str, Any], + data: Mapping[str, object], user_api_key_dict: UserAPIKeyAuth, ) -> None: vector_store_ids: Final = _collect_response_file_search_vector_store_ids(data) @@ -700,7 +741,7 @@ async def create_response( # Preserve status code from HTTPException (e.g., guardrail blocks) error_status: Final = getattr(e, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR) - raw_detail: Final = getattr(e, "detail", "Error processing stream start") + raw_detail: Final = _getattr_object(e, "detail", "Error processing stream start") message, structured_fields = _serialize_http_exception_detail(raw_detail) existing_fields: Final = getattr(e, "provider_specific_fields", None) or {} @@ -711,7 +752,7 @@ async def create_response( # Match ProxyException.to_dict() shape so streaming and non-streaming # error frames are byte-identical. - error_obj: Final[dict[str, Any]] = { + error_obj: Final[dict[str, object]] = { "message": message, "type": getattr(e, "type", "None"), "param": getattr(e, "param", "None"), @@ -777,7 +818,7 @@ def _is_azure_model_router_request(model: str) -> bool: def _override_openai_response_model( *, - response_obj: Any, + response_obj: object, requested_model: str, log_context: str, return_raw_model_name: bool = False, @@ -972,7 +1013,7 @@ def _log_llm_api_exception(e: Exception) -> None: async def _cancel_llm_call_on_client_disconnect( request: Request, - llm_api_call: "asyncio.Future[Any]", + llm_api_call: "asyncio.Future[object]", disconnect_event: asyncio.Event, ) -> None: try: @@ -1023,7 +1064,7 @@ class ProxyBaseLLMRequestProcessing: version: str | None = None, model_region: str | None = None, response_cost: float | str | None = None, - hidden_params: dict | None = None, + hidden_params: Mapping[str, object] | None = None, fastest_response_batch_completion: bool | None = None, request_data: dict | None = {}, timeout: float | httpx.Timeout | None = None, @@ -1115,7 +1156,7 @@ class ProxyBaseLLMRequestProcessing: @staticmethod async def build_litellm_proxy_success_headers_from_llm_response( *, - response: Any, + response: object, request_data: dict, request: Request, user_api_key_dict: UserAPIKeyAuth, @@ -1906,7 +1947,7 @@ class ProxyBaseLLMRequestProcessing: _captured_user_api_key_dict: Final = user_api_key_dict _captured_logging_obj: Final = logging_obj - async def _on_deferred_stream_complete(assembled_response, cache_hit): + async def _on_deferred_stream_complete(assembled_response: object, cache_hit: object) -> None: await ProxyBaseLLMRequestProcessing._run_deferred_stream_guardrails( captured_data=_captured_data, captured_user_api_key_dict=_captured_user_api_key_dict, @@ -2157,7 +2198,7 @@ class ProxyBaseLLMRequestProcessing: @staticmethod async def _record_container_owners_from_responses_if_needed( - response: Any, + response: object, user_api_key_dict: UserAPIKeyAuth, ) -> None: """Register code-interpreter containers so follow-up file APIs pass ownership checks.""" @@ -2180,7 +2221,7 @@ class ProxyBaseLLMRequestProcessing: ) @staticmethod - def _extract_completed_responses_response(stream_response: Any) -> Any: + def _extract_completed_responses_response(stream_response: object) -> object: """Pull the assembled ``ResponsesAPIResponse`` off a streaming iterator. ``ResponsesAPIStreamingIterator`` stores the terminal stream event @@ -2190,17 +2231,17 @@ class ProxyBaseLLMRequestProcessing: ``ResponsesAPIResponse`` directly. Handle both shapes so the container-ownership recording path can walk ``.output`` either way. """ - completed: Final = getattr(stream_response, "completed_response", None) + completed: Final = _getattr_object(stream_response, "completed_response") if completed is None: return None - response_obj: Final = getattr(completed, "response", None) + response_obj: Final = _getattr_object(completed, "response") if response_obj is not None: return response_obj return completed @staticmethod async def _wrap_responses_stream_for_container_ownership( - original_stream_response: Any, + original_stream_response: object, wrapped_generator: Any, user_api_key_dict: UserAPIKeyAuth, ): @@ -2299,12 +2340,13 @@ class ProxyBaseLLMRequestProcessing: if isinstance(result, Response): return result - content: Final = await result.aread() + upstream: Final = _as_upstream_response(result) + content: Final = await upstream.aread() return Response( content=content, - status_code=result.status_code, + status_code=upstream.status_code, headers=HttpPassThroughEndpointHelpers.get_response_headers( - headers=result.headers, + headers=upstream.headers, custom_headers=dict(fastapi_response.headers), ), ) @@ -2435,9 +2477,10 @@ class ProxyBaseLLMRequestProcessing: HttpPassThroughEndpointHelpers, ) + upstream: Final = _as_upstream_response(response) try: - response_status: Final[int] = response.status_code - content_type: Final[str] = response.headers.get("content-type", "") + response_status: Final[int] = upstream.status_code + content_type: Final[str] = _as_header_reader(upstream.headers).get("content-type", "") except AttributeError: return None @@ -2451,20 +2494,20 @@ class ProxyBaseLLMRequestProcessing: return None response_headers: Final = HttpPassThroughEndpointHelpers.get_response_headers( - headers=response.headers, + headers=upstream.headers, custom_headers=custom_headers, ) callback_headers: Final = await proxy_logging_obj.post_call_response_headers_hook( data=self.data, user_api_key_dict=user_api_key_dict, - response=response, + response=upstream, request_headers=request_headers, ) if callback_headers: response_headers.update(callback_headers) if is_event_stream: - body_bytes = await response.aread() + body_bytes = await upstream.aread() modified_bytes: Final = await self._handle_event_stream_allm_passthrough_route( body_bytes=body_bytes, proxy_logging_obj=proxy_logging_obj, @@ -2477,7 +2520,7 @@ class ProxyBaseLLMRequestProcessing: headers=response_headers, ) - body_bytes = await response.aread() + body_bytes = await upstream.aread() try: parsed: Final = _json.loads(body_bytes) except (_json.JSONDecodeError, UnicodeDecodeError): @@ -2566,9 +2609,9 @@ class ProxyBaseLLMRequestProcessing: async def _run_deferred_stream_guardrails( captured_data: dict, captured_user_api_key_dict: "UserAPIKeyAuth", - captured_logging_obj: Any, + captured_logging_obj: LiteLLMLoggingObj, assembled_response: Any, - cache_hit: Any, + cache_hit: object, ) -> None: """ Run non-streaming post-call guardrail hooks on an assembled streaming @@ -2646,7 +2689,7 @@ class ProxyBaseLLMRequestProcessing: # _is_sync_litellm_request (which only recognizes a subset of # async markers stored in litellm_params). asyncio.create_task( - captured_logging_obj.dispatch_success_handlers( + _as_success_dispatcher(captured_logging_obj).dispatch_success_handlers( _response, cache_hit=cache_hit, start_time=None, @@ -2717,7 +2760,7 @@ class ProxyBaseLLMRequestProcessing: headers = getattr(e, "headers", None) or {} if not headers: # Try to get headers from e.response.headers (httpx.Response) - _response: Final = getattr(e, "response", None) + _response: Final = _getattr_object(e, "response") if _response is not None: _response_headers: Final = getattr(_response, "headers", None) if _response_headers: @@ -2749,7 +2792,7 @@ class ProxyBaseLLMRequestProcessing: raise e if isinstance(e, HTTPException): - raw_detail: Final = getattr(e, "detail", str(e)) + raw_detail: Final = _getattr_object(e, "detail", str(e)) message, structured_fields = _serialize_http_exception_detail(raw_detail) existing_fields: Final = getattr(e, "provider_specific_fields", None) or {} if structured_fields: @@ -3042,8 +3085,16 @@ class ProxyBaseLLMRequestProcessing: request=request, ) + @overload @staticmethod - def _process_chunk_with_cost_injection(chunk: Any, model_name: str) -> Any: + def _process_chunk_with_cost_injection(chunk: bytes, model_name: str) -> bytes: ... + + @overload + @staticmethod + def _process_chunk_with_cost_injection(chunk: object, model_name: str) -> object: ... + + @staticmethod + def _process_chunk_with_cost_injection(chunk: object, model_name: str) -> object: """ Process a streaming chunk and inject cost information if enabled. diff --git a/litellm/proxy/db/spend_log_tool_index.py b/litellm/proxy/db/spend_log_tool_index.py index fc605dca257..93bbc567430 100644 --- a/litellm/proxy/db/spend_log_tool_index.py +++ b/litellm/proxy/db/spend_log_tool_index.py @@ -35,7 +35,7 @@ class ToolUsageTransaction: total_tokens: int -def response_tool_call_names(completion_response: Any) -> tuple[str, ...]: +def response_tool_call_names(completion_response: object) -> tuple[str, ...]: """Tool names invoked in a completion response, in call order, for any response surface get_tool_calls_from_response understands (chat completions, Responses API output items, Anthropic Messages tool_use blocks). Reads every choice of @@ -59,7 +59,7 @@ def build_tool_usage_transaction( mcp_namespaced_tool_name: str | None, spend: float, total_tokens: int, - completion_response: Any, + completion_response: object, realtime_tool_calls: Any = None, ) -> ToolUsageTransaction | None: """None when the request invoked no tools. Realtime sessions carry invoked diff --git a/litellm/proxy/hooks/parallel_request_limiter.py b/litellm/proxy/hooks/parallel_request_limiter.py index 79c85571fc9..b313cb64c3f 100644 --- a/litellm/proxy/hooks/parallel_request_limiter.py +++ b/litellm/proxy/hooks/parallel_request_limiter.py @@ -20,6 +20,7 @@ from litellm.proxy.auth.auth_utils import ( from litellm.proxy.auth.budget_throttle import throttled_limit from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitError from litellm.proxy.hooks.rate_limiter_utils import resolve_llm_provider_for_rate_limit +from litellm.types.utils import Usage if TYPE_CHECKING: from opentelemetry.trace import Span as _Span @@ -33,6 +34,13 @@ else: InternalUsageCache = Any +def _response_total_tokens(response_obj: object) -> int: + if not isinstance(response_obj, (ModelResponse, EmbeddingResponse, TextCompletionResponse)): + return 0 + response_usage: Final = getattr(response_obj, "usage", None) + return response_usage.total_tokens if isinstance(response_usage, Usage) else 0 + + class CacheObject(TypedDict): current_global_requests: dict | None request_count_api_key: dict | None @@ -480,7 +488,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): ) # don't block execution for cache updates ) - async def async_log_success_event(self, kwargs, response_obj, start_time, end_time): + async def async_log_success_event(self, kwargs, response_obj: object, start_time, end_time): from litellm.proxy.common_utils.callback_utils import ( get_model_group_from_litellm_kwargs, ) @@ -529,21 +537,18 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): current_minute: Final = datetime.now().strftime("%M") precise_minute: Final = f"{current_date}-{current_hour}-{current_minute}" - total_tokens = 0 - - if isinstance(response_obj, (ModelResponse, EmbeddingResponse, TextCompletionResponse)): - total_tokens = response_obj.usage.total_tokens + total_tokens: int = _response_total_tokens(response_obj) # ------------ # Update usage - API Key # ------------ - values_to_update_in_cache: Final = [] + values_to_update_in_cache: Final[list[tuple[str, object]]] = [] if user_api_key is not None: request_count_api_key = f"{user_api_key}::{precise_minute}::request_count" - current = await self.internal_usage_cache.async_get_cache( + current: dict[str, int] = await self.internal_usage_cache.async_get_cache( key=request_count_api_key, litellm_parent_otel_span=litellm_parent_otel_span, ) or { @@ -606,13 +611,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): # Update usage - User # ------------ if user_api_key_user_id is not None: - total_tokens = 0 - - if isinstance( - response_obj, - (ModelResponse, EmbeddingResponse, TextCompletionResponse), - ): - total_tokens = response_obj.usage.total_tokens + total_tokens = _response_total_tokens(response_obj) request_count_api_key = f"{user_api_key_user_id}::{precise_minute}::request_count" @@ -638,13 +637,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): # Update usage - Team # ------------ if user_api_key_team_id is not None: - total_tokens = 0 - - if isinstance( - response_obj, - (ModelResponse, EmbeddingResponse, TextCompletionResponse), - ): - total_tokens = response_obj.usage.total_tokens + total_tokens = _response_total_tokens(response_obj) request_count_api_key = f"{user_api_key_team_id}::{precise_minute}::request_count" @@ -670,13 +663,7 @@ class _PROXY_MaxParallelRequestsHandler(CustomLogger): # Update usage - End User # ------------ if user_api_key_end_user_id is not None: - total_tokens = 0 - - if isinstance( - response_obj, - (ModelResponse, EmbeddingResponse, TextCompletionResponse), - ): - total_tokens = response_obj.usage.total_tokens + total_tokens = _response_total_tokens(response_obj) request_count_api_key = f"{user_api_key_end_user_id}::{precise_minute}::request_count" diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 38b5d755535..b578fb40ef2 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -189,6 +189,16 @@ class _PrismaTableActions(Protocol[_PrismaRowT]): ) -> _PrismaRowT | None: ... +class _UserRowLike(Protocol): + user_id: str | None + user_email: str | None + user_alias: str | None + + def model_dump(self) -> Mapping[str, object]: ... + + def dict(self) -> Mapping[str, object]: ... + + class _TxTables(Protocol): litellm_proxymodeltable: _PrismaTableActions[object] @@ -4222,7 +4232,7 @@ def _transform_verification_tokens_to_deleted_records( record = deleted_record.model_dump() # Map org_id to organization_id (model uses org_id, but schema expects organization_id) - org_id_value = record.pop("org_id", None) + org_id_value: object = record.pop("org_id", None) if org_id_value is not None: record["organization_id"] = org_id_value @@ -4691,7 +4701,7 @@ async def _execute_virtual_key_regeneration( grace_period=data.grace_period if data else None, ) - updated_token: Final = await VerificationTokenRepository(prisma_client).table.update( + updated_token: Final[Mapping[str, object] | None] = await VerificationTokenRepository(prisma_client).table.update( where={"token": hashed_api_key}, data=jsonified_update_data, ) @@ -5988,7 +5998,9 @@ async def _list_key_helper( created_by_ids: Final = [key.created_by for key in keys if key.created_by] all_ids: Final = list(set(user_ids + created_by_ids)) # Remove duplicates if all_ids: - users: Final = await UserRepository(prisma_client).table.find_many(where={"user_id": {"in": all_ids}}) + users: Final[Sequence[_UserRowLike]] = await UserRepository(prisma_client).table.find_many( + where={"user_id": {"in": all_ids}} + ) user_map = {user.user_id: user for user in users} # Prepare response diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py index dfa58b182b6..64d8b2929b6 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/gemini_passthrough_logging_handler.py @@ -92,7 +92,7 @@ class GeminiPassthroughLoggingHandler: litellm_params={}, api_key="", request_data={}, - encoding=litellm.encoding, + encoding=getattr(litellm, "encoding", None), ) kwargs = GeminiPassthroughLoggingHandler._create_gemini_response_logging_payload_for_generate_content( litellm_model_response=litellm_model_response, diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py index f79b589f6b3..cd6dee3f473 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/openai_passthrough_logging_handler.py @@ -327,7 +327,7 @@ class OpenAIPassthroughLoggingHandler(BasePassthroughLoggingHandler): optional_params=request_body.get("optional_params", {}), api_key="", request_data=request_body, - encoding=litellm.encoding, + encoding=getattr(litellm, "encoding", None), json_mode=request_body.get("response_format", {}).get("type") == "json_object", litellm_params=existing_litellm_params, ) diff --git a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py index 9c5b7dc563e..7dee0e4a364 100644 --- a/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py +++ b/litellm/proxy/pass_through_endpoints/llm_provider_handlers/vertex_passthrough_logging_handler.py @@ -133,7 +133,7 @@ class VertexPassthroughLoggingHandler: litellm_params={}, api_key="", request_data={}, - encoding=litellm.encoding, + encoding=getattr(litellm, "encoding", None), ) kwargs = VertexPassthroughLoggingHandler._create_vertex_response_logging_payload_for_generate_content( litellm_model_response=litellm_model_response, diff --git a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py index 8a526fcd6cb..fc5e0e48dc3 100644 --- a/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/pass_through_endpoints.py @@ -5,10 +5,10 @@ import json import posixpath import traceback from base64 import b64encode -from collections.abc import AsyncGenerator, Mapping +from collections.abc import AsyncGenerator, Callable, Mapping from datetime import datetime from itertools import groupby -from typing import Any, Final, cast +from typing import Any, Final, TypedDict, cast from urllib.parse import urlencode, urlparse import httpx @@ -92,7 +92,7 @@ router: Final = APIRouter() pass_through_endpoint_logging: Final = PassThroughEndpointLogging() # Global registry to track registered pass-through routes and prevent memory leaks -_registered_pass_through_routes: Final[dict[str, dict[str, str | bool | list[str] | dict[str, Any]]]] = {} +_registered_pass_through_routes: Final[dict[str, dict[str, str | bool | list[str] | Mapping[str, object]]]] = {} def get_response_body(response: httpx.Response) -> dict | None: @@ -1128,15 +1128,22 @@ async def pass_through_request( else: # SigV4-signed callers (Bedrock) supply the exact pre-signed bytes; # otherwise httpx encodes the parsed JSON dict as before. - body_kwargs: Final[dict[str, Any]] = ( - {"content": state_raw_body} if state_raw_body is not None else {"json": _parsed_body} - ) - req: Final = async_client.build_request( - request.method, - url, - params=requested_query_params, - headers=headers, - **body_kwargs, + req: Final = ( + async_client.build_request( + request.method, + url, + params=requested_query_params, + headers=headers, + content=state_raw_body, + ) + if state_raw_body is not None + else async_client.build_request( + request.method, + url, + params=requested_query_params, + headers=headers, + json=_parsed_body, + ) ) response = await async_client.send(req, stream=stream) @@ -1584,9 +1591,15 @@ def _update_metadata_with_tags_in_header(request: Request, metadata: dict) -> di return metadata +class _PassThroughRequestEnvelope(TypedDict, total=False): + query_params: Mapping[str, object] | None + custom_body: Mapping[str, object] | None + stream: bool | None + + async def _parse_request_data_by_content_type( request: Request, -) -> tuple[Any | None, Any | None, Any | None, Any | None]: +) -> tuple[object, object, None, bool | None]: """ Parse request data based on content type. @@ -1605,7 +1618,7 @@ async def _parse_request_data_by_content_type( if "application/json" in content_type: # ✅ Handle JSON try: - body = await request.json() + body: _PassThroughRequestEnvelope = await request.json() query_params_data = body.get("query_params") custom_body_data = body.get("custom_body") stream = body.get("stream") @@ -1646,7 +1659,7 @@ async def _parse_request_data_by_content_type( def create_pass_through_route( endpoint, target: str, - custom_headers: Mapping[str, Any] | None = None, + custom_headers: Mapping[str, object] | None = None, _forward_headers: bool | None = False, _merge_query_params: bool | None = False, dependencies: list | None = None, @@ -1656,7 +1669,7 @@ def create_pass_through_route( is_streaming_request: bool | None = False, query_params: dict | None = None, default_query_params: dict | None = None, - guardrails: dict[str, Any] | None = None, + guardrails: dict[str, object] | None = None, config_file_path: str | None = None, timeout: float | None = None, ): @@ -1887,7 +1900,7 @@ async def websocket_passthrough_request( # Initialize tracking variables start_time: Final = datetime.now() - websocket_messages: Final[list[dict[str, Any]]] = [] + websocket_messages: Final[list[dict[str, object]]] = [] litellm_call_id: Final = str(uuid.uuid4()) verbose_proxy_logger.info("WebSocket passthrough (%s): Starting WebSocket connection to %s", endpoint, target) @@ -1980,7 +1993,7 @@ async def websocket_passthrough_request( ) ### CALL HOOKS ### - modify incoming data / reject request before calling the model - websocket_data: dict[str, Any] = {} + websocket_data: dict[str, object] = {} websocket_data = await proxy_logging_obj.pre_call_hook( user_api_key_dict=user_api_key_dict, data=websocket_data, @@ -2009,8 +2022,8 @@ async def websocket_passthrough_request( await upstream_ws.close() break - text_data = message.get("text") - bytes_data = message.get("bytes") + text_data: str | None = message.get("text") + bytes_data: bytes | None = message.get("bytes") if text_data is not None: # Try to extract model from client setup message for Vertex AI Live @@ -2086,7 +2099,7 @@ async def websocket_passthrough_request( # Ensure raw_response is bytes before decoding if isinstance(raw_response, str): raw_response = raw_response.encode("ascii") - setup_response: Final = json.loads(raw_response.decode("ascii")) + setup_response: Final[Mapping[str, object]] = json.loads(raw_response.decode("ascii")) verbose_proxy_logger.debug("Setup response: %s", setup_response) # Extract model and provider from setup response for Vertex AI Live @@ -2129,7 +2142,7 @@ async def websocket_passthrough_request( await websocket.send_bytes(upstream_message) # Parse and collect for cost tracking try: - message_data = json.loads(upstream_message.decode()) + message_data: dict[str, object] = json.loads(upstream_message.decode()) websocket_messages.append(message_data) except (json.JSONDecodeError, UnicodeDecodeError): pass @@ -2315,7 +2328,8 @@ def _should_buffer_passthrough_response(response: httpx.Response) -> bool: """ if response.status_code >= 400: return True - media_type: Final = response.headers.get("content-type", "").split(";")[0].strip().lower() + content_type_header: Final[str] = response.headers.get("content-type", "") + media_type: Final = content_type_header.split(";")[0].strip().lower() return media_type in ("", "application/json") or media_type.endswith("+json") @@ -2368,7 +2382,7 @@ async def _relay_passthrough_response_bytes( ) -def _extract_model_from_vertex_ai_setup(setup_response: dict) -> str | None: +def _extract_model_from_vertex_ai_setup(setup_response: Mapping[str, object]) -> str | None: """ Extract the model name from Vertex AI Live setup response. @@ -2434,7 +2448,7 @@ class SafeRouteAdder: def add_api_route_if_not_exists( app: FastAPI, path: str, - endpoint: Any, + endpoint: Callable[..., object], methods: list[str], dependencies: list | None = None, ) -> bool: @@ -2767,7 +2781,7 @@ def _get_combined_pass_through_endpoints( async def _register_pass_through_endpoint( - endpoint: dict[str, Any] | PassThroughGenericEndpoint, + endpoint: dict[str, object] | PassThroughGenericEndpoint, app: FastAPI, premium_user: bool, visited_endpoints: set[str], @@ -2783,8 +2797,8 @@ async def _register_pass_through_endpoint( endpoint_data["id"] = str(uuid.uuid4()) endpoint_id: Final = cast(str, endpoint_data["id"]) - target: Final = endpoint_data.get("target") - path: Final = endpoint_data.get("path") + target: Final[str | None] = endpoint_data.get("target") + path: Final[str | None] = endpoint_data.get("path") if path is None: raise ValueError("Path is required for pass-through endpoint") @@ -2792,7 +2806,7 @@ async def _register_pass_through_endpoint( forward_headers: Final = endpoint_data.get("forward_headers") merge_query_params: Final = endpoint_data.get("merge_query_params") default_query_params: Final = endpoint_data.get("default_query_params") - auth: Final = endpoint_data.get("auth") + auth: Final[bool | str | None] = endpoint_data.get("auth") dependencies = None auth_enforced: Final = auth is not None and str(auth).lower() == "true" @@ -2951,12 +2965,12 @@ def _get_pass_through_endpoints_from_config() -> list[PassThroughGenericEndpoint if isinstance(endpoint, dict): endpoint_dict = dict(endpoint) endpoint_dict["is_from_config"] = True - returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict)) + returned_endpoints.append(PassThroughGenericEndpoint.model_validate(endpoint_dict)) elif isinstance(endpoint, PassThroughGenericEndpoint): # Create a copy with is_from_config=True endpoint_dict = endpoint.model_dump() endpoint_dict["is_from_config"] = True - returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict)) + returned_endpoints.append(PassThroughGenericEndpoint.model_validate(endpoint_dict)) except ValidationError as e: verbose_proxy_logger.warning( "Skipping malformed pass-through endpoint from config: %s", @@ -2994,11 +3008,11 @@ async def _get_pass_through_endpoints_from_db( if isinstance(endpoint, dict): endpoint_dict = dict(endpoint) endpoint_dict["is_from_config"] = False - returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict)) + returned_endpoints.append(PassThroughGenericEndpoint.model_validate(endpoint_dict)) elif isinstance(endpoint, PassThroughGenericEndpoint): endpoint_dict = endpoint.model_dump() endpoint_dict["is_from_config"] = False - returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict)) + returned_endpoints.append(PassThroughGenericEndpoint.model_validate(endpoint_dict)) else: # Find specific endpoint by ID found_endpoint: Final = _find_endpoint_by_id(pass_through_endpoint_data, endpoint_id) @@ -3009,7 +3023,7 @@ async def _get_pass_through_endpoints_from_db( else dict(found_endpoint) ) endpoint_dict["is_from_config"] = False - returned_endpoints.append(PassThroughGenericEndpoint(**endpoint_dict)) + returned_endpoints.append(PassThroughGenericEndpoint.model_validate(endpoint_dict)) return returned_endpoints @@ -3191,7 +3205,7 @@ async def update_pass_through_endpoints( endpoint_dict.pop("is_from_config", None) # Create updated endpoint object - updated_endpoint: Final = PassThroughGenericEndpoint(**endpoint_dict) + updated_endpoint: Final = PassThroughGenericEndpoint.model_validate(endpoint_dict) # Update the list pass_through_endpoint_data[endpoint_index] = endpoint_dict @@ -3212,9 +3226,10 @@ async def update_pass_through_endpoints( _custom_headers: dict | None = updated_endpoint.headers or {} _custom_headers = await set_env_variables_in_header(custom_headers=_custom_headers) + route_app: Final[FastAPI] = request.app if updated_endpoint.include_subpath: InitPassThroughEndpointHelpers.add_subpath_route( - app=request.app, + app=route_app, path=updated_endpoint.path, target=updated_endpoint.target, custom_headers=_custom_headers, @@ -3231,7 +3246,7 @@ async def update_pass_through_endpoints( ) else: InitPassThroughEndpointHelpers.add_exact_path_route( - app=request.app, + app=route_app, path=updated_endpoint.path, target=updated_endpoint.target, custom_headers=_custom_headers, @@ -3297,15 +3312,16 @@ async def create_pass_through_endpoints( await update_config_general_settings(data=updated_data, user_api_key_dict=user_api_key_dict) # Return the created endpoint with the generated ID - created_endpoint: Final = PassThroughGenericEndpoint(**data_dict) + created_endpoint: Final = PassThroughGenericEndpoint.model_validate(data_dict) # Register the new route _custom_headers: dict | None = created_endpoint.headers or {} _custom_headers = await set_env_variables_in_header(custom_headers=_custom_headers) + route_app: Final[FastAPI] = request.app if created_endpoint.include_subpath: InitPassThroughEndpointHelpers.add_subpath_route( - app=request.app, + app=route_app, path=created_endpoint.path, target=created_endpoint.target, custom_headers=_custom_headers, @@ -3322,7 +3338,7 @@ async def create_pass_through_endpoints( ) else: InitPassThroughEndpointHelpers.add_exact_path_route( - app=request.app, + app=route_app, path=created_endpoint.path, target=created_endpoint.target, custom_headers=_custom_headers, diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index ff1c12d08d7..de320586d54 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -76,17 +76,13 @@ class PassThroughStreamingHandler: PassThroughStreamingHandler._stamp_first_chunk_if_needed(litellm_logging_obj) if endpoint_type == EndpointType.VERTEX_AI: if "streamRawPredict" in url_route or "rawPredict" in url_route: - modified_chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( + chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( chunk, resolved_model_name ) - if modified_chunk is not None: - chunk = modified_chunk else: # EndpointType.ANTHROPIC - modified_chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( + chunk = ProxyBaseLLMRequestProcessing._process_chunk_with_cost_injection( chunk, resolved_model_name ) - if modified_chunk is not None: - chunk = modified_chunk yield chunk except Exception as e: diff --git a/litellm/proxy/pass_through_endpoints/success_handler.py b/litellm/proxy/pass_through_endpoints/success_handler.py index 5b816dc24b3..3dcc8257b82 100644 --- a/litellm/proxy/pass_through_endpoints/success_handler.py +++ b/litellm/proxy/pass_through_endpoints/success_handler.py @@ -122,7 +122,7 @@ class PassThroughEndpointLogging: def normalize_llm_passthrough_logging_payload( self, httpx_response: httpx.Response, - response_body: dict | None, + response_body: dict | list[dict[str, object]] | None, request_body: dict, logging_obj: LiteLLMLoggingObj, url_route: str, @@ -142,7 +142,7 @@ class PassThroughEndpointLogging: if self.is_gemini_route(url_route, custom_llm_provider): gemini_passthrough_logging_handler_result = GeminiPassthroughLoggingHandler.gemini_passthrough_handler( httpx_response=httpx_response, - response_body=response_body or {}, + response_body=response_body if isinstance(response_body, dict) else {}, logging_obj=logging_obj, url_route=url_route, result=result, @@ -172,7 +172,7 @@ class PassThroughEndpointLogging: anthropic_passthrough_logging_handler_result: Final = ( AnthropicPassthroughLoggingHandler.anthropic_passthrough_handler( httpx_response=httpx_response, - response_body=response_body or {}, + response_body=response_body if isinstance(response_body, dict) else {}, logging_obj=logging_obj, url_route=url_route, result=result, @@ -189,7 +189,7 @@ class PassThroughEndpointLogging: elif self.is_cohere_route(url_route): cohere_passthrough_logging_handler_result = cohere_passthrough_logging_handler.cohere_passthrough_handler( httpx_response=httpx_response, - response_body=response_body or {}, + response_body=response_body if isinstance(response_body, dict) else {}, logging_obj=logging_obj, url_route=url_route, result=result, @@ -208,7 +208,7 @@ class PassThroughEndpointLogging: openai_passthrough_logging_handler_result = OpenAIPassthroughLoggingHandler.openai_passthrough_handler( httpx_response=httpx_response, - response_body=response_body or {}, + response_body=response_body if isinstance(response_body, dict) else {}, logging_obj=logging_obj, url_route=url_route, result=result, @@ -224,7 +224,7 @@ class PassThroughEndpointLogging: elif self.is_cursor_route(url_route, custom_llm_provider): cursor_passthrough_logging_handler_result = CursorPassthroughLoggingHandler.cursor_passthrough_handler( httpx_response=httpx_response, - response_body=response_body or {}, + response_body=response_body if isinstance(response_body, dict) else {}, logging_obj=logging_obj, url_route=url_route, result=result, @@ -266,7 +266,7 @@ class PassThroughEndpointLogging: async def pass_through_async_success_handler( self, httpx_response: httpx.Response, - response_body: dict | None, + response_body: dict | list[dict[str, object]] | None, logging_obj: LiteLLMLoggingObj, url_route: str, result: str, @@ -285,7 +285,7 @@ class PassThroughEndpointLogging: return self.assemblyai_passthrough_logging_handler.assemblyai_passthrough_logging_handler( httpx_response=httpx_response, - response_body=response_body or {}, + response_body=response_body if isinstance(response_body, dict) else {}, logging_obj=logging_obj, url_route=url_route, result=result, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f79819c76d3..91526ad70d0 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -9,13 +9,14 @@ import random import re import secrets import shutil +import socket import subprocess import sys import threading import time import traceback import warnings -from collections.abc import AsyncGenerator, AsyncIterator, Callable, Mapping +from collections.abc import AsyncGenerator, AsyncIterator, Callable, Mapping, MutableMapping, Sequence from datetime import datetime, timedelta, timezone from types import MappingProxyType, UnionType from typing import ( @@ -25,6 +26,8 @@ from typing import ( Literal, NamedTuple, Optional, + Protocol, + TypeAlias, TypedDict, Union, cast, @@ -128,6 +131,7 @@ from litellm.utils import ( if TYPE_CHECKING: from aiohttp import ClientSession + from fastapi.routing import APIRoute from opentelemetry.trace import Span as _Span from litellm.integrations.opentelemetry import OpenTelemetry @@ -137,7 +141,7 @@ else: Span = Any OpenTelemetry = Any -REALTIME_REQUEST_SCOPE_TEMPLATE: Final[dict[str, Any]] = { +REALTIME_REQUEST_SCOPE_TEMPLATE: Final[dict[str, object]] = { "type": "http", "method": "POST", "path": "/v1/realtime", @@ -255,6 +259,7 @@ from litellm.litellm_core_utils.sensitive_data_masker import ( SensitiveDataMasker, mask_sensitive_keys, ) +from litellm.litellm_core_utils.streaming_handler import validated_stream_logging_obj from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.proxy._lazy_features import attach_lazy_features @@ -596,6 +601,7 @@ from litellm.proxy.utils import ( update_spend, ) from litellm.proxy.video_endpoints.endpoints import router as video_router +from litellm.repositories.base_repository import SupportsModelDump from litellm.repositories.credentials_repository import CredentialsRepository from litellm.router import ( AssistantsTypedDict, @@ -870,6 +876,18 @@ async def proxy_shutdown_event(): cleanup_router_config_variables() +_AiohttpAddrInfo: TypeAlias = tuple[int | socket.AddressFamily, int | socket.SocketKind, int, str, tuple[object, ...]] + + +class _AiohttpConnectorKwargs(TypedDict, total=False): + keepalive_timeout: float + ttl_dns_cache: int + enable_cleanup_closed: bool + limit: int + limit_per_host: int + socket_factory: Callable[[_AiohttpAddrInfo], socket.socket] + + async def _initialize_shared_aiohttp_session(): """Initialize shared aiohttp session for connection reuse with connection limits.""" try: @@ -879,7 +897,7 @@ async def _initialize_shared_aiohttp_session(): _build_aiohttp_keepalive_socket_factory, ) - connector_kwargs: Final[dict[str, Any]] = { + connector_kwargs: Final[_AiohttpConnectorKwargs] = { "keepalive_timeout": AIOHTTP_KEEPALIVE_TIMEOUT, "ttl_dns_cache": AIOHTTP_TTL_DNS_CACHE, } @@ -1234,7 +1252,7 @@ async def proxy_startup_event(app: FastAPI): await proxy_shutdown_event() -def _generate_stable_operation_id(route: Any) -> str: +def _generate_stable_operation_id(route: "APIRoute") -> str: operation_id = re.sub(r"\W", "_", f"{route.name}{route.path_format}") route_methods: Final = sorted(route.methods or []) if len(route_methods) == 1: @@ -1493,7 +1511,7 @@ async def openai_exception_handler(request: Request, exc: ProxyException): def _close_dangling_otel_server_span(request: Request, status_code: int, exc: Exception | None = None) -> None: - parent_otel_span: Final = getattr(request.state, "parent_otel_span", None) + parent_otel_span: Final[_Span | None] = getattr(request.state, "parent_otel_span", None) if parent_otel_span is None: return if open_telemetry_logger is None: @@ -1536,17 +1554,80 @@ async def management_problem_exception_handler(request: Request, exc: Management return problem_response(exc.problem) +class _ConfigParamRow(Protocol): + param_name: str + param_value: Mapping[str, JsonValue] | None + + +class _ConfigOverridesRow(Protocol): + config_value: Mapping[str, JsonValue] | None + + +class _SSOConfigRow(Protocol): + sso_settings: MutableMapping[str, object] + + +class _UISettingsRow(Protocol): + ui_settings: Mapping[str, object] | str | None + + +class _InvitationLinkRow(Protocol): + user_id: str + expires_at: datetime + is_accepted: bool + accepted_at: datetime | None + created_by: str + + +class _UserTableRow(Protocol): + user_id: str + user_email: str | None + user_role: str + + +class _ModelTableRow(Protocol): + model_id: str | None + created_by: str | None + + +class _TTFTRow(TypedDict): + api_base: str + model: str + time_to_first_token: float + request_id: str + day: str + + +class _LatencyRow(TypedDict): + api_base: str | None + model: str + day: str + avg_latency_per_token: float + + +class _ExceptionRow(TypedDict, total=False): + combined_model_api_base: str + total_exceptions: int + exception_counts: Mapping[str, int] + + +class _ValidationErrorDetail(TypedDict): + loc: tuple[int | str, ...] + msg: str + + @app.exception_handler(RequestValidationError) async def otel_request_validation_exception_handler(request: Request, exc: RequestValidationError): if request.url.path.startswith(MANAGEMENT_V1_PREFIX): _close_dangling_otel_server_span(request, 400, exc=exc) + validation_errors: Final[Sequence[_ValidationErrorDetail]] = exc.errors() return problem_response( ProblemDetail( type=f"{PROBLEM_TYPE_BASE}invalid-query-parameter", title="Invalid query parameter", status=400, detail="; ".join( - f"{'.'.join(str(part) for part in error['loc'][1:])}: {error['msg']}" for error in exc.errors() + f"{'.'.join(str(part) for part in error['loc'][1:])}: {error['msg']}" for error in validation_errors ) or "The request query parameters are invalid.", ) @@ -2150,13 +2231,14 @@ db_writer_client: AsyncHTTPHandler | None = None ### logger ### -def _resolve_typed_dict_type(typ): +def _resolve_typed_dict_type(typ: object): """Resolve the actual TypedDict class from a potentially wrapped type.""" from typing_extensions import _TypedDictMeta - origin: Final = get_origin(typ) + origin: Final[object] = get_origin(typ) if origin is Union or origin is UnionType: # Check if it's a Union (like Optional) - for arg in get_args(typ): + union_args: Final[tuple[object, ...]] = get_args(typ) + for arg in union_args: if isinstance(arg, _TypedDictMeta): return arg elif isinstance(typ, type) and isinstance(typ, dict): @@ -2164,12 +2246,13 @@ def _resolve_typed_dict_type(typ): return None -def _resolve_pydantic_type(typ) -> list: +def _resolve_pydantic_type(typ: object) -> list: """Resolve the actual TypedDict class from a potentially wrapped type.""" - origin: Final = get_origin(typ) + origin: Final[object] = get_origin(typ) typs: Final = [] if origin is Union or origin is UnionType: # Check if it's a Union (like Optional) - for arg in get_args(typ): + union_args: Final[tuple[object, ...]] = get_args(typ) + for arg in union_args: if arg is not None and "NoneType" not in str(arg): typs.append(arg) elif isinstance(typ, type) and isinstance(typ, BaseModel): @@ -2499,7 +2582,7 @@ async def increment_spend_counters( increment=cost, ) - key_obj: Final = await user_api_key_cache.async_get_cache(key=hashed_token) + key_obj: Final[object] = await user_api_key_cache.async_get_cache(key=hashed_token) if key_obj is None: return key_budget_limits = getattr(key_obj, "budget_limits", None) or ( @@ -2530,7 +2613,7 @@ async def increment_spend_counters( increment=cost, ) - team_obj: Final = await user_api_key_cache.async_get_cache(key=f"team_id:{scope_team_id}") + team_obj: Final[object] = await user_api_key_cache.async_get_cache(key=f"team_id:{scope_team_id}") if team_obj is None: return team_budget_limits = getattr(team_obj, "budget_limits", None) or ( @@ -2824,7 +2907,7 @@ async def _ensure_window_spend_counter_initialized( async def _is_spend_counter_cache_warm(counter_key: str) -> bool: if spend_counter_cache.redis_cache is not None: try: - current_value: Final = await spend_counter_cache.redis_cache.async_get_cache( + current_value: Final[object] = await spend_counter_cache.redis_cache.async_get_cache( key=counter_key, ) if current_value is None: @@ -2896,7 +2979,7 @@ async def update_cache( Put any alerting logic in here. """ - values_to_update_in_cache: Final[list[tuple[Any, Any]]] = [] + values_to_update_in_cache: Final[list[tuple[str, object]]] = [] ### UPDATE KEY SPEND ### async def _update_key_cache(token: str, response_cost: float): @@ -4108,7 +4191,9 @@ class ProxyConfig: if prisma_client is None or not (general_settings.get("store_model_in_db", False) is True or store_model_in_db): return - row = await ConfigRepository(prisma_client).table.find_first(where={"param_name": "environment_variables"}) + row: Final[_ConfigParamRow | None] = await ConfigRepository(prisma_client).table.find_first( + where={"param_name": "environment_variables"} + ) existing: Final[dict] = dict(row.param_value) if row is not None and row.param_value is not None else {} to_set: Final = {k: v for k, v in updates.items() if v is not None} @@ -5906,7 +5991,7 @@ class ProxyConfig: 4. Update router settings """ if llm_router is not None and prisma_client is not None: - db_router_settings: Final = await ConfigRepository(prisma_client).table.find_first( + db_router_settings: Final[_ConfigParamRow | None] = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "router_settings"} ) @@ -6391,7 +6476,9 @@ class ProxyConfig: new_models=new_models, proxy_logging_obj=proxy_logging_obj ) - db_general_settings: Final = await get_config_param(prisma_client, "general_settings") + db_general_settings: Final[_ConfigParamRow | None] = await get_config_param( + prisma_client, "general_settings" + ) # update general settings if db_general_settings is not None: @@ -6587,7 +6674,7 @@ class ProxyConfig: """ try: - sso_settings: Final = await call_with_db_reconnect_retry( + sso_settings: Final[_SSOConfigRow | None] = await call_with_db_reconnect_retry( prisma_client, lambda: SSOConfigRepository(prisma_client).table.find_unique(where={"id": "sso_config"}), reason="init_sso_settings_in_db_lookup_failure", @@ -6618,7 +6705,7 @@ class ProxyConfig: ) try: - db_record: Final = await call_with_db_reconnect_retry( + db_record: Final[_ConfigOverridesRow | None] = await call_with_db_reconnect_retry( prisma_client, lambda: ConfigOverridesRepository(prisma_client).table.find_unique( where={"config_type": "hashicorp_vault"} @@ -6836,7 +6923,7 @@ class ProxyConfig: from litellm.types.prompts.init_prompts import PromptSpec try: - prompts_in_db: Final = await PromptRepository(prisma_client).table.find_many() + prompts_in_db: Final[Sequence[object]] = await PromptRepository(prisma_client).table.find_many() for prompt in prompts_in_db: # Convert DB object to dict and create versioned prompt_id prompt_spec = self._get_prompt_spec_for_db_prompt(db_prompt=prompt) @@ -8448,7 +8535,9 @@ class ProxyStartupEvent: if prisma_client is None: return - db_record: Final = await UISettingsRepository(prisma_client).table.find_unique(where={"id": "ui_settings"}) + db_record: Final[_UISettingsRow | None] = await UISettingsRepository(prisma_client).table.find_unique( + where={"id": "ui_settings"} + ) if db_record and db_record.ui_settings: raw: Final = db_record.ui_settings ui_settings: Final = json.loads(raw) if isinstance(raw, str) else dict(raw) @@ -8597,7 +8686,7 @@ class ProxyStartupEvent: # but YAML config has False. if store_model_in_db is not True and prisma_client is not None: try: - _db_gs_record: Final = await ConfigRepository(prisma_client).table.find_first( + _db_gs_record: Final[_ConfigParamRow | None] = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "general_settings"} ) if _db_gs_record is not None and isinstance(_db_gs_record.param_value, dict): @@ -9660,7 +9749,7 @@ async def chat_completion( completion_stream=_iterator, model=e.model, custom_llm_provider="cached_response", - logging_obj=_logging_obj, + logging_obj=validated_stream_logging_obj(_logging_obj), ) selected_data_generator = select_data_generator( response=_streaming_response, @@ -9677,6 +9766,7 @@ async def chat_completion( return _chat_response except RejectedRequestError as e: _data = e.request_data + _rejected_request_data: Final[dict[str, object]] = e.request_data await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, @@ -9691,7 +9781,7 @@ async def chat_completion( completion_stream=_iterator, model=data.get("model", ""), custom_llm_provider="cached_response", - logging_obj=_data.get("litellm_logging_obj", None), + logging_obj=validated_stream_logging_obj(_rejected_request_data.get("litellm_logging_obj")), ) selected_data_generator = select_data_generator( response=_streaming_response, @@ -10441,7 +10531,7 @@ async def vertex_ai_live_passthrough_endpoint( None, description="Override the Vertex AI region (for example, 'us-central1').", ), - user_api_key_dict=Depends(user_api_key_auth_websocket), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth_websocket), ): """ Vertex AI Live API WebSocket Pass-through Endpoint @@ -10489,7 +10579,7 @@ async def realtime_websocket_endpoint( None, description="Comma-separated list of guardrail names to apply to this request.", ), - user_api_key_dict=Depends(user_api_key_auth_websocket), + user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth_websocket), ): requested_protocols: Final = [ p.strip() for p in (websocket.headers.get("sec-websocket-protocol") or "").split(",") if p.strip() @@ -10521,7 +10611,7 @@ async def realtime_websocket_endpoint( # Only use explicit parameters, not all query params query_params: Final = cast(RealtimeQueryParams, dict(_realtime_query_params_template(model, intent))) - data: dict[str, Any] = { + data: dict[str, object] = { "model": route_model, "websocket": websocket, "query_params": query_params, # Only explicit params @@ -11654,7 +11744,7 @@ async def _check_if_model_is_user_added( id = model.get("model_info", {}).get("id", None) if id is None: continue - db_model = await ModelRepository(prisma_client).table.find_unique(where={"model_id": id}) + db_model: _ModelTableRow | None = await ModelRepository(prisma_client).table.find_unique(where={"model_id": id}) if db_model is not None: if db_model.created_by == user_api_key_dict.user_id: filtered_models.append(model) @@ -11838,7 +11928,7 @@ async def get_all_team_models( team_db_objects_typed: list[LiteLLM_TeamTable] = [] if user_teams == "*": - team_db_objects = await TeamRepository(prisma_client).table.find_many() + team_db_objects: Sequence[SupportsModelDump] = await TeamRepository(prisma_client).table.find_many() team_db_objects_typed = [ LiteLLM_TeamTable.model_validate(team_db_object.model_dump()) for team_db_object in team_db_objects ] @@ -11917,7 +12007,7 @@ async def _populate_team_access_on_models( user_teams = "*" direct_access_models = llm_router.get_model_ids(exclude_team_models=True) # has access to all models elif user_api_key_dict.user_id is not None: - user_db_object: Final = await UserRepository(prisma_client).table.find_unique( + user_db_object: Final[SupportsModelDump | None] = await UserRepository(prisma_client).table.find_unique( where={"user_id": user_api_key_dict.user_id} ) if user_db_object is not None: @@ -12473,7 +12563,9 @@ def _team_models_resolve_to_names(team_models: list[str], access_groups: dict[st async def _load_team_object_for_model_filter(team_id: str, prisma_client: PrismaClient) -> LiteLLM_TeamTable | None: """Load team row from DB; returns None if missing or on error.""" try: - team_db_object: Final = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id}) + team_db_object: Final[SupportsModelDump | None] = await TeamRepository(prisma_client).table.find_unique( + where={"team_id": team_id} + ) if team_db_object is None: verbose_proxy_logger.warning("Team %s not found in database", team_id) return None @@ -12522,7 +12614,7 @@ async def _gather_team_accessible_model_ids( try: if team_object.models and SpecialModelNames.all_proxy_models.value not in team_object.models: _resolved_names: Final = _team_models_resolve_to_names(team_object.models, access_groups) - db_models: Final = await ModelRepository(prisma_client).table.find_many( + db_models: Final[Sequence[_ModelTableRow]] = await ModelRepository(prisma_client).table.find_many( where={"model_name": {"in": _resolved_names}} ) for db_model in db_models: @@ -12984,7 +13076,9 @@ async def model_streaming_metrics( """ _all_api_bases: Final = set() - db_response: Final = await prisma_client.db.query_raw(sql_query, _selected_model_group, startTime, endTime) + db_response: Final[Sequence[_TTFTRow] | None] = await prisma_client.db.query_raw( + sql_query, _selected_model_group, startTime, endTime + ) _daily_entries: dict = {} # {"Jun 23": {"model1": 0.002, "model2": 0.003}} if db_response is not None: for model_data in db_response: @@ -13106,7 +13200,7 @@ async def model_metrics( avg_latency_per_token DESC; """ _all_api_bases: Final = set() - db_response: Final = await prisma_client.db.query_raw( + db_response: Final[Sequence[_LatencyRow] | None] = await prisma_client.db.query_raw( sql_query, _selected_model_group, startTime, endTime, api_key, customer ) _daily_entries: dict = {} # {"Jun 23": {"model1": 0.002, "model2": 0.003}} @@ -13297,7 +13391,9 @@ async def model_metrics_exceptions( ORDER BY total_exceptions DESC LIMIT 200; """ - db_response: Final = await prisma_client.db.query_raw(sql_query, startTime, endTime, _selected_model_group, api_key) + db_response: Final[Sequence[_ExceptionRow] | None] = await prisma_client.db.query_raw( + sql_query, startTime, endTime, _selected_model_group, api_key + ) response: Final[list[dict]] = [] exception_types: Final = set() @@ -14459,7 +14555,9 @@ async def onboarding(invite_link: str, request: Request): detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - invite_obj: Final = await InvitationLinkRepository(prisma_client).table.find_unique(where={"id": invite_link}) + invite_obj: Final[_InvitationLinkRow | None] = await InvitationLinkRepository(prisma_client).table.find_unique( + where={"id": invite_link} + ) if invite_obj is None: raise HTTPException(status_code=401, detail={"error": "Invitation link does not exist in db."}) #### CHECK IF EXPIRED @@ -14477,7 +14575,9 @@ async def onboarding(invite_link: str, request: Request): ) ### GET USER OBJECT ### - user_obj: Final = await UserRepository(prisma_client).table.find_unique(where={"user_id": invite_obj.user_id}) + user_obj: Final[_UserTableRow | None] = await UserRepository(prisma_client).table.find_unique( + where={"user_id": invite_obj.user_id} + ) if user_obj is None: raise HTTPException(status_code=401, detail={"error": "User does not exist in db."}) @@ -14649,7 +14749,9 @@ async def claim_onboarding_link(data: InvitationClaim, request: Request): detail={"error": CommonProxyErrors.db_not_connected_error.value}, ) - invite_obj = await InvitationLinkRepository(prisma_client).table.find_unique(where={"id": data.invitation_link}) + invite_obj: Final[_InvitationLinkRow | None] = await InvitationLinkRepository(prisma_client).table.find_unique( + where={"id": data.invitation_link} + ) if invite_obj is None: raise HTTPException(status_code=401, detail={"error": "Invitation link does not exist in db."}) #### CHECK IF EXPIRED @@ -14704,7 +14806,7 @@ async def claim_onboarding_link(data: InvitationClaim, request: Request): ) ### UPDATE USER OBJECT ### - user_obj: Final = await tx.litellm_usertable.update( + user_obj: Final[_UserTableRow | None] = await tx.litellm_usertable.update( where={"user_id": invite_obj.user_id}, data={"password": hashed_pw} ) @@ -14930,7 +15032,7 @@ async def new_invitation(data: InvitationNew, user_api_key_dict: UserAPIKeyAuth detail={"error": "You can only create invitations for users in your organization or team."}, ) - response: Final = await create_invitation_for_user( + response: Final[object] = await create_invitation_for_user( data=data, user_api_key_dict=user_api_key_dict, ) @@ -14972,7 +15074,9 @@ async def invitation_info(invitation_id: str, user_api_key_dict: UserAPIKeyAuth detail={"error": f"{CommonProxyErrors.not_allowed_access.value}, your role={user_api_key_dict.user_role}"}, ) - response: Final = await InvitationLinkRepository(prisma_client).table.find_unique(where={"id": invitation_id}) + response: Final[object] = await InvitationLinkRepository(prisma_client).table.find_unique( + where={"id": invitation_id} + ) if response is None: raise HTTPException( @@ -15020,7 +15124,7 @@ async def invitation_update( ) current_time: Final = litellm.utils.get_utc_datetime() - response: Final = await InvitationLinkRepository(prisma_client).table.update( + response: Final[object] = await InvitationLinkRepository(prisma_client).table.update( where={"id": data.invitation_id}, data={ "id": data.invitation_id, @@ -15086,7 +15190,9 @@ async def invitation_delete( # Org admins can only delete invitations they created if is_other_admin and not is_proxy_admin: - invitation = await InvitationLinkRepository(prisma_client).table.find_unique(where={"id": data.invitation_id}) + invitation: Final[_InvitationLinkRow | None] = await InvitationLinkRepository(prisma_client).table.find_unique( + where={"id": data.invitation_id} + ) if invitation is None: raise HTTPException( status_code=400, @@ -15098,7 +15204,9 @@ async def invitation_delete( detail={"error": "Organization admins can only delete invitations they created."}, ) - response: Final = await InvitationLinkRepository(prisma_client).table.delete(where={"id": data.invitation_id}) + response: Final[object] = await InvitationLinkRepository(prisma_client).table.delete( + where={"id": data.invitation_id} + ) if response is None: raise HTTPException( @@ -15136,7 +15244,9 @@ async def update_config( raise Exception("No DB Connected") async def _read_section(param_name: str) -> dict: - row: Final = await ConfigRepository(prisma_client).table.find_first(where={"param_name": param_name}) + row: Final[_ConfigParamRow | None] = await ConfigRepository(prisma_client).table.find_first( + where={"param_name": param_name} + ) if row is None or row.param_value is None: return {} return dict(row.param_value) @@ -15159,7 +15269,7 @@ async def update_config( if config_info.general_settings is not None: existing = await _read_section("general_settings") before_general_settings: Final = copy.deepcopy(existing) - updates = config_info.general_settings.dict(exclude_none=True) + updates: Mapping[str, JsonValue] = config_info.general_settings.dict(exclude_none=True) for k, v in updates.items(): if k == "alert_to_webhook_url": if "alerting" not in existing: @@ -15599,7 +15709,7 @@ async def get_config_general_settings( ) ## get general settings from db - db_general_settings: Final = await ConfigRepository(prisma_client).table.find_first( + db_general_settings: Final[_ConfigParamRow | None] = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "general_settings"} ) ### pop the value @@ -15788,12 +15898,12 @@ async def get_config_list( is_full_admin: Final = user_api_key_dict.user_role == LitellmUserRoles.PROXY_ADMIN ## get general settings from db - db_general_settings: Final = await ConfigRepository(prisma_client).table.find_first( + db_general_settings: Final[_ConfigParamRow | None] = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "general_settings"} ) if db_general_settings is not None and db_general_settings.param_value is not None: - db_general_settings_dict = dict(db_general_settings.param_value) + db_general_settings_dict: Mapping[str, JsonValue] = dict(db_general_settings.param_value) else: db_general_settings_dict = {} @@ -15884,7 +15994,7 @@ async def get_config_list( ) return_val.append(_response_obj) - db_litellm_settings_row: Final = await ConfigRepository(prisma_client).table.find_first( + db_litellm_settings_row: Final[_ConfigParamRow | None] = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "litellm_settings"} ) db_litellm_settings: Final[dict] = ( @@ -15961,7 +16071,7 @@ async def delete_config_general_settings( ) ## get general settings from db - db_general_settings: Final = await ConfigRepository(prisma_client).table.find_first( + db_general_settings: Final[_ConfigParamRow | None] = await ConfigRepository(prisma_client).table.find_first( where={"param_name": "general_settings"} ) ### pop the value @@ -16528,7 +16638,7 @@ async def reload_anthropic_beta_headers( last_anthropic_beta_headers_reload = current_time.isoformat() # Set force reload flag in database for other pods, preserving existing interval_hours - existing_beta_config: Final = await ConfigRepository(prisma_client).table.find_unique( + existing_beta_config: Final[_ConfigParamRow | None] = await ConfigRepository(prisma_client).table.find_unique( where={"param_name": "anthropic_beta_headers_reload_config"} ) existing_beta_interval = None @@ -17116,7 +17226,7 @@ async def _is_mcp_access_group_cached(name: str) -> bool: ) cache_key: Final = f"mcp_access_group_exists:{name}" - cached: Final = await user_api_key_cache.async_get_cache(key=cache_key) + cached: Final[object] = await user_api_key_cache.async_get_cache(key=cache_key) if cached is not None: return bool(cached) result: Final = bool(await MCPRequestHandler._get_mcp_servers_from_access_groups([name])) diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 3e5a9f2fb3b..ee430a41b42 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -497,7 +497,10 @@ async def cursor_chat_completions( from litellm.completion_extras.litellm_responses_transformation.handler import ( responses_api_bridge, ) - from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.litellm_core_utils.streaming_handler import ( + CustomStreamWrapper, + validated_stream_logging_obj, + ) from litellm.proxy.proxy_server import ( async_data_generator, chat_completion, @@ -545,7 +548,7 @@ async def cursor_chat_completions( processor: Final = ProxyBaseLLMRequestProcessing(data=data) - def cursor_data_generator(response, user_api_key_dict, request_data, request=None): + def cursor_data_generator(response, user_api_key_dict, request_data: dict[str, object], request=None): """ Custom generator that transforms Responses API streaming chunks to chat completion chunks. @@ -579,7 +582,7 @@ async def cursor_chat_completions( completion_stream=completion_stream, model=request_data.get("model", ""), custom_llm_provider=None, - logging_obj=logging_obj, + logging_obj=validated_stream_logging_obj(logging_obj), ) # Use async_data_generator to format as SSE return async_data_generator( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index dd0c57aa911..fa24846268a 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -15,7 +15,7 @@ from dataclasses import dataclass, field from datetime import date, datetime, timedelta, timezone from email.mime.multipart import MIMEMultipart from email.mime.text import MIMEText -from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, Union, cast, overload +from typing import TYPE_CHECKING, Any, ClassVar, Final, Literal, Optional, TypeVar, Union, cast, overload from litellm import _custom_logger_compatible_callbacks_literal from litellm.constants import ( @@ -163,23 +163,26 @@ if TYPE_CHECKING: from mcp.types import CallToolResult from opentelemetry.trace import Span as _Span from prisma.client import TransactionManager + from prisma.types import HttpConfig from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.models.team import LiteLLM_TeamTableCachedObj from litellm.proxy.db.autorouter_session_rollup import AutoRouterTurnTransaction from litellm.proxy.db.spend_log_tool_index import ToolUsageTransaction - Span = _Span | Any + Span = _Span | object else: Span = Any +_T: Final = TypeVar("_T") + unified_guardrail: Final = UnifiedLLMGuardrails() NON_OPENAI_STREAM_GUARDRAIL_TRANSLATION_CALL_TYPES: "frozenset[CallTypes]" = frozenset({CallTypes.anthropic_messages}) -def print_verbose(print_statement): +def print_verbose(print_statement: object): """ Prints the given `print_statement` to the console if `litellm.set_verbose` is True. Also logs the `print_statement` at the debug level using `verbose_proxy_logger`. @@ -227,10 +230,10 @@ class InternalUsageCache: async def async_get_cache( self, - key, + key: str, litellm_parent_otel_span: Span | None, local_only: bool = False, - **kwargs, + **kwargs: object, ) -> Any: return await self.dual_cache.async_get_cache( key=key, @@ -241,11 +244,11 @@ class InternalUsageCache: async def async_set_cache( self, - key, - value, + key: str, + value: object, litellm_parent_otel_span: Span | None, local_only: bool = False, - **kwargs, + **kwargs: object, ) -> None: return await self.dual_cache.async_set_cache( key=key, @@ -257,10 +260,10 @@ class InternalUsageCache: async def async_batch_set_cache( self, - cache_list: list, + cache_list: list[tuple[str, object]], litellm_parent_otel_span: Span | None, local_only: bool = False, - **kwargs, + **kwargs: object, ) -> None: return await self.dual_cache.async_set_cache_pipeline( cache_list=cache_list, @@ -271,19 +274,19 @@ class InternalUsageCache: async def async_batch_get_cache( self, - keys: list, + keys: Sequence[str | None], parent_otel_span: Span | None = None, local_only: bool = False, ): return await self.dual_cache.async_batch_get_cache( - keys=keys, + keys=list(keys), parent_otel_span=parent_otel_span, local_only=local_only, ) async def async_increment_cache( self, - key, + key: str, value: float, litellm_parent_otel_span: Span | None, local_only: bool = False, @@ -299,10 +302,10 @@ class InternalUsageCache: def set_cache( self, - key, - value, + key: str, + value: object, local_only: bool = False, - **kwargs, + **kwargs: object, ) -> None: return self.dual_cache.set_cache( key=key, @@ -313,9 +316,9 @@ class InternalUsageCache: def get_cache( self, - key, + key: str, local_only: bool = False, - **kwargs, + **kwargs: object, ) -> Any: return self.dual_cache.get_cache( key=key, @@ -338,7 +341,7 @@ def _accepts_litellm_call_info(cb: CustomLogger) -> bool: return _CALLBACK_ACCEPTS_CALL_INFO[key] -def _enrich_http_exception_with_guardrail_context(exc: BaseException, callback: Any) -> None: +def _enrich_http_exception_with_guardrail_context(exc: BaseException, callback: object) -> None: """ If `exc` is an HTTPException with a dict `detail`, mutate it in place to add `guardrail_name` and `guardrail_mode` taken from the callback instance. @@ -391,7 +394,7 @@ class _CallbackCapabilities: # Resolved CustomLogger callbacks in original order. Pre-resolving once # avoids the per-request ``get_custom_logger_compatible_class`` walk for # every string entry in ``litellm.callbacks``. - resolved_callbacks: tuple[Any, ...] = field(default_factory=tuple) + resolved_callbacks: tuple[object, ...] = field(default_factory=tuple) class ProxyLogging: @@ -674,7 +677,7 @@ class ProxyLogging: return synthetic_data - def _convert_llm_result_to_mcp_response(self, llm_result, request_obj) -> Any | None: + def _convert_llm_result_to_mcp_response(self, llm_result, request_obj) -> MCPPreCallResponseObject | None: """ Convert LLM guardrail result back to MCP response format. """ @@ -800,7 +803,7 @@ class ProxyLogging: verbose_proxy_logger.error("Error in manual argument parsing: %s", e) return None - def _convert_llm_result_to_mcp_during_response(self, llm_result, request_obj) -> Any | None: + def _convert_llm_result_to_mcp_during_response(self, llm_result, request_obj) -> MCPDuringCallResponseObject | None: """ Convert LLM guardrail result back to MCP during call response format. """ @@ -846,7 +849,7 @@ class ProxyLogging: self, response: MCPPreCallResponseObject, original_request: MCPPreCallRequestObject, - ) -> dict[str, Any]: + ) -> Mapping[str, object]: """ Parse the response from the pre_mcp_tool_call_hook @@ -949,8 +952,8 @@ class ProxyLogging: data: dict, user_api_key_dict: UserAPIKeyAuth | None, call_type: CallTypesLiteral, - response: Any | None = None, - ) -> Any: + response: LLMResponseTypes | None = None, + ) -> object: """ Execute a single guardrail's hook. @@ -1004,8 +1007,8 @@ class ProxyLogging: data: dict, user_api_key_dict: UserAPIKeyAuth | None, call_type: CallTypesLiteral, - response: Any | None = None, - ) -> Any: + response: LLMResponseTypes | None = None, + ) -> object: """ Execute a guardrail using the router's load balancing. @@ -1140,8 +1143,8 @@ class ProxyLogging: self, data: dict, litellm_logging_obj: Any, - prompt_id: Any, - prompt_version: Any, + prompt_id: str, + prompt_version: int | None, call_type: CallTypesLiteral, ) -> None: """Process prompt template if applicable.""" @@ -1362,8 +1365,8 @@ class ProxyLogging: return None litellm_logging_obj: Final = cast(Optional["LiteLLMLoggingObj"], data.get("litellm_logging_obj", None)) - prompt_id: Final = data.get("prompt_id", None) - prompt_version: Final = data.get("prompt_version", None) + prompt_id: Final[str | None] = data.get("prompt_id", None) + prompt_version: Final[int | None] = data.get("prompt_version", None) ## PROMPT TEMPLATE CHECK ## @@ -1444,7 +1447,7 @@ class ProxyLogging: if call_type == "call_mcp_tool" and user_api_key_dict is None: continue - response = await _callback.async_pre_call_hook( + response: Exception | str | Mapping[str, object] | None = await _callback.async_pre_call_hook( user_api_key_dict=user_api_key_dict, cache=self.call_details["user_api_key_cache"], data=data, @@ -1612,7 +1615,7 @@ class ProxyLogging: break @staticmethod - async def _run_guardrail_with_metrics(callback: Any, coro: Awaitable[Any], hook_type: str) -> Any: + async def _run_guardrail_with_metrics(callback: object, coro: Awaitable[_T], hook_type: str) -> _T: """ Await `coro`, recording its latency and status to the `litellm_guardrail_latency_seconds` metric under `hook_type`, and @@ -1644,8 +1647,8 @@ class ProxyLogging: @staticmethod async def _wrap_streaming_iterator_with_enrichment( - callback: Any, gen: AsyncGenerator[Any, None] - ) -> AsyncGenerator[Any, None]: + callback: object, gen: AsyncGenerator[_T, None] + ) -> AsyncGenerator[_T, None]: """ Yield from `gen`; if iteration raises an HTTPException with dict detail, enrich the detail with the originating callback's `guardrail_name` and @@ -1690,11 +1693,11 @@ class ProxyLogging: has_guardrail = False has_pre_call_override = False iterator_overrides: Final[list[tuple[Any, str]]] = [] # (callback, kind) - resolved_callbacks: Final[list[Any]] = [] + resolved_callbacks: Final[list[CustomLogger]] = [] for callback in callbacks: if isinstance(callback, str): - resolved: Any = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( + resolved = litellm.litellm_core_utils.litellm_logging.get_custom_logger_compatible_class( cast(_custom_logger_compatible_callbacks_literal, callback) ) else: @@ -2539,7 +2542,7 @@ class ProxyLogging: self, data: dict, user_api_key_dict: UserAPIKeyAuth, - response: Any, + response: object, request_headers: dict[str, str] | None = None, ) -> dict[str, str]: """ @@ -2595,7 +2598,7 @@ class ProxyLogging: return merged_headers @staticmethod - def _build_litellm_call_info(data: dict, response: Any) -> dict[str, Any]: + def _build_litellm_call_info(data: dict, response: object) -> dict[str, object]: """ Build a normalized dict of routing metadata from response._hidden_params and data, abstracting away the metadata vs litellm_metadata split. @@ -2872,7 +2875,7 @@ _DEPRECATED_KEY_CACHE_TTL_SECONDS: Final = 60 async def _lookup_deprecated_key( - db: Any, + db: PrismaWrapper | RoutingPrismaWrapper, hashed_token: str, ) -> str | None: """ @@ -2940,7 +2943,7 @@ def _config_cache_key(param_name: str) -> str: return f"litellm_config:param:{param_name}" -def _pack_config_row(row: Any) -> dict[str, Any]: +def _pack_config_row(row: Any) -> dict[str, object]: return {"param_name": row.param_name, "param_value": row.param_value} @@ -2952,7 +2955,7 @@ def _unpack_config_row(cached: Any) -> _ConfigRow | None: return None -async def get_config_param(prisma_client: Any, param_name: str) -> Any | None: +async def get_config_param(prisma_client: "PrismaClient", param_name: str) -> Any | None: """Cached read of a LiteLLM_Config row; returns row, _ConfigRow shim, or None.""" cache_key: Final = _config_cache_key(param_name) cached: Final = await litellm_config_cache.async_get_cache(cache_key) @@ -2960,7 +2963,7 @@ async def get_config_param(prisma_client: Any, param_name: str) -> Any | None: return _unpack_config_row(cached) row: Final = await prisma_client.get_generic_data(key="param_name", value=param_name, table_name="config") - cache_value: Final[Any] = _pack_config_row(row) if row is not None else _CONFIG_CACHE_MISS + cache_value: Final[Mapping[str, object] | str] = _pack_config_row(row) if row is not None else _CONFIG_CACHE_MISS await litellm_config_cache.async_set_cache(cache_key, cache_value, ttl=LITELLM_CONFIG_CACHE_TTL_SECONDS) return row @@ -2975,7 +2978,7 @@ async def invalidate_config_param(param_name: str) -> None: await publish_config_param_change(param_name) -async def prefetch_config_params(prisma_client: Any, param_names: list[str]) -> None: +async def prefetch_config_params(prisma_client: "PrismaClient | None", param_names: list[str]) -> None: """Batch-load LiteLLM_Config rows into the cache with one find_many.""" if not param_names: return @@ -2990,7 +2993,7 @@ async def prefetch_config_params(prisma_client: Any, param_names: list[str]) -> by_name: Final = {row.param_name: row for row in rows} for name in param_names: row = by_name.get(name) - cache_value: Any = _pack_config_row(row) if row is not None else _CONFIG_CACHE_MISS + cache_value: Mapping[str, object] | str = _pack_config_row(row) if row is not None else _CONFIG_CACHE_MISS await litellm_config_cache.async_set_cache( _config_cache_key(name), cache_value, ttl=LITELLM_CONFIG_CACHE_TTL_SECONDS ) @@ -3017,7 +3020,7 @@ class PrismaClient: self, database_url: str, proxy_logging_obj: ProxyLogging, - http_client: Any | None = None, + http_client: "HttpConfig | None" = None, ): ## init logging object self.proxy_logging_obj = proxy_logging_obj @@ -3309,7 +3312,7 @@ class PrismaClient: async def get_generic_data( self, key: str, - value: Any, + value: object, table_name: Literal["users", "keys", "config", "spend"], ): """ @@ -5494,7 +5497,7 @@ class ProxyUpdateSpend: prisma_client: PrismaClient, db_writer_client: AsyncHTTPHandler | None, proxy_logging_obj: ProxyLogging, - logs_to_process: list[dict[str, Any]] | None = None, + logs_to_process: list[dict[str, object]] | None = None, ): BATCH_SIZE: Final = 1000 # Preferred size of each batch to write to the database MAX_LOGS_PER_INTERVAL: Final = 10000 # Maximum number of logs to flush in a single interval @@ -6725,7 +6728,7 @@ def model_dump_with_preserved_fields( obj: Any, preserve_fields: list[str] | None = None, exclude_unset: bool = True, -) -> dict[str, Any]: +) -> dict[str, object]: """ Serialize a Pydantic model to a dictionary while preserving specific fields even if they are None. diff --git a/litellm/responses/main.py b/litellm/responses/main.py index 7b02c1b8023..e0af363b1a5 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -1,6 +1,6 @@ import asyncio import contextvars -from collections.abc import Coroutine, Iterable +from collections.abc import Coroutine, Iterable, Mapping from functools import partial from typing import TYPE_CHECKING, Any, Final, Literal, Optional, cast @@ -53,6 +53,7 @@ from litellm.utils import ( ) if TYPE_CHECKING: + from fastapi import WebSocket from mcp.types import Tool as MCPTool else: MCPTool = Any @@ -66,7 +67,7 @@ litellm_completion_transformation_handler: Final = LiteLLMCompletionTransformati ################################################# -def _has_file_search_tool(tools: Any | None) -> bool: +def _has_file_search_tool(tools: Iterable[Mapping[str, object]] | None) -> bool: """Return True if any tool in the list has type 'file_search'.""" if not tools: return False @@ -132,7 +133,7 @@ async def aresponses_api_with_mcp( instructions: str | None = None, max_output_tokens: int | None = None, prompt: PromptObject | None = None, - metadata: dict[str, Any] | None = None, + metadata: dict[str, object] | None = None, parallel_tool_calls: bool | None = None, previous_response_id: str | None = None, reasoning: Reasoning | None = None, @@ -148,9 +149,9 @@ async def aresponses_api_with_mcp( user: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, @@ -397,7 +398,7 @@ async def aresponses( instructions: str | None = None, max_output_tokens: int | None = None, prompt: PromptObject | None = None, - metadata: dict[str, Any] | None = None, + metadata: dict[str, object] | None = None, parallel_tool_calls: bool | None = None, previous_response_id: str | None = None, reasoning: Reasoning | None = None, @@ -416,9 +417,9 @@ async def aresponses( safety_identifier: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, @@ -564,9 +565,9 @@ def _apply_prompt_management_to_responses_call( custom_llm_provider: str | None, litellm_logging_obj: LiteLLMLoggingObj | None, kwargs: dict[str, Any], - local_vars: dict[str, Any], + local_vars: dict[str, object], ) -> tuple[str | ResponseInputParam, str, str | None]: - async_merged: Final = kwargs.pop("_async_prompt_merged_params", None) + async_merged: Final[Mapping[str, object] | None] = kwargs.pop("_async_prompt_merged_params", None) if async_merged is not None: for key, value in async_merged.items(): local_vars[key] = value @@ -633,7 +634,7 @@ def _normalize_openai_chat_completions_responses_model(model: str) -> tuple[str, return f"openai/{remainder}", True -def _pop_use_chat_completions_api_kw(kwargs: dict[str, Any]) -> bool: +def _pop_use_chat_completions_api_kw(kwargs: dict[str, object]) -> bool: """Pop use_chat_completions_api; True when the chat-completions bridge is requested.""" use_cc: Final = kwargs.pop("use_chat_completions_api", None) return bool(use_cc) @@ -643,7 +644,7 @@ def _resolve_model_provider_for_responses( model: str, custom_llm_provider: str | None, litellm_params: GenericLiteLLMParams, - local_vars: dict[str, Any], + local_vars: dict[str, object], ) -> tuple[str, str | None]: if custom_llm_provider is not None and not litellm_params.custom_llm_provider: litellm_params.custom_llm_provider = custom_llm_provider @@ -668,7 +669,7 @@ def _apply_managed_file_id_mapping( input: str | ResponseInputParam, tools: Iterable[ToolParam] | None, kwargs: dict[str, Any], - local_vars: dict[str, Any], + local_vars: dict[str, object], ) -> tuple[str | ResponseInputParam, Iterable[ToolParam] | None]: model_file_id_mapping: Final = kwargs.get("model_file_id_mapping") model_info_id = kwargs.get("model_info", {}).get("id") if isinstance(kwargs.get("model_info"), dict) else None @@ -706,7 +707,7 @@ def _responses_try_dispatch_mcp_gateway( instructions: str | None, max_output_tokens: int | None, prompt: PromptObject | None, - metadata: dict[str, Any] | None, + metadata: dict[str, object] | None, parallel_tool_calls: bool | None, previous_response_id: str | None, reasoning: Reasoning | None, @@ -719,9 +720,9 @@ def _responses_try_dispatch_mcp_gateway( top_p: float | None, truncation: Literal["auto", "disabled"] | None, user: str | None, - extra_headers: dict[str, Any] | None, - extra_query: dict[str, Any] | None, - extra_body: dict[str, Any] | None, + extra_headers: dict[str, object] | None, + extra_query: dict[str, object] | None, + extra_body: dict[str, object] | None, timeout: float | httpx.Timeout | None, custom_llm_provider: str | None, kwargs: dict[str, Any], @@ -778,7 +779,7 @@ def _responses_try_dispatch_emulated_file_search( instructions: str | None, max_output_tokens: int | None, prompt: PromptObject | None, - metadata: dict[str, Any] | None, + metadata: dict[str, object] | None, parallel_tool_calls: bool | None, previous_response_id: str | None, reasoning: Reasoning | None, @@ -795,14 +796,14 @@ def _responses_try_dispatch_emulated_file_search( safety_identifier: str | None, text_format: type[BaseModel] | dict | None, allowed_openai_params: list[str] | None, - extra_headers: dict[str, Any] | None, - extra_query: dict[str, Any] | None, - extra_body: dict[str, Any] | None, + extra_headers: dict[str, object] | None, + extra_query: dict[str, object] | None, + extra_body: dict[str, object] | None, timeout: float | httpx.Timeout | None, custom_llm_provider: str | None, kwargs: dict[str, Any], _is_async: bool, -) -> Any | None: +) -> ResponsesAPIResponse | Coroutine[object, object, ResponsesAPIResponse] | None: """Return a response when emulated file_search handles the call; otherwise None.""" if not _has_file_search_tool(tools) or not ( responses_api_provider_config is None @@ -864,7 +865,7 @@ def responses( instructions: str | None = None, max_output_tokens: int | None = None, prompt: PromptObject | None = None, - metadata: dict[str, Any] | None = None, + metadata: dict[str, object] | None = None, parallel_tool_calls: bool | None = None, previous_response_id: str | None = None, reasoning: Reasoning | None = None, @@ -883,9 +884,9 @@ def responses( safety_identifier: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, allowed_openai_params: list[str] | None = None, @@ -1148,9 +1149,9 @@ async def adelete_responses( response_id: str, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, @@ -1209,14 +1210,14 @@ def delete_responses( response_id: str, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, **kwargs, -) -> DeleteResponseResult | Coroutine[Any, Any, DeleteResponseResult]: +) -> DeleteResponseResult | Coroutine[object, object, DeleteResponseResult]: """ Synchronous version of the DELETE Responses API @@ -1299,9 +1300,9 @@ async def aget_responses( response_id: str, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, @@ -1374,14 +1375,14 @@ def get_responses( response_id: str, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, **kwargs, -) -> ResponsesAPIResponse | Coroutine[Any, Any, ResponsesAPIResponse]: +) -> ResponsesAPIResponse | Coroutine[object, object, ResponsesAPIResponse]: """ Fetch a response by its ID. @@ -1481,7 +1482,7 @@ async def alist_input_items( include: list[str] | None = None, limit: int = 20, order: Literal["asc", "desc"] = "desc", - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, @@ -1537,11 +1538,11 @@ def list_input_items( include: list[str] | None = None, limit: int = 20, order: Literal["asc", "desc"] = "desc", - extra_headers: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, custom_llm_provider: str | None = None, **kwargs, -) -> dict | Coroutine[Any, Any, dict]: +) -> dict | Coroutine[object, object, dict]: """List input items for a response""" local_vars: Final = locals() try: @@ -1612,9 +1613,9 @@ async def acancel_responses( response_id: str, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, @@ -1673,14 +1674,14 @@ def cancel_responses( response_id: str, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, **kwargs, -) -> ResponsesAPIResponse | Coroutine[Any, Any, ResponsesAPIResponse]: +) -> ResponsesAPIResponse | Coroutine[object, object, ResponsesAPIResponse]: """ Synchronous version of the POST Responses API @@ -1766,9 +1767,9 @@ async def acompact_responses( previous_response_id: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, @@ -1844,14 +1845,14 @@ def compact_responses( previous_response_id: str | None = None, # Use the following arguments if you need to pass additional parameters to the API that aren't available via kwargs. # The extra values given here take precedence over values defined on the client or passed to this method. - extra_headers: dict[str, Any] | None = None, - extra_query: dict[str, Any] | None = None, - extra_body: dict[str, Any] | None = None, + extra_headers: dict[str, object] | None = None, + extra_query: dict[str, object] | None = None, + extra_body: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, # LiteLLM specific params, custom_llm_provider: str | None = None, **kwargs, -) -> ResponsesAPIResponse | Coroutine[Any, Any, ResponsesAPIResponse]: +) -> ResponsesAPIResponse | Coroutine[object, object, ResponsesAPIResponse]: """ Synchronous version of the POST Compact Responses API @@ -1975,7 +1976,7 @@ def _build_litellm_metadata_for_ws(kwargs: dict) -> dict: @client async def _aresponses_websocket( model: str, - websocket: Any, + websocket: "WebSocket", api_base: str | None = None, api_key: str | None = None, timeout: float | None = None, diff --git a/litellm/responses/mcp/chat_completions_handler.py b/litellm/responses/mcp/chat_completions_handler.py index 49564dc7f07..38e6d07c626 100644 --- a/litellm/responses/mcp/chat_completions_handler.py +++ b/litellm/responses/mcp/chat_completions_handler.py @@ -233,7 +233,7 @@ async def acompletion_with_mcp( self.follow_up_iterator = None self.follow_up_exhausted = False - async def __aiter__(self): + def __aiter__(self): return self def _add_mcp_list_tools_to_chunk(self, chunk: ModelResponseStream) -> ModelResponseStream: @@ -497,12 +497,12 @@ async def acompletion_with_mcp( # Create a wrapper class that delegates to our custom iterator # We'll use a simple approach: just replace the __aiter__ method class MCPStreamWrapper(CustomStreamWrapper): - def __init__(self, original_wrapper, custom_iterator): + def __init__(self, original_wrapper: CustomStreamWrapper, custom_iterator: MCPStreamingIterator): # Initialize with the same parameters as original wrapper super().__init__( completion_stream=None, model=getattr(original_wrapper, "model", "unknown"), - logging_obj=getattr(original_wrapper, "logging_obj", None), + logging_obj=original_wrapper.logging_obj, custom_llm_provider=getattr(original_wrapper, "custom_llm_provider", None), stream_options=getattr(original_wrapper, "stream_options", None), make_call=getattr(original_wrapper, "make_call", None), diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 2e1e1a44594..d7f6ece5cd1 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -9,7 +9,7 @@ from collections.abc import Awaitable, Callable, Mapping from datetime import datetime from functools import lru_cache from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, runtime_checkable import httpx from openai._streaming import SSEDecoder @@ -313,8 +313,10 @@ class BaseResponsesAPIStreamingIterator: if encrypted_content and isinstance(encrypted_content, str): model_id: Final = _model_id_from_metadata(self.litellm_metadata) if model_id: - wrapped_content = ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( - encrypted_content, model_id + wrapped_content: Final = ( + ResponsesAPIRequestUtils._wrap_encrypted_content_with_model_id( + encrypted_content, model_id + ) ) setattr(item, "encrypted_content", wrapped_content) @@ -336,7 +338,9 @@ class BaseResponsesAPIStreamingIterator: usage_obj: Final[ResponseAPIUsage | None] = getattr(response_obj, "usage", None) if usage_obj is not None: try: - cost: float | None = self.logging_obj._response_cost_calculator(result=response_obj) + cost: Final[float | None] = self.logging_obj._response_cost_calculator( + result=response_obj + ) if cost is not None: setattr(usage_obj, "cost", cost) except Exception: @@ -1029,6 +1033,16 @@ class CachedResponsesAPIStreamingIterator(BaseResponsesAPIStreamingIterator): return evt +@runtime_checkable +class _HasModelDump(Protocol): + def model_dump(self, *, exclude_none: bool = ...) -> Mapping[str, object]: ... + + +@runtime_checkable +class _HasModelDumpJson(Protocol): + def model_dump_json(self, *, exclude_none: bool = ...) -> str: ... + + def _dump_response_object(obj: Any) -> dict[str, Any]: if hasattr(obj, "model_dump"): return obj.model_dump() @@ -1358,7 +1372,7 @@ class ResponsesWebSocketStreaming: # response.create frame to prevent deployment-substitution attacks. self.authorized_model: str | None = authorized_model - def _should_store_event(self, event_obj: dict[str, object]) -> bool: + def _should_store_event(self, event_obj: Mapping[str, object]) -> bool: return event_obj.get("type") in RESPONSES_WS_LOGGED_EVENT_TYPES def _store_event(self, event: str | bytes | dict[str, object]) -> None: @@ -1636,7 +1650,7 @@ class ResponsesWebSocketStreaming: metadata: Final = self.request_data.get("metadata") raw_pii_tokens: Final = metadata.get("pii_tokens") if _is_json_object(metadata) else None - pii_tokens: Final[dict[str, str]] = raw_pii_tokens if _is_str_mapping(raw_pii_tokens) else {} + pii_tokens: Final[Mapping[str, str]] = raw_pii_tokens if _is_str_mapping(raw_pii_tokens) else {} if not pii_tokens: return response_str @@ -1883,11 +1897,11 @@ class ManagedResponsesWebSocketHandler: def _serialize_chunk(chunk: Any) -> str | None: """Serialize a streaming chunk to a JSON string for WebSocket transmission.""" try: - if hasattr(chunk, "model_dump_json"): + if isinstance(chunk, _HasModelDumpJson): return chunk.model_dump_json(exclude_none=True) - if hasattr(chunk, "model_dump"): + if isinstance(chunk, _HasModelDump): return json.dumps(chunk.model_dump(exclude_none=True), default=str) - if isinstance(chunk, dict): + if _is_json_object(chunk): return json.dumps(chunk, default=str) return json.dumps(str(chunk)) except Exception as exc: diff --git a/litellm/router.py b/litellm/router.py index 9cece2014af..acab3ccf54e 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -23,7 +23,7 @@ from collections import defaultdict from collections.abc import AsyncGenerator, Callable, Generator, Mapping, Sequence from functools import lru_cache from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeVar, Union, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, TypeAlias, TypeVar, Union, cast import anyio import httpx @@ -258,6 +258,14 @@ else: QualityRouter = Any PreRoutingHookResponse = Any +RouterStrategySelector: TypeAlias = ( + LeastBusyLoggingHandler + | LowestCostLoggingHandler + | LowestLatencyLoggingHandler + | LowestTPMLoggingHandler + | LowestTPMLoggingHandler_v2 +) + def _cost_value_as_float(value: str | float | None) -> float | None: if value is None: @@ -728,7 +736,7 @@ class Router: routing_strategy_args=routing_strategy_args, ) self._init_routing_groups(self._routing_groups_input) - self._override_selectors: dict[str, Any] = {} + self._override_selectors: dict[str, RouterStrategySelector | None] = {} self._override_selectors_lock = threading.Lock() self.access_groups = None ## USAGE TRACKING ## @@ -921,13 +929,13 @@ class Router: strategy: RoutingStrategy | str, routing_strategy_args: dict, register_callbacks: bool = True, - ) -> Any | None: + ) -> RouterStrategySelector | None: """ Constructs a strategy selector for a given strategy. Returns None for `simple-shuffle` (no selector needed) and unknown strategies. """ - selector: Any | None = None + selector: RouterStrategySelector | None = None match self._normalize_strategy(strategy): case RoutingStrategy.LEAST_BUSY.value: selector = LeastBusyLoggingHandler(router_cache=self.cache) @@ -964,7 +972,7 @@ class Router: return selector - def _unregister_router_selectors(self, selectors: list[Any]) -> None: + def _unregister_router_selectors(self, selectors: Sequence[object]) -> None: """ Drop router-owned strategy selectors from litellm's global callback lists by identity. Used before re-init (`routing_strategy_init` / @@ -1021,13 +1029,14 @@ class Router: `"default"` group, whose selectors are the `self._logger` attributes set up in `routing_strategy_init`. """ - self._unregister_router_selectors( - [sel for selectors in getattr(self, "_group_selectors", {}).values() for sel in selectors.values()] + group_selectors: Final[Mapping[str, Mapping[str, RouterStrategySelector]]] = getattr( + self, "_group_selectors", {} ) + self._unregister_router_selectors([sel for selectors in group_selectors.values() for sel in selectors.values()]) self._routing_groups: dict[str, RoutingGroup] = {} self._model_to_group: dict[str, str] = {} - self._group_selectors: dict[str, dict[str, Any]] = {} + self._group_selectors: dict[str, dict[str, RouterStrategySelector]] = {} if not groups_input: return @@ -1105,7 +1114,7 @@ class Router: return None return strategy - def _get_override_strategy_selector(self, strategy: str) -> Any | None: + def _get_override_strategy_selector(self, strategy: str) -> RouterStrategySelector | None: """ Returns the selector for a per-request strategy override. @@ -1126,7 +1135,9 @@ class Router: ) return self._override_selectors[strategy] - def _get_routing_context(self, model: str, request_kwargs: dict | None = None) -> tuple[str | None, Any | None]: + def _get_routing_context( + self, model: str, request_kwargs: dict | None = None + ) -> tuple[str | None, RouterStrategySelector | None]: """ Resolves the routing strategy and selector to use for the given model. @@ -1949,7 +1960,7 @@ class Router: return silent_kwargs - def _silent_experiment_completion(self, silent_model: str, messages: list[Any], **kwargs): + def _silent_experiment_completion(self, silent_model: str, messages: Sequence[Mapping[str, str]], **kwargs): """ Run a silent experiment in the background (thread). """ @@ -2299,7 +2310,7 @@ class Router: # in __init__ rather than declaring it as a class field, so # static narrowing doesn't expose it. Mirror the sync path # (_completion_streaming_iterator) and pull via getattr. - chat: Final = getattr(built, "usage", None) if built is not None else None + chat: Final[object | None] = getattr(built, "usage", None) if built is not None else None if chat is not None: # getattr-with-default because the test path may # substitute a SimpleNamespace lacking some fields; @@ -2393,7 +2404,7 @@ class Router: # ResponseOutputMessageParam, ...) — annotating as List[Dict[str, Any]] # rejects the list() spread of input_val. We cast the combined list to # ResponseInputParam at the return. - base: list[Any] + base: list[object] if isinstance(input_val, str): base = [ { @@ -2406,7 +2417,7 @@ class Router: base = list(input_val) else: base = [] - continuation: Final[list[Any]] = [ + continuation: Final[list[object]] = [ { "type": "message", "role": "developer", @@ -2783,7 +2794,7 @@ class Router: return SyncFallbackStreamWrapper(stream_with_fallbacks()) - async def _silent_experiment_acompletion(self, silent_model: str, messages: list[Any], **kwargs): + async def _silent_experiment_acompletion(self, silent_model: str, messages: Sequence[Mapping[str, str]], **kwargs): """ Run a silent experiment in the background. """ @@ -3039,7 +3050,7 @@ class Router: pass def _stamp_failed_deployment_id_with_effective_model_info( - self, exception: Exception, deployment: Mapping[str, Any], kwargs: Mapping[str, Any] + self, exception: Exception, deployment: Mapping[str, object], kwargs: Mapping[str, object] ) -> None: # A client-side-credential call gets a dynamic deployment id generated inside # _update_kwargs_with_deployment and stamped into kwargs["model_info"]; stamping @@ -3562,8 +3573,8 @@ class Router: model: str, priority: int, original_function: Callable, - args: tuple[Any, ...], - kwargs: dict[str, Any], + args: tuple[object, ...], + kwargs: dict[str, object], ): parent_otel_span: Final = _get_parent_otel_span_from_kwargs(kwargs) ### FLOW ITEM ### @@ -4651,7 +4662,7 @@ class Router: # fallback to the original reference for any non-picklable value. # The original_generic_function is preserved so the per-attempt # helper knows which underlying API to call on fallback. - fallback_kwargs: Final[dict[str, Any]] = kwargs.copy() + fallback_kwargs: Final[dict[str, object]] = kwargs.copy() if isinstance(fallback_kwargs.get("litellm_metadata"), dict): fallback_kwargs["litellm_metadata"] = safe_deep_copy(fallback_kwargs["litellm_metadata"]) if isinstance(fallback_kwargs.get("metadata"), dict): @@ -5698,7 +5709,7 @@ class Router: def sync_wrapper( custom_llm_provider: str | None = None, - client: Any | None = None, + client: object | None = None, **kwargs, ): return self._generic_api_call_with_fallbacks(original_function=original_function, **kwargs) @@ -5714,7 +5725,7 @@ class Router: def vector_store_sync_wrapper( custom_llm_provider: str | None = None, - client: Any | None = None, + client: object | None = None, **kwargs, ): if custom_llm_provider and "custom_llm_provider" not in kwargs: @@ -5736,7 +5747,7 @@ class Router: def vector_store_file_sync_wrapper( custom_llm_provider: str | None = None, - client: Any | None = None, + client: object | None = None, **kwargs, ): return original_function( @@ -5757,7 +5768,7 @@ class Router: def managed_agents_sync_wrapper( custom_llm_provider: str | None = None, - client: Any | None = None, + client: object | None = None, **kwargs, ): if custom_llm_provider and "custom_llm_provider" not in kwargs: @@ -7142,7 +7153,9 @@ class Router: except Exception as e: raise e - async def async_deployment_callback_on_failure(self, kwargs, completion_response: Any | None, start_time, end_time): + async def async_deployment_callback_on_failure( + self, kwargs, completion_response: object | None, start_time, end_time + ): """ Update RPM usage for a deployment """ @@ -7843,7 +7856,7 @@ class Router: continue if self._has_registered_strategy(self.adaptive_routers, model_name, tagged.tags): continue - adaptive_router = complexity_router._ensure_adaptive_router() + adaptive_router: AdaptiveRouter | None = complexity_router._ensure_adaptive_router() if adaptive_router is not None: self.adaptive_routers[model_name] = [ *self.adaptive_routers.get(model_name, []), @@ -8131,7 +8144,7 @@ class Router: self.provider_default_deployment_ids.append(deployment.model_info.id) _team_id: Final = deployment.model_info.get("team_id") - _team_public_model_name: Final = deployment.model_info.get("team_public_model_name") + _team_public_model_name: Final[str | None] = deployment.model_info.get("team_public_model_name") if _team_id is not None and _team_public_model_name is not None and "*" in _team_public_model_name: if _team_id not in self.team_pattern_routers: self.team_pattern_routers[_team_id] = PatternMatchRouter() @@ -9485,7 +9498,7 @@ class Router: async def set_response_headers( self, - response: Any, + response: object, model_group: str | None = None, request_kwargs: dict | None = None, ) -> Any: @@ -11318,7 +11331,7 @@ class Router: @staticmethod def _redact_prompt_text_if_needed( - request_kwargs: Mapping[str, Any], + request_kwargs: Mapping[str, object], routing_decision: StandardLoggingRoutingDecision, ) -> StandardLoggingRoutingDecision: """Drop verbatim prompt text from the record when message logging is redacted. @@ -11680,7 +11693,7 @@ class Router: flag. Used by credential-lookup helpers so passthrough file / batch endpoints cannot bypass the pause by resolving credentials directly. """ - model_info: Final = getattr(deployment, "model_info", None) + model_info: Final[object | None] = getattr(deployment, "model_info", None) if model_info is None: return False return getattr(model_info, "blocked", None) is True diff --git a/litellm/types/completion.py b/litellm/types/completion.py index 84c804e9910..c1c6cc9ed1c 100644 --- a/litellm/types/completion.py +++ b/litellm/types/completion.py @@ -217,7 +217,7 @@ class _CompletionDispatchContext: headers: dict hf_model_name: str | None kwargs: dict - litellm_params: dict + litellm_params: dict[str, object] logger_fn: Callable | None logging: LiteLLMLoggingObj max_retries: int | None diff --git a/litellm/utils.py b/litellm/utils.py index 87937c99a0c..e33c8c847d8 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -234,9 +234,11 @@ except (ImportError, AttributeError, TypeError): # Convert to str (if necessary) claude_json_str = json.dumps(json_data) import importlib.metadata -from collections.abc import Callable, Iterable, Mapping +from collections.abc import Callable, Iterable, Mapping, Sequence from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast, get_args +from litellm import utils as litellm_utils + # These are lazy loaded via __getattr__ from litellm.llms.base_llm.base_utils import ( BaseLLMModelInfo, @@ -263,6 +265,7 @@ if TYPE_CHECKING: map_finish_reason, process_response_headers, ) + from litellm.litellm_core_utils.credential_accessor import CredentialAccessor from litellm.litellm_core_utils.dot_notation_indexing import ( delete_nested_value, is_nested_path, @@ -351,6 +354,24 @@ if TYPE_CHECKING: ) from litellm.llms.base_llm.videos.transformation import BaseVideoConfig from litellm.llms.bedrock.common_utils import BedrockModelInfo + from litellm.llms.bedrock.embed.amazon_nova_transformation import ( + AmazonNovaEmbeddingConfig, + ) + from litellm.llms.bedrock.embed.amazon_titan_g1_transformation import ( + AmazonTitanG1Config, + ) + from litellm.llms.bedrock.embed.amazon_titan_multimodal_transformation import ( + AmazonTitanMultimodalEmbeddingG1Config, + ) + from litellm.llms.bedrock.embed.amazon_titan_v2_transformation import ( + AmazonTitanV2Config, + ) + from litellm.llms.bedrock.embed.cohere_transformation import ( + BedrockCohereEmbeddingConfig, + ) + from litellm.llms.bedrock.embed.twelvelabs_marengo_transformation import ( + TwelveLabsMarengoEmbeddingConfig, + ) from litellm.llms.cohere.common_utils import CohereModelInfo from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.mistral.ocr.transformation import MistralOCRConfig @@ -574,7 +595,7 @@ def get_request_guardrails(kwargs: dict[str, Any]) -> list[str]: return applied_guardrails -def get_applied_guardrails(kwargs: dict[str, Any]) -> list[str]: +def get_applied_guardrails(kwargs: dict[str, object]) -> list[str]: """ - Add 'default_on' guardrails to the list - Add request guardrails to the list @@ -601,7 +622,7 @@ def load_credentials_from_list(kwargs: dict): credential_name: Final = kwargs.get("litellm_credential_name") if credential_name and litellm.credential_list: - credential_accessor: Final = CredentialAccessor.get_credential_values(credential_name) + credential_accessor: Final[Mapping[str, object]] = CredentialAccessor.get_credential_values(credential_name) for key, value in credential_accessor.items(): if key not in kwargs: kwargs[key] = value @@ -789,7 +810,7 @@ def function_setup( function_id: Final[str | None] = kwargs["id"] if "id" in kwargs else None ## LAZY LOAD COROUTINE CHECKER ## - get_coroutine_checker_fn: Final = getattr(sys.modules[__name__], "get_coroutine_checker") + get_coroutine_checker_fn: Final = litellm_utils.get_coroutine_checker coroutine_checker: Final = get_coroutine_checker_fn() ## DYNAMIC CALLBACKS ## @@ -925,7 +946,7 @@ def function_setup( elif kwargs.get("messages", None): messages = kwargs["messages"] ### PRE-CALL RULES ### - Rules: Final = getattr(sys.modules[__name__], "Rules") + Rules: Final = litellm_utils.Rules if ( Rules.has_pre_call_rules() and isinstance(messages, list) @@ -1033,7 +1054,7 @@ def function_setup( ) contents_param: Final = args[1] if len(args) > 1 else kwargs.get("contents") - model_param: Final = args[0] if len(args) > 0 else kwargs.get("model", "") + model_param: Final[str] = args[0] if len(args) > 0 else kwargs.get("model", "") if contents_param: adapter: Final = GoogleGenAIAdapter() @@ -1078,7 +1099,7 @@ def function_setup( ) ## check if metadata is passed in - litellm_params: Final[dict[str, Any]] = {"api_base": ""} + litellm_params: Final[dict[str, object]] = {"api_base": ""} if "metadata" in kwargs: litellm_params["metadata"] = kwargs["metadata"] if "litellm_metadata" in kwargs and isinstance(kwargs["litellm_metadata"], dict): @@ -1154,9 +1175,11 @@ def _get_wrapper_num_retries(kwargs: dict[str, Any], exception: Exception) -> tu if num_retries is None: num_retries = litellm.num_retries if kwargs.get("retry_policy", None): - get_num_retries_from_retry_policy: Final = getattr(sys.modules[__name__], "get_num_retries_from_retry_policy") - reset_retry_policy: Final = getattr(sys.modules[__name__], "reset_retry_policy") - retry_policy_num_retries: Final = get_num_retries_from_retry_policy( + get_num_retries_from_retry_policy: Final[Callable[..., int | None]] = getattr( + sys.modules[__name__], "get_num_retries_from_retry_policy" + ) + reset_retry_policy: Final = litellm_utils.reset_retry_policy + retry_policy_num_retries: Final[int | None] = get_num_retries_from_retry_policy( exception=exception, retry_policy=kwargs.get("retry_policy"), ) @@ -1167,7 +1190,7 @@ def _get_wrapper_num_retries(kwargs: dict[str, Any], exception: Exception) -> tu return num_retries, kwargs -def _get_wrapper_timeout(kwargs: dict[str, Any], exception: Exception) -> float | int | httpx.Timeout | None: +def _get_wrapper_timeout(kwargs: dict[str, object], exception: Exception) -> float | int | httpx.Timeout | None: """ Get the timeout from the kwargs Used for the wrapper functions. @@ -1179,7 +1202,7 @@ def _get_wrapper_timeout(kwargs: dict[str, Any], exception: Exception) -> float def check_coroutine(value) -> bool: - get_coroutine_checker: Final = getattr(sys.modules[__name__], "get_coroutine_checker") + get_coroutine_checker: Final = litellm_utils.get_coroutine_checker return get_coroutine_checker().is_async_callable(value) @@ -1207,7 +1230,7 @@ async def async_pre_call_deployment_hook(kwargs: dict[str, Any], call_type: str) async def async_post_call_success_deployment_hook( - request_data: dict, response: Any, call_type: CallTypes | None + request_data: dict, response: object, call_type: CallTypes | None ) -> Any | None: """ Allow modifying / reviewing the response just after it's received from the deployment. @@ -1317,7 +1340,7 @@ def post_call_processing( def client(original_function): - Rules: Final = getattr(sys.modules[__name__], "Rules") + Rules: Final = litellm_utils.Rules rules_obj: Final = Rules() @wraps(original_function) @@ -1551,10 +1574,10 @@ def client(original_function): if call_type == CallTypes.completion.value: num_retries = kwargs.get("num_retries", None) or litellm.num_retries or None if kwargs.get("retry_policy", None): - get_num_retries_from_retry_policy = getattr( + get_num_retries_from_retry_policy: Callable[..., int | None] = getattr( sys.modules[__name__], "get_num_retries_from_retry_policy" ) - reset_retry_policy = getattr(sys.modules[__name__], "reset_retry_policy") + reset_retry_policy = litellm_utils.reset_retry_policy num_retries = get_num_retries_from_retry_policy( exception=e, retry_policy=kwargs.get("retry_policy"), @@ -1593,7 +1616,7 @@ def client(original_function): get_num_retries_from_retry_policy = getattr( sys.modules[__name__], "get_num_retries_from_retry_policy" ) - reset_retry_policy = getattr(sys.modules[__name__], "reset_retry_policy") + reset_retry_policy = litellm_utils.reset_retry_policy num_retries = get_num_retries_from_retry_policy( exception=e, retry_policy=kwargs.get("retry_policy"), @@ -1939,7 +1962,7 @@ def client(original_function): if not _is_streaming_response_for_correlation(result): _restore_correlation_context_if_supported(logging_obj) - get_coroutine_checker: Final = getattr(sys.modules[__name__], "get_coroutine_checker") + get_coroutine_checker: Final = litellm_utils.get_coroutine_checker is_coroutine: Final = get_coroutine_checker().is_async_callable(original_function) # Return the appropriate wrapper based on the original function type @@ -1992,7 +2015,7 @@ _STREAMING_CALL_TYPES: Final = frozenset( def _is_streaming_request( - kwargs: dict[str, Any], + kwargs: dict[str, object], call_type: CallTypes | str, ) -> bool: """ @@ -2323,7 +2346,7 @@ def supports_response_schema(model: str, custom_llm_provider: str | None = None) """ ## GET LLM PROVIDER ## try: - get_llm_provider: Final = getattr(sys.modules[__name__], "get_llm_provider") + get_llm_provider: Final = litellm_utils.get_llm_provider model, custom_llm_provider, _, _ = get_llm_provider(model=model, custom_llm_provider=custom_llm_provider) except Exception as e: verbose_logger.debug( @@ -2700,7 +2723,7 @@ _CACHE_PRICING_FIELDS: Final = ( ) -def _resolve_builtin_model_cost_entry(key: str, provider: str) -> dict[str, Any] | None: +def _resolve_builtin_model_cost_entry(key: str, provider: str) -> dict[str, object] | None: """Best-effort lookup of a built-in ``model_cost`` entry for a custom key whose shape ``get_model_info`` cannot resolve (repeated provider prefixes like ``bedrock/bedrock/bedrock/us.anthropic.claude-sonnet-4-6`` or region @@ -2992,7 +3015,7 @@ def get_optional_params_transcription( passed_params.pop("OPENAI_TRANSCRIPTION_PARAMS") custom_llm_provider = passed_params.pop("custom_llm_provider") drop_params = passed_params.pop("drop_params") - special_params: Final = passed_params.pop("kwargs") + special_params: Final[Mapping[str, object]] = passed_params.pop("kwargs") for k, v in special_params.items(): passed_params[k] = v @@ -3101,7 +3124,7 @@ def get_optional_params_image_gen( provider_config = passed_params.pop("provider_config", None) drop_params = passed_params.pop("drop_params", None) additional_drop_params = passed_params.pop("additional_drop_params", None) - special_params: Final = passed_params.pop("kwargs") + special_params: Final[Mapping[str, object]] = passed_params.pop("kwargs") for k, v in special_params.items(): if ( k.startswith("aws_") @@ -3133,7 +3156,7 @@ def get_optional_params_image_gen( default_params=default_params, additional_drop_params=additional_drop_params, ) - optional_params: dict[str, Any] = {} + optional_params: dict[str, object] = {} ## raise exception if non-default value passed for non-openai/azure embedding calls def _check_valid_arg(supported_params): @@ -3365,7 +3388,14 @@ def get_optional_params_embeddings( elif custom_llm_provider == "bedrock": # if dimensions is in non_default_params -> pass it for model=bedrock/amazon.titan-embed-text-v2 if "amazon.titan-embed-text-v1" in model: - object: Any = litellm.AmazonTitanG1Config() + object: ( + AmazonTitanG1Config + | AmazonTitanMultimodalEmbeddingG1Config + | AmazonTitanV2Config + | BedrockCohereEmbeddingConfig + | TwelveLabsMarengoEmbeddingConfig + | AmazonNovaEmbeddingConfig + ) = litellm.AmazonTitanG1Config() elif "amazon.titan-embed-image-v1" in model: object = litellm.AmazonTitanMultimodalEmbeddingG1Config() elif "amazon.titan-embed-text-v2:0" in model: @@ -4949,7 +4979,7 @@ def get_max_tokens(model: str) -> int | None: response.raise_for_status() # Raise an exception for bad responses (4xx or 5xx) # Parse the JSON response - config_json: Final = response.json() + config_json: Final[Mapping[str, int]] = response.json() # Extract and return the max_position_embeddings max_position_embeddings: Final = config_json.get("max_position_embeddings") if max_position_embeddings is not None: @@ -4965,7 +4995,7 @@ def get_max_tokens(model: str) -> int | None: return litellm.model_cost[model]["max_output_tokens"] elif "max_tokens" in litellm.model_cost[model]: return litellm.model_cost[model]["max_tokens"] - get_llm_provider: Final = getattr(sys.modules[__name__], "get_llm_provider") + get_llm_provider: Final = litellm_utils.get_llm_provider model, custom_llm_provider, _, _ = get_llm_provider(model=model) if custom_llm_provider == "huggingface": max_tokens: Final = _get_max_position_embeddings(model_name=model) @@ -5253,7 +5283,7 @@ def _get_potential_model_names(model: str, custom_llm_provider: str | None) -> P if custom_llm_provider is None: # Get custom_llm_provider try: - get_llm_provider: Final = getattr(sys.modules[__name__], "get_llm_provider") + get_llm_provider: Final = litellm_utils.get_llm_provider split_model, custom_llm_provider, _, _ = get_llm_provider(model=model) except Exception: split_model = model @@ -5297,7 +5327,7 @@ def _get_max_position_embeddings(model_name: str) -> int | None: response.raise_for_status() # Raise an exception for bad responses (4xx or 5xx) # Parse the JSON response - config_json: Final = response.json() + config_json: Final[Mapping[str, int]] = response.json() # Extract and return the max_position_embeddings max_position_embeddings: Final = config_json.get("max_position_embeddings") @@ -6066,7 +6096,7 @@ def validate_environment( } ## EXTRACT LLM PROVIDER - if model name provided try: - get_llm_provider: Final = getattr(sys.modules[__name__], "get_llm_provider") + get_llm_provider: Final = litellm_utils.get_llm_provider _, custom_llm_provider, _, _ = get_llm_provider(model=model) except Exception: custom_llm_provider = None @@ -6543,7 +6573,7 @@ def _get_retry_after_from_exception_header( # ". See https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Retry-After#syntax for # details. if response_headers is not None: - retry_header: Final = response_headers.get("retry-after") + retry_header: Final[str] = response_headers.get("retry-after") try: retry_after = int(retry_header) except Exception: @@ -6634,7 +6664,7 @@ def register_prompt_template( complete_model: Final = model potential_models: Final = [complete_model] try: - get_llm_provider: Final = getattr(sys.modules[__name__], "get_llm_provider") + get_llm_provider: Final = litellm_utils.get_llm_provider model = get_llm_provider(model=model)[0] potential_models.append(model) except Exception: @@ -7276,7 +7306,7 @@ def _get_base_model_from_metadata(model_call_details=None): return _base_model metadata: Final = litellm_params.get("metadata") or {} - _get_base_model_from_litellm_call_metadata = getattr( + _get_base_model_from_litellm_call_metadata: Callable[..., str | None] = getattr( sys.modules[__name__], "_get_base_model_from_litellm_call_metadata" ) base_model_from_metadata: Final = _get_base_model_from_litellm_call_metadata(metadata=metadata) @@ -7969,7 +7999,7 @@ class ProviderConfigManager: @staticmethod def _get_cohere_config(model: str) -> BaseConfig: """Get Cohere config based on route.""" - CohereModelInfo: Final = getattr(sys.modules[__name__], "CohereModelInfo") + CohereModelInfo: Final = litellm_utils.CohereModelInfo route: Final = CohereModelInfo.get_cohere_route(model) if route == "v2": return litellm.CohereV2ChatConfig() @@ -9006,7 +9036,7 @@ class ProviderConfigManager: return ReductoParseLegacyConfig() return None - MistralOCRConfig: Final = getattr(sys.modules[__name__], "MistralOCRConfig") + MistralOCRConfig: Final = litellm_utils.MistralOCRConfig PROVIDER_TO_CONFIG_MAP: Final = { litellm.LlmProviders.MISTRAL: MistralOCRConfig, } @@ -9285,13 +9315,14 @@ def extract_duration_from_srt_or_vtt(srt_or_vtt_content: str) -> float | None: # Regular expression to match timestamps in the format "hh:mm:ss,ms" or "hh:mm:ss.ms" timestamp_pattern: Final = r"(\d{2}):(\d{2}):(\d{2})[.,](\d{3})" - timestamps: Final = re.findall(timestamp_pattern, srt_or_vtt_content) + timestamps: Final[Sequence[tuple[str, str, str, str]]] = re.findall(timestamp_pattern, srt_or_vtt_content) if not timestamps: return None # Convert timestamps to seconds and find the max (end time) durations: Final = [] + match: tuple[str, str, str, str] for match in timestamps: hours, minutes, seconds, milliseconds = map(int, match) total_seconds = hours * 3600 + minutes * 60 + seconds + milliseconds / 1000.0 @@ -9338,11 +9369,11 @@ def _add_path_to_api_base(api_base: str, ending_path: str) -> str: return str(modified_url.copy_with(params=original_url.params)) -def get_standard_openai_params(params: dict) -> dict: +def get_standard_openai_params(params: Mapping[str, object]) -> dict: return {k: v for k, v in params.items() if k in litellm.OPENAI_CHAT_COMPLETION_PARAMS and v is not None} -def get_non_default_completion_params(kwargs: dict) -> dict: +def get_non_default_completion_params(kwargs: Mapping[str, object]) -> dict: openai_params: Final = litellm.OPENAI_CHAT_COMPLETION_PARAMS default_params: Final = openai_params + all_litellm_params non_default_params: Final = { @@ -9352,7 +9383,7 @@ def get_non_default_completion_params(kwargs: dict) -> dict: return non_default_params -def peek_reasoning_summary_aliases(optional_params: dict) -> Any | None: +def peek_reasoning_summary_aliases(optional_params: dict) -> object | None: """Read AI-SDK-style reasoning summary from optional_params or nested extra_body. Uses key membership (not ``or`` chains) so falsy values like ``""`` are not skipped. @@ -9372,7 +9403,7 @@ def peek_reasoning_summary_aliases(optional_params: dict) -> Any | None: def strip_reasoning_summary_aliases_from_optional_params( optional_params: dict, -) -> tuple[dict, Any | None]: +) -> tuple[dict, object | None]: """Copy optional_params; remove reasoningSummary aliases from top-level and extra_body.""" op: Final = dict(optional_params) rs_val = op.pop("reasoningSummary", None) @@ -9404,7 +9435,7 @@ def get_non_default_transcription_params(kwargs: dict) -> dict: def add_openai_metadata( - metadata: Mapping[str, Any] | None, + metadata: Mapping[str, object] | None, ) -> dict[str, str] | None: """ Add metadata to openai optional parameters, excluding hidden params. @@ -9438,7 +9469,7 @@ def add_openai_metadata( return visible_metadata.copy() -def get_requester_metadata(metadata: dict): +def get_requester_metadata(metadata: Mapping[str, object]): if not metadata: return None @@ -9498,7 +9529,7 @@ def return_raw_request(endpoint: CallTypes, kwargs: dict) -> RawRequestTypedDict ) -def jsonify_tools(tools: list[Any]) -> list[dict]: +def jsonify_tools(tools: Sequence[object]) -> list[dict]: """ Fixes https://github.com/BerriAI/litellm/issues/9321 @@ -9524,9 +9555,9 @@ def get_empty_usage() -> Usage: def should_run_mock_completion( - mock_response: Any | None, - mock_tool_calls: Any | None, - mock_timeout: Any | None, + mock_response: object | None, + mock_tool_calls: object | None, + mock_timeout: object | None, ) -> bool: if mock_response or mock_tool_calls or mock_timeout: return True diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index da0f608fdb5..9b97dcdb43c 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,21 +1,21 @@ { "ANN001": { - "limit": 3106 + "limit": 3056 }, "ANN002": { "limit": 71 }, "ANN003": { - "limit": 832 + "limit": 827 }, "ANN201": { - "limit": 2023 + "limit": 2022 }, "ANN202": { - "limit": 860 + "limit": 855 }, "ANN204": { - "limit": 713 + "limit": 712 }, "ANN205": { "limit": 114 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 1495 + "limit": 1388 }, "ASYNC230": { "limit": 11 @@ -39,7 +39,7 @@ "limit": 505 }, "B009": { - "limit": 79 + "limit": 64 }, "B010": { "limit": 190 @@ -171,7 +171,7 @@ "limit": 3 }, "RET504": { - "limit": 177 + "limit": 176 }, "RUF012": { "limit": 241 @@ -201,7 +201,7 @@ "limit": 58 }, "SIM102": { - "limit": 322 + "limit": 321 }, "SIM103": { "limit": 119 @@ -234,7 +234,7 @@ "limit": 5 }, "TID251": { - "limit": 1226 + "limit": 1224 }, "TRY002": { "limit": 524 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index c990ae52ff2..ad3f56c7bdf 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23057 + "limit": 23023 }, "LIT002": { - "limit": 27156 + "limit": 27148 }, "LIT003": { "limit": 269 @@ -15,19 +15,19 @@ "limit": 0 }, "LIT006": { - "limit": 1078 + "limit": 1077 }, "LIT007": { "limit": 0 }, "LIT008": { - "limit": 951 + "limit": 950 }, "LIT009": { "limit": 0 }, "LIT010": { - "limit": 16744 + "limit": 16733 }, "LIT011": { "limit": 5596 From df5425675e2877fa3100e4d2f1b4c1e98a3b4afb Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 11 Aug 2026 15:16:31 +0000 Subject: [PATCH 07/24] fix(schema): declare supports_tool_search in the model prices schemas --- model_prices_and_context_window.schema.json | 3 +++ tests/test_litellm/test_utils.py | 1 + 2 files changed, 4 insertions(+) diff --git a/model_prices_and_context_window.schema.json b/model_prices_and_context_window.schema.json index 882f514b199..56400e0666b 100644 --- a/model_prices_and_context_window.schema.json +++ b/model_prices_and_context_window.schema.json @@ -671,6 +671,9 @@ "supports_tool_choice": { "type": "boolean" }, + "supports_tool_search": { + "type": "boolean" + }, "supports_url_context": { "type": "boolean" }, diff --git a/tests/test_litellm/test_utils.py b/tests/test_litellm/test_utils.py index ed5a9f1dd63..acd8ef4c96c 100644 --- a/tests/test_litellm/test_utils.py +++ b/tests/test_litellm/test_utils.py @@ -921,6 +921,7 @@ def test_aaamodel_prices_and_context_window_json_is_valid(): "supports_response_schema": {"type": "boolean"}, "supports_system_messages": {"type": "boolean"}, "supports_tool_choice": {"type": "boolean"}, + "supports_tool_search": {"type": "boolean"}, "supports_video_input": {"type": "boolean"}, "supports_vision": {"type": "boolean"}, "supports_web_search": {"type": "boolean"}, From 262d1b4ca070ddbec219155d87dbf0e846dcc4fb Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 11 Aug 2026 08:22:57 -0700 Subject: [PATCH 08/24] fix: remove over-strict stream logging validator, restore base seam behavior --- basedpyright-code-budget.json | 16 ++++++++-------- litellm/litellm_core_utils/streaming_handler.py | 8 -------- litellm/main.py | 9 ++++----- litellm/proxy/image_endpoints/endpoints.py | 9 +++------ litellm/proxy/proxy_server.py | 6 ++---- litellm/proxy/rerank_endpoints/endpoints.py | 5 +++-- .../proxy/response_api_endpoints/endpoints.py | 11 ++++------- ruff-strict-budget.json | 2 +- type-discipline-budget.json | 2 +- 9 files changed, 26 insertions(+), 42 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 14fdd3ed0bc..1b2563a59a4 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -3,13 +3,13 @@ "limit": 25153 }, "reportArgumentType": { - "limit": 2596 + "limit": 2597 }, "reportAssignmentType": { - "limit": 327 + "limit": 325 }, "reportAttributeAccessIssue": { - "limit": 510 + "limit": 501 }, "reportCallIssue": { "limit": 114 @@ -54,7 +54,7 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5770 + "limit": 5772 }, "reportMissingTypeArgument": { "limit": 15676 @@ -99,16 +99,16 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44911 + "limit": 44914 }, "reportUnknownLambdaType": { "limit": 113 }, "reportUnknownMemberType": { - "limit": 39464 + "limit": 39456 }, "reportUnknownParameterType": { - "limit": 20058 + "limit": 20060 }, "reportUnknownVariableType": { "limit": 31038 @@ -129,7 +129,7 @@ "limit": 0 }, "reportUntypedFunctionDecorator": { - "limit": 33 + "limit": 30 }, "reportUnusedClass": { "limit": 23 diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 570c34169d0..99b1c1a2ab7 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -155,14 +155,6 @@ class _TextCompletionChoiceLike(Protocol): finish_reason: str | None -def validated_stream_logging_obj(candidate: object) -> LiteLLMLoggingObject: - from litellm.litellm_core_utils.litellm_logging import Logging - - if isinstance(candidate, Logging): - return candidate - raise TypeError("CustomStreamWrapper requires a LiteLLMLoggingObject") - - class CustomStreamWrapper: def __init__( self, diff --git a/litellm/main.py b/litellm/main.py index 5d9144674de..27c20159322 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -96,7 +96,6 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( from litellm.litellm_core_utils.request_timeout_resolver import ( get_configured_request_timeout, ) -from litellm.litellm_core_utils.streaming_handler import validated_stream_logging_obj from litellm.llms.base_llm import BaseConfig, BaseImageGenerationConfig from litellm.llms.base_llm.base_model_iterator import ( convert_model_response_to_streaming, @@ -597,7 +596,7 @@ async def acompletion( _, custom_llm_provider, _, _ = get_llm_provider( model=model, custom_llm_provider=custom_llm_provider, - api_base=completion_kwargs.get("base_url", None), + api_base=base_url, ) fallbacks = fallbacks or litellm.model_fallbacks @@ -827,7 +826,7 @@ def mock_completion( mock_response: MOCK_RESPONSE_TYPE | None = "This is a mock request", mock_tool_calls: list | None = None, mock_timeout: bool | None = False, - logging: LiteLLMLoggingObj | None = None, + logging=None, custom_llm_provider=None, timeout: float | str | httpx.Timeout | None = None, **kwargs, @@ -911,7 +910,7 @@ def mock_completion( ), model=model, custom_llm_provider="openai", - logging_obj=validated_stream_logging_obj(logging), + logging_obj=logging, ) return CustomStreamWrapper( completion_stream=mock_completion_streaming_obj( @@ -919,7 +918,7 @@ def mock_completion( ), model=model, custom_llm_provider="openai", - logging_obj=validated_stream_logging_obj(logging), + logging_obj=logging, ) if isinstance(mock_response, litellm.MockException): raise mock_response diff --git a/litellm/proxy/image_endpoints/endpoints.py b/litellm/proxy/image_endpoints/endpoints.py index 24ee1d96a0d..414beabe014 100644 --- a/litellm/proxy/image_endpoints/endpoints.py +++ b/litellm/proxy/image_endpoints/endpoints.py @@ -1,8 +1,10 @@ import asyncio +import io import traceback +from typing import Final import orjson -from fastapi import APIRouter, Depends, File, HTTPException, Request, Response, status +from fastapi import APIRouter, Depends, File, HTTPException, Request, Response, UploadFile, status from fastapi.responses import ORJSONResponse import litellm @@ -18,11 +20,6 @@ from litellm.types.llms.openai import ChatCompletionUserMessage router: Final = APIRouter() -import io -from typing import Final - -from fastapi import UploadFile - async def uploadfile_to_bytesio(upload: UploadFile) -> io.BytesIO: """ diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 91526ad70d0..78171bd3cec 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -259,7 +259,6 @@ from litellm.litellm_core_utils.sensitive_data_masker import ( SensitiveDataMasker, mask_sensitive_keys, ) -from litellm.litellm_core_utils.streaming_handler import validated_stream_logging_obj from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, HTTPHandler from litellm.llms.vertex_ai.vertex_llm_base import VertexBase from litellm.proxy._lazy_features import attach_lazy_features @@ -9749,7 +9748,7 @@ async def chat_completion( completion_stream=_iterator, model=e.model, custom_llm_provider="cached_response", - logging_obj=validated_stream_logging_obj(_logging_obj), + logging_obj=_logging_obj, ) selected_data_generator = select_data_generator( response=_streaming_response, @@ -9766,7 +9765,6 @@ async def chat_completion( return _chat_response except RejectedRequestError as e: _data = e.request_data - _rejected_request_data: Final[dict[str, object]] = e.request_data await proxy_logging_obj.post_call_failure_hook( user_api_key_dict=user_api_key_dict, original_exception=e, @@ -9781,7 +9779,7 @@ async def chat_completion( completion_stream=_iterator, model=data.get("model", ""), custom_llm_provider="cached_response", - logging_obj=validated_stream_logging_obj(_rejected_request_data.get("litellm_logging_obj")), + logging_obj=_data.get("litellm_logging_obj", None), ) selected_data_generator = select_data_generator( response=_streaming_response, diff --git a/litellm/proxy/rerank_endpoints/endpoints.py b/litellm/proxy/rerank_endpoints/endpoints.py index fab97f0bdab..45b190c1f9d 100644 --- a/litellm/proxy/rerank_endpoints/endpoints.py +++ b/litellm/proxy/rerank_endpoints/endpoints.py @@ -1,5 +1,8 @@ #### Rerank Endpoints ##### +import asyncio +from typing import Final + import orjson from fastapi import APIRouter, Depends, HTTPException, Request, Response, status from fastapi.responses import ORJSONResponse @@ -10,8 +13,6 @@ from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing router: Final = APIRouter() -import asyncio -from typing import Final @router.post( diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index ee430a41b42..e5ba5182bed 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -95,7 +95,7 @@ def _normalize_tool_dialect( def _is_chat_completions_body(data: Mapping[str, Any]) -> bool: messages: Final = data.get("messages") - if isinstance(messages, list) and len(messages) > 0: + if isinstance(messages, list) and messages: return True return "messages" in data and "input" not in data @@ -497,10 +497,7 @@ async def cursor_chat_completions( from litellm.completion_extras.litellm_responses_transformation.handler import ( responses_api_bridge, ) - from litellm.litellm_core_utils.streaming_handler import ( - CustomStreamWrapper, - validated_stream_logging_obj, - ) + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.proxy.proxy_server import ( async_data_generator, chat_completion, @@ -548,7 +545,7 @@ async def cursor_chat_completions( processor: Final = ProxyBaseLLMRequestProcessing(data=data) - def cursor_data_generator(response, user_api_key_dict, request_data: dict[str, object], request=None): + def cursor_data_generator(response, user_api_key_dict, request_data, request=None): """ Custom generator that transforms Responses API streaming chunks to chat completion chunks. @@ -582,7 +579,7 @@ async def cursor_chat_completions( completion_stream=completion_stream, model=request_data.get("model", ""), custom_llm_provider=None, - logging_obj=validated_stream_logging_obj(logging_obj), + logging_obj=logging_obj, ) # Use async_data_generator to format as SSE return async_data_generator( diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index 9b97dcdb43c..5c2ae134b94 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 3056 + "limit": 3058 }, "ANN002": { "limit": 71 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index ad3f56c7bdf..e33cd34f609 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,6 +1,6 @@ { "LIT001": { - "limit": 23023 + "limit": 23021 }, "LIT002": { "limit": 27148 From c8f385483354a85168d159b90e31c4cd05f58824 Mon Sep 17 00:00:00 2001 From: mateo Date: Tue, 11 Aug 2026 02:44:47 +0000 Subject: [PATCH 09/24] docs: require a user flow and a stuck-at proof in feature requests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/ISSUE_TEMPLATE/feature_request.yml | 57 ++++++++++++++++++++-- CLAUDE.md | 2 +- 2 files changed, 55 insertions(+), 4 deletions(-) diff --git a/.github/ISSUE_TEMPLATE/feature_request.yml b/.github/ISSUE_TEMPLATE/feature_request.yml index 4cc42901897..d1c31e4ff6d 100644 --- a/.github/ISSUE_TEMPLATE/feature_request.yml +++ b/.github/ISSUE_TEMPLATE/feature_request.yml @@ -24,10 +24,61 @@ body: validations: required: true - type: textarea - id: motivation + id: user-flow attributes: - label: Motivation, pitch - description: Please outline the motivation for the proposal. Is your feature request related to a specific problem? e.g., "I'm working on X and would like Y to be possible". If this is related to another GitHub issue, please link here too. + label: User Flow + description: Two numbered lists walking the same end user through the same task, one today without the feature and one with it. Keep the guidance comments in the box while you fill it in, they explain every rule. + value: | + + + Today (without the feature): + + 1. + 2. + 3. + + With the feature: + + 1. + 2. + 3. + validations: + required: true + - type: textarea + id: how-far-you-got + attributes: + label: How far you got + description: Walk the "With the feature" list against a live proxy and paste the commands and output up to the step where you get stuck. Keep the guidance comments in the box while you fill it in, they explain every rule. + value: | + + validations: required: true - type: dropdown diff --git a/CLAUDE.md b/CLAUDE.md index 0e63b86036d..02cc8e024df 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -31,7 +31,7 @@ When creating PRs, don't set base to `main`. `litellm_internal_staging` is the d When writing a PR body, treat the comments and imperative instructions inside @.github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule -Same applies for filing bug reports and .github/ISSUE_TEMPLATE/bug_report.yml +Same applies for filing bug reports and feature requests, and .github/ISSUE_TEMPLATE/bug_report.yml and .github/ISSUE_TEMPLATE/feature_request.yml If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just leave the section blank From 471bb834ba9962452e446eda34308b9c7bbf7b60 Mon Sep 17 00:00:00 2001 From: mateo Date: Tue, 11 Aug 2026 02:45:44 +0000 Subject: [PATCH 10/24] docs: rename the feature request flow lists to before/after this feature Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/ISSUE_TEMPLATE/feature_request.yml | 18 +++++++++--------- 1 file changed, 9 insertions(+), 9 deletions(-) diff --git a/.github/ISSUE_TEMPLATE/feature_request.yml b/.github/ISSUE_TEMPLATE/feature_request.yml index d1c31e4ff6d..4ffa491883c 100644 --- a/.github/ISSUE_TEMPLATE/feature_request.yml +++ b/.github/ISSUE_TEMPLATE/feature_request.yml @@ -27,25 +27,25 @@ body: id: user-flow attributes: label: User Flow - description: Two numbered lists walking the same end user through the same task, one today without the feature and one with it. Keep the guidance comments in the box while you fill it in, they explain every rule. + description: Two numbered lists walking the same end user through the same task, one as it goes today and one as it would ideally go with the feature. Keep the guidance comments in the box while you fill it in, they explain every rule. value: | - - Today (without the feature): + Before this feature (today): 1. 2. 3. - With the feature: + After this feature (ideal user flow): 1. 2. @@ -70,9 +70,9 @@ body: id: how-far-you-got attributes: label: How far you got - description: Walk the "With the feature" list against a live proxy and paste the commands and output up to the step where you get stuck. Keep the guidance comments in the box while you fill it in, they explain every rule. + description: Walk the "After this feature (ideal user flow)" list against a live proxy and paste the commands and output up to the step where you get stuck. Keep the guidance comments in the box while you fill it in, they explain every rule. value: | - + Config / setup the proxy ran with: + + Version or commit: + + Commands and their full output, up to the step that dead-ends: + + What stopped me there: + validations: required: true + - type: checkboxes + id: attempt-attestation + attributes: + label: About that attempt + options: + - label: I ran it against a live proxy myself, with no mocks, and the output above is what I actually got back + required: true - type: dropdown id: component attributes: From 38a1bec5c2de09e722faaf8a498e204d1f24ecd5 Mon Sep 17 00:00:00 2001 From: mateo Date: Tue, 11 Aug 2026 02:56:38 +0000 Subject: [PATCH 12/24] fix(triage): require evidence of the dead-end in feature requests Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/scripts/triage_with_llm.py | 5 +++++ tests/test_litellm/test_github_triage_with_llm.py | 13 +++++++++++++ 2 files changed, 18 insertions(+) diff --git a/.github/scripts/triage_with_llm.py b/.github/scripts/triage_with_llm.py index a7bd145dbc8..0908c3d76cc 100644 --- a/.github/scripts/triage_with_llm.py +++ b/.github/scripts/triage_with_llm.py @@ -597,6 +597,11 @@ def build_issue_prompt(*, title: str, body: str) -> str: that it does not today). - Motivation / use case with a concrete example (config, API call, UI flow, or scenario showing what's blocked today). + - END-TO-END EVIDENCE OF THE DEAD-END: a video, a screenshot, or the + exact command(s) actually run paired with their real output, + showing the point where the flow stops today. Mocked or stubbed + dependencies do NOT count, and an unfilled template scaffold + (bare headings, empty numbered lists) counts as absent. For an issue that is neither a bug report nor a feature request (a question, support request, or discussion), PASS as long as it has a diff --git a/tests/test_litellm/test_github_triage_with_llm.py b/tests/test_litellm/test_github_triage_with_llm.py index 300fd7c0710..3172dc93e82 100644 --- a/tests/test_litellm/test_github_triage_with_llm.py +++ b/tests/test_litellm/test_github_triage_with_llm.py @@ -678,6 +678,19 @@ class TestBuildPrompts: assert "unfilled template scaffold" in normalized assert "counts as absent, not as evidence" in normalized + def test_issue_feature_rubric_requires_evidence_of_the_dead_end( + self, triage_module + ): + # The feature form asks the requester to walk the ideal flow against a + # live proxy and paste output up to the step that dead-ends, so the + # judge has to demand that evidence, and must not accept an unedited + # scaffold of bare headings as if it were a real attempt. + prompt = triage_module.build_issue_prompt(title="t", body="x") + normalized = " ".join(prompt.split()) + assert "END-TO-END EVIDENCE OF THE DEAD-END" in normalized + assert "showing the point where the flow stops today" in normalized + assert "unfilled template scaffold" in normalized + def test_should_not_crash_when_pr_body_contains_curly_braces(self, triage_module): """User-supplied content with `{` / `}` must NOT be re-parsed by `str.format()`. `format` only scans the template literal for From 1d9500066e31fc2c69012b809f8154978dea197b Mon Sep 17 00:00:00 2001 From: mateo Date: Tue, 11 Aug 2026 02:59:10 +0000 Subject: [PATCH 13/24] fix: move feature request guidance out of prefilled values so required means filled Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/ISSUE_TEMPLATE/feature_request.yml | 52 ++++++++-------------- 1 file changed, 18 insertions(+), 34 deletions(-) diff --git a/.github/ISSUE_TEMPLATE/feature_request.yml b/.github/ISSUE_TEMPLATE/feature_request.yml index fc79f5af6d2..5fbed194223 100644 --- a/.github/ISSUE_TEMPLATE/feature_request.yml +++ b/.github/ISSUE_TEMPLATE/feature_request.yml @@ -27,18 +27,16 @@ body: id: user-flow attributes: label: User Flow - description: Two numbered lists walking the same end user through the same task, one as it goes today and one as it would ideally go with the feature. Keep the guidance comments in the box while you fill it in, they explain every rule. - value: | - - - Before this feature (today): - - 1. - 2. - 3. - - After this feature (ideal user flow): - - 1. - 2. - 3. validations: required: true - type: textarea id: how-far-you-got attributes: label: How far you got - description: Walk the "After this feature (ideal user flow)" list against a live proxy and paste the commands and output up to the step where you get stuck. Keep the guidance comments in the box while you fill it in, they explain every rule. - value: | - + description: | + Run as many steps of the "After this feature (ideal user flow)" list as you can against a live proxy you ran yourself (e.g., `litellm --config config.yaml --detailed_debug` on localhost:4000), then paste the commands (e.g., curl) and their full output, ending at the step that dead-ends. Every rule below applies. + - Say plainly what stopped you there, in user terms: the option you passed came back ignored, the response 400'd naming an unsupported field, there is no button on the page for it. This is what proves the feature is genuinely missing rather than undocumented, and it is the single most useful thing you can give us + - No mocks. Where the flow involves a provider call, hit the real provider API, even though that costs real $. `pytest` commands are not enough + - Include the config.yaml (or SDK setup) and env vars the proxy ran with, plus the version or commit you were on. Keep the real values for env vars that aren't sensitive, and redact only the secrets: never paste a real API key, virtual key, database URL, or other credential, here or anywhere else in the issue + - If the provider already supports this, link their API docs and paste a direct call to them succeeding, so we can see the shape LiteLLM should be sending + - For UI asks: include screenshots of the page you got stuck on and its URL. Scrub keys and tokens out of screenshots too, they show up in headers, request panels, and the Admin UI + placeholder: | Config / setup the proxy ran with: Version or commit: @@ -86,7 +71,6 @@ body: Commands and their full output, up to the step that dead-ends: What stopped me there: - validations: required: true - type: checkboxes From 4eb4511f20267c392b5386d4f37d4a63169fc7f1 Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 11 Aug 2026 08:35:22 -0700 Subject: [PATCH 14/24] chore: make it clearer --- CLAUDE.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/CLAUDE.md b/CLAUDE.md index 02cc8e024df..a3c24b84ea8 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -31,7 +31,7 @@ When creating PRs, don't set base to `main`. `litellm_internal_staging` is the d When writing a PR body, treat the comments and imperative instructions inside @.github/pull_request_template.md as rules to follow, not just layout. Agent harnesses may strip HTML comments from copies of that file injected into context, so read .github/pull_request_template.md from disk before writing a PR body to make sure you see every comment rule -Same applies for filing bug reports and feature requests, and .github/ISSUE_TEMPLATE/bug_report.yml and .github/ISSUE_TEMPLATE/feature_request.yml +Same applies for filing bug reports and feature requests, with .github/ISSUE_TEMPLATE/bug_report.yml and .github/ISSUE_TEMPLATE/feature_request.yml, respectively If you're resolving a linear ticket, in the "## Linear ticket" section of the PR, say "Resolves LIT-1234", replacing "LIT-1234" with the actual ticket id that you're resolving. If you don't have the ticket id, don't make one up or search for it. Just leave the section blank From cf1766be1bb49f61c020ed2ba8def74f9be46454 Mon Sep 17 00:00:00 2001 From: mateo Date: Tue, 11 Aug 2026 16:26:10 +0000 Subject: [PATCH 15/24] chore: drop the feature request attempt attestation and tighten its wording Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/ISSUE_TEMPLATE/feature_request.yml | 13 +++---------- 1 file changed, 3 insertions(+), 10 deletions(-) diff --git a/.github/ISSUE_TEMPLATE/feature_request.yml b/.github/ISSUE_TEMPLATE/feature_request.yml index 5fbed194223..30a7e839eb3 100644 --- a/.github/ISSUE_TEMPLATE/feature_request.yml +++ b/.github/ISSUE_TEMPLATE/feature_request.yml @@ -58,11 +58,11 @@ body: description: | Run as many steps of the "After this feature (ideal user flow)" list as you can against a live proxy you ran yourself (e.g., `litellm --config config.yaml --detailed_debug` on localhost:4000), then paste the commands (e.g., curl) and their full output, ending at the step that dead-ends. Every rule below applies. - - Say plainly what stopped you there, in user terms: the option you passed came back ignored, the response 400'd naming an unsupported field, there is no button on the page for it. This is what proves the feature is genuinely missing rather than undocumented, and it is the single most useful thing you can give us - - No mocks. Where the flow involves a provider call, hit the real provider API, even though that costs real $. `pytest` commands are not enough + - Say plainly what stopped you there, in user terms: the option you passed came back ignored, the response 400'd naming an unsupported field, there is no button on the page for it. This is what proves the feature is genuinely missing rather than undocumented + - No mocks. Where the flow involves a provider call, hit the real provider API, even if it costs real $$$. `pytest` commands are not enough - Include the config.yaml (or SDK setup) and env vars the proxy ran with, plus the version or commit you were on. Keep the real values for env vars that aren't sensitive, and redact only the secrets: never paste a real API key, virtual key, database URL, or other credential, here or anywhere else in the issue - If the provider already supports this, link their API docs and paste a direct call to them succeeding, so we can see the shape LiteLLM should be sending - - For UI asks: include screenshots of the page you got stuck on and its URL. Scrub keys and tokens out of screenshots too, they show up in headers, request panels, and the Admin UI + - For UI asks: include screenshots of the page you got stuck on and its URL. Scrub keys and tokens out of screenshots too (for example, the virtual key is briefly shown in the panel right after you create a virtual key) placeholder: | Config / setup the proxy ran with: @@ -73,13 +73,6 @@ body: What stopped me there: validations: required: true - - type: checkboxes - id: attempt-attestation - attributes: - label: About that attempt - options: - - label: I ran it against a live proxy myself, with no mocks, and the output above is what I actually got back - required: true - type: dropdown id: component attributes: From d4468ba63ab3c2a5048935ad5b5d53327ac684ba Mon Sep 17 00:00:00 2001 From: mateo Date: Tue, 11 Aug 2026 16:31:27 +0000 Subject: [PATCH 16/24] chore: say real $ instead of $$$ in the feature request attempt rules Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .github/ISSUE_TEMPLATE/feature_request.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/ISSUE_TEMPLATE/feature_request.yml b/.github/ISSUE_TEMPLATE/feature_request.yml index 30a7e839eb3..41b097041f1 100644 --- a/.github/ISSUE_TEMPLATE/feature_request.yml +++ b/.github/ISSUE_TEMPLATE/feature_request.yml @@ -59,7 +59,7 @@ body: Run as many steps of the "After this feature (ideal user flow)" list as you can against a live proxy you ran yourself (e.g., `litellm --config config.yaml --detailed_debug` on localhost:4000), then paste the commands (e.g., curl) and their full output, ending at the step that dead-ends. Every rule below applies. - Say plainly what stopped you there, in user terms: the option you passed came back ignored, the response 400'd naming an unsupported field, there is no button on the page for it. This is what proves the feature is genuinely missing rather than undocumented - - No mocks. Where the flow involves a provider call, hit the real provider API, even if it costs real $$$. `pytest` commands are not enough + - No mocks. Where the flow involves a provider call, hit the real provider API, even if it costs real $. `pytest` commands are not enough - Include the config.yaml (or SDK setup) and env vars the proxy ran with, plus the version or commit you were on. Keep the real values for env vars that aren't sensitive, and redact only the secrets: never paste a real API key, virtual key, database URL, or other credential, here or anywhere else in the issue - If the provider already supports this, link their API docs and paste a direct call to them succeeding, so we can see the shape LiteLLM should be sending - For UI asks: include screenshots of the page you got stuck on and its URL. Scrub keys and tokens out of screenshots too (for example, the virtual key is briefly shown in the panel right after you create a virtual key) From e59da407503145a2087253f89cb1955d645a309b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 11 Aug 2026 10:43:16 -0700 Subject: [PATCH 17/24] fix(triage): ask for dead-end evidence in feature request recovery comments --- .github/scripts/triage_with_llm.py | 13 +++++++++---- .../test_litellm/test_github_triage_with_llm.py | 17 +++++++++++++++++ 2 files changed, 26 insertions(+), 4 deletions(-) diff --git a/.github/scripts/triage_with_llm.py b/.github/scripts/triage_with_llm.py index 0908c3d76cc..6499a063b39 100644 --- a/.github/scripts/triage_with_llm.py +++ b/.github/scripts/triage_with_llm.py @@ -843,8 +843,11 @@ def format_issue_close_comment(verdict: dict) -> str: "video, a screenshot, or the exact commands you ran with their real output / " "traceback) plus expected vs. actual behavior. Written steps with no run output, " "video, or screenshot don't count, and mocked or stubbed runs don't count.\n" - " - For **feature requests**: a concrete description of what should change, plus a " - "use case and example (config / API call / UI flow).\n" + " - For **feature requests**: a concrete description of what should change, a " + "use case and example (config / API call / UI flow), plus end-to-end evidence of " + "the dead-end (a video, a screenshot, or the exact commands you ran with their " + "real output showing where the flow stops today). Mocked or stubbed runs don't " + "count.\n" "2. Comment `@agent-shin reconsider`. I'll re-run triage and reopen the issue if it " "now meets the bar. (GitHub doesn't let external authors reopen an issue a maintainer " "or bot closed, so the comment-based reconsider is the reliable path.)\n" @@ -950,8 +953,10 @@ def format_grace_warning_issue_comment(verdict: dict) -> str: "screenshot, or the exact commands you ran with their real output / traceback) plus " "expected vs. actual behavior. Written steps with no run output don't count, and " "mocked or stubbed runs don't count.\n" - "- For **feature requests**: a concrete description of what should change, plus a use " - "case and example (config / API call / UI flow).\n" + "- For **feature requests**: a concrete description of what should change, a use " + "case and example (config / API call / UI flow), plus end-to-end evidence of the " + "dead-end (a video, a screenshot, or the exact commands you ran with their real " + "output showing where the flow stops today). Mocked or stubbed runs don't count.\n" "\n" "**If the issue does get auto-closed in 2 hours**, comment `@agent-shin reconsider` " "and I'll re-evaluate. If it now meets the bar, I'll reopen the issue.\n" diff --git a/tests/test_litellm/test_github_triage_with_llm.py b/tests/test_litellm/test_github_triage_with_llm.py index 3172dc93e82..50b198f62c5 100644 --- a/tests/test_litellm/test_github_triage_with_llm.py +++ b/tests/test_litellm/test_github_triage_with_llm.py @@ -207,6 +207,23 @@ class TestCloseCommentText: assert "end-to-end qa proof" in body.lower() assert "mock" in body.lower() + def test_issue_recovery_comments_should_name_feature_dead_end_evidence( + self, triage_module + ): + # The feature-request pass bar demands end-to-end evidence of the + # dead-end, so the close and grace-warning recovery bullets must ask + # for it too — otherwise a requester follows those exact instructions + # (description + use case only) and fails `reconsider` again with no + # hint of what else was needed. + verdict = {"verdict": "fail", "missing": [], "explanation": ""} + for body in ( + triage_module.format_issue_close_comment(verdict), + triage_module.format_grace_warning_issue_comment(verdict), + ): + normalized = " ".join(body.split()) + assert "end-to-end evidence of the dead-end" in normalized + assert "showing where the flow stops today" in normalized + def test_all_agent_shin_comments_should_use_bullet_train_emoji(self, triage_module): # The bullet train (🚅) is Agent Shin's symbol, matching the LiteLLM # logo; the previous wave (👋) was generic and didn't match the bot's From 812abcc7f55153cf5ac11a49d76c0b3a6775f4e3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 11 Aug 2026 10:57:46 -0700 Subject: [PATCH 18/24] fix(triage): track and credit feature dead-end evidence in the verdict --- .github/scripts/triage_with_llm.py | 17 ++++++++---- .../test_github_triage_with_llm.py | 26 +++++++++++++++++++ 2 files changed, 38 insertions(+), 5 deletions(-) diff --git a/.github/scripts/triage_with_llm.py b/.github/scripts/triage_with_llm.py index 6499a063b39..e23a012425a 100644 --- a/.github/scripts/triage_with_llm.py +++ b/.github/scripts/triage_with_llm.py @@ -597,11 +597,13 @@ def build_issue_prompt(*, title: str, body: str) -> str: that it does not today). - Motivation / use case with a concrete example (config, API call, UI flow, or scenario showing what's blocked today). - - END-TO-END EVIDENCE OF THE DEAD-END: a video, a screenshot, or the - exact command(s) actually run paired with their real output, - showing the point where the flow stops today. Mocked or stubbed - dependencies do NOT count, and an unfilled template scaffold - (bare headings, empty numbered lists) counts as absent. + - END-TO-END EVIDENCE OF THE DEAD-END (set + `has_dead_end_evidence=true` only when this is present): a video, + a screenshot, or the exact command(s) actually run paired with + their real output, showing the point where the flow stops today. + Mocked or stubbed dependencies do NOT count, and an unfilled + template scaffold (bare headings, empty numbered lists) counts as + absent. For an issue that is neither a bug report nor a feature request (a question, support request, or discussion), PASS as long as it has a @@ -615,6 +617,7 @@ def build_issue_prompt(*, title: str, body: str) -> str: "has_repro": boolean, "has_expected_vs_actual": boolean, "has_motivation_example": boolean, + "has_dead_end_evidence": boolean, "missing": ["plain-english strings naming what is missing"], "explanation": "1-2 sentence reasoning for the team to skim" }} @@ -712,6 +715,10 @@ _ISSUE_BUG_LABELS: tuple[tuple[str, str], ...] = ( ) _ISSUE_FEATURE_LABELS: tuple[tuple[str, str], ...] = ( ("has_motivation_example", "Motivation and concrete example"), + ( + "has_dead_end_evidence", + "End-to-end evidence of the dead-end (video, screenshot, or command + real output)", + ), ) diff --git a/tests/test_litellm/test_github_triage_with_llm.py b/tests/test_litellm/test_github_triage_with_llm.py index 50b198f62c5..96b77e80457 100644 --- a/tests/test_litellm/test_github_triage_with_llm.py +++ b/tests/test_litellm/test_github_triage_with_llm.py @@ -306,6 +306,27 @@ class TestCloseCommentText: assert "Expected vs. actual behavior" in body assert "- ✅ End-to-end evidence of the bug" not in body + def test_issue_close_comment_should_credit_feature_dead_end_evidence( + self, triage_module + ): + # A feature requester who pasted their dead-end run but skipped the + # motivation must see the evidence credited and only the motivation + # listed as a gap — without a dedicated verdict field the praise + # block could never acknowledge the work they did do. + body = triage_module.format_issue_close_comment( + { + "verdict": "fail", + "kind": "feature", + "has_motivation_example": False, + "has_dead_end_evidence": True, + "missing": ["motivation / use case"], + "explanation": "no use case given", + } + ) + assert "What you got right" in body + assert "- ✅ End-to-end evidence of the dead-end" in body + assert "- ✅ Motivation and concrete example" not in body + def test_close_comments_should_use_softer_park_for_later_framing( self, triage_module ): @@ -707,6 +728,11 @@ class TestBuildPrompts: assert "END-TO-END EVIDENCE OF THE DEAD-END" in normalized assert "showing the point where the flow stops today" in normalized assert "unfilled template scaffold" in normalized + # The evidence has its own verdict field so feature requesters who + # provided it get credited in "What you got right", exactly like + # `has_repro` credits bug evidence. + assert "`has_dead_end_evidence=true` only when this is present" in normalized + assert '"has_dead_end_evidence": boolean' in normalized def test_should_not_crash_when_pr_body_contains_curly_braces(self, triage_module): """User-supplied content with `{` / `}` must NOT be re-parsed by From b144b15d48ddf78e9096488bad531aee8709dfd3 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 11 Aug 2026 11:02:57 -0700 Subject: [PATCH 19/24] fix(proxy): add config_updated_at audit timestamp for virtual keys (#36488) * fix(proxy): add config_updated_at audit timestamp for virtual keys updated_at carries Prisma's @updatedAt, so every batched spend flush rewrites it and it cannot distinguish config changes from usage. Add an additive config_updated_at column stamped only by key management writes (update, bulk update, regenerate, block, unblock) via a shared helper, expose it on key responses, and switch the key page's Last Updated to it with a created_at fallback. * test(proxy): assert config_updated_at survives key archival * refactor(proxy): rename config_updated_at to settings_updated_at --- .../migration.sql | 5 + .../litellm_proxy_extras/schema.prisma | 2 + litellm/models/verification_token.py | 1 + .../key_management_endpoints.py | 7 +- .../management_helpers/key_settings_audit.py | 14 +++ litellm/proxy/schema.prisma | 2 + litellm/proxy/utils.py | 3 +- schema.prisma | 2 + .../proxy/db/test_db_spend_update_writer.py | 64 ++++++++++ .../test_key_management_endpoints.py | 110 ++++++++++++++++-- tests/test_litellm/proxy/test_proxy_utils.py | 22 ++++ .../components/key_team_helpers/key_list.tsx | 1 + .../templates/key_info_view.test.tsx | 40 +++++++ .../components/templates/key_info_view.tsx | 4 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 6 + 15 files changed, 269 insertions(+), 14 deletions(-) create mode 100644 litellm-proxy-extras/litellm_proxy_extras/migrations/20260810000000_add_verificationtoken_settings_updated_at/migration.sql create mode 100644 litellm/proxy/management_helpers/key_settings_audit.py diff --git a/litellm-proxy-extras/litellm_proxy_extras/migrations/20260810000000_add_verificationtoken_settings_updated_at/migration.sql b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260810000000_add_verificationtoken_settings_updated_at/migration.sql new file mode 100644 index 00000000000..fa12f4eb138 --- /dev/null +++ b/litellm-proxy-extras/litellm_proxy_extras/migrations/20260810000000_add_verificationtoken_settings_updated_at/migration.sql @@ -0,0 +1,5 @@ +-- AlterTable +ALTER TABLE "LiteLLM_VerificationToken" ADD COLUMN IF NOT EXISTS "settings_updated_at" TIMESTAMP(3); + +-- AlterTable +ALTER TABLE "LiteLLM_DeletedVerificationToken" ADD COLUMN IF NOT EXISTS "settings_updated_at" TIMESTAMP(3); diff --git a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma index 33fd9389b63..854602f5380 100644 --- a/litellm-proxy-extras/litellm_proxy_extras/schema.prisma +++ b/litellm-proxy-extras/litellm_proxy_extras/schema.prisma @@ -452,6 +452,7 @@ model LiteLLM_VerificationToken { created_by String? updated_at DateTime? @default(now()) @updatedAt @map("updated_at") updated_by String? + settings_updated_at DateTime? @map("settings_updated_at") last_active DateTime? // When this key was last used rotation_count Int? @default(0) // Number of times key has been rotated auto_rotate Boolean? @default(false) // Whether this key should be auto-rotated @@ -548,6 +549,7 @@ model LiteLLM_DeletedVerificationToken { created_by String? // Original creator updated_at DateTime? // Last update timestamp before deletion updated_by String? // Last user who updated before deletion + settings_updated_at DateTime? // Last configuration change before deletion last_active DateTime? // When this key was last used before deletion rotation_count Int? @default(0) auto_rotate Boolean? @default(false) diff --git a/litellm/models/verification_token.py b/litellm/models/verification_token.py index ea822c2dab0..fec3caec457 100644 --- a/litellm/models/verification_token.py +++ b/litellm/models/verification_token.py @@ -49,6 +49,7 @@ class LiteLLM_VerificationToken(LiteLLMPydanticObjectBase): created_by: str | None = None updated_at: datetime | None = None updated_by: str | None = None + settings_updated_at: datetime | None = None last_active: datetime | None = None object_permission_id: str | None = None object_permission: LiteLLM_ObjectPermissionTable | None = None diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 38b5d755535..836b223c6ad 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -88,6 +88,7 @@ from litellm.proxy.management_endpoints.common_utils import ( from litellm.proxy.management_endpoints.model_management_endpoints import ( _add_model_to_db, ) +from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at from litellm.proxy.management_helpers.object_permission_utils import ( _set_object_permission, attach_object_permission_to_dict, @@ -4693,7 +4694,7 @@ async def _execute_virtual_key_regeneration( updated_token: Final = await VerificationTokenRepository(prisma_client).table.update( where={"token": hashed_api_key}, - data=jsonified_update_data, + data=with_settings_updated_at(jsonified_update_data), ) updated_token_dict: Final[dict[str, object]] = dict(updated_token) if updated_token is not None else {} updated_token_dict["key"] = new_token @@ -6203,7 +6204,7 @@ async def block_key( record: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).update( where={"token": hashed_token}, - data={"blocked": True}, + data=with_settings_updated_at({"blocked": True}), ) ## UPDATE KEY CACHE - invalidate so next read re-fetches from DB @@ -6316,7 +6317,7 @@ async def unblock_key( record: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).update( where={"token": hashed_token}, - data={"blocked": False}, + data=with_settings_updated_at({"blocked": False}), ) ## UPDATE KEY CACHE - invalidate so next read re-fetches from DB diff --git a/litellm/proxy/management_helpers/key_settings_audit.py b/litellm/proxy/management_helpers/key_settings_audit.py new file mode 100644 index 00000000000..a2c4bd8cac0 --- /dev/null +++ b/litellm/proxy/management_helpers/key_settings_audit.py @@ -0,0 +1,14 @@ +"""Audit stamping for virtual key configuration changes.""" + +from collections.abc import Mapping +from datetime import datetime, timezone + + +def with_settings_updated_at(data: Mapping[str, object]) -> dict[str, object]: + """Stamp a key update payload with the time its configuration changed. + + ``updated_at`` carries Prisma's ``@updatedAt`` and so is rewritten by every + spend flush, which makes it useless for auditing; ``settings_updated_at`` is + written only from key-management write paths. + """ + return {**data, "settings_updated_at": datetime.now(timezone.utc)} diff --git a/litellm/proxy/schema.prisma b/litellm/proxy/schema.prisma index 33fd9389b63..854602f5380 100644 --- a/litellm/proxy/schema.prisma +++ b/litellm/proxy/schema.prisma @@ -452,6 +452,7 @@ model LiteLLM_VerificationToken { created_by String? updated_at DateTime? @default(now()) @updatedAt @map("updated_at") updated_by String? + settings_updated_at DateTime? @map("settings_updated_at") last_active DateTime? // When this key was last used rotation_count Int? @default(0) // Number of times key has been rotated auto_rotate Boolean? @default(false) // Whether this key should be auto-rotated @@ -548,6 +549,7 @@ model LiteLLM_DeletedVerificationToken { created_by String? // Original creator updated_at DateTime? // Last update timestamp before deletion updated_by String? // Last user who updated before deletion + settings_updated_at DateTime? // Last configuration change before deletion last_active DateTime? // When this key was last used before deletion rotation_count Int? @default(0) auto_rotate Boolean? @default(false) diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index dd0c57aa911..b5972935806 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -135,6 +135,7 @@ from litellm.proxy.hooks.sensitive_data_routing import ( _PROXY_SensitiveDataRoutingHandler, ) from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup +from litellm.proxy.management_helpers.key_settings_audit import with_settings_updated_at from litellm.proxy.policy_engine.pipeline_executor import PipelineExecutor from litellm.repositories.budget_repository import BudgetRepository from litellm.repositories.config_repository import ConfigRepository @@ -3996,7 +3997,7 @@ class PrismaClient: db_data["token"] = token response: Final = await VerificationTokenRepository(self).table.update( where={"token": token}, - data={**db_data}, + data=with_settings_updated_at(db_data), ) verbose_proxy_logger.debug("\033[91m" + f"DB Token Table update succeeded {response}" + "\033[0m") _data: dict = {} diff --git a/schema.prisma b/schema.prisma index 33fd9389b63..854602f5380 100644 --- a/schema.prisma +++ b/schema.prisma @@ -452,6 +452,7 @@ model LiteLLM_VerificationToken { created_by String? updated_at DateTime? @default(now()) @updatedAt @map("updated_at") updated_by String? + settings_updated_at DateTime? @map("settings_updated_at") last_active DateTime? // When this key was last used rotation_count Int? @default(0) // Number of times key has been rotated auto_rotate Boolean? @default(false) // Whether this key should be auto-rotated @@ -548,6 +549,7 @@ model LiteLLM_DeletedVerificationToken { created_by String? // Original creator updated_at DateTime? // Last update timestamp before deletion updated_by String? // Last user who updated before deletion + settings_updated_at DateTime? // Last configuration change before deletion last_active DateTime? // When this key was last used before deletion rotation_count Int? @default(0) auto_rotate Boolean? @default(false) diff --git a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py index 51810a28cdd..2c426e0f071 100644 --- a/tests/test_litellm/proxy/db/test_db_spend_update_writer.py +++ b/tests/test_litellm/proxy/db/test_db_spend_update_writer.py @@ -2157,3 +2157,67 @@ async def test_daily_transaction_compression_saved_tokens_zero_when_absent(): assert transaction["compression_saved_tokens"] == 0 assert transaction["compression_savings_spend"] == 0 assert transaction["prompt_caching_savings_spend"] == 0 + + +@pytest.mark.asyncio +async def test_commit_spend_updates_to_db_does_not_stamp_key_settings_updated_at(): + """Spend flushes must leave settings_updated_at alone, or it decays into + another `updated_at` and stops being an audit signal.""" + db_writer = DBSpendUpdateWriter() + + mock_batcher = MagicMock() + mock_batcher.litellm_verificationtoken = MagicMock() + mock_batcher.litellm_verificationtoken.update_many = MagicMock() + mock_batcher.litellm_usertable = MagicMock() + mock_batcher.litellm_usertable.update_many = MagicMock() + mock_batcher.litellm_teamtable = MagicMock() + mock_batcher.litellm_teamtable.update_many = MagicMock() + mock_batcher.litellm_teammembership = MagicMock() + mock_batcher.litellm_teammembership.update_many = MagicMock() + mock_batcher.litellm_organizationtable = MagicMock() + mock_batcher.litellm_organizationtable.update_many = MagicMock() + mock_batcher.litellm_tagtable = MagicMock() + mock_batcher.litellm_tagtable.update_many = MagicMock() + mock_batcher.litellm_agentstable = MagicMock() + mock_batcher.litellm_agentstable.update_many = MagicMock() + + mock_transaction = AsyncMock() + mock_transaction.__aenter__ = AsyncMock(return_value=mock_transaction) + mock_transaction.__aexit__ = AsyncMock(return_value=False) + mock_transaction.batch_ = MagicMock( + return_value=AsyncMock( + __aenter__=AsyncMock(return_value=mock_batcher), + __aexit__=AsyncMock(return_value=False), + ) + ) + + mock_prisma_client = MagicMock() + mock_prisma_client.db = MagicMock() + mock_prisma_client.db.tx = MagicMock(return_value=mock_transaction) + + token = "hashed-token-abc" + response_cost = 0.25 + db_spend_update_transactions = { + "user_list_transactions": {}, + "end_user_list_transactions": {}, + "key_list_transactions": {token: response_cost}, + "team_list_transactions": {}, + "team_member_list_transactions": {}, + "org_list_transactions": {}, + "tag_list_transactions": {}, + "agent_list_transactions": {}, + } + + with patch("litellm.proxy.utils._raise_failed_update_spend_exception"): + await db_writer._commit_spend_updates_to_db( + prisma_client=mock_prisma_client, + n_retry_times=0, + proxy_logging_obj=MagicMock(), + db_spend_update_transactions=db_spend_update_transactions, + ) + + mock_batcher.litellm_verificationtoken.update_many.assert_called_once() + call_kwargs = mock_batcher.litellm_verificationtoken.update_many.call_args[1] + assert call_kwargs["where"] == {"token": token} + assert set(call_kwargs["data"]) == {"spend", "last_active"} + assert call_kwargs["data"]["spend"] == {"increment": response_cost} diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 8f151ed882c..4c13a3367d9 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -2293,9 +2293,9 @@ async def test_unblock_key_supports_both_sk_and_hashed_tokens(monkeypatch): ) # Verify that the database update was called with hashed token - mock_prisma_client.db.litellm_verificationtoken.update.assert_called_with( - where={"token": test_hashed_token}, data={"blocked": False} - ) + sk_token_call = mock_prisma_client.db.litellm_verificationtoken.update.call_args.kwargs + assert sk_token_call["where"] == {"token": test_hashed_token} + assert sk_token_call["data"]["blocked"] is False assert result == mock_key_record @@ -2313,9 +2313,9 @@ async def test_unblock_key_supports_both_sk_and_hashed_tokens(monkeypatch): ) # Verify that the database update was called with the same hashed token - mock_prisma_client.db.litellm_verificationtoken.update.assert_called_with( - where={"token": test_hashed_token}, data={"blocked": False} - ) + hashed_token_call = mock_prisma_client.db.litellm_verificationtoken.update.call_args.kwargs + assert hashed_token_call["where"] == {"token": test_hashed_token} + assert hashed_token_call["data"]["blocked"] is False assert result == mock_key_record @@ -2849,9 +2849,10 @@ async def test_block_key_existing_key_succeeds(monkeypatch): mock_prisma_client.db.litellm_verificationtoken.find_unique.assert_called_once_with( where={"token": test_hashed_token} ) - mock_prisma_client.db.litellm_verificationtoken.update.assert_called_once_with( - where={"token": test_hashed_token}, data={"blocked": True} - ) + mock_prisma_client.db.litellm_verificationtoken.update.assert_called_once() + block_call = mock_prisma_client.db.litellm_verificationtoken.update.call_args.kwargs + assert block_call["where"] == {"token": test_hashed_token} + assert block_call["data"]["blocked"] is True assert result == mock_updated_record @@ -4717,6 +4718,7 @@ def test_transform_verification_tokens_to_deleted_records(): user_role=LitellmUserRoles.PROXY_ADMIN.value, ) + config_stamp = datetime(2026, 8, 10, 12, 30, 45, tzinfo=timezone.utc) key1 = LiteLLM_VerificationToken( token="hashed-token-1", user_id="user-123", @@ -4733,6 +4735,7 @@ def test_transform_verification_tokens_to_deleted_records(): model_spend={}, soft_budget_cooldown=False, allowed_routes=[], + settings_updated_at=config_stamp, ) key2 = LiteLLM_VerificationToken( @@ -4775,6 +4778,7 @@ def test_transform_verification_tokens_to_deleted_records(): assert record1["token"] == "hashed-token-1" assert record1["user_id"] == "user-123" assert record1["team_id"] == "team-456" + assert record1["settings_updated_at"] == config_stamp assert isinstance(record1["aliases"], str) assert isinstance(record1["config"], str) assert isinstance(record1["permissions"], str) @@ -15817,3 +15821,91 @@ async def test_regenerate_key_output_token_estimate_lowered_rejected_for_non_adm assert exc.value.status_code == 403 assert "Only proxy admins can set" in str(exc.value.detail) + + +@pytest.mark.asyncio +async def test_execute_virtual_key_regeneration_stamps_settings_updated_at(): + """Regenerate rewrites the key's config, so it must move settings_updated_at.""" + from datetime import datetime, timezone + + from litellm.proxy._types import RegenerateKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import ( + _execute_virtual_key_regeneration, + ) + + mock_prisma_client = _make_regenerate_mock_prisma() + + with _patch_regenerate_side_effects(): + before = datetime.now(timezone.utc) + await _execute_virtual_key_regeneration( + prisma_client=mock_prisma_client, + key_in_db=_make_regenerate_existing_key(), + hashed_api_key="abc123", + key="abc123", + data=RegenerateKeyRequest(max_budget=42.0), + user_api_key_dict=_make_regenerate_user_api_key_dict(), + litellm_changed_by=None, + user_api_key_cache=MagicMock(), + proxy_logging_obj=MagicMock(), + ) + after = datetime.now(timezone.utc) + + sent = mock_prisma_client.db.litellm_verificationtoken.update.call_args.kwargs["data"] + assert sent["max_budget"] == 42.0 + assert before <= sent["settings_updated_at"] <= after + + +@pytest.mark.asyncio +async def test_block_key_stamps_settings_updated_at(monkeypatch): + """Blocking a key is a config change, not spend activity.""" + from datetime import datetime, timezone + + from litellm.proxy._types import BlockKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import block_key + + mock_prisma_client, _ = _setup_block_unblock_mocks(monkeypatch) + + before = datetime.now(timezone.utc) + await block_key( + data=BlockKeyRequest(key="sk-test123456789"), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin_user", + ), + litellm_changed_by=None, + ) + after = datetime.now(timezone.utc) + + sent = mock_prisma_client.db.litellm_verificationtoken.update.call_args.kwargs["data"] + assert sent["blocked"] is True + assert before <= sent["settings_updated_at"] <= after + + +@pytest.mark.asyncio +async def test_unblock_key_stamps_settings_updated_at(monkeypatch): + """Unblocking a key is a config change, not spend activity.""" + from datetime import datetime, timezone + + from litellm.proxy._types import BlockKeyRequest + from litellm.proxy.management_endpoints.key_management_endpoints import unblock_key + + mock_prisma_client, _ = _setup_block_unblock_mocks(monkeypatch) + + before = datetime.now(timezone.utc) + await unblock_key( + data=BlockKeyRequest(key="sk-test123456789"), + http_request=MagicMock(), + user_api_key_dict=UserAPIKeyAuth( + user_role=LitellmUserRoles.PROXY_ADMIN, + api_key="sk-admin", + user_id="admin_user", + ), + litellm_changed_by=None, + ) + after = datetime.now(timezone.utc) + + sent = mock_prisma_client.db.litellm_verificationtoken.update.call_args.kwargs["data"] + assert sent["blocked"] is False + assert before <= sent["settings_updated_at"] <= after diff --git a/tests/test_litellm/proxy/test_proxy_utils.py b/tests/test_litellm/proxy/test_proxy_utils.py index a8e81e92ebd..a4f93e90673 100644 --- a/tests/test_litellm/proxy/test_proxy_utils.py +++ b/tests/test_litellm/proxy/test_proxy_utils.py @@ -1169,3 +1169,25 @@ async def test_prisma_health_check_failure_redacts_database_credentials(caplog): assert emitted assert all("hunter2" not in message for message in emitted) assert any("postgresql://REDACTED@db.internal" in message for message in emitted) + + +@pytest.mark.asyncio +async def test_update_data_key_branch_stamps_settings_updated_at(): + """`updated_at` carries Prisma's @updatedAt and is rewritten by every spend + flush, so key config edits need their own audit column.""" + from datetime import datetime, timezone + from unittest.mock import AsyncMock + + from litellm.proxy.utils import PrismaClient + + client = MagicMock() + client.jsonify_object = MagicMock(side_effect=lambda data: dict(data)) + client.db.litellm_verificationtoken.update = AsyncMock(return_value=None) + + before = datetime.now(timezone.utc) + await PrismaClient.update_data(client, token="sk-test-key", data={"models": ["gpt-4"]}) + after = datetime.now(timezone.utc) + + sent = client.db.litellm_verificationtoken.update.call_args.kwargs["data"] + assert sent["models"] == ["gpt-4"] + assert before <= sent["settings_updated_at"] <= after diff --git a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx index 4b446b0c283..920d1f4af5a 100644 --- a/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx +++ b/ui/litellm-dashboard/src/components/key_team_helpers/key_list.tsx @@ -60,6 +60,7 @@ export interface KeyResponse { created_at: string; created_by?: string; updated_at: string; + settings_updated_at?: string | null; last_active: string | null; team_spend: number; team_alias: string; diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx index 8e9dd2fa975..680fbe9ff1e 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.test.tsx @@ -149,6 +149,46 @@ describe("KeyInfoView", () => { await userEvent.click(await screen.findByRole("button", { name: /more key actions/i })); }; + describe("last updated", () => { + const renderWithTimestamps = (overrides: Partial) => { + vi.mocked(useAuthorized).mockReturnValue(baseUseAuthorizedMock); + + return renderWithProviders( + {}} + keyId={"test-key-id"} + onKeyDataUpdate={() => {}} + teams={[]} + />, + ); + }; + + const findLastUpdatedText = async () => { + const label = await screen.findByText("Last Updated"); + return label.closest("div")?.parentElement?.parentElement?.textContent ?? ""; + }; + + it("should show when the key was last configured, not when it last recorded spend", async () => { + renderWithTimestamps({ settings_updated_at: "2022-06-15T12:00:00Z" }); + + expect(await findLastUpdatedText()).toMatch(/Jun \d+, 2022/); + expect(screen.queryByText(/Jun \d+, 2023/)).not.toBeInTheDocument(); + }); + + it("should fall back to creation time for a key that was never reconfigured", async () => { + renderWithTimestamps({ settings_updated_at: null }); + + expect(await findLastUpdatedText()).toMatch(/Jun \d+, 2021/); + expect(screen.queryByText(/Jun \d+, 2023/)).not.toBeInTheDocument(); + }); + }); + it("should render tags", async () => { vi.mocked(useAuthorized).mockReturnValue(baseUseAuthorizedMock); diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index b1ada186bb8..6a2a4a88425 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -444,6 +444,8 @@ export default function KeyInfoView({ ); }; + const lastConfiguredAt = currentKeyData.settings_updated_at || currentKeyData.created_at; + const parentTeam = currentKeyData.team_id ? teamsData?.find((team) => team.team_id === currentKeyData.team_id) : null; const budgetDisplay = @@ -468,7 +470,7 @@ export default function KeyInfoView({ currentKeyData.created_by || "", createdAt: currentKeyData.created_at ? formatTimestamp(currentKeyData.created_at) : "", - lastUpdated: currentKeyData.updated_at ? formatTimestamp(currentKeyData.updated_at) : "", + lastUpdated: lastConfiguredAt ? formatTimestamp(lastConfiguredAt) : "", lastActive: currentKeyData.last_active ? formatTimestamp(currentKeyData.last_active) : "Never", expires: currentKeyData.expires ? formatTimestamp(currentKeyData.expires) : "Never", }} diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 8989e098f54..35adc9e92cf 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -26280,6 +26280,8 @@ export interface components { } | null; /** Rpm Limit */ rpm_limit?: number | null; + /** Settings Updated At */ + settings_updated_at?: string | null; /** * Soft Budget Cooldown * @default false @@ -27719,6 +27721,8 @@ export interface components { } | null; /** Rpm Limit */ rpm_limit?: number | null; + /** Settings Updated At */ + settings_updated_at?: string | null; /** * Soft Budget Cooldown * @default false @@ -34887,6 +34891,8 @@ export interface components { rpm_limit_per_model?: { [key: string]: number; } | null; + /** Settings Updated At */ + settings_updated_at?: string | null; /** Soft Budget */ soft_budget?: number | null; /** From e9156afdd03ff972dd79791f75369b8ab37fed1e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Tue, 11 Aug 2026 11:34:26 -0700 Subject: [PATCH 20/24] ci: retry transient network fetch failures in lint workflow --- .github/workflows/test-linting.yml | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/.github/workflows/test-linting.yml b/.github/workflows/test-linting.yml index 280ec476cdf..69495cff896 100644 --- a/.github/workflows/test-linting.yml +++ b/.github/workflows/test-linting.yml @@ -43,9 +43,10 @@ jobs: BASE_SHA: ${{ github.event.pull_request.base.sha }} HEAD_SHA: ${{ github.event.pull_request.head.sha }} run: | - MERGE_BASE=$(gh api "repos/${{ github.repository }}/compare/${BASE_SHA}...${HEAD_SHA}?per_page=1" --jq '.merge_base_commit.sha') + retry() { "$@" || { sleep 15; "$@"; } || { sleep 30; "$@"; }; } + MERGE_BASE=$(retry gh api "repos/${{ github.repository }}/compare/${BASE_SHA}...${HEAD_SHA}?per_page=1" --jq '.merge_base_commit.sha') test -n "$MERGE_BASE" - git fetch --no-tags --depth=1 origin "$MERGE_BASE" + retry git fetch --no-tags --depth=1 origin "$MERGE_BASE" echo "GATE_BASE_SHA=$MERGE_BASE" >> "$GITHUB_ENV" - name: Set up Python @@ -161,7 +162,8 @@ jobs: env: BASE_SHA: ${{ github.event.pull_request.base.sha }} run: | - git fetch --no-tags --depth=1 origin "$BASE_SHA" + retry() { "$@" || { sleep 15; "$@"; } || { sleep 30; "$@"; }; } + retry git fetch --no-tags --depth=1 origin "$BASE_SHA" - name: Set up Python uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 @@ -205,7 +207,8 @@ jobs: GITGUARDIAN_API_KEY: ${{ secrets.GITGUARDIAN_API_KEY }} run: | if [ -n "$GITGUARDIAN_API_KEY" ]; then - git fetch --no-tags --unshallow origin + retry() { "$@" || { sleep 15; "$@"; } || { sleep 30; "$@"; }; } + retry git fetch --no-tags --unshallow origin uv tool run --from 'ggshield==1.48.0' ggshield secret scan repo . else echo "GITGUARDIAN_API_KEY not set, skipping ggshield scan" From 84d6666a594e7ddbf95ba51db8c74f547f5f38ee Mon Sep 17 00:00:00 2001 From: Deepanshu Lulla Date: Tue, 11 Aug 2026 14:49:18 -0400 Subject: [PATCH 21/24] feat(router): add required-AND (&) tag prefix and allow_fail_open flag (#36193) * feat(router): add required-AND (&) tag prefix and allow_fail_open flag Tag routing supported inclusion-OR and independent "!" negation, but had no way to express a hard "must match all of these" constraint per request, and no way for a model group to opt into degrading gracefully instead of raising when a constraint eliminates every deployment. Adds a "&tag" prefix for required-AND inclusion, composing with existing plain (OR) and "!" (negate) tags: negation still applies first, then required tags narrow the survivors, then plain tags apply today's OR/AND preference logic unchanged. Adds model_info.allow_fail_open (default false) so a chain can opt into falling back to the default-tagged pool instead of raising no_deployments_with_tag_routing when "!" or "&" empties the candidate set; existing chains without the flag keep today's fail-closed behavior exactly. * fix(router): gate mixed negation on allow_fail_open and stop diluting required-only requests Two gaps in the initial required-AND/allow_fail_open change: a "!" exclusion combined with a plain positive tag that emptied the candidate set raised unconditionally, bypassing allow_fail_open entirely, since the fail-open check only looked at required-AND exhaustion. And a request using only "&" tags could get narrowed down to just the deployment matching an incidental tag_regex/User-Agent preference, silently dropping other deployments that satisfied the required tags but had no tag_regex at all. Fixes both: the fail-open check now fires whenever either "!" or "&" leaves the candidate set empty, not just "&". And regex/header preference no longer counts as a positive filter when a required-AND ask is present, so a required-only request returns every deployment satisfying the required tags regardless of regex/header matching. Also regenerates ui/litellm-dashboard/src/lib/http/schema.d.ts for the new model_info.allow_fail_open field, and removes source comments explaining the router logic per repository convention. * fix(router): let allow_fail_open cover a non-empty !/& survivor set that fails the plain-tag preference The unconditional raise inside the has_positive_filter loop was the one remaining path a chain could hit despite setting allow_fail_open: when "!" or "&" leaves a non-empty candidate set but none of the survivors match the request's plain preference tag or carry "default", the request still raised instead of degrading. Routes that raise through the same allow_fail_open check used everywhere else, so it now falls back to the default-tagged pool for opted-in chains and keeps raising unconditionally for everyone else. This also let the now-redundant pre-loop empty-candidates shortcut be removed, since the loop reaches the same outcome on its own. * fix(router): deny allow_fail_open when an unrecognized required tag is masking a satisfiable answer A caller could add a single "&" tag no deployment in the group has ever carried to force an empty required-AND set on demand. On a chain with allow_fail_open, that emptied set fell back to the default-tagged pool unconditionally, discarding every other constraint merged into the same request, including ones inherited from key/team policy, even when the rest of those constraints were still individually satisfiable. Before falling back, drop any required tag not carried by any deployment in the group and recompute: if a specific, non-empty answer exists using only the recognized tags, the unrecognized tag was the actual cause of the exhaustion, and fail-open must not paper over it. If every required tag is already recognized, or none are, there's nothing hidden behind an invented tag, and fail-open proceeds exactly as before; this keeps a single opted-in deployment's legitimate catch-all behavior working when a caller's tag simply doesn't exist anywhere in that group. Ratchets ANN401 and LIT001 budgets down to reflect fixes already earned in this branch. * test(router): cover required-AND, allow_fail_open, and unknown-tag denial across fallback chains and model groups Extends coverage beyond single-hop scenarios: & exhausting a primary group falls through to a fallback group exactly like ! already does; !, &, and allow_fail_open composed together across three chained model groups each raise or fall back independently per-hop; and the unknown-tag denial from the previous commit is evaluated fresh per hop rather than leaking state across groups in a fallback chain. * feat(router): add model_info.enable_tag_filtering per-model-group override enable_tag_filtering was router-wide only: an operator turning it on for one model group that needs tag-driven routing exposed every other model group on the same proxy to the same tag evaluation, even ones that never use tags. Adds model_info.enable_tag_filtering, checked against any deployment sharing a model_name, so a chain can flip the router-wide default in either direction for itself alone: opt a specific group into filtering while the rest of the proxy stays off, or opt a group out (e.g. an incident-response catch-all) while the rest of the proxy enforces it. Precedence, low to high: router-wide default, then the chain override if set, then the existing request-level escalation (from key/team settings), which still only ever turns filtering on, never off, over whatever the router and chain already decided. Also regenerates ui/litellm-dashboard/src/lib/http/schema.d.ts for the new field. * fix(router): gate plain-tag exhaustion on allow_fail_open when the tag is known to the group A model group where every deployment is tagged "default" (a legitimate cross-cutting safety-net pattern) never has an empty default_deployments list, so the existing exhaustion check (len(new)==0 and len(default)==0) never fired for a plain positive tag that matched nothing among the currently healthy candidates. The request silently fell through to whatever "default"-tagged deployment happened to survive, even when allow_fail_open was never set and the caller's intent (e.g. quality:high) was never honored. Adds a check for whether the requested tag is part of the group's real vocabulary at all: if some deployment configured under this model_name (regardless of current health) genuinely carries the tag, and nothing healthy currently matches it, the request now raises by default or falls back per allow_fail_open, through the same _resolve_or_fail_open gate every other exhaustion path already uses. A tag that's foreign to the group entirely (e.g. one meant for an unrelated mechanism sharing the same request-tags list) keeps falling back to the default pool unconditionally, unchanged, since there's nothing this group's own routing intent could be violating. * fix(router): preserve inherited tag constraints when allow_fail_open discards a caller-caused exhaustion Adds metadata.caller_tags in litellm_pre_call_utils.py, populated only from what the request itself supplied (header, body tags, body metadata.tags), never from key/team metadata merged into the same metadata.tags list. get_deployments_for_tag now uses it to compute a trusted-only pool before falling open: a required/excluded tag attributable to the caller can be discarded on fail-open, one inherited from key/team policy cannot. If the trusted-only pool is itself empty, allow_fail_open raises instead of silently routing around an unsatisfiable inherited constraint. When caller_tags carries no information at all (direct SDK Router usage, bypassing the proxy layer), behavior is unchanged: unconditional fall-open to the default pool, exactly as before this fix. * feat(router): add opt-in tag_routing_prefix for collision-proof tag disambiguation router_settings.tag_routing_prefix lets a caller explicitly mark which x-litellm-tags/metadata.tags values are routing directives, exempting them from the known-tag-vocabulary heuristic used to guard fail-open against caller-invented "&"/"!" tags. Unprefixed tags keep going through today's existing handling unchanged (hybrid, no migration required); default "" is a full no-op. Fixes a bug caught during live-proxy verification: the prefix-stripped "confirmed" set kept the "&"/"!" marker character, so it never matched required_set/excluded_set (which _split_tags always strips bare) -- the entire trusted-required/excluded-tag mechanism silently no-opped for its primary use case. Adds regression tests for the bare-value mismatch and updates existing _chain_allows_fail_open/_tag_known_to_group/ _caller_constraint_sets call sites for the new routing_confirmed/ routing_prefix parameters. * fix(router): resolve model_info.enable_tag_filtering override from the full model group, not just healthy deployments Cooldown filtering runs before get_deployments_for_tag, so _chain_tag_filtering_override only saw the survivors of that filter. A model group whose only enable_tag_filtering-carrying deployment goes into cooldown lost the override entirely, silently falling back to the router-wide default and letting any !/&/tag constraint on that chain be bypassed by driving the one overriding deployment into cooldown. Resolve the override from every deployment configured for the model instead, mirroring _tag_known_to_group's existing pattern. Verified live: with a bad-key deployment carrying the override forced into real cooldown via allowed_fails=1, an explicit "!provider:openai" ban on the remaining deployment reproducibly returned 200 via OpenAI before this fix and 401 (tag filtering still enforced) after it. * fix(router): avoid Final-reassignment lint error and a MagicMock router fixture gap from tag_routing_prefix _chain_tag_filtering_override's try/except reassigned a Final-annotated name across branches, which basedpyright flags as illegal; extracted the lookup-with-fallback into its own helper so the binding is assigned once. Also sets tag_routing_prefix on the bare MagicMock router used by test_router_tag_regex_routing.py's fixture, which otherwise returns an auto-generated MagicMock (truthy, non-string) for the new attribute and crashes _strip_routing_prefix's removeprefix() call. * fix(router): key inherited-tag protection off provenance, not value subtraction allow_fail_open's trusted-only pool computed "not caller-attributable" as required_set - caller_required_set. A caller who resubmits the exact value of an inherited "&"/"!" tag (e.g. an inherited "®ion:eu" alongside a caller-supplied "®ion:eu" plus a conflicting "!region:eu") collapses both origins to the same set value, so the subtraction zeroes out the inherited requirement's protection too, letting fail-open route outside a key/team-enforced constraint. Adds metadata.inherited_tags in litellm_pre_call_utils.py: a snapshot of "tags" taken after key/team/project policy is merged in but before this request's own caller-supplied tags are merged on top. A required or excluded tag is now protected from fail-open discard if it has ANY inherited backing (set intersection with inherited_tags), regardless of whether the caller also happens to submit the identical value -- this is what set membership alone could never tell apart under the old subtraction-based approach. caller_tags is kept (documented as the complementary record) but no longer consulted for this decision. Verified live: a virtual key with metadata.tags=["®ion:eu"] hit with header x-litellm-tags: ®ion:eu,!region:eu (the exact value-collision attack) reproducibly routed to the OpenAI/us deployment before this fix and stayed on the Anthropic/eu deployment after it. * fix(lint): re-ratchet budget ceilings after rebasing onto litellm_internal_staging Regenerated ruff-strict-budget.json and type-discipline-budget.json via make lint-ruff-budget-update / lint-type-discipline-budget-update against the post-rebase merge-base. * fix(proxy): compute inherited_tags from key/team/project sources directly, not a tags-list snapshot apply_client_tag_policy_pre_auth (run from user_api_key_auth, for _tag_max_budget_check) merges the caller's x-litellm-tags header into the same metadata.tags list before add_litellm_data_to_request ever runs. The previous inherited_tags snapshot ("whatever's in tags before this function's own caller-tag merge") therefore misattributed that caller-controlled value as policy-backed whenever a request arrived with the header set -- Greptile flagged this as a P1 security finding. inherited_tags is now built directly from key_metadata/team_metadata/ project_metadata's own "tags" fields, independent of the shared, pipeline-position-dependent "tags" list's mutation history. Verified with a direct reproduction mirroring the real pipeline (calling apply_client_tag_policy_pre_auth on the same data dict before add_litellm_data_to_request, as user_api_key_auth actually does): the caller's header tag no longer appears in inherited_tags. Added a regression test exercising that same call order; confirmed it fails against the pre-fix snapshot approach and passes against this fix. * fix(router): make tag_routing_prefix configurable through update_settings/get_settings and UpdateRouterConfig router_settings.tag_routing_prefix was only ever applied via the Router() constructor. Router.update_settings's _allowed_settings (used directly by proxy_server.py's _add_router_settings_from_db_config for the DB-backed router_settings path) and get_settings's vars_to_include both omitted it, so an operator relying on that path had the value silently ignored -- flagged by veria-ai. Also adds it to UpdateRouterConfig (the pydantic schema behind POST /config/update), the same bug shape LIT-3152 previously fixed for retry_policy: a field missing from that schema gets silently dropped by model_dump(exclude_none=True) before update_settings is ever called. * chore(ui): regenerate schema.d.ts for UpdateRouterConfig.tag_routing_prefix Adding tag_routing_prefix to UpdateRouterConfig changed the proxy's OpenAPI spec; regenerate the dashboard's generated API types to match. * fix(lint): re-ratchet budget ceilings after rebasing onto litellm_internal_staging Regenerated ruff-strict-budget.json and type-discipline-budget.json against the post-rebase merge-base. LIT002/LIT011 ceilings reflect this branch's true current counts (confirmed unchanged across the rebase by diffing against the pre-rebase commit); the base's own counts moved independently. * fix(lint): replace mutable-collection fallbacks with immutable ones in inherited_tags computation key_metadata/team_metadata/project_metadata's "tags" fallbacks used `or {}` / `or []` literals, each a LIT002 mutable-collection-construction violation that pushed the branch 4 over its ratchet ceiling relative to a moved base. Swapped to MappingProxyType({}) / () to match the immutable idiom the rest of tag_based_routing.py already uses; no behavior change, since both are falsy and only ever read via .get()/ unpacking. Tightens type-discipline-budget.json's LIT002 ceiling back down to match, fully closing that gap (LIT011 keeps a genuine 1-count gap from pre-existing, untouched lines in this file, non-gating). * fix(lint): suppress LIT011 on the two new data[...] mutation sites Both new lines follow this file's established data[...] mutation idiom for add_litellm_data_to_request, matching the existing suppression already on the inherited_tags line. * test(router): lock in fallback + tag-filtering interaction Cover the router-level fallbacks mechanism composing with tag-based routing: a plain negation exhausting a group correctly advances to the fallback group, the same exclusion tag exhausting every hop correctly raises, and allow_fail_open resolving locally must not spuriously trigger an unrelated external fallback. * chore: retrigger CI now that litellm-docs#814 is merged --------- Co-authored-by: Deepanshu --- litellm/proxy/litellm_pre_call_utils.py | 49 +- litellm/router.py | 4 + litellm/router_strategy/tag_based_routing.py | 376 +++- litellm/types/router.py | 15 + ruff-strict-budget.json | 2 +- .../proxy/test_litellm_pre_call_utils.py | 236 +++ .../test_router_tag_regex_routing.py | 1 + .../test_router_tag_routing.py | 1717 ++++++++++++++++- type-discipline-budget.json | 6 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 6 + 10 files changed, 2346 insertions(+), 66 deletions(-) diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 10142a894a1..0924b6aebea 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -1862,6 +1862,24 @@ async def add_litellm_data_to_request( tags_to_add=project_metadata["tags"], ) + # inherited_tags: every tag key/team/project policy contributed, read + # directly from those three sources rather than snapshotted off the shared + # "tags" list. A pre-auth pass (apply_client_tag_policy_pre_auth, run from + # user_api_key_auth for _tag_max_budget_check) may already have merged the + # caller's own header tags into that same list before this function ever + # runs, so a snapshot taken here -- at any point in this function -- would + # misattribute caller-supplied tags as policy-backed. tag_based_routing.py's + # allow_fail_open reads this (rather than subtracting caller_tags from the + # final merged set) so a caller can't strip an inherited "!"/"&" + # constraint's protection just by resubmitting its exact value alongside a + # conflicting one. + _key_tags: Final = (key_metadata or MappingProxyType({})).get("tags") or () + _team_tags: Final = team_metadata.get("tags") or () + _project_tags: Final = project_metadata.get("tags") or () + data[_metadata_variable_name]["inherited_tags"] = tuple( # rebind-ok: matches this file's data[...] mutation idiom + dict.fromkeys((*_key_tags, *_team_tags, *_project_tags)) + ) + ## TEAM-LEVEL METADATA data = LiteLLMProxyRequestSetup.add_management_endpoint_metadata_to_request_metadata( data=data, @@ -1958,15 +1976,28 @@ async def add_litellm_data_to_request( tags_to_add=tags, ) - if _metadata_variable_name != "metadata": - _user_metadata = data.get("metadata") - if isinstance(_user_metadata, dict): - _user_tags: Final = _user_metadata.get("tags") - if isinstance(_user_tags, list) and _user_tags: - data[_metadata_variable_name]["tags"] = LiteLLMProxyRequestSetup._merge_tags( - request_tags=data[_metadata_variable_name].get("tags"), - tags_to_add=_user_tags, - ) + _caller_body_metadata: Final = data.get("metadata") if _metadata_variable_name != "metadata" else None + _caller_body_tags: Final = ( + _caller_body_metadata.get("tags") + if isinstance(_caller_body_metadata, dict) and isinstance(_caller_body_metadata.get("tags"), list) + else None + ) + if _caller_body_tags: + data[_metadata_variable_name]["tags"] = LiteLLMProxyRequestSetup._merge_tags( # rebind-ok: matches file idiom + request_tags=data[_metadata_variable_name].get("tags"), + tags_to_add=_caller_body_tags, + ) + + # caller_tags: exactly what this request itself supplied (x-litellm-tags header, + # body "tags", or body "metadata.tags" on litellm_metadata routes), never + # anything from key/team/project metadata. Read directly from the header and + # body values here, the same way inherited_tags above is read directly from + # key/team/project metadata -- neither is derived by inspecting the shared + # "tags" list, which a pre-auth pass (apply_client_tag_policy_pre_auth) may + # have already merged caller header tags into before this function runs. + data[_metadata_variable_name]["caller_tags"] = tuple( # rebind-ok: matches file idiom + dict.fromkeys((*(tags or ()), *(_caller_body_tags or ()))) + ) # Team Callbacks controls callback_settings_obj: Final = _get_dynamic_logging_metadata( diff --git a/litellm/router.py b/litellm/router.py index 9cece2014af..6af1c1a78e7 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -402,6 +402,7 @@ class Router: enable_pre_call_checks: bool = False, enable_tag_filtering: bool = False, tag_filtering_match_any: bool = True, + tag_routing_prefix: str = "", plugins: list[RoutingPlugin] | None = None, retry_after: int = 0, # min time to wait before retrying a failed request retry_policy: RetryPolicy | dict | None = None, # set custom retries for different exceptions @@ -511,6 +512,7 @@ class Router: self.enable_pre_call_checks = enable_pre_call_checks self.enable_tag_filtering = enable_tag_filtering self.tag_filtering_match_any = tag_filtering_match_any + self.tag_routing_prefix = tag_routing_prefix from litellm._service_logger import ServiceLogging self.service_logger_obj: ServiceLogging = ServiceLogging() @@ -10109,6 +10111,7 @@ class Router: "model_group_alias", "enable_weighted_failover", "enable_tag_filtering", + "tag_routing_prefix", ] for var in vars_to_include: @@ -10146,6 +10149,7 @@ class Router: "model_group_alias", "enable_weighted_failover", "enable_tag_filtering", + "tag_routing_prefix", ] _int_settings: Final = [ diff --git a/litellm/router_strategy/tag_based_routing.py b/litellm/router_strategy/tag_based_routing.py index c952b54e672..bbe97613c57 100644 --- a/litellm/router_strategy/tag_based_routing.py +++ b/litellm/router_strategy/tag_based_routing.py @@ -4,9 +4,12 @@ Use this to route requests between Teams - If tags in request is a subset of tags in deployment, return deployment - if deployments are set with default tags, return all default deployment - If no default_deployments are set, return all deployments +- A "!tag" excludes deployments carrying that tag; a "&tag" requires it """ import re +from collections.abc import Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal from litellm._logging import verbose_logger @@ -114,14 +117,52 @@ def _match_deployment( return None -def _split_tags(tags: list[str]) -> tuple[list[str], list[str]]: - positive: Final = [t for t in tags if not t.startswith("!")] - excluded: Final = [tag[1:] for tag in tags if tag.startswith("!") and len(tag) > 1] - return positive, excluded +def _bare_tag_value(tag: str) -> str | None: + # Mirrors _split_tags' own stripping rule exactly, so a confirmed value + # compares equal to whatever required_set/excluded_set/positive_tags end up + # holding for the same tag: a "&"/"!" marker is stripped only when something + # follows it; a lone marker with nothing after it parses to nothing in any + # of the three sets, so it must not become a confirmed value either. + if tag.startswith(("&", "!")): + return tag[1:] if len(tag) > 1 else None + return tag + + +def _strip_routing_prefix(tags: Sequence[str], prefix: str) -> tuple[tuple[str, ...], frozenset[str]]: + # Strips the configured routing-prefix marker from any tag carrying it, used + # exactly as configured with no delimiter auto-appended, and separately + # tracks the post-strip, post-marker-strip values that arrived prefixed: tags + # whose routing intent the caller declared explicitly, exempt from the "maybe + # foreign to this group" heuristics in _unknown_required_tag_hides_an_answer + # and _tag_known_to_group below. Confirmed values are compared against + # required_set/excluded_set downstream, which are themselves already stripped + # of their "&"/"!" marker by _split_tags -- confirmed must match that same + # bare form, not the raw post-prefix-strip value that still carries the + # marker character. An empty prefix must return every tag unconfirmed, not + # run every tag through str.startswith(""), which is trivially True for + # every string and would mark everything confirmed. + if not prefix: + return tuple(tags), frozenset() + rewritten: Final = tuple(t.removeprefix(prefix) for t in tags) + confirmed: Final = frozenset( + bare + for bare in (_bare_tag_value(t.removeprefix(prefix)) for t in tags if t.startswith(prefix)) + if bare is not None + ) + return rewritten, confirmed + + +def _split_tags(tags: Sequence[str]) -> tuple[tuple[str, ...], list[str], tuple[str, ...]]: + required: Final = tuple(tag[1:] for tag in tags if tag.startswith("&") and len(tag) > 1) + positive: Final = [ + t for t in tags if not t.startswith("!") and not t.startswith("&") + ] # mutable-ok: feeds _match_deployment's existing list[str]-typed request_tags param + excluded: Final = tuple(tag[1:] for tag in tags if tag.startswith("!") and len(tag) > 1) + return required, positive, excluded def _exclude_deployments( - deployments: list[Any] | dict[Any, Any], + deployments: Sequence[Any] | Mapping[Any, Any], excluded_set: frozenset[str], ) -> list[Any]: if not excluded_set: @@ -129,24 +170,223 @@ def _exclude_deployments( return [d for d in deployments if not excluded_set.intersection(d.get("litellm_params", {}).get("tags") or [])] -def _require_candidates( - candidates: list[Any], +def _require_all_tags( + deployments: Sequence[Any] | Mapping[Any, Any], + required_set: frozenset[str], +) -> tuple[Any, ...]: + if not required_set: + return tuple(deployments) + return tuple(d for d in deployments if required_set.issubset(d.get("litellm_params", {}).get("tags") or [])) + + +def _default_tagged_pool( + deployments: Sequence[Any] | Mapping[Any, Any], +) -> tuple[Any, ...]: + defaults: Final = tuple(d for d in deployments if "default" in (d.get("litellm_params", {}).get("tags") or [])) + return defaults if defaults else tuple(deployments) + + +def _known_tag_values(deployments: Sequence[Any] | Mapping[Any, Any]) -> frozenset[str]: + return frozenset( + tag for d in deployments for tag in (d.get("litellm_params", MappingProxyType({})).get("tags") or ()) + ) + + +def _unknown_required_tag_hides_an_answer( + healthy_deployments: Sequence[Any] | Mapping[Any, Any], + excluded_set: frozenset[str], + required_set: frozenset[str], + routing_confirmed: frozenset[str], +) -> bool: + # A caller-invented "&" tag (one no deployment in this group has ever carried) + # guarantees an empty required-AND result on its own, regardless of whether the + # rest of the request's required tags were satisfiable. Dropping the unknown + # tags and recomputing: if that reveals a specific, non-empty answer, the invented + # tag was the actual cause of the exhaustion, and fail-open must not paper over + # it. If every required tag is known, or none are, there's nothing hidden to + # protect: either the caller made a real, honestly-unsatisfiable ask (fail-open + # proceeds normally), or the whole required set is unrecognized noise with no + # narrower answer to hide behind it. routing_confirmed (tag_routing_prefix) + # counts as known too: the caller explicitly declared it a routing directive, + # so it is never treated as invented noise regardless of deployment vocabulary. + known_required: Final = required_set & (_known_tag_values(healthy_deployments) | routing_confirmed) + if not known_required or known_required == required_set: + return False + allowed: Final = _exclude_deployments(healthy_deployments, excluded_set) + return bool(_require_all_tags(allowed, known_required)) + + +def _chain_allows_fail_open( + healthy_deployments: Sequence[Any] | Mapping[Any, Any], + excluded_set: frozenset[str], + required_set: frozenset[str], + routing_confirmed: frozenset[str], +) -> bool: + if _unknown_required_tag_hides_an_answer(healthy_deployments, excluded_set, required_set, routing_confirmed): + return False + return any((d.get("model_info") or {}).get("allow_fail_open") is True for d in healthy_deployments) + + +def _trusted_only_pool( + healthy_deployments: Sequence[Any] | Mapping[Any, Any], + excluded_set: frozenset[str], + required_set: frozenset[str], + inherited_excluded_set: frozenset[str] | None, + inherited_required_set: frozenset[str] | None, +) -> tuple[Any, ...]: + # inherited_*_set is None only when this request carries no origin information + # at all (e.g. direct SDK Router usage, bypassing the proxy layer that + # populates metadata.inherited_tags) -- treat every constraint as + # caller-controlled in that case (protected == empty), reproducing this + # function's pre-provenance behavior exactly: an unconditional fall-open to the + # full default-tagged pool, constraints discarded entirely. Otherwise, a tag + # value is protected the moment it has ANY inherited backing, even when the + # caller also happens to submit the identical value themselves -- set + # membership can't distinguish "this value came from policy" from "this value + # coincidentally matches policy," so presence in the inherited set (not + # absence from a caller-supplied set) is what must gate discardability. This + # is deliberately intersection with inherited_*_set, not subtraction of a + # caller-supplied set: subtraction would let a caller strip an inherited + # requirement's protection just by resubmitting its exact value alongside a + # conflicting one (e.g. inherited "®ion:eu" plus caller "®ion:eu" + # and "!region:eu" would otherwise cancel the inherited requirement out). + trusted_excluded: Final = ( + frozenset[str]() if inherited_excluded_set is None else inherited_excluded_set & excluded_set + ) + trusted_required: Final = ( + frozenset[str]() if inherited_required_set is None else inherited_required_set & required_set + ) + return _require_all_tags(_exclude_deployments(healthy_deployments, trusted_excluded), trusted_required) + + +def _resolve_or_fail_open( + pool: Sequence[Any], + healthy_deployments: Sequence[Any] | Mapping[Any, Any], + excluded_set: frozenset[str], + required_set: frozenset[str], + inherited_excluded_set: frozenset[str] | None, + inherited_required_set: frozenset[str] | None, + routing_confirmed: frozenset[str], model: str, - request_tags: Any, -) -> list[Any]: - if not candidates: - raise ValueError( - f"{RouterErrors.no_deployments_with_tag_routing.value}. Passed model={model} and tags={request_tags}" + request_tags: object, +) -> tuple[Any, ...]: + if pool: + return tuple(pool) + if _chain_allows_fail_open(healthy_deployments, excluded_set, required_set, routing_confirmed): + # Fall open only within whatever still satisfies whichever constraints + # trace back to key/team policy. A constraint with no inherited backing at + # all (or, when inherited_tags is unavailable, any constraint at all) can + # be discarded; one inherited from key/team policy cannot -- if that alone + # is unsatisfiable, raise instead of silently routing around it. + trusted_pool: Final = _trusted_only_pool( + healthy_deployments, excluded_set, required_set, inherited_excluded_set, inherited_required_set ) - return candidates + if trusted_pool: + return _default_tagged_pool(trusted_pool) + raise ValueError( + f"{RouterErrors.no_deployments_with_tag_routing.value}. Passed model={model} and tags={request_tags}" + ) -def _ban_only_base_pool( - deployments: list[Any] | dict[Any, Any], -) -> list[Any]: - # Mirrors untagged-request semantics so callers can't use !tags to escape the default pool. - defaults: Final = [d for d in deployments if "default" in (d.get("litellm_params", {}).get("tags") or [])] - return defaults if defaults else list(deployments) +def _resolve_constraint_only_pool( + healthy_deployments: Sequence[Any] | Mapping[Any, Any], + excluded_set: frozenset[str], + required_set: frozenset[str], + inherited_excluded_set: frozenset[str] | None, + inherited_required_set: frozenset[str] | None, + routing_confirmed: frozenset[str], + model: str, + request_tags: object, +) -> tuple[Any, ...]: + pool: Final = ( + _require_all_tags(_exclude_deployments(healthy_deployments, excluded_set), required_set) + if required_set + else _exclude_deployments(_default_tagged_pool(healthy_deployments), excluded_set) + ) + return _resolve_or_fail_open( + pool, + healthy_deployments, + excluded_set, + required_set, + inherited_excluded_set, + inherited_required_set, + routing_confirmed, + model, + request_tags, + ) + + +def _all_deployments_or_fallback( + llm_router_instance: LitellmRouter, + model: str, + fallback: Sequence[Any] | Mapping[Any, Any], +) -> Sequence[Any] | Mapping[Any, Any]: + try: + return llm_router_instance._get_all_deployments(model_name=model) + except Exception: # noqa: BLE001 # fail safe toward today's healthy-only behavior on lookup errors + return fallback + + +def _chain_tag_filtering_override( + llm_router_instance: LitellmRouter, + model: str, + healthy_deployments: Sequence[Any] | Mapping[Any, Any], +) -> bool | None: + # Resolved from every deployment configured for this model group, not just the + # ones that survived cooldown/health filtering (async_get_healthy_deployments + # filters cooldowns before calling get_deployments_for_tag) -- otherwise the + # sole deployment carrying this group's only explicit override loses its effect + # the moment it's transiently unhealthy, silently falling back to the + # router-wide default and letting an attacker disable a chain's tag policy by + # repeatedly failing that one deployment into cooldown. Falls back to + # healthy_deployments on a lookup error, preserving today's behavior rather + # than crashing the request. + all_deployments: Final = _all_deployments_or_fallback(llm_router_instance, model, healthy_deployments) + for d in all_deployments: + value = (d.get("model_info") or MappingProxyType({})).get("enable_tag_filtering") + if value is not None: + return value + return None + + +def _inherited_constraint_sets( + inherited_tags: object, routing_prefix: str +) -> tuple[frozenset[str] | None, frozenset[str] | None]: + # None means no origin information is available at all (e.g. this request + # bypassed the proxy layer that populates metadata.inherited_tags, as direct + # SDK Router usage does) -- callers of this must treat that as "nothing is + # protected," not "nothing is inherited," see _trusted_only_pool. + # metadata.inherited_tags is a snapshot of whatever key/team/project policy + # merged into "tags" *before* this request's own caller-supplied tags were + # merged in on top (see litellm_pre_call_utils.py), so a value present here is + # policy-backed regardless of whether the caller also happens to submit the + # identical value. inherited_tags is stripped through the same routing_prefix + # as the main request tags so a policy-inherited prefixed tag still matches + # correctly against the (already-stripped) required_set/excluded_set computed + # from request_tags. + if not isinstance(inherited_tags, (list, tuple)): + return None, None + rewritten_inherited_tags: Final = _strip_routing_prefix(inherited_tags, routing_prefix)[0] + inherited_required, _inherited_positive, inherited_excluded = _split_tags(rewritten_inherited_tags) + return frozenset(inherited_required), frozenset(inherited_excluded) + + +def _tag_known_to_group( + llm_router_instance: LitellmRouter, + model: str, + positive_tags: Sequence[str], + routing_confirmed: frozenset[str], +) -> bool: + tag_set: Final = frozenset(positive_tags) + if tag_set & routing_confirmed: + return True + try: + all_deployments: Final = llm_router_instance._get_all_deployments(model_name=model) + except Exception: # noqa: BLE001 # fail safe toward "unrecognized" so lookup errors preserve the existing silent-fallback behavior + return False + return any( + tag_set.intersection(d.get("litellm_params", MappingProxyType({})).get("tags") or ()) for d in all_deployments + ) async def get_deployments_for_tag( @@ -161,24 +401,29 @@ async def get_deployments_for_tag( Executes tag based filtering based on the tags in request metadata and the tags on the deployments - Runs when the router-level `enable_tag_filtering` is True or the request carries - `enable_tag_filtering=True` (set from key/team router_settings by the proxy). - A request-level False never disables a router-level True, so per-request settings - cannot escape an operator's global tag-routing policy. + Runs when the effective enable_tag_filtering is True. Effective value: a + request-level enable_tag_filtering=True (set from key/team router_settings by + the proxy) always wins; otherwise model_info.enable_tag_filtering on this model + group, if set on any of its deployments, overrides the router-wide default. + A request-level False never disables either of those, so per-request settings + cannot escape an operator's or a chain owner's tag-routing policy. """ - request_enable_tag_filtering: Final = request_kwargs.get("enable_tag_filtering") if request_kwargs else None - if request_enable_tag_filtering is not True and llm_router_instance.enable_tag_filtering is not True: - return healthy_deployments - - if request_kwargs is None: + if request_kwargs is None or not healthy_deployments: verbose_logger.debug( - "get_deployments_for_tag: request_kwargs is None returning healthy_deployments: %s", + "get_deployments_for_tag: skipping tag filter (request_kwargs=%s, healthy_deployments=%s)", + request_kwargs, healthy_deployments, ) return healthy_deployments - if not healthy_deployments: - verbose_logger.debug("get_deployments_for_tag: empty or None healthy_deployments; skipping tag filter") + request_enable_tag_filtering: Final = request_kwargs.get("enable_tag_filtering") + chain_enable_tag_filtering: Final = _chain_tag_filtering_override(llm_router_instance, model, healthy_deployments) + chain_default: Final = ( + chain_enable_tag_filtering + if chain_enable_tag_filtering is not None + else llm_router_instance.enable_tag_filtering + ) + if request_enable_tag_filtering is not True and chain_default is not True: return healthy_deployments verbose_logger.debug("request metadata: %s", request_kwargs.get(metadata_variable_name)) @@ -186,29 +431,52 @@ async def get_deployments_for_tag( metadata: Final = request_kwargs[metadata_variable_name] request_tags: Final = metadata.get("tags") match_any: Final = llm_router_instance.tag_filtering_match_any + routing_prefix: Final = llm_router_instance.tag_routing_prefix or "" # Build header strings for regex matching from what the proxy already stores. # Currently we match against User-Agent; format matches "^User-Agent: claude-code/..." user_agent: Final = metadata.get("user_agent", "") header_strings: Final[list[str]] = [f"User-Agent: {user_agent}"] if user_agent else [] - positive_tags, excluded_patterns = _split_tags(request_tags or []) + # A tag_routing_prefix-marked tag is stripped before matching -- everything + # downstream (_split_tags, deployment matching) works off the unprefixed + # value, exactly as if the caller had sent it unprefixed -- and its + # post-strip value is remembered in routing_confirmed as an explicit, + # caller-declared routing directive, exempt from the "maybe foreign to this + # group" heuristics that unprefixed tags still go through unchanged below. + rewritten_tags, routing_confirmed = _strip_routing_prefix(request_tags or [], routing_prefix) + required_tags, positive_tags, excluded_patterns = _split_tags(rewritten_tags) + inherited_required_set, inherited_excluded_set = _inherited_constraint_sets( + metadata.get("inherited_tags"), routing_prefix + ) excluded_set: Final = frozenset(excluded_patterns) - candidates: Final = _exclude_deployments(healthy_deployments, excluded_set) + required_set: Final = frozenset(required_tags) + allowed_deployments: Final = _exclude_deployments(healthy_deployments, excluded_set) + candidates: Final = _require_all_tags(allowed_deployments, required_set) has_regex_deployments: Final = any(d.get("litellm_params", {}).get("tag_regex") for d in candidates) - has_tag_filter: Final = bool(positive_tags) or (bool(header_strings) and has_regex_deployments) - ban_only: Final = bool(excluded_set) and not has_tag_filter + has_positive_filter: Final = bool(positive_tags) or ( + bool(header_strings) and has_regex_deployments and not required_set + ) + constraint_only: Final = (bool(excluded_set) or bool(required_set)) and not has_positive_filter - if ban_only: - pool: Final = _exclude_deployments(_ban_only_base_pool(healthy_deployments), excluded_set) - return _require_candidates(pool, model, request_tags) + if constraint_only: + return _resolve_constraint_only_pool( + healthy_deployments, + excluded_set, + required_set, + inherited_excluded_set, + inherited_required_set, + routing_confirmed, + model, + request_tags, + ) new_healthy_deployments: Final[list[Any]] = [] default_deployments: Final[list[Any]] = [] - if has_tag_filter: + if has_positive_filter: verbose_logger.debug( "get_deployments_for_tag routing: request_tags=%s user_agent=%s", request_tags, @@ -245,9 +513,33 @@ async def get_deployments_for_tag( default_deployments.append(deployment) if len(new_healthy_deployments) == 0 and len(default_deployments) == 0: - raise ValueError( - f"{RouterErrors.no_deployments_with_tag_routing.value}." - f" Passed model={model} and tags={request_tags}" + return _resolve_or_fail_open( + (), + healthy_deployments, + excluded_set, + required_set, + inherited_excluded_set, + inherited_required_set, + routing_confirmed, + model, + request_tags, + ) + + if ( + len(new_healthy_deployments) == 0 + and positive_tags + and _tag_known_to_group(llm_router_instance, model, positive_tags, routing_confirmed) + ): + return _resolve_or_fail_open( + (), + healthy_deployments, + excluded_set, + required_set, + inherited_excluded_set, + inherited_required_set, + routing_confirmed, + model, + request_tags, ) return new_healthy_deployments if len(new_healthy_deployments) > 0 else default_deployments diff --git a/litellm/types/router.py b/litellm/types/router.py index 3ac59c0f581..d7ff8d12aa6 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -123,6 +123,7 @@ class UpdateRouterConfig(BaseModel): context_window_fallbacks: list[dict] | None = None model_group_alias: dict[str, str | dict] | None = {} enable_tag_filtering: bool | None = None + tag_routing_prefix: str | None = None model_config = ConfigDict(protected_namespaces=()) @@ -170,6 +171,20 @@ class ModelInfo(MirroredPricingParams): ptu_effective_from: datetime.datetime | None = None ptu_effective_to: datetime.datetime | None = None + # when tag-based routing's "!" or "&" constraints eliminate every deployment + # in this model group, fall back to the default-tagged pool instead of + # raising no_deployments_with_tag_routing. Defaults to False (raise), so + # existing "!" negation behavior is unchanged unless explicitly opted in. + allow_fail_open: bool | None = None + + # per-model-group override for router_settings.enable_tag_filtering; unset + # defers to the router-wide default. Checked against any deployment in the + # group, so set it consistently across every deployment sharing this + # model_name. A request-level enable_tag_filtering=True (from key/team + # settings) still wins over this, exactly as it already does over the + # router-wide default. + enable_tag_filtering: bool | None = None + def __init__(self, id: str | int | None = None, **params) -> None: if id is None: id = str(uuid.uuid4()) # Generate a UUID if id is None or not provided diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index da0f608fdb5..12a98dee7ed 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 1495 + "limit": 1491 }, "ASYNC230": { "limit": 11 diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index f48e1dba601..293cce5fa5d 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -6309,3 +6309,239 @@ class TestPromotedTraceControlFields: assert "litellm_metadata" not in updated assert updated["metadata"]["trace_id"] == "trace-1" assert updated["metadata"]["session_id"] == "session-1" + + +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_inherited_tags_excludes_caller_tags(): + """inherited_tags must carry only what key/team/project policy contributed, + never anything the caller's own request (header/body) supplied, even when the + caller resubmits the identical value -- it's a snapshot taken before the + caller's own tags are merged in, not a set difference against caller_tags. + tag_based_routing.py's allow_fail_open relies on this so a caller can't strip + an inherited constraint's protection by resubmitting its exact value.""" + request_mock = MagicMock(spec=Request) + request_mock.url.path = "/v1/chat/completions" + request_mock.url = MagicMock() + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + + data = { + "model": "gpt-3.5-turbo", + # Caller resubmits the exact value the key policy also contributes. + "tags": ["key-supplied"], + } + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + user_id="real-user", + metadata={"tags": ["key-supplied"]}, + team_metadata={"tags": ["team-supplied"]}, + spend=0.0, + max_budget=100.0, + model_max_budget={}, + team_spend=0.0, + team_max_budget=200.0, + ) + + updated = await add_litellm_data_to_request( + data=data, + request=request_mock, + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert set(updated["metadata"]["tags"]) == {"key-supplied", "team-supplied"} + assert set(updated["metadata"]["inherited_tags"]) == {"key-supplied", "team-supplied"} + assert tuple(updated["metadata"]["caller_tags"]) == ("key-supplied",) + + +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_inherited_tags_survives_pre_auth_header_merge(): + """Regression: apply_client_tag_policy_pre_auth (run from user_api_key_auth, + for _tag_max_budget_check) merges the caller's x-litellm-tags header into the + same metadata.tags list this function later reads from -- before this + function ever runs. A snapshot-based inherited_tags would misattribute that + caller-controlled value as policy-backed; inherited_tags must instead be read + directly from key/team/project metadata, immune to whatever the pre-auth pass + already merged into "tags".""" + request_mock = MagicMock(spec=Request) + request_mock.url.path = "/v1/chat/completions" + request_mock.url = MagicMock() + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json", "x-litellm-tags": "caller-invented-tag"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + + data: dict = {"model": "gpt-3.5-turbo"} + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + user_id="real-user", + metadata={"tags": ["key-supplied"]}, + team_metadata={}, + spend=0.0, + max_budget=100.0, + model_max_budget={}, + team_spend=0.0, + team_max_budget=200.0, + ) + + # Simulate the real request pipeline: the pre-auth merge runs first, on the + # same data dict, before add_litellm_data_to_request is ever called. + LiteLLMProxyRequestSetup.apply_client_tag_policy_pre_auth( + request=request_mock, + request_data=data, + user_api_key_dict=user_api_key_dict, + ) + assert data["metadata"]["tags"] == ["caller-invented-tag"] + + updated = await add_litellm_data_to_request( + data=data, + request=request_mock, + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert set(updated["metadata"]["tags"]) == {"caller-invented-tag", "key-supplied"} + assert updated["metadata"]["inherited_tags"] == ("key-supplied",) + assert updated["metadata"]["caller_tags"] == ("caller-invented-tag",) + + +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_caller_tags_excludes_key_and_team_tags(): + """caller_tags must carry only what the caller itself sent (header + body + tags), never anything merged in from key/team metadata, even though the + merged "tags" field (used for matching) legitimately contains all three.""" + request_mock = MagicMock(spec=Request) + request_mock.url.path = "/v1/chat/completions" + request_mock.url = MagicMock() + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + + data = { + "model": "gpt-3.5-turbo", + "tags": ["caller-supplied"], + } + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + user_id="real-user", + metadata={"tags": ["key-supplied"]}, + team_metadata={"tags": ["team-supplied"]}, + spend=0.0, + max_budget=100.0, + model_max_budget={}, + team_spend=0.0, + team_max_budget=200.0, + ) + + updated = await add_litellm_data_to_request( + data=data, + request=request_mock, + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert set(updated["metadata"]["tags"]) == {"caller-supplied", "key-supplied", "team-supplied"} + assert tuple(updated["metadata"]["caller_tags"]) == ("caller-supplied",) + + +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_caller_tags_includes_header_tags(): + """The x-litellm-tags header is as much a caller-controlled input as the + body's "tags" field; both must land in caller_tags.""" + request_mock = MagicMock(spec=Request) + request_mock.url.path = "/v1/chat/completions" + request_mock.url = MagicMock() + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json", "x-litellm-tags": "header-tag"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + + data = {"model": "gpt-3.5-turbo"} + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + user_id="real-user", + metadata={"tags": ["key-supplied"]}, + team_metadata={}, + spend=0.0, + max_budget=100.0, + model_max_budget={}, + team_spend=0.0, + team_max_budget=200.0, + ) + + updated = await add_litellm_data_to_request( + data=data, + request=request_mock, + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert set(updated["metadata"]["tags"]) == {"header-tag", "key-supplied"} + assert tuple(updated["metadata"]["caller_tags"]) == ("header-tag",) + + +@pytest.mark.asyncio +async def test_add_litellm_data_to_request_caller_tags_empty_when_caller_sends_nothing(): + """caller_tags must be present (an empty tuple), not absent, when the caller + supplied no tags of their own -- an empty-but-present value tells + tag_based_routing.py's allow_fail_open that any required/excluded tag on the + request is entirely inherited, not that no origin information is available. + """ + request_mock = MagicMock(spec=Request) + request_mock.url.path = "/v1/chat/completions" + request_mock.url = MagicMock() + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + + data = {"model": "gpt-3.5-turbo"} + + user_api_key_dict = UserAPIKeyAuth( + api_key="hashed-key", + user_id="real-user", + metadata={"tags": ["key-supplied"]}, + team_metadata={}, + spend=0.0, + max_budget=100.0, + model_max_budget={}, + team_spend=0.0, + team_max_budget=200.0, + ) + + updated = await add_litellm_data_to_request( + data=data, + request=request_mock, + user_api_key_dict=user_api_key_dict, + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated["metadata"]["tags"] == ["key-supplied"] + assert updated["metadata"]["caller_tags"] == () diff --git a/tests/test_litellm/router_strategy/test_router_tag_regex_routing.py b/tests/test_litellm/router_strategy/test_router_tag_regex_routing.py index dca2bd84f92..6591478a4e7 100644 --- a/tests/test_litellm/router_strategy/test_router_tag_regex_routing.py +++ b/tests/test_litellm/router_strategy/test_router_tag_regex_routing.py @@ -112,6 +112,7 @@ def _make_router_mock(enable_tag_filtering=True, match_any=True): mock = MagicMock() mock.enable_tag_filtering = enable_tag_filtering mock.tag_filtering_match_any = match_any + mock.tag_routing_prefix = "" return mock diff --git a/tests/test_litellm/router_strategy/test_router_tag_routing.py b/tests/test_litellm/router_strategy/test_router_tag_routing.py index 98506aad594..9e19e981f80 100644 --- a/tests/test_litellm/router_strategy/test_router_tag_routing.py +++ b/tests/test_litellm/router_strategy/test_router_tag_routing.py @@ -429,42 +429,58 @@ def test_get_tags_from_request_kwargs_various_inputs(): def test_split_tags_positive_only(): from litellm.router_strategy.tag_based_routing import _split_tags - positive, excluded = _split_tags(["paid", "teamA"]) + required, positive, excluded = _split_tags(["paid", "teamA"]) + assert required == () assert positive == ["paid", "teamA"] - assert excluded == [] + assert excluded == () def test_split_tags_negation_only(): from litellm.router_strategy.tag_based_routing import _split_tags - positive, excluded = _split_tags(["!provider:anthropic"]) + required, positive, excluded = _split_tags(["!provider:anthropic"]) + assert required == () assert positive == [] - assert excluded == ["provider:anthropic"] + assert excluded == ("provider:anthropic",) + + +def test_split_tags_required_only(): + from litellm.router_strategy.tag_based_routing import _split_tags + + required, positive, excluded = _split_tags(["&reasoning_type:high", "&provider:anthropic"]) + assert required == ("reasoning_type:high", "provider:anthropic") + assert positive == [] + assert excluded == () def test_split_tags_mixed(): from litellm.router_strategy.tag_based_routing import _split_tags - positive, excluded = _split_tags(["paid", "!provider:anthropic", "!inference:cerebras"]) + required, positive, excluded = _split_tags( + ["paid", "!provider:anthropic", "!inference:cerebras", "&reasoning_type:high"] + ) + assert required == ("reasoning_type:high",) assert positive == ["paid"] assert len(excluded) == 2 -def test_split_tags_bare_bang_skipped(): +def test_split_tags_bare_bang_and_amp_skipped(): from litellm.router_strategy.tag_based_routing import _split_tags - # A bare "!" with nothing after it is not a valid negation tag; skip it - positive, excluded = _split_tags(["paid", "!"]) + # A bare "!" or "&" with nothing after it is not a valid tag; skip it + required, positive, excluded = _split_tags(["paid", "!", "&"]) + assert required == () assert positive == ["paid"] - assert excluded == [] + assert excluded == () def test_split_tags_empty(): from litellm.router_strategy.tag_based_routing import _split_tags - positive, excluded = _split_tags([]) + required, positive, excluded = _split_tags([]) + assert required == () assert positive == [] - assert excluded == [] + assert excluded == () # --- get_deployments_for_tag negation integration tests --- @@ -1115,3 +1131,1682 @@ async def test_request_level_enable_tag_filtering_false_cannot_disable_global(): mock_response="hi", ) assert response._hidden_params["model_id"] == "team-a-deployment" + + +# --- model_info.enable_tag_filtering per-chain override --- + + +class _FakeRouterForChainOverride: + def __init__(self, all_deployments): + self._all_deployments = all_deployments + + def _get_all_deployments(self, model_name): + return self._all_deployments + + +def test_chain_tag_filtering_override_reads_any_member(): + from litellm.router_strategy.tag_based_routing import _chain_tag_filtering_override + + deployments = [ + {"model_info": {}}, + {"model_info": {"enable_tag_filtering": False}}, + ] + router = _FakeRouterForChainOverride(deployments) + assert _chain_tag_filtering_override(router, "gpt-4", deployments) is False + + +def test_chain_tag_filtering_override_none_when_unset_anywhere(): + from litellm.router_strategy.tag_based_routing import _chain_tag_filtering_override + + deployments = [{"model_info": {}}, {}] + router = _FakeRouterForChainOverride(deployments) + assert _chain_tag_filtering_override(router, "gpt-4", deployments) is None + + +def test_chain_tag_filtering_override_survives_the_overriding_member_going_unhealthy(): + # Regression: the per-group override must be resolved from every deployment + # configured for the model, not just the ones that survived cooldown/health + # filtering. async_get_healthy_deployments filters cooldowns before calling + # into get_deployments_for_tag, so healthy_deployments alone can be missing + # the one deployment that carries the group's only explicit override. + from litellm.router_strategy.tag_based_routing import _chain_tag_filtering_override + + all_deployments = [ + {"model_info": {"enable_tag_filtering": True}}, + {"model_info": {}}, + ] + router = _FakeRouterForChainOverride(all_deployments) + # The overriding deployment (index 0) is cooled down and absent from + # healthy_deployments -- the override must still be found via the full-group + # lookup, not silently lost. + healthy_deployments = [all_deployments[1]] + assert _chain_tag_filtering_override(router, "gpt-4", healthy_deployments) is True + + +def test_chain_tag_filtering_override_falls_back_to_healthy_deployments_on_lookup_error(): + from litellm.router_strategy.tag_based_routing import _chain_tag_filtering_override + + class _BrokenRouter: + def _get_all_deployments(self, model_name): + raise RuntimeError("model group not found") + + healthy_deployments = [{"model_info": {"enable_tag_filtering": False}}] + assert _chain_tag_filtering_override(_BrokenRouter(), "gpt-4", healthy_deployments) is False + + +@pytest.mark.asyncio() +async def test_chain_enable_tag_filtering_true_overrides_router_level_false(): + # Router-wide tag filtering is off; this model group opts in on its own via + # model_info.enable_tag_filtering, so tags still apply to requests for it. + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["teamA"], + }, + "model_info": {"id": "team-a-deployment", "enable_tag_filtering": True}, + }, + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["teamB"], + }, + "model_info": {"id": "team-b-deployment", "enable_tag_filtering": True}, + }, + ], + enable_tag_filtering=False, + ) + + for _ in range(5): + response = await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["teamA"]}, + mock_response="hi", + ) + assert response._hidden_params["model_id"] == "team-a-deployment" + + +@pytest.mark.asyncio() +async def test_chain_enable_tag_filtering_false_overrides_router_level_true(): + # Router-wide tag filtering is on, but this model group opts itself out via + # model_info.enable_tag_filtering: tags are ignored for requests to this group, + # so an untagged-style request just gets ordinary load-balanced routing. + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["teamA"], + }, + "model_info": {"id": "team-a-deployment", "enable_tag_filtering": False}, + }, + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["teamB"], + }, + "model_info": {"id": "team-b-deployment", "enable_tag_filtering": False}, + }, + ], + enable_tag_filtering=True, + ) + + seen_ids = set() + for _ in range(10): + response = await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["teamA"]}, + mock_response="hi", + ) + seen_ids.add(response._hidden_params["model_id"]) + + assert seen_ids == {"team-a-deployment", "team-b-deployment"} + + +@pytest.mark.asyncio() +async def test_request_level_enable_tag_filtering_still_wins_over_chain_level_false(): + # A key/team's own request-level enable_tag_filtering=True must still win over + # a chain that opted itself out, exactly as it already wins over the router + # default: request-level escalation is the highest-precedence layer. + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["teamA"], + }, + "model_info": {"id": "team-a-deployment", "enable_tag_filtering": False}, + }, + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["teamB"], + }, + "model_info": {"id": "team-b-deployment", "enable_tag_filtering": False}, + }, + ], + enable_tag_filtering=False, + ) + + for _ in range(5): + response = await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["teamA"]}, + enable_tag_filtering=True, + mock_response="hi", + ) + assert response._hidden_params["model_id"] == "team-a-deployment" + + +# --- _require_all_tags / _chain_allows_fail_open unit tests --- + + +def test_require_all_tags_empty_required_set_is_noop(): + from litellm.router_strategy.tag_based_routing import _require_all_tags + + deployments = [{"litellm_params": {"tags": ["a"]}}, {"litellm_params": {"tags": []}}] + assert _require_all_tags(deployments, frozenset()) == tuple(deployments) + + +def test_require_all_tags_keeps_only_deployments_with_every_required_tag(): + from litellm.router_strategy.tag_based_routing import _require_all_tags + + has_both = {"litellm_params": {"tags": ["reasoning_type:high", "provider:anthropic"]}} + has_one = {"litellm_params": {"tags": ["reasoning_type:high"]}} + has_neither = {"litellm_params": {"tags": ["provider:openai"]}} + + result = _require_all_tags( + [has_both, has_one, has_neither], frozenset({"reasoning_type:high", "provider:anthropic"}) + ) + assert result == (has_both,) + + +def test_chain_allows_fail_open_true_when_any_member_sets_flag(): + from litellm.router_strategy.tag_based_routing import _chain_allows_fail_open + + deployments = [ + {"model_info": {}, "litellm_params": {"tags": ["provider:anthropic"]}}, + {"model_info": {"allow_fail_open": True}, "litellm_params": {"tags": ["provider:openai"]}}, + ] + assert _chain_allows_fail_open(deployments, frozenset(), frozenset({"provider:anthropic"}), frozenset()) is True + + +def test_chain_allows_fail_open_false_by_default(): + from litellm.router_strategy.tag_based_routing import _chain_allows_fail_open + + deployments = [{"model_info": {}}, {}] + assert _chain_allows_fail_open(deployments, frozenset(), frozenset(), frozenset()) is False + + +def test_chain_allows_fail_open_true_when_no_required_tag_is_known_at_all(): + # An entirely-invented required tag with nothing else known to compare against + # has no narrower answer to hide; a single-deployment catch-all fallback is a + # legitimate use of allow_fail_open, not something to deny. + from litellm.router_strategy.tag_based_routing import _chain_allows_fail_open + + deployments = [ + {"model_info": {"allow_fail_open": True}, "litellm_params": {"tags": ["default", "reasoning_type:low"]}}, + ] + assert _chain_allows_fail_open(deployments, frozenset(), frozenset({"reasoning_type:high"}), frozenset()) is True + + +def test_unknown_required_tag_hides_an_answer_denies_fail_open(): + from litellm.router_strategy.tag_based_routing import _chain_allows_fail_open + + deployments = [ + { + "model_info": {}, + "litellm_params": {"tags": ["provider:anthropic", "region:us-east"]}, + }, + { + "model_info": {"allow_fail_open": True}, + "litellm_params": {"tags": ["default", "provider:openai"]}, + }, + ] + # region:us-east is real and satisfiable on the first deployment; the invented tag + # alone forces emptiness. Dropping it reveals a specific, non-default answer, so + # fail-open must be denied even though the flag is set on the group. + assert ( + _chain_allows_fail_open( + deployments, frozenset(), frozenset({"region:us-east", "totally-invented-tag-nobody-has"}), frozenset() + ) + is False + ) + + +def test_unknown_required_tag_allows_fail_open_when_no_answer_is_hidden(): + from litellm.router_strategy.tag_based_routing import _chain_allows_fail_open + + deployments = [ + { + "model_info": {"allow_fail_open": True}, + "litellm_params": {"tags": ["provider:anthropic", "region:us-east"]}, + }, + { + "model_info": {"allow_fail_open": True}, + "litellm_params": {"tags": ["provider:eu", "region:eu"]}, + }, + { + "model_info": {"allow_fail_open": True}, + "litellm_params": {"tags": ["default", "provider:openai"]}, + }, + ] + # region:us-east and region:eu are both real, known tags; no single deployment + # carries both, so this is a genuinely unsatisfiable combination, not an invented + # tag masking a narrower answer. Fail-open must proceed normally. + assert ( + _chain_allows_fail_open(deployments, frozenset(), frozenset({"region:us-east", "region:eu"}), frozenset()) + is True + ) + + +# --- _strip_routing_prefix / _bare_tag_value unit tests --- + + +def test_strip_routing_prefix_empty_prefix_is_noop(): + from litellm.router_strategy.tag_based_routing import _strip_routing_prefix + + tags = ["provider:anthropic", "®ion:eu", "!region:us"] + rewritten, confirmed = _strip_routing_prefix(tags, "") + assert rewritten == tuple(tags) + assert confirmed == frozenset() + + +def test_strip_routing_prefix_splits_routed_from_other(): + from litellm.router_strategy.tag_based_routing import _strip_routing_prefix + + rewritten, confirmed = _strip_routing_prefix(["feature:demo", "route:!provider:openai"], "route:") + assert rewritten == ("feature:demo", "!provider:openai") + assert confirmed == frozenset({"provider:openai"}) + + +def test_strip_routing_prefix_confirmed_matches_bare_required_and_excluded_values(): + # Regression: confirmed must carry the same bare (marker-stripped) form that + # _split_tags produces for required_set/excluded_set downstream. A prior bug + # left the "&"/"!" marker in `confirmed`, so `required_set & routing_confirmed` + # never intersected for any prefixed "&"/"!" tag -- the entire "trusted, + # caller-declared required/excluded tag" mechanism silently no-opped. + from litellm.router_strategy.tag_based_routing import _strip_routing_prefix + + _, confirmed = _strip_routing_prefix(["route:&provider:anthropic", "route:!region:eu"], "route:") + assert confirmed == frozenset({"provider:anthropic", "region:eu"}) + + +def test_strip_routing_prefix_lone_marker_confirms_nothing(): + from litellm.router_strategy.tag_based_routing import _strip_routing_prefix + + # A lone "&"/"!" with nothing after it parses to nothing in required_set, + # excluded_set, or positive_tags (see test_split_tags_bare_bang_and_amp_skipped); + # confirmed must not invent a value for it either. + _, confirmed = _strip_routing_prefix(["route:&", "route:!"], "route:") + assert confirmed == frozenset() + + +def test_chain_allows_fail_open_true_when_prefixed_unknown_required_tag_is_confirmed(): + # Regression for the same bug: a required tag no deployment carries is normally + # treated as invented noise that can hide a narrower answer (see + # test_unknown_required_tag_hides_an_answer_denies_fail_open) -- but once the + # caller has explicitly marked it via the routing prefix, it counts as a known, + # honest ask, and fail-open must proceed rather than get denied. + from litellm.router_strategy.tag_based_routing import _chain_allows_fail_open + + deployments = [ + { + "model_info": {"allow_fail_open": True}, + "litellm_params": {"tags": ["default", "provider:anthropic"]}, + }, + ] + required_set = frozenset({"provider:anthropic", "typo-tag"}) + assert _chain_allows_fail_open(deployments, frozenset(), required_set, frozenset()) is False + assert _chain_allows_fail_open(deployments, frozenset(), required_set, required_set) is True + + +# --- get_deployments_for_tag required-AND ("&") integration tests --- + + +@pytest.mark.asyncio() +async def test_required_and_matches_deployment_with_all_tags(): + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["reasoning_type:high", "provider:anthropic"], + }, + "model_info": {"id": "high-reasoning-anthropic"}, + }, + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["reasoning_type:high", "provider:openai"], + }, + "model_info": {"id": "high-reasoning-openai"}, + }, + ], + enable_tag_filtering=True, + ) + + for _ in range(5): + response = await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["&reasoning_type:high", "&provider:anthropic"]}, + mock_response="hi", + ) + assert response._hidden_params["model_id"] == "high-reasoning-anthropic" + + +@pytest.mark.asyncio() +async def test_required_and_excludes_deployment_missing_one_tag(): + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["reasoning_type:high", "provider:anthropic"], + }, + "model_info": {"id": "high-reasoning-anthropic"}, + }, + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["reasoning_type:low", "provider:anthropic"], + }, + "model_info": {"id": "low-reasoning-anthropic"}, + }, + ], + enable_tag_filtering=True, + ) + + for _ in range(5): + response = await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["&reasoning_type:high", "&provider:anthropic"]}, + mock_response="hi", + ) + assert response._hidden_params["model_id"] == "high-reasoning-anthropic" + + +@pytest.mark.asyncio() +async def test_required_and_composes_with_negation(): + # &reasoning_type:high requires the tag; !provider:anthropic bans that provider. + # Negation applies first, so the anthropic deployment is excluded even though + # it satisfies the required tag. + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["reasoning_type:high", "provider:anthropic"], + }, + "model_info": {"id": "high-reasoning-anthropic"}, + }, + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["reasoning_type:high", "provider:openai"], + }, + "model_info": {"id": "high-reasoning-openai"}, + }, + ], + enable_tag_filtering=True, + ) + + for _ in range(5): + response = await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["&reasoning_type:high", "!provider:anthropic"]}, + mock_response="hi", + ) + assert response._hidden_params["model_id"] == "high-reasoning-openai" + + +@pytest.mark.asyncio() +async def test_required_and_combines_with_positive_or_preference(): + # &reasoning_type:high is a hard requirement; provider:anthropic/provider:openai + # is a preference (OR) applied on top of the survivors. + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["reasoning_type:high", "provider:anthropic"], + }, + "model_info": {"id": "high-reasoning-anthropic"}, + }, + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["reasoning_type:high", "provider:vertex"], + }, + "model_info": {"id": "high-reasoning-vertex"}, + }, + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["reasoning_type:low", "provider:anthropic"], + }, + "model_info": {"id": "low-reasoning-anthropic"}, + }, + ], + enable_tag_filtering=True, + ) + + for _ in range(5): + response = await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["&reasoning_type:high", "provider:anthropic", "provider:openai"]}, + mock_response="hi", + ) + assert response._hidden_params["model_id"] == "high-reasoning-anthropic" + + +@pytest.mark.asyncio() +async def test_required_and_single_tag_matches_trivially(): + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["reasoning_type:high"], + }, + "model_info": {"id": "high-reasoning"}, + }, + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["reasoning_type:low"], + }, + "model_info": {"id": "low-reasoning"}, + }, + ], + enable_tag_filtering=True, + ) + + for _ in range(5): + response = await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["&reasoning_type:high"]}, + mock_response="hi", + ) + assert response._hidden_params["model_id"] == "high-reasoning" + + +@pytest.mark.asyncio() +async def test_required_and_unmatched_raises_by_default(): + # allow_fail_open unset -> unmatched required-AND raises, same as today's "!" behavior. + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["reasoning_type:low"], + }, + "model_info": {"id": "low-reasoning"}, + }, + ], + enable_tag_filtering=True, + ) + + with pytest.raises(Exception) as exc_info: + await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["&reasoning_type:high"]}, + mock_response="hi", + ) + + from litellm.types.router import RouterErrors + + assert RouterErrors.no_deployments_with_tag_routing.value in str(exc_info.value) + + +@pytest.mark.asyncio() +async def test_required_and_combined_with_positive_unmatched_raises_by_default(): + # &A eliminates every candidate before the positive-tag preference even runs; + # this must be gated by allow_fail_open too, not just the required-AND-only path. + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["reasoning_type:low", "provider:anthropic"], + }, + "model_info": {"id": "low-reasoning-anthropic"}, + }, + ], + enable_tag_filtering=True, + ) + + with pytest.raises(Exception) as exc_info: + await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["&reasoning_type:high", "provider:anthropic"]}, + mock_response="hi", + ) + + from litellm.types.router import RouterErrors + + assert RouterErrors.no_deployments_with_tag_routing.value in str(exc_info.value) + + +# --- get_deployments_for_tag allow_fail_open integration tests --- + + +@pytest.mark.asyncio() +async def test_allow_fail_open_required_and_unmatched_falls_back_to_default_pool(): + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["default", "reasoning_type:low"], + }, + "model_info": {"id": "default-model", "allow_fail_open": True}, + }, + ], + enable_tag_filtering=True, + ) + + for _ in range(5): + response = await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["&reasoning_type:high"]}, + mock_response="hi", + ) + assert response._hidden_params["model_id"] == "default-model" + + +@pytest.mark.asyncio() +async def test_allow_fail_open_negation_eliminates_everything_includes_banned_deployment(): + # The core backwards-compatibility risk: once allow_fail_open opts a chain in, + # a "!" ban that eliminates every deployment falls back to the full default + # pool, INCLUDING the deployment the request tried to ban. This must never + # silently disappear (still raise) nor silently reappear on chains without + # the flag set (see test_negation_all_excluded_raises). + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["provider:anthropic"], + }, + "model_info": {"id": "anthropic-model", "allow_fail_open": True}, + }, + ], + enable_tag_filtering=True, + ) + + for _ in range(5): + response = await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["!provider:anthropic"]}, + mock_response="hi", + ) + assert response._hidden_params["model_id"] == "anthropic-model" + + +@pytest.mark.asyncio() +async def test_allow_fail_open_prefers_default_tagged_deployment_on_fallback(): + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["provider:anthropic"], + }, + "model_info": {"id": "anthropic-model", "allow_fail_open": True}, + }, + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["provider:anthropic", "default"], + }, + "model_info": {"id": "anthropic-default-model", "allow_fail_open": True}, + }, + ], + enable_tag_filtering=True, + ) + + for _ in range(5): + response = await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["!provider:anthropic"]}, + mock_response="hi", + ) + assert response._hidden_params["model_id"] == "anthropic-default-model" + + +@pytest.mark.asyncio() +async def test_allow_fail_open_per_hop_across_fallback_chain(): + # required-AND fail-open must be re-evaluated fresh on every hop, the same + # per-hop guarantee the negation feature already established. + router = litellm.Router( + model_list=[ + { + "model_name": "primary", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["reasoning_type:low"], + }, + "model_info": {"id": "primary-low-reasoning"}, + }, + { + "model_name": "fallback", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["default", "reasoning_type:low"], + }, + "model_info": {"id": "fallback-model", "allow_fail_open": True}, + }, + ], + fallbacks=[{"primary": ["fallback"]}], + enable_tag_filtering=True, + ) + + response = await router.acompletion( + model="primary", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["&reasoning_type:high"]}, + mock_response="hi", + ) + assert response._hidden_params["model_id"] == "fallback-model" + + +@pytest.mark.asyncio() +async def test_allow_fail_open_resolves_locally_without_triggering_external_fallback(): + # allow_fail_open on the primary group's own default deployment absorbs the + # exhaustion internally (_resolve_or_fail_open returns a non-empty pool, so + # get_deployments_for_tag never raises); router.async_function_with_fallbacks + # only invokes the configured "fallbacks" chain on an exception, so a + # separate, unrelated fallback group must never be touched even though one is + # configured. A fallback deployment that would trivially satisfy the request + # tag if it were ever consulted makes this a meaningful negative assertion, + # not a vacuous one. + router = litellm.Router( + model_list=[ + { + "model_name": "primary", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["reasoning_type:high"], + }, + "model_info": {"id": "primary-high-reasoning"}, + }, + { + "model_name": "primary", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["default", "reasoning_type:low"], + }, + "model_info": {"id": "primary-default", "allow_fail_open": True}, + }, + { + "model_name": "fallback", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["region:eu"], + }, + "model_info": {"id": "fallback-should-never-be-used"}, + }, + ], + fallbacks=[{"primary": ["fallback"]}], + enable_tag_filtering=True, + ) + + response = await router.acompletion( + model="primary", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["®ion:eu"]}, + mock_response="hi", + ) + assert response._hidden_params["model_id"] == "primary-default" + + +# --- allow_fail_open must also gate "!" exhaustion combined with a plain positive tag --- + + +@pytest.mark.asyncio() +async def test_negation_combined_with_positive_unmatched_raises_by_default(): + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["provider:anthropic", "paid"], + }, + "model_info": {"id": "anthropic-paid"}, + }, + ], + enable_tag_filtering=True, + ) + + with pytest.raises(Exception) as exc_info: + await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["!provider:anthropic", "paid"]}, + mock_response="hi", + ) + + from litellm.types.router import RouterErrors + + assert RouterErrors.no_deployments_with_tag_routing.value in str(exc_info.value) + + +@pytest.mark.asyncio() +async def test_negation_combined_with_positive_unmatched_falls_open_when_allowed(): + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["provider:anthropic", "paid", "default"], + }, + "model_info": {"id": "anthropic-paid", "allow_fail_open": True}, + }, + ], + enable_tag_filtering=True, + ) + + for _ in range(5): + response = await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["!provider:anthropic", "paid"]}, + mock_response="hi", + ) + assert response._hidden_params["model_id"] == "anthropic-paid" + + +# --- a required-AND-only request must not be diluted by incidental regex/header preference --- + + +@pytest.mark.asyncio() +async def test_required_and_only_returns_every_matching_deployment_despite_regex_header(): + # Deployment A satisfies &reasoning_type:high and also happens to carry a tag_regex + # that matches the caller's User-Agent. Deployment B also satisfies the required tag + # but has no tag_regex at all. A required-AND-only request (no plain positive tags) + # must be free to route to either survivor, not be narrowed down to only the one + # that happens to match the incidental regex/header preference. + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["reasoning_type:high"], + "tag_regex": ["^User-Agent: claude-code\\/"], + }, + "model_info": {"id": "high-reasoning-with-regex"}, + }, + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["reasoning_type:high"], + }, + "model_info": {"id": "high-reasoning-no-regex"}, + }, + ], + enable_tag_filtering=True, + ) + + seen_ids = set() + for _ in range(30): + response = await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["&reasoning_type:high"], "user_agent": "claude-code/1.2.3"}, + mock_response="hi", + ) + seen_ids.add(response._hidden_params["model_id"]) + + assert seen_ids == {"high-reasoning-with-regex", "high-reasoning-no-regex"} + + +@pytest.mark.asyncio() +async def test_required_and_only_excludes_regex_deployment_missing_the_required_tag(): + # The tag_regex deployment matches the caller's User-Agent but does NOT carry the + # required tag; a required-AND-only request must not let it through on the strength + # of the regex/header match alone. + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["reasoning_type:low"], + "tag_regex": ["^User-Agent: claude-code\\/"], + }, + "model_info": {"id": "low-reasoning-with-regex"}, + }, + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["reasoning_type:high"], + }, + "model_info": {"id": "high-reasoning-no-regex"}, + }, + ], + enable_tag_filtering=True, + ) + + for _ in range(5): + response = await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["&reasoning_type:high"], "user_agent": "claude-code/1.2.3"}, + mock_response="hi", + ) + assert response._hidden_params["model_id"] == "high-reasoning-no-regex" + + +# --- allow_fail_open must also gate exhaustion after a non-empty required-AND survivor +# set fails to match a plain preference tag, not just full !/& exhaustion --- + + +@pytest.mark.asyncio() +async def test_mixed_constraint_survivor_unmatched_by_positive_tag_raises_by_default(): + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["reasoning_type:high", "provider:anthropic"], + }, + "model_info": {"id": "high-reasoning-anthropic"}, + }, + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["default", "reasoning_type:low"], + }, + "model_info": {"id": "default-fallback"}, + }, + ], + enable_tag_filtering=True, + ) + + with pytest.raises(Exception) as exc_info: + await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["&reasoning_type:high", "provider:openai"]}, + mock_response="hi", + ) + + from litellm.types.router import RouterErrors + + assert RouterErrors.no_deployments_with_tag_routing.value in str(exc_info.value) + + +@pytest.mark.asyncio() +async def test_mixed_constraint_survivor_unmatched_by_positive_tag_falls_open_when_allowed(): + # &reasoning_type:high survives to a non-empty candidate set (the anthropic + # deployment), but the plain preference tag provider:openai matches none of the + # survivors, and the surviving deployment itself is not "default"-tagged (so the + # pre-existing in-loop default-collection escape hatch can't mask the fix). Greptile + # flagged this exact path as bypassing allow_fail_open by raising unconditionally; + # it must instead fall back to the group's actual default-tagged deployment, which + # is a different deployment than the one &reasoning_type:high matched. + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["reasoning_type:high", "provider:anthropic"], + }, + "model_info": {"id": "high-reasoning-anthropic", "allow_fail_open": True}, + }, + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["default", "reasoning_type:low"], + }, + "model_info": {"id": "default-fallback", "allow_fail_open": True}, + }, + ], + enable_tag_filtering=True, + ) + + for _ in range(5): + response = await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["&reasoning_type:high", "provider:openai"]}, + mock_response="hi", + ) + assert response._hidden_params["model_id"] == "default-fallback" + + +# --- allow_fail_open must not be triggerable by an invented tag the chain has never +# carried; a caller-supplied garbage tag must not be able to force an otherwise- +# satisfiable constraint (e.g. one inherited from the key/team) to be discarded --- + + +@pytest.mark.asyncio() +async def test_allow_fail_open_denied_when_request_includes_unknown_tag(): + # region:us-east is a real, satisfiable constraint on anthropic-deployment. Adding + # a single invented tag no deployment in this group has ever carried empties the + # required-AND set regardless of region:us-east's own satisfiability. allow_fail_open + # is set on the default deployment, but must not fire here: none of the *other* + # deployments carry the invented tag either, so it is unknown to the chain, and + # falling back would silently discard the still-satisfiable region:us-east ask. + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["provider:anthropic", "region:us-east"], + }, + "model_info": {"id": "anthropic-deployment"}, + }, + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["default", "provider:openai"], + }, + "model_info": {"id": "openai-default", "allow_fail_open": True}, + }, + ], + enable_tag_filtering=True, + ) + + with pytest.raises(Exception) as exc_info: + await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["®ion:us-east", "&totally-invented-tag-nobody-has"]}, + mock_response="hi", + ) + + from litellm.types.router import RouterErrors + + assert RouterErrors.no_deployments_with_tag_routing.value in str(exc_info.value) + + +@pytest.mark.asyncio() +async def test_allow_fail_open_still_fires_when_every_requested_tag_is_known(): + # region:us-east and region:eu are both real tags this chain uses; no single + # deployment carries both, so the combination is genuinely unsatisfiable, not + # invented. allow_fail_open must still fall back normally in this case. + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["provider:anthropic", "region:us-east"], + }, + "model_info": {"id": "anthropic-deployment", "allow_fail_open": True}, + }, + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["provider:eu", "region:eu"], + }, + "model_info": {"id": "eu-deployment", "allow_fail_open": True}, + }, + { + "model_name": "gpt-4", + "litellm_params": { + "model": "openai/gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["default", "provider:openai"], + }, + "model_info": {"id": "openai-default", "allow_fail_open": True}, + }, + ], + enable_tag_filtering=True, + ) + + for _ in range(5): + response = await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["®ion:us-east", "®ion:eu"]}, + mock_response="hi", + ) + assert response._hidden_params["model_id"] == "openai-default" + + +# --- required-AND, allow_fail_open, and the unknown-tag denial across fallback +# chains spanning multiple model groups --- + + +@pytest.mark.asyncio() +async def test_required_and_exhausts_primary_group_falls_through_to_fallback_group(): + # &reasoning_type:high matches nothing on "primary" (raises internally, same as + # negation's own fallback-chain behavior), so the router advances to "fallback" + # where the tag is satisfiable. No allow_fail_open involved; this is the plain + # fallback-chain mechanics already established for "!" extended to "&". + router = litellm.Router( + model_list=[ + { + "model_name": "primary", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["reasoning_type:low"], + }, + "model_info": {"id": "primary-low-reasoning"}, + }, + { + "model_name": "fallback", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["reasoning_type:high"], + }, + "model_info": {"id": "fallback-high-reasoning"}, + }, + ], + fallbacks=[{"primary": ["fallback"]}], + enable_tag_filtering=True, + ) + + response = await router.acompletion( + model="primary", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["&reasoning_type:high"]}, + mock_response="hi", + ) + assert response._hidden_params["model_id"] == "fallback-high-reasoning" + + +@pytest.mark.asyncio() +async def test_required_and_negation_and_allow_fail_open_combine_across_three_model_groups(): + # A single request routes through three independent model groups via two + # fallback hops, exercising "!", "&", and allow_fail_open together at each hop: + # - "primary" is banned outright by "!provider:anthropic" -> raises, advances. + # - "secondary" satisfies the negation but not &reasoning_type:high, and has no + # allow_fail_open -> raises exactly as today, advances. + # - "tertiary" has reasoning_type:high, but only on the deployment the same + # "!provider:anthropic" also bans; the tag is known to the chain but its only + # carrier is legitimately excluded, not hidden behind an invented tag, so the + # opted-in allow_fail_open falls back to the group's own default deployment. + router = litellm.Router( + model_list=[ + { + "model_name": "primary", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["provider:anthropic", "reasoning_type:high"], + }, + "model_info": {"id": "primary-anthropic"}, + }, + { + "model_name": "secondary", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["provider:openai", "reasoning_type:low"], + }, + "model_info": {"id": "secondary-openai"}, + }, + { + "model_name": "tertiary", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["provider:anthropic", "reasoning_type:high", "region:eu"], + }, + "model_info": {"id": "tertiary-anthropic-high-reasoning"}, + }, + { + "model_name": "tertiary", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["default", "provider:openai", "reasoning_type:low"], + }, + "model_info": {"id": "tertiary-default", "allow_fail_open": True}, + }, + ], + fallbacks=[{"primary": ["secondary"]}, {"secondary": ["tertiary"]}], + enable_tag_filtering=True, + ) + + response = await router.acompletion( + model="primary", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["!provider:anthropic", "&reasoning_type:high"]}, + mock_response="hi", + ) + assert response._hidden_params["model_id"] == "tertiary-default" + + +@pytest.mark.asyncio() +async def test_unknown_tag_denial_is_scoped_per_hop_not_leaked_across_fallback_groups(): + # On "primary": region:us-east is real and satisfiable there, but the invented + # tag masks it -> denies fail-open -> raises -> advances to "fallback". + # On "fallback": neither region:us-east nor the invented tag is known to this + # entirely different, unrelated group at all, so there's no answer for the + # invented tag to hide -> falls open normally. Each hop must independently + # discover what its own group knows; a deny decision from a prior hop's group + # must not leak forward and block a later hop that has no relevant knowledge. + router = litellm.Router( + model_list=[ + { + "model_name": "primary", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["region:us-east"], + }, + "model_info": {"id": "primary-us-east", "allow_fail_open": True}, + }, + { + "model_name": "fallback", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["default", "provider:openai"], + }, + "model_info": {"id": "fallback-default", "allow_fail_open": True}, + }, + ], + fallbacks=[{"primary": ["fallback"]}], + enable_tag_filtering=True, + ) + + response = await router.acompletion( + model="primary", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["®ion:us-east", "&totally-invented-tag-nobody-has"]}, + mock_response="hi", + ) + assert response._hidden_params["model_id"] == "fallback-default" + + +@pytest.mark.asyncio() +async def test_required_and_only_finds_compliant_non_default_deployment_over_noncompliant_default(): + # A required-AND-only request must be checked against every deployment in the + # group, not just the one tagged "default". A compliant, healthy deployment that + # simply isn't the operator's default must win over routing to a noncompliant + # default just because allow_fail_open happened to be set. + router = litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["provider:anthropic", "region:us-east"], + }, + "model_info": {"id": "anthropic-us-east"}, + }, + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["default", "provider:openai"], + }, + "model_info": {"id": "openai-default", "allow_fail_open": True}, + }, + ], + enable_tag_filtering=True, + ) + + for _ in range(5): + response = await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["®ion:us-east"]}, + mock_response="hi", + ) + assert response._hidden_params["model_id"] == "anthropic-us-east" + + +# --- plain positive-tag exhaustion must not be masked by a universally-applied +# "default" tag; allow_fail_open must still be consulted (or hard-fail without it) --- + + +def _quality_high_cost_low_router(): + return litellm.Router( + model_list=[ + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["default", "quality:high"], + }, + "model_info": {"id": "quality-high-1"}, + }, + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["default", "quality:high"], + }, + "model_info": {"id": "quality-high-2"}, + }, + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["default", "cost:low"], + }, + "model_info": {"id": "cost-low-1"}, + }, + { + "model_name": "gpt-4", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["default", "cost:low"], + }, + "model_info": {"id": "cost-low-2"}, + }, + ], + enable_tag_filtering=True, + ) + + +@pytest.mark.asyncio() +async def test_plain_tag_exhaustion_with_universal_default_tag_raises_by_default(): + # Every deployment in the group is tagged "default" (a legitimate cross-cutting + # safety-net pattern), so default_deployments is never empty on its own. With + # the quality:high deployments unhealthy, a request asking for quality:high + # must still hard-fail, not silently get served by a cost:low deployment just + # because it happens to also carry "default". + from unittest.mock import AsyncMock, patch + + router = _quality_high_cost_low_router() + + with patch( + "litellm.router._async_get_cooldown_deployments", + new=AsyncMock(return_value=["quality-high-1", "quality-high-2"]), + ): + with pytest.raises(Exception) as exc_info: + await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["quality:high"]}, + mock_response="hi", + ) + + from litellm.types.router import RouterErrors + + assert RouterErrors.no_deployments_with_tag_routing.value in str(exc_info.value) + + +@pytest.mark.asyncio() +async def test_plain_tag_exhaustion_with_universal_default_tag_falls_open_when_allowed(): + router = _quality_high_cost_low_router() + for deployment in router.model_list: + deployment["model_info"]["allow_fail_open"] = True + + from unittest.mock import AsyncMock, patch + + with patch( + "litellm.router._async_get_cooldown_deployments", + new=AsyncMock(return_value=["quality-high-1", "quality-high-2"]), + ): + response = await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["quality:high"]}, + mock_response="hi", + ) + + assert response._hidden_params["model_id"] in ("cost-low-1", "cost-low-2") + + +@pytest.mark.asyncio() +async def test_plain_tag_unknown_to_group_still_falls_back_silently_unconditionally(): + # A tag that no deployment in this group has ever carried (foreign to this + # group entirely, e.g. an attribution tag meant for an unrelated mechanism + # sharing the same request-tags list) must keep falling back to the + # "default"-tagged pool unconditionally, exactly like today, regardless of + # allow_fail_open. Only a tag that IS part of this group's real vocabulary + # triggers the new hard-fail/fail-open gate. + router = _quality_high_cost_low_router() + + for _ in range(5): + response = await router.acompletion( + model="gpt-4", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["llm-preference-include:some-unrelated-mechanism"]}, + mock_response="hi", + ) + assert response._hidden_params["model_id"] in ( + "quality-high-1", + "quality-high-2", + "cost-low-1", + "cost-low-2", + ) + + +def test_tag_known_to_group_true_for_real_tag(): + from litellm.router_strategy.tag_based_routing import _tag_known_to_group + + router = _quality_high_cost_low_router() + assert _tag_known_to_group(router, "gpt-4", ["quality:high"], frozenset()) is True + + +def test_tag_known_to_group_false_for_foreign_tag(): + from litellm.router_strategy.tag_based_routing import _tag_known_to_group + + router = _quality_high_cost_low_router() + assert _tag_known_to_group(router, "gpt-4", ["llm-preference-include:unrelated"], frozenset()) is False + + +def test_inherited_constraint_sets_none_when_inherited_tags_absent(): + from litellm.router_strategy.tag_based_routing import _inherited_constraint_sets + + assert _inherited_constraint_sets(None, "") == (None, None) + + +def test_inherited_constraint_sets_splits_required_and_excluded(): + from litellm.router_strategy.tag_based_routing import _inherited_constraint_sets + + inherited_required_set, inherited_excluded_set = _inherited_constraint_sets( + ["®ion:eu", "!region:us", "plain"], "" + ) + assert inherited_required_set == frozenset({"region:eu"}) + assert inherited_excluded_set == frozenset({"region:us"}) + + +def test_inherited_constraint_sets_none_for_non_sequence_value(): + from litellm.router_strategy.tag_based_routing import _inherited_constraint_sets + + # A malformed/unexpected inherited_tags value (anything but a list/tuple) must + # be treated the same as "no origin information", never as "nothing is + # inherited" -- the two are not interchangeable, see _trusted_only_pool. + assert _inherited_constraint_sets("not-a-sequence", "") == (None, None) + + +def test_trusted_only_pool_discards_everything_when_inherited_sets_are_none(): + from litellm.router_strategy.tag_based_routing import _trusted_only_pool + + deployments = ({"litellm_params": {"tags": ["region:us"]}},) + # No origin info at all -> reproduce the pre-provenance unconditional + # fall-open: the trusted-only pool ignores excluded_set/required_set entirely. + assert _trusted_only_pool(deployments, frozenset({"region:eu"}), frozenset({"region:apac"}), None, None) == deployments + + +def test_trusted_only_pool_keeps_constraint_backed_by_inherited_tags(): + from litellm.router_strategy.tag_based_routing import _trusted_only_pool + + eu = {"litellm_params": {"tags": ["region:eu"]}} + us = {"litellm_params": {"tags": ["region:us"]}} + # required_set={"region:eu"} IS in inherited_required_set -> protected, kept. + result = _trusted_only_pool( + (eu, us), frozenset(), frozenset({"region:eu"}), frozenset(), frozenset({"region:eu"}) + ) + assert result == (eu,) + + +def test_trusted_only_pool_discards_a_value_with_no_inherited_backing_even_if_the_caller_also_sent_it(): + # Regression for the value-collision bypass Greptile and veria-ai both + # flagged: a value with zero inherited backing is discardable even when it + # happens to be the exact value the caller submitted -- there is nothing here + # to distinguish "caller-only" from "caller happened to guess a real policy + # value" at this function's level, which is exactly why protection must be + # keyed off presence in inherited_required_set, never absence from a + # caller-supplied set (see the router-level regression below for the full + # bypass this replaces). + from litellm.router_strategy.tag_based_routing import _trusted_only_pool + + eu = {"litellm_params": {"tags": ["region:eu"]}} + us = {"litellm_params": {"tags": ["region:us"]}} + result = _trusted_only_pool((eu, us), frozenset(), frozenset({"region:eu"}), frozenset(), frozenset()) + assert result == (eu, us) + + +def _eu_region_router(): + # eu-1 deliberately carries no "default" tag, and us-default is the only + # "default"-tagged deployment -- this keeps _default_tagged_pool's outcome a + # single, deterministic deployment id in every scenario below, regardless of + # which of the two candidate pools (trusted-only vs fully-unconstrained) a + # given code path resolves to. + return litellm.Router( + model_list=[ + { + "model_name": "chat", + "litellm_params": { + "model": "gpt-4o", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["region:eu"], + }, + "model_info": {"id": "eu-1", "allow_fail_open": True}, + }, + { + "model_name": "chat", + "litellm_params": { + "model": "gpt-4o-mini", + "api_base": "https://exampleopenaiendpoint-production.up.railway.app/", + "tags": ["region:us", "default"], + }, + "model_info": {"id": "us-default", "allow_fail_open": True}, + }, + ], + enable_tag_filtering=True, + ) + + +@pytest.mark.asyncio() +async def test_allow_fail_open_preserves_inherited_constraint_when_caller_tag_causes_exhaustion(): + # ®ion:eu simulates a key/team-inherited hard requirement, captured in + # inherited_tags (a snapshot taken before the caller's own tags are merged + # in); !region:eu simulates the caller's own tag. Combined they exhaust the + # pool (nothing can both carry and not carry region:eu), but allow_fail_open + # must fall back to what still satisfies the inherited requirement, not the + # fully-unconstrained default pool (us-default), and not raise either. + router = _eu_region_router() + + response = await router.acompletion( + model="chat", + messages=[{"role": "user", "content": "hi"}], + metadata={ + "tags": ["®ion:eu", "!region:eu"], + "inherited_tags": ["®ion:eu"], + "caller_tags": ["!region:eu"], + }, + mock_response="hi", + ) + + assert response._hidden_params["model_id"] == "eu-1" + + +@pytest.mark.asyncio() +async def test_allow_fail_open_stays_protected_when_caller_duplicates_the_inherited_tag(): + # Regression for the value-collision bypass Greptile and veria-ai both + # flagged: a caller who resubmits the exact value of an inherited "&" tag + # (here alongside a conflicting "!" for the same value) must not be able to + # strip that value's protection just because it now also appears in + # caller_tags. Protection is keyed off presence in inherited_tags, not + # absence from caller_tags -- if it were the latter, subtracting + # caller_required_set={"region:eu"} from required_set would zero out the + # inherited requirement entirely and this would incorrectly resolve to + # us-default instead of eu-1. + router = _eu_region_router() + + response = await router.acompletion( + model="chat", + messages=[{"role": "user", "content": "hi"}], + metadata={ + "tags": ["®ion:eu", "!region:eu"], + "inherited_tags": ["®ion:eu"], + "caller_tags": ["®ion:eu", "!region:eu"], + }, + mock_response="hi", + ) + + assert response._hidden_params["model_id"] == "eu-1" + + +@pytest.mark.asyncio() +async def test_allow_fail_open_raises_when_inherited_constraint_alone_is_unsatisfiable(): + # Both region:eu and region:us are known to the group (so the unknown-tag + # masking guard does not apply), but no single deployment carries both, and + # inherited_tags confirms the entire required-AND set traces back to policy. + # allow_fail_open must not paper over an inherited requirement that is + # unsatisfiable on its own; it should raise exactly as it would with + # allow_fail_open unset. + router = _eu_region_router() + + with pytest.raises(Exception) as exc_info: + await router.acompletion( + model="chat", + messages=[{"role": "user", "content": "hi"}], + metadata={ + "tags": ["®ion:eu", "®ion:us"], + "inherited_tags": ["®ion:eu", "®ion:us"], + "caller_tags": [], + }, + mock_response="hi", + ) + + from litellm.types.router import RouterErrors + + assert RouterErrors.no_deployments_with_tag_routing.value in str(exc_info.value) + + +@pytest.mark.asyncio() +async def test_allow_fail_open_unconditional_discard_when_inherited_tags_key_absent(): + # No "inherited_tags" key at all (e.g. a direct SDK Router call that never + # went through the proxy's litellm_pre_call_utils.py) must reproduce the exact + # pre-provenance behavior: unconditional fall-open to the default pool, even + # though region:eu here would otherwise look like an inherited requirement. + router = _eu_region_router() + + response = await router.acompletion( + model="chat", + messages=[{"role": "user", "content": "hi"}], + metadata={"tags": ["®ion:eu", "!region:eu"]}, + mock_response="hi", + ) + + assert response._hidden_params["model_id"] == "us-default" + + +# --- tag_routing_prefix must be configurable through every settings-update +# path the router already supports for its sibling enable_tag_filtering, not +# just the config.yaml constructor argument --- + + +def test_router_update_settings_applies_tag_routing_prefix(): + # Regression: tag_routing_prefix was missing from Router.update_settings's + # _allowed_settings, so an operator configuring it via the DB-backed + # router_settings path (proxy_server.py's _add_router_settings_from_db_config, + # which calls update_settings directly) had the value silently ignored. + router = litellm.Router(model_list=[{"model_name": "x", "litellm_params": {"model": "openai/gpt-4o-mini"}}]) + assert router.tag_routing_prefix == "" + + router.update_settings(tag_routing_prefix="route:") + + assert router.tag_routing_prefix == "route:" + + +def test_router_get_settings_includes_tag_routing_prefix(): + router = litellm.Router(model_list=[{"model_name": "x", "litellm_params": {"model": "openai/gpt-4o-mini"}}]) + router.update_settings(tag_routing_prefix="route:") + + assert router.get_settings()["tag_routing_prefix"] == "route:" + + +def test_update_router_config_schema_includes_tag_routing_prefix(): + # The Admin UI's POST /config/update path validates through + # UpdateRouterConfig before calling update_settings; a field missing here + # causes model_dump(exclude_none=True) to silently drop it before + # update_settings is ever called -- the same bug shape LIT-3152 fixed for + # retry_policy (see tests/test_litellm/test_router_retry_policy_update.py). + from litellm.types.router import UpdateRouterConfig + + config = UpdateRouterConfig(tag_routing_prefix="route:") + assert config.model_dump(exclude_none=True)["tag_routing_prefix"] == "route:" diff --git a/type-discipline-budget.json b/type-discipline-budget.json index c990ae52ff2..eee067815d2 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23057 + "limit": 23039 }, "LIT002": { - "limit": 27156 + "limit": 27154 }, "LIT003": { "limit": 269 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16744 + "limit": 16742 }, "LIT011": { "limit": 5596 diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 35adc9e92cf..05cba547f23 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -34222,6 +34222,8 @@ export interface components { routing_strategy_args?: { [key: string]: unknown; } | null; + /** Tag Routing Prefix */ + tag_routing_prefix?: string | null; /** Timeout */ timeout?: number | null; }; @@ -35391,6 +35393,8 @@ export interface components { }; /** ModelInfo */ litellm__types__router__ModelInfo: { + /** Allow Fail Open */ + allow_fail_open?: boolean | null; /** Base Model */ base_model?: string | null; /** Blocked */ @@ -35410,6 +35414,8 @@ export interface components { * @default false */ db_model: boolean; + /** Enable Tag Filtering */ + enable_tag_filtering?: boolean | null; /** Id */ id: string | null; /** Input Cost Per Character */ From cbf85a015fd1a12f86fa0ef7c1bdbbf82e7007dc Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 11 Aug 2026 11:53:11 -0700 Subject: [PATCH 22/24] feat(proxy): per-key prompt caching toggle via enable_prompt_caching (#36466) * feat(proxy): per-key prompt caching auto-injection via enable_prompt_caching Adds a key-level enable_prompt_caching toggle that auto-injects Anthropic cache_control breakpoints on requests made with that key, without requiring the gateway-wide enable_anthropic_prompt_caching flag. The flag lives in key metadata, is stamped onto the request root by add_key_level_controls, rides kwargs into both the /chat/completions seeding path and the native /v1/messages path, and reuses every existing gate (anthropic/bedrock only, supports_prompt_caching, client markers win). Client-supplied body values are stripped as an untrusted root control field. Includes the Admin UI switch on key create and key edit plus a read-only settings row, and dedupes the key edit view's drifted initial-values objects. * fix(proxy): drop section comment and suppress LIT011 on key-level prompt caching stamp --- .../anthropic_cache_control_hook.py | 30 ++++--- litellm/main.py | 2 + litellm/proxy/_types.py | 2 + litellm/proxy/litellm_pre_call_utils.py | 4 + .../key_management_endpoints.py | 2 + litellm/types/utils.py | 1 + .../test_anthropic_cache_control_hook.py | 81 +++++++++++++++++++ .../test_key_management_endpoints.py | 15 ++++ .../proxy/test_litellm_pre_call_utils.py | 38 +++++++++ .../organisms/create_key_button.tsx | 15 ++++ .../templates/key_edit_view.test.tsx | 30 +++++++ .../components/templates/key_edit_view.tsx | 48 +++++------ .../components/templates/key_info_view.tsx | 7 ++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 12 +++ 14 files changed, 248 insertions(+), 39 deletions(-) diff --git a/litellm/integrations/anthropic_cache_control_hook.py b/litellm/integrations/anthropic_cache_control_hook.py index f2ef8d63a07..4df6fce74c0 100644 --- a/litellm/integrations/anthropic_cache_control_hook.py +++ b/litellm/integrations/anthropic_cache_control_hook.py @@ -382,19 +382,23 @@ class AnthropicCacheControlHook(CustomPromptManagement): model: str, custom_llm_provider: str | None, tools: list | None = None, + enable_prompt_caching: bool | None = None, ) -> list[CacheControlInjectionPoint]: """Default breakpoints when ``litellm.enable_anthropic_prompt_caching`` is on. - Caches the system prompt and the trailing turn, so the stable prefix - (system + tools + history) is reused while the breakpoint advances with - the conversation. Returns [] (stand down) when the flag is off, the - provider does not consume cache_control breakpoints (only anthropic / - bedrock do), the model lacks prompt-caching support, or the request - already carries client-supplied cache_control. + ``enable_prompt_caching`` is the per-request override (stamped from key + metadata by the proxy); True turns auto-injection on for this request + even when the global flag is off. Caches the system prompt and the + trailing turn, so the stable prefix (system + tools + history) is + reused while the breakpoint advances with the conversation. Returns [] + (stand down) when neither flag is on, the provider does not consume + cache_control breakpoints (only anthropic / bedrock do), the model + lacks prompt-caching support, or the request already carries + client-supplied cache_control. """ import litellm - if litellm.enable_anthropic_prompt_caching is not True: + if litellm.enable_anthropic_prompt_caching is not True and enable_prompt_caching is not True: return [] provider = custom_llm_provider @@ -433,6 +437,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): model: str, custom_llm_provider: str | None, tools: list | None = None, + enable_prompt_caching: bool | None = None, ) -> None: """For /chat/completions: resolve the injection points the request should carry. @@ -458,6 +463,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): model=model, custom_llm_provider=custom_llm_provider, tools=tools, + enable_prompt_caching=enable_prompt_caching, ) if points: non_default_params["cache_control_injection_points"] = points @@ -478,12 +484,17 @@ class AnthropicCacheControlHook(CustomPromptManagement): judgment happens once per request; points a prior pass wrote back carry the judged stamp and are never re-judged (see ``_should_stand_down``). When none are configured but - ``litellm.enable_anthropic_prompt_caching`` is on, synthesize default - breakpoints for the native /v1/messages path. Pops the key from kwargs; + ``litellm.enable_anthropic_prompt_caching`` or the per-request + ``enable_prompt_caching`` kwarg (stamped from key metadata) is on, + synthesize default breakpoints for the native /v1/messages path. Pops + both keys from kwargs; if remaining (non-message) points exist they are written back so downstream transforms can handle them. """ typed_messages = cast(list[AllMessageValues], messages) # cast-ok: Anthropic-shaped dicts from v1/messages + enable_prompt_caching: Final = cast( # cast-ok: kwargs is untyped; key stamped as bool by the proxy + bool | None, kwargs.pop("enable_prompt_caching", None) + ) configured: Final = cast( # cast-ok: kwargs is untyped; this key only holds the documented injection-point list list[CacheControlInjectionPoint] | None, kwargs.pop("cache_control_injection_points", None) ) @@ -497,6 +508,7 @@ class AnthropicCacheControlHook(CustomPromptManagement): tools=tools, model=model, custom_llm_provider=custom_llm_provider, + enable_prompt_caching=enable_prompt_caching, ) if not injection_points: return messages, system diff --git a/litellm/main.py b/litellm/main.py index c70a41c891a..bf6ec8004ce 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -504,6 +504,7 @@ async def acompletion( model=model, custom_llm_provider=cast(str | None, custom_llm_provider), # cast-ok: read from untyped kwargs tools=tools, + enable_prompt_caching=cast(bool | None, kwargs.get("enable_prompt_caching")), # cast-ok: untyped kwargs ) if isinstance(litellm_logging_obj, LiteLLMLoggingObj) and ( @@ -5105,6 +5106,7 @@ def completion( model=model, custom_llm_provider=cast(str | None, kwargs.get("custom_llm_provider")), # cast-ok: untyped kwargs tools=tools, + enable_prompt_caching=cast(bool | None, kwargs.get("enable_prompt_caching")), # cast-ok: untyped kwargs ) if isinstance(litellm_logging_obj, LiteLLMLoggingObj) and ( diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index dc4f17c7b31..08348187645 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -1108,6 +1108,7 @@ class KeyRequestBase(GenerateRequestBase): budget_id: str | None = None tags: list[str] | None = None disable_global_guardrails: bool | None = None + enable_prompt_caching: bool | None = None throttle_on_budget_exceeded: bool | None = None enforced_params: list[str] | None = None allowed_routes: list | None = [] @@ -4124,6 +4125,7 @@ LiteLLM_ManagementEndpoint_MetadataFields: Final = [ "enforced_batch_output_expires_after", "enforced_file_expires_after", "throttle_on_budget_exceeded", + "enable_prompt_caching", ] LiteLLM_ManagementEndpoint_MetadataFields_Premium: Final = [ diff --git a/litellm/proxy/litellm_pre_call_utils.py b/litellm/proxy/litellm_pre_call_utils.py index 0924b6aebea..f83061a15ce 100644 --- a/litellm/proxy/litellm_pre_call_utils.py +++ b/litellm/proxy/litellm_pre_call_utils.py @@ -201,6 +201,7 @@ _UNTRUSTED_ROOT_CONTROL_FIELDS: Final = ( "mock_tool_calls", "disable_global_guardrails", "disable_global_guardrail", + "enable_prompt_caching", "opted_out_global_guardrails", "applied_guardrails", "applied_policies", @@ -1333,6 +1334,9 @@ class LiteLLMProxyRequestSetup: if "disable_fallbacks" in key_metadata and isinstance(key_metadata["disable_fallbacks"], bool): data["disable_fallbacks"] = key_metadata["disable_fallbacks"] + if isinstance(key_metadata.get("enable_prompt_caching"), bool): + data["enable_prompt_caching"] = key_metadata["enable_prompt_caching"] # rebind-ok: data is an out-param + ## KEY-LEVEL METADATA data = LiteLLMProxyRequestSetup.add_management_endpoint_metadata_to_request_metadata( data=data, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 836b223c6ad..ebf44497658 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -1593,6 +1593,7 @@ async def generate_key_fn( - policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules. - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. - throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely. + - enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Anthropic and Bedrock Claude models only. - permissions: Optional[dict] - key-specific permissions. Currently just used for turning off pii masking (if connected). Example - {"pii": false} - model_max_budget: Optional[Dict[str, BudgetConfig]] - Model-specific budgets {"gpt-4": {"budget_limit": 0.0005, "time_period": "30d"}}}. IF null or {} then no model specific budget. - budget_fallbacks: Optional[Dict[str, List[str]]] - Per-model fallback chain tried in order when that model's own `model_max_budget` is exceeded, e.g. {"gpt-4o": ["gpt-4o-mini"]}. @@ -2693,6 +2694,7 @@ async def update_key_fn( - policies: Optional[List[str]] - List of policy names to apply to the key. Policies define guardrails, conditions, and inheritance rules. - disable_global_guardrails: Optional[bool] - Whether to disable global guardrails for the key. - throttle_on_budget_exceeded: Optional[bool] - When the key exceeds its max_budget, throttle its tpm/rpm to the global budget_exceeded_throttle_percentage instead of blocking the key entirely. + - enable_prompt_caching: Optional[bool] - Auto-inject prompt caching breakpoints (Anthropic cache_control markers) on requests made with this key. Anthropic and Bedrock Claude models only. - prompts: Optional[List[str]] - List of prompts that the key is allowed to use. - blocked: Optional[bool] - Whether the key is blocked - aliases: Optional[dict] - Model aliases for the key - [Docs](https://litellm.vercel.app/docs/proxy/virtual_keys#model-aliases) diff --git a/litellm/types/utils.py b/litellm/types/utils.py index 74311d59d8e..d4c26df54d8 100644 --- a/litellm/types/utils.py +++ b/litellm/types/utils.py @@ -3467,6 +3467,7 @@ all_litellm_params = ( "caching_groups", "ttl", "cache", + "enable_prompt_caching", "no-log", "base_model", "stream_timeout", diff --git a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py index 47baacd61d7..cc43a424419 100644 --- a/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py +++ b/tests/test_litellm/integrations/test_anthropic_cache_control_hook.py @@ -1728,6 +1728,87 @@ class TestEnableAnthropicPromptCaching: assert result_msgs[-1]["content"][-1]["cache_control"] == {"type": "ephemeral"} assert "cache_control" not in result_msgs[0]["content"][-1] + +class TestPerKeyEnablePromptCaching: + """Per-request enable_prompt_caching override (stamped from key metadata) with the global flag off.""" + + MESSAGES: List[AllMessageValues] = [ + {"role": "system", "content": "a long system prompt"}, + {"role": "user", "content": "latest turn"}, + ] + + def _points(self, enable_prompt_caching, model="claude-sonnet-4-5", provider="anthropic", messages=None): + return AnthropicCacheControlHook.get_default_injection_points( + messages=copy.deepcopy(self.MESSAGES) if messages is None else messages, + system=None, + model=model, + custom_llm_provider=provider, + enable_prompt_caching=enable_prompt_caching, + ) + + def test_true_injects_with_global_flag_off(self): + assert litellm.enable_anthropic_prompt_caching is False + assert self._points(True) == [ + {"location": "message", "role": "system", "index": None, "control": {"type": "ephemeral"}}, + {"location": "message", "role": None, "index": -1, "control": {"type": "ephemeral"}}, + ] + + @pytest.mark.parametrize("enable_prompt_caching", [False, None]) + def test_false_and_none_fall_back_to_global_flag(self, enable_prompt_caching): + assert self._points(enable_prompt_caching) == [] + + def test_false_does_not_suppress_global_flag(self, monkeypatch): + monkeypatch.setattr(litellm, "enable_anthropic_prompt_caching", True) + assert [p["index"] for p in self._points(False)] == [None, -1] + + def test_provider_gate_still_applies(self): + assert self._points(True, model="gpt-4o", provider="openai") == [] + + def test_unsupported_model_gate_still_applies(self): + assert self._points(True, model="anthropic.claude-3-5-sonnet-20240620-v1:0", provider="bedrock") == [] + + def test_client_markers_still_win(self): + messages = [ + {"role": "system", "content": [{"type": "text", "text": "s", "cache_control": {"type": "ephemeral"}}]}, + {"role": "user", "content": "latest turn"}, + ] + assert self._points(True, messages=messages) == [] + + def test_seed_injects_with_global_flag_off(self): + params: dict = {} + AnthropicCacheControlHook.maybe_seed_default_injection_points( + non_default_params=params, + messages=copy.deepcopy(self.MESSAGES), + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + enable_prompt_caching=True, + ) + assert [p["index"] for p in params["cache_control_injection_points"]] == [None, -1] + + def test_v1_messages_injects_and_pops_flag_from_kwargs(self): + kwargs: dict = {"enable_prompt_caching": True} + result_msgs, result_sys = AnthropicCacheControlHook.maybe_inject_cache_control( + [{"role": "user", "content": [{"type": "text", "text": "latest"}]}], + "a system prompt", + kwargs, + model="claude-sonnet-4-5", + custom_llm_provider="anthropic", + ) + assert result_sys == [{"type": "text", "text": "a system prompt", "cache_control": {"type": "ephemeral"}}] + assert result_msgs[-1]["content"][-1]["cache_control"] == {"type": "ephemeral"} + assert "enable_prompt_caching" not in kwargs + + def test_v1_messages_pops_flag_even_when_noop(self): + kwargs: dict = {"enable_prompt_caching": True} + AnthropicCacheControlHook.maybe_inject_cache_control( + [{"role": "user", "content": [{"type": "text", "text": "hi"}]}], + None, + kwargs, + model="gpt-4o", + custom_llm_provider="openai", + ) + assert "enable_prompt_caching" not in kwargs + def test_v1_messages_is_noop_when_disabled(self): messages = [{"role": "user", "content": [{"type": "text", "text": "hi"}]}] result_msgs, result_sys = AnthropicCacheControlHook.maybe_inject_cache_control( diff --git a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py index 4c13a3367d9..0a88f59f677 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_key_management_endpoints.py @@ -1733,6 +1733,21 @@ async def test_update_service_account_works_with_team_id(): await prepare_key_update_data(data=data, existing_key_row=existing_key) +@pytest.mark.asyncio +@pytest.mark.parametrize("flag_value", [True, False]) +async def test_update_key_enable_prompt_caching_folds_into_metadata(flag_value): + """Top-level enable_prompt_caching on /key/update lands in key metadata, including flipping back to False.""" + data = UpdateKeyRequest(key="sk-1", enable_prompt_caching=flag_value) + existing_key = LiteLLM_VerificationToken( + token="hashed", metadata={"enable_prompt_caching": not flag_value} + ) + + updated = await prepare_key_update_data(data=data, existing_key_row=existing_key) + + assert updated["metadata"]["enable_prompt_caching"] is flag_value + assert "enable_prompt_caching" not in {k for k in updated if k != "metadata"} + + @pytest.mark.asyncio async def test_update_preserves_service_account_id_when_metadata_replaced(): """ diff --git a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py index 293cce5fa5d..e31058f402e 100644 --- a/tests/test_litellm/proxy/test_litellm_pre_call_utils.py +++ b/tests/test_litellm/proxy/test_litellm_pre_call_utils.py @@ -688,6 +688,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): "mock_response": "free response", "mock_tool_calls": [{"id": "call_1"}], "disable_global_guardrails": True, + "enable_prompt_caching": True, "routing_decision": {"cause": "forged", "routed_model": "spoofed"}, "metadata": copy.deepcopy(malicious_metadata), "litellm_metadata": copy.deepcopy(malicious_metadata), @@ -705,6 +706,7 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): assert "mock_response" not in updated assert "mock_tool_calls" not in updated assert "disable_global_guardrails" not in updated + assert "enable_prompt_caching" not in updated assert "routing_decision" not in updated stripped_keys = { @@ -741,6 +743,42 @@ async def test_add_litellm_data_to_request_strips_user_control_fields(): assert "pillar_response_headers" not in snapshot_body["metadata"] +@pytest.mark.asyncio +@pytest.mark.parametrize( + "key_value, expected", + [(True, True), (False, False), ("yes", None), (None, None)], +) +async def test_key_metadata_enable_prompt_caching_promoted_to_request_root(key_value, expected): + """Key metadata enable_prompt_caching is stamped onto the request root (bools only), even when the client spoofs it.""" + request_mock = MagicMock(spec=Request) + request_mock.url.path = "/v1/chat/completions" + request_mock.url = MagicMock() + request_mock.url.__str__.return_value = "http://localhost/v1/chat/completions" + request_mock.method = "POST" + request_mock.query_params = {} + request_mock.headers = {"Content-Type": "application/json"} + request_mock.client = MagicMock() + request_mock.client.host = "127.0.0.1" + + data = { + "model": "claude-sonnet-4-5", + "messages": [{"role": "user", "content": "hello"}], + "enable_prompt_caching": "spoofed-by-client", + } + key_metadata = {} if key_value is None else {"enable_prompt_caching": key_value} + + updated = await add_litellm_data_to_request( + data=data, + request=request_mock, + user_api_key_dict=UserAPIKeyAuth(api_key="hashed-key", metadata=key_metadata), + proxy_config=MagicMock(), + general_settings={}, + version="test-version", + ) + + assert updated.get("enable_prompt_caching") == expected + + @pytest.mark.asyncio @pytest.mark.parametrize( "control_field", diff --git a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx index 19abae73549..8951a471f84 100644 --- a/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx +++ b/ui/litellm-dashboard/src/components/organisms/create_key_button.tsx @@ -1191,6 +1191,21 @@ const CreateKey: React.FC = ({ team, teams, data, addKey, autoOp > + + Enable Prompt Caching{" "} + + + + + } + name="enable_prompt_caching" + valuePropName="checked" + > + + diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx index 29b32c1f7fc..fb8ab87e45f 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.test.tsx @@ -410,6 +410,36 @@ describe("KeyEditView", () => { }); }); + it("should initialize and submit enable_prompt_caching from key metadata", async () => { + const onSubmitMock = vi.fn().mockResolvedValue(undefined); + const keyDataWithPromptCaching = { + ...MOCK_KEY_DATA, + metadata: { ...MOCK_KEY_DATA.metadata, enable_prompt_caching: true }, + }; + + renderWithProviders( + {}} + onSubmit={onSubmitMock} + accessToken={"test-token"} + userID={"test-user"} + userRole={"admin"} + premiumUser={false} + />, + ); + + await waitFor(() => { + expect(screen.getByText("Enable Prompt Caching")).toBeInTheDocument(); + }); + + await userEvent.click(screen.getByRole("button", { name: /save changes/i })); + + await waitFor(() => { + expect(onSubmitMock).toHaveBeenCalledWith(expect.objectContaining({ enable_prompt_caching: true })); + }); + }); + it("should disable models field when management routes are selected", async () => { const keyDataWithManagementRoutes = { ...MOCK_KEY_DATA, diff --git a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx index e407527f562..36dc02f7bd0 100644 --- a/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_edit_view.tsx @@ -150,6 +150,7 @@ export function KeyEditView({ guardrails: keyData.metadata?.guardrails, disable_global_guardrails: keyData.metadata?.disable_global_guardrails || false, throttle_on_budget_exceeded: keyData.metadata?.throttle_on_budget_exceeded || false, + enable_prompt_caching: keyData.metadata?.enable_prompt_caching || false, ...estimateFields(keyData.metadata), prompts: keyData.metadata?.prompts, tags: keyData.metadata?.tags, @@ -178,36 +179,8 @@ export function KeyEditView({ }; useEffect(() => { - form.setFieldsValue({ - ...keyData, - token: keyData.token || keyData.token_id, - budget_duration: canonicalBudgetDuration(keyData.budget_duration), - metadata: formatMetadataForDisplay(stripTagsFromMetadata(keyData.metadata)), - guardrails: keyData.metadata?.guardrails, - disable_global_guardrails: keyData.metadata?.disable_global_guardrails || false, - prompts: keyData.metadata?.prompts, - tags: keyData.metadata?.tags, - vector_stores: keyData.object_permission?.vector_stores || [], - mcp_servers_and_groups: { - servers: keyData.object_permission?.mcp_servers || [], - accessGroups: keyData.object_permission?.mcp_access_groups || [], - toolsets: keyData.object_permission?.mcp_toolsets || [], - }, - mcp_tool_permissions: keyData.object_permission?.mcp_tool_permissions || {}, - throttle_on_budget_exceeded: keyData.metadata?.throttle_on_budget_exceeded || false, - ...estimateFields(keyData.metadata), - logging_settings: extractLoggingSettings(keyData.metadata), - disabled_callbacks: Array.isArray(keyData.metadata?.litellm_disabled_callbacks) - ? mapInternalToDisplayNames(keyData.metadata.litellm_disabled_callbacks) - : [], - access_group_ids: keyData.access_group_ids || [], - auto_rotate: keyData.auto_rotate || false, - ...(keyData.rotation_interval && { rotation_interval: keyData.rotation_interval }), - allowed_routes: - Array.isArray(keyData.allowed_routes) && keyData.allowed_routes.length > 0 - ? keyData.allowed_routes.join(", ") - : "", - }); + form.setFieldsValue(initialValues); + // eslint-disable-next-line react-hooks/exhaustive-deps -- initialValues is rebuilt from keyData every render; depending on it would re-run each render }, [keyData, form]); // Sync auto-rotation state with form values @@ -532,6 +505,21 @@ export function KeyEditView({ + + Enable Prompt Caching{" "} + + + + + } + name="enable_prompt_caching" + valuePropName="checked" + > + + + diff --git a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx index 6a2a4a88425..6fd6547a995 100644 --- a/ui/litellm-dashboard/src/components/templates/key_info_view.tsx +++ b/ui/litellm-dashboard/src/components/templates/key_info_view.tsx @@ -782,6 +782,13 @@ export default function KeyInfoView({ {currentKeyData.expires ? formatTimestamp(currentKeyData.expires) : "Never"} + {Boolean(currentKeyData.metadata?.enable_prompt_caching) && ( +
+ Prompt Caching + Enabled (auto-injects cache_control markers on Anthropic and Bedrock Claude requests) +
+ )} + Date: Tue, 11 Aug 2026 12:31:43 -0700 Subject: [PATCH 23/24] fix(ui): stub useIsOrgAdmin in UsageTab tests so useCan needs no QueryClient (#36565) 255d65192e added useCan to UsageTab, whose useIsOrgAdmin leg calls useOrganizations (react-query), so every UsageTab test died with 'No QueryClient set'. Stub the org-admin leg; role gating still flows through the real hasCapability with the varied userRole. --- .../cost-optimization/_components/UsageTab.test.tsx | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx index ad68111bba7..25be956f18b 100644 --- a/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx +++ b/ui/litellm-dashboard/src/app/(dashboard)/cost-optimization/_components/UsageTab.test.tsx @@ -13,6 +13,12 @@ vi.mock("@/app/(dashboard)/hooks/useAuthorized", () => ({ default: useAuthorizedMock, })); +// useCan reaches useOrganizations (react-query) through useIsOrgAdmin; stub the +// org-admin leg so role gating flows through hasCapability without a QueryClient +vi.mock("@/app/(dashboard)/hooks/useIsOrgAdmin", () => ({ + default: () => false, +})); + vi.mock("@/components/networking", () => ({ getToolSpend: (...args: unknown[]) => mockGetToolSpend(...args), })); From be71a8fdbf46f0335a8ba71daa0ffb9e39568c02 Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Tue, 11 Aug 2026 12:41:11 -0700 Subject: [PATCH 24/24] fix(alerting): dedupe scheduled Slack spend reports across pods (#36489) * fix(alerting): dedupe scheduled Slack spend reports across pods Every pod ran its own weekly/monthly spend report jobs, prometheus fallback stats cron, and daily report loop, so deployments with multiple replicas or uvicorn workers received one copy per pod. Gate each scheduled send behind the shared PodLockManager redis lock. The lock is never released: its TTL (the full reporting window for the weekly interval job, whose per-pod anchors drift by boot time and jitter) doubles as a sent-this-window marker. acquire_lock returning None (no redis wired) proceeds, preserving single-pod behavior. Also generalize the pod lock could-not-acquire log line, which claimed to be about spend tracking for every consumer. Fixes #14809 * fix(alerting): harden spend report locks after adversarial review Weekly lock TTL gets an hour haircut: with ttl equal to the interval, the winner re-fires just before its own key expires, reacquires without a TTL refresh, and the key then lapses in time for a trailing pod to re-send. Job/lock ids move to litellm/constants.py per convention, and spend_report_frequency now rejects non-positive day counts, which previously coerced to an every-second schedule and would now compute a negative lock TTL that silently never sends. Adds the missing test coverage the review flagged: startup_event's pod_lock_manager wiring (identity-asserted), the prometheus closure's positive path, and the ungated immediate prometheus send pinned to exactly one await. * test(alerting): consolidate spend_report_frequency validator coverage Drops a duplicate non-positive-days test and parametrizes the survivor over the suffix half of the validator too * fix(alerting): route the startup prometheus fallback send through the pod lock Greptile caught that the boot-time send still ran once per pod when PROMETHEUS_URL is set, the same duplication class this PR removes * fix(alerting): make report lock acquisition non-reentrant Greptile caught that a pod booting within an hour of the fallback stats cron sent twice: the startup send takes the lock, then the cron fire hits acquire_lock's reacquire branch, which returns True for the holder. Window-marker gates now pass allow_reentrant=False so a live lock blocks everyone including its holder; leader-election consumers keep the reentrant default * test(proxy): give spec'd ProxyLogging mocks a db_spend_update_writer _initialize_slack_alerting_jobs now reads it for the pod lock manager, and spec=ProxyLogging blocks instance-only attributes --- litellm/constants.py | 4 + .../SlackAlerting/slack_alerting.py | 27 +- .../db_transaction_queue/pod_lock_manager.py | 21 +- litellm/proxy/proxy_server.py | 58 +- litellm/proxy/utils.py | 5 +- .../SlackAlerting/test_slack_alerting.py | 170 ++- .../test_pod_lock_manager.py | 35 +- .../proxy/proxy_server/test_lifecycle.py | 184 ++- tests/test_litellm/proxy/test_proxy_server.py | 1241 +++++------------ .../utils/proxy_logging/test_lifecycle.py | 36 + 10 files changed, 807 insertions(+), 974 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index 87d6fa1a744..c9d9ff155ff 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1478,6 +1478,10 @@ CLOUDZERO_MAX_FETCHED_DATA_RECORDS: Final = int(os.getenv("CLOUDZERO_MAX_FETCHED SPEND_LOG_CLEANUP_JOB_NAME: Final = "spend_log_cleanup" KEY_ROTATION_JOB_NAME: Final = "litellm_key_rotation_job" EXPIRED_UI_SESSION_KEY_CLEANUP_JOB_NAME: Final = "litellm_expired_ui_session_key_cleanup_job" +WEEKLY_SPEND_REPORT_JOB_ID: Final = "weekly_spend_report_job" +MONTHLY_SPEND_REPORT_JOB_ID: Final = "monthly_spend_report_job" +PROMETHEUS_FALLBACK_STATS_JOB_ID: Final = "prometheus_fallback_stats_job" +SLACK_DAILY_REPORT_LOCK_ID: Final = "slack_daily_report" SPEND_LOG_RUN_LOOPS: Final = int(os.getenv("SPEND_LOG_RUN_LOOPS", 500)) SPEND_LOG_CLEANUP_BATCH_SIZE: Final = int(os.getenv("SPEND_LOG_CLEANUP_BATCH_SIZE", 1000)) SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES = int(os.getenv("SPEND_LOG_CLEANUP_MAX_CONSECUTIVE_BATCH_FAILURES", 3)) diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index 7edfb93e581..f3cd937599c 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -17,7 +17,7 @@ import litellm.litellm_core_utils.litellm_logging import litellm.types from litellm._logging import verbose_logger, verbose_proxy_logger from litellm.caching.caching import DualCache -from litellm.constants import HOURS_IN_A_DAY +from litellm.constants import HOURS_IN_A_DAY, SLACK_DAILY_REPORT_LOCK_ID from litellm.integrations.custom_batch_logger import CustomBatchLogger from litellm.integrations.SlackAlerting.budget_alert_types import get_budget_alert_type from litellm.integrations.SlackAlerting.hanging_request_check import ( @@ -51,6 +51,7 @@ from .batching_handler import send_to_webhook, squash_payloads from .utils import process_slack_alerting_variables if TYPE_CHECKING: + from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager from litellm.router import Router as _Router Router = _Router @@ -1576,7 +1577,11 @@ Model Info: except Exception: pass - async def _run_scheduler_helper(self, llm_router) -> bool: + async def _run_scheduler_helper( + self, + llm_router, + pod_lock_manager: "PodLockManager | None" = None, + ) -> bool: """ Returns: - True -> report sent @@ -1601,6 +1606,16 @@ Model Info: interval_seconds: Final = self.alerting_args.daily_report_frequency if current_time - report_sent >= interval_seconds: + if ( + pod_lock_manager is not None + and ( + await pod_lock_manager.acquire_lock( + cronjob_id=SLACK_DAILY_REPORT_LOCK_ID, ttl=interval_seconds, allow_reentrant=False + ) + ) + is False + ): + return False # Sneak in the reporting logic here await self.send_daily_reports(router=llm_router) # Also, don't forget to update the report_sent time after sending the report! @@ -1612,7 +1627,11 @@ Model Info: return report_sent_bool - async def _run_scheduled_daily_report(self, llm_router: Any | None = None): + async def _run_scheduled_daily_report( + self, + llm_router: Any | None = None, + pod_lock_manager: "PodLockManager | None" = None, + ): """ If 'daily_reports' enabled @@ -1625,7 +1644,7 @@ Model Info: if "daily_reports" in self.alert_types: while True: - await self._run_scheduler_helper(llm_router=llm_router) + await self._run_scheduler_helper(llm_router=llm_router, pod_lock_manager=pod_lock_manager) interval = random.randint( self.alerting_args.report_check_interval - 3, self.alerting_args.report_check_interval + 3, diff --git a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py index c74cb412c68..4be1331e955 100644 --- a/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py +++ b/litellm/proxy/db/db_transaction_queue/pod_lock_manager.py @@ -43,6 +43,7 @@ end self, cronjob_id: str, ttl: int | None = None, + allow_reentrant: bool = True, ) -> bool | None: """ Attempt to acquire the lock for a specific cron job using Redis. @@ -53,6 +54,10 @@ end ttl: Optional custom TTL in seconds. Defaults to DEFAULT_CRON_JOB_LOCK_TTL_SECONDS. Use a longer TTL for jobs that may take longer than the default 60s (e.g. key rotation with many keys). + allow_reentrant: With the default True, a pod that already holds the lock + acquires it again (leader election semantics). Pass False when the live + lock marks work as already done for this window, so not even the holder + may redo it before the TTL expires. """ if self.redis_cache is None: verbose_proxy_logger.debug("redis_cache is None, skipping acquire_lock") @@ -88,7 +93,7 @@ end if current_value is not None: if isinstance(current_value, bytes): current_value = current_value.decode("utf-8") - if current_value == self.pod_id: + if current_value == self.pod_id and allow_reentrant: verbose_proxy_logger.info( "Pod %s already holds the Redis lock for cronjob_id=%s", self.pod_id, @@ -96,14 +101,12 @@ end ) self._emit_acquired_lock_event(cronjob_id, self.pod_id) return True - else: - verbose_proxy_logger.info( - "Spend tracking - pod %s could not acquire lock for cronjob_id=%s, " - "held by pod %s. Spend updates in Redis will wait for the leader pod to commit.", - self.pod_id, - cronjob_id, - current_value, - ) + verbose_proxy_logger.info( + "Pod %s could not acquire lock for cronjob_id=%s, held by pod %s.", + self.pod_id, + cronjob_id, + current_value, + ) return False except Exception as e: verbose_proxy_logger.error("Error acquiring Redis lock for %s: %s", cronjob_id, e) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index f79819c76d3..19030d50110 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -233,6 +233,8 @@ from litellm.constants import ( GLOBAL_PROXY_SPEND_CACHE_KEY, LITELLM_PROXY_ADMIN_NAME, LITELLM_PROXY_BUDGET_NAME, + MONTHLY_SPEND_REPORT_JOB_ID, + PROMETHEUS_FALLBACK_STATS_JOB_ID, PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS, PROXY_BATCH_POLLING_ENABLED, PROXY_BATCH_POLLING_INTERVAL, @@ -240,6 +242,7 @@ from litellm.constants import ( PROXY_BUDGET_RESCHEDULER_MAX_TIME, PROXY_BUDGET_RESCHEDULER_MIN_TIME, PROXY_CONFIG_RELOAD_INTERVAL_SECONDS, + WEEKLY_SPEND_REPORT_JOB_ID, ) from litellm.exceptions import RejectedRequestError from litellm.integrations.custom_guardrail import ModifyResponseException @@ -9043,41 +9046,76 @@ class ProxyStartupEvent: spend_report_frequency: Final[str] = general_settings.get("spend_report_frequency", "7d") or "7d" days: Final = int(spend_report_frequency[:-1]) - if spend_report_frequency[-1].lower() != "d": - raise ValueError("spend_report_frequency must be specified in days, e.g., '1d', '7d'") + if spend_report_frequency[-1].lower() != "d" or days <= 0: + raise ValueError("spend_report_frequency must be a positive number of days, e.g., '1d', '7d'") + + pod_lock_manager: Final = proxy_logging_obj.db_spend_update_writer.pod_lock_manager + weekly_lock_ttl: Final = duration_in_seconds(spend_report_frequency) - 3600 + + async def _scheduled_weekly_spend_report() -> None: + # TTL spans the whole reporting window: each pod's interval anchor is its own + # boot time + jitter, so a shorter lock would let a later pod re-send the report. + # Minus an hour so the next window's first firer finds a free key + if ( + await pod_lock_manager.acquire_lock( + cronjob_id=WEEKLY_SPEND_REPORT_JOB_ID, ttl=weekly_lock_ttl, allow_reentrant=False + ) + is False + ): + return + await proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report(spend_report_frequency) + + async def _scheduled_monthly_spend_report() -> None: + if ( + await pod_lock_manager.acquire_lock( + cronjob_id=MONTHLY_SPEND_REPORT_JOB_ID, ttl=3600, allow_reentrant=False + ) + is False + ): + return + await proxy_logging_obj.slack_alerting_instance.send_monthly_spend_report() scheduler.add_job( - proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report, + _scheduled_weekly_spend_report, "interval", days=days, next_run_time=datetime.now() + timedelta(seconds=10 + random.randint(0, 300)), - args=[spend_report_frequency], - id="weekly_spend_report_job", + id=WEEKLY_SPEND_REPORT_JOB_ID, replace_existing=True, misfire_grace_time=APSCHEDULER_MISFIRE_GRACE_TIME, ) scheduler.add_job( - proxy_logging_obj.slack_alerting_instance.send_monthly_spend_report, + _scheduled_monthly_spend_report, "cron", day=1, - id="monthly_spend_report_job", + id=MONTHLY_SPEND_REPORT_JOB_ID, replace_existing=True, ) if os.getenv("PROMETHEUS_URL"): from zoneinfo import ZoneInfo + async def _scheduled_fallback_stats() -> None: + if ( + await pod_lock_manager.acquire_lock( + cronjob_id=PROMETHEUS_FALLBACK_STATS_JOB_ID, ttl=3600, allow_reentrant=False + ) + is False + ): + return + await proxy_logging_obj.slack_alerting_instance.send_fallback_stats_from_prometheus() + scheduler.add_job( - proxy_logging_obj.slack_alerting_instance.send_fallback_stats_from_prometheus, + _scheduled_fallback_stats, "cron", hour=PROMETHEUS_FALLBACK_STATS_SEND_TIME_HOURS, minute=0, timezone=ZoneInfo("America/Los_Angeles"), - id="prometheus_fallback_stats_job", + id=PROMETHEUS_FALLBACK_STATS_JOB_ID, replace_existing=True, ) - await proxy_logging_obj.slack_alerting_instance.send_fallback_stats_from_prometheus() + await _scheduled_fallback_stats() @classmethod async def _setup_prisma_client( diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index b5972935806..9653106f7e0 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -468,7 +468,10 @@ class ProxyLogging: and not self.daily_report_started ): asyncio.create_task( - self.slack_alerting_instance._run_scheduled_daily_report(llm_router=llm_router) + self.slack_alerting_instance._run_scheduled_daily_report( + llm_router=llm_router, + pod_lock_manager=self.db_spend_update_writer.pod_lock_manager, + ) ) # RUN DAILY REPORT (if scheduled) self.daily_report_started = True diff --git a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py index 1ea4795207d..23a35098697 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py +++ b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py @@ -1,17 +1,21 @@ +import asyncio import datetime import json import os import sys +import time import unittest -from typing import List, Optional, Tuple +from typing import Final, List, Optional, Tuple from unittest.mock import ANY, AsyncMock, MagicMock, Mock, patch -sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system-path +import pytest + +sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system-path import litellm +from litellm.caching.caching import DualCache from litellm.integrations.SlackAlerting.slack_alerting import SlackAlerting from litellm.proxy._types import CallInfo, Litellm_EntityType +from litellm.types.integrations.slack_alerting import SlackAlertingCacheKeys class TestSlackAlerting(unittest.TestCase): @@ -20,37 +24,27 @@ class TestSlackAlerting(unittest.TestCase): def test_get_percent_of_max_budget_left(self): # Test case 1: When max_budget is None - user_info = CallInfo( - max_budget=None, spend=50.0, event_group=Litellm_EntityType.KEY - ) + user_info = CallInfo(max_budget=None, spend=50.0, event_group=Litellm_EntityType.KEY) result = self.slack_alerting._get_percent_of_max_budget_left(user_info) self.assertEqual(result, 0.0) # Test case 2: When max_budget is 0 - user_info = CallInfo( - max_budget=0.0, spend=50.0, event_group=Litellm_EntityType.KEY - ) + user_info = CallInfo(max_budget=0.0, spend=50.0, event_group=Litellm_EntityType.KEY) result = self.slack_alerting._get_percent_of_max_budget_left(user_info) self.assertEqual(result, 0.0) # Test case 3: When spend is less than max_budget - user_info = CallInfo( - max_budget=100.0, spend=75.0, event_group=Litellm_EntityType.KEY - ) + user_info = CallInfo(max_budget=100.0, spend=75.0, event_group=Litellm_EntityType.KEY) result = self.slack_alerting._get_percent_of_max_budget_left(user_info) self.assertEqual(result, 0.25) # Test case 4: When spend equals max_budget - user_info = CallInfo( - max_budget=100.0, spend=100.0, event_group=Litellm_EntityType.KEY - ) + user_info = CallInfo(max_budget=100.0, spend=100.0, event_group=Litellm_EntityType.KEY) result = self.slack_alerting._get_percent_of_max_budget_left(user_info) self.assertEqual(result, 0.0) # Test case 5: When spend exceeds max_budget - user_info = CallInfo( - max_budget=100.0, spend=120.0, event_group=Litellm_EntityType.KEY - ) + user_info = CallInfo(max_budget=100.0, spend=120.0, event_group=Litellm_EntityType.KEY) result = self.slack_alerting._get_percent_of_max_budget_left(user_info) self.assertEqual(result, -0.2) @@ -189,7 +183,9 @@ class TestSlackAlerting(unittest.TestCase): # Test the specific formatting logic we're interested in alert_type_formatted = f"Alert type: `{alert_type.name}`\n" - formatted_message = f"{alert_type_formatted}\n Level: `{level}`\nTimestamp: `{current_time}`\n\nMessage: {message}" + formatted_message = ( + f"{alert_type_formatted}\n Level: `{level}`\nTimestamp: `{current_time}`\n\nMessage: {message}" + ) # Verify alert_type is in the formatted message as expected self.assertIn("Alert type: `llm_exceptions`", formatted_message) @@ -214,9 +210,7 @@ class TestSlackAlerting(unittest.TestCase): json.dumps(outage_value) # Verify the specific error message - self.assertIn( - "Object of type set is not JSON serializable", str(context.exception) - ) + self.assertIn("Object of type set is not JSON serializable", str(context.exception)) def test_fixed_redis_serialization(self): """Test that our fix resolves the Redis serialization error.""" @@ -245,3 +239,133 @@ class TestSlackAlerting(unittest.TestCase): ) self.assertEqual(parsed_data["alerts"], [408]) self.assertEqual(parsed_data["provider_region_id"], "vertex_aius-east1") + + +_REPORT_SENT_KEY: Final = SlackAlertingCacheKeys.report_sent_key.value +_DAILY_REPORT_FREQUENCY: Final = 900 + + +async def _slack_alerting_with_due_daily_report() -> SlackAlerting: + slack_alerting: Final = SlackAlerting( + internal_usage_cache=DualCache(), + alerting_args={"daily_report_frequency": _DAILY_REPORT_FREQUENCY}, + ) + await slack_alerting.internal_usage_cache.async_set_cache( + key=_REPORT_SENT_KEY, + value=time.time() - _DAILY_REPORT_FREQUENCY - 1, + ) + slack_alerting.send_daily_reports = AsyncMock() + return slack_alerting + + +async def _read_report_sent(slack_alerting: SlackAlerting) -> float: + return await slack_alerting.internal_usage_cache.async_get_cache( + key=_REPORT_SENT_KEY, + parent_otel_span=None, + ) + + +@pytest.mark.asyncio +async def test_daily_report_skipped_when_another_pod_holds_the_lock(): + """regression: issue #14809 - every pod sent its own copy of the daily report. + + The losing pod must also leave report_sent untouched so the winner's window still counts. + """ + slack_alerting: Final = await _slack_alerting_with_due_daily_report() + report_sent_before: Final = await _read_report_sent(slack_alerting) + pod_lock_manager: Final = AsyncMock() + pod_lock_manager.acquire_lock.return_value = False + + result: Final = await slack_alerting._run_scheduler_helper( + llm_router=MagicMock(), + pod_lock_manager=pod_lock_manager, + ) + + assert result is False + slack_alerting.send_daily_reports.assert_not_awaited() + assert await _read_report_sent(slack_alerting) == report_sent_before + pod_lock_manager.acquire_lock.assert_awaited_once_with( + cronjob_id="slack_daily_report", + ttl=_DAILY_REPORT_FREQUENCY, + allow_reentrant=False, + ) + + +@pytest.mark.asyncio +async def test_daily_report_sent_by_the_pod_that_wins_the_lock(): + slack_alerting: Final = await _slack_alerting_with_due_daily_report() + report_sent_before: Final = await _read_report_sent(slack_alerting) + llm_router: Final = MagicMock() + pod_lock_manager: Final = AsyncMock() + pod_lock_manager.acquire_lock.return_value = True + + result: Final = await slack_alerting._run_scheduler_helper( + llm_router=llm_router, + pod_lock_manager=pod_lock_manager, + ) + + assert result is True + slack_alerting.send_daily_reports.assert_awaited_once_with(router=llm_router) + assert await _read_report_sent(slack_alerting) > report_sent_before + pod_lock_manager.acquire_lock.assert_awaited_once_with( + cronjob_id="slack_daily_report", + ttl=_DAILY_REPORT_FREQUENCY, + allow_reentrant=False, + ) + + +@pytest.mark.parametrize("lock_state", ["no_pod_lock_manager", "no_redis_configured"]) +@pytest.mark.asyncio +async def test_daily_report_still_sent_without_a_working_lock(lock_state: str): + """Single-pod parity: a missing lock manager, or one whose acquire_lock returns None + because redis isn't configured, must not suppress the report.""" + slack_alerting: Final = await _slack_alerting_with_due_daily_report() + report_sent_before: Final = await _read_report_sent(slack_alerting) + llm_router: Final = MagicMock() + pod_lock_manager: Final = ( + None if lock_state == "no_pod_lock_manager" else AsyncMock(acquire_lock=AsyncMock(return_value=None)) + ) + + result: Final = await slack_alerting._run_scheduler_helper( + llm_router=llm_router, + pod_lock_manager=pod_lock_manager, + ) + + assert result is True + slack_alerting.send_daily_reports.assert_awaited_once_with(router=llm_router) + assert await _read_report_sent(slack_alerting) > report_sent_before + + +@pytest.mark.asyncio +async def test_daily_report_lock_not_attempted_before_the_interval_elapses(): + """The lock is a per-window marker, so a pod must not burn it on a check that isn't due yet.""" + slack_alerting: Final = await _slack_alerting_with_due_daily_report() + await slack_alerting.internal_usage_cache.async_set_cache(key=_REPORT_SENT_KEY, value=time.time()) + pod_lock_manager: Final = AsyncMock() + pod_lock_manager.acquire_lock.return_value = True + + result: Final = await slack_alerting._run_scheduler_helper( + llm_router=MagicMock(), + pod_lock_manager=pod_lock_manager, + ) + + assert result is False + pod_lock_manager.acquire_lock.assert_not_awaited() + slack_alerting.send_daily_reports.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_scheduled_daily_report_threads_the_pod_lock_manager_through(): + """The loop in _run_scheduled_daily_report is where the lock manager reaches the gate.""" + slack_alerting: Final = SlackAlerting(alert_types=["daily_reports"]) + pod_lock_manager: Final = AsyncMock() + slack_alerting._run_scheduler_helper = AsyncMock(side_effect=asyncio.CancelledError) + + with pytest.raises(asyncio.CancelledError): + await slack_alerting._run_scheduled_daily_report( + llm_router=MagicMock(), + pod_lock_manager=pod_lock_manager, + ) + + _, kwargs = slack_alerting._run_scheduler_helper.await_args + assert kwargs["pod_lock_manager"] is pod_lock_manager diff --git a/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py b/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py index f2745052faa..7a1ab60c547 100644 --- a/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py +++ b/tests/test_litellm/proxy/db/db_transaction_queue/test_pod_lock_manager.py @@ -7,9 +7,7 @@ from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi.testclient import TestClient -sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system path +sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system path from litellm.constants import DEFAULT_CRON_JOB_LOCK_TTL_SECONDS from litellm.proxy.db.db_transaction_queue.pod_lock_manager import PodLockManager @@ -310,9 +308,7 @@ async def test_lock_takeover_race_condition(mock_redis): @pytest.mark.asyncio -async def test_release_lock_uses_atomic_compare_delete_script_when_available( - pod_lock_manager, mock_redis -): +async def test_release_lock_uses_atomic_compare_delete_script_when_available(pod_lock_manager, mock_redis): """ Test that release_lock prefers atomic compare-and-delete Lua script when redis cache exposes script registration. @@ -323,12 +319,8 @@ async def test_release_lock_uses_atomic_compare_delete_script_when_available( await pod_lock_manager.release_lock(cronjob_id="test_job") lock_key = pod_lock_manager.get_redis_lock_key(cronjob_id="test_job") - mock_redis.async_register_script.assert_called_once_with( - PodLockManager._COMPARE_AND_DELETE_LOCK_SCRIPT - ) - script_callable.assert_called_once_with( - keys=[lock_key], args=[json.dumps(pod_lock_manager.pod_id)] - ) + mock_redis.async_register_script.assert_called_once_with(PodLockManager._COMPARE_AND_DELETE_LOCK_SCRIPT) + script_callable.assert_called_once_with(keys=[lock_key], args=[json.dumps(pod_lock_manager.pod_id)]) mock_redis.async_get_cache.assert_not_called() mock_redis.async_delete_cache.assert_not_called() @@ -359,9 +351,7 @@ async def test_release_lock_lua_path_emits_released_event(pod_lock_manager, mock with patch.object(pod_lock_manager, "_emit_released_lock_event") as mock_emit: await pod_lock_manager.release_lock(cronjob_id="test_job") - mock_emit.assert_called_once_with( - cronjob_id="test_job", pod_id=pod_lock_manager.pod_id - ) + mock_emit.assert_called_once_with(cronjob_id="test_job", pod_id=pod_lock_manager.pod_id) class FakeRedisLockStore: @@ -437,9 +427,7 @@ async def test_release_lock_preserves_lock_held_by_other_pod(): @pytest.mark.asyncio -async def test_release_lock_falls_back_to_get_del_when_lua_execution_fails( - pod_lock_manager, mock_redis -): +async def test_release_lock_falls_back_to_get_del_when_lua_execution_fails(pod_lock_manager, mock_redis): """ Test that release_lock falls back to GET+DEL when Lua script execution raises (e.g. Redis restart cleared loaded scripts). @@ -457,3 +445,14 @@ async def test_release_lock_falls_back_to_get_del_when_lua_execution_fails( mock_redis.async_delete_cache.assert_called_once_with(lock_key) # Cached script handle should be reset so next call re-registers assert pod_lock_manager._release_lock_script is None + + +@pytest.mark.asyncio +async def test_acquire_lock_own_lock_not_reentrant(pod_lock_manager, mock_redis): + """With allow_reentrant=False a live lock means the window's work is done, so even + the holder gets False; the default stays reentrant for leader-election callers.""" + mock_redis.async_set_cache.return_value = False + mock_redis.async_get_cache.return_value = pod_lock_manager.pod_id + + assert await pod_lock_manager.acquire_lock(cronjob_id="test_job", allow_reentrant=False) is False + assert await pod_lock_manager.acquire_lock(cronjob_id="test_job") is True diff --git a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py index 6ac1e15e7b5..40ca7e3a64e 100644 --- a/tests/test_litellm/proxy/proxy_server/test_lifecycle.py +++ b/tests/test_litellm/proxy/proxy_server/test_lifecycle.py @@ -22,6 +22,7 @@ import inspect import json import logging import os +from collections.abc import Awaitable, Callable from typing import List, Optional, Union from unittest.mock import AsyncMock, MagicMock, patch @@ -397,9 +398,7 @@ def test__redact_worker_config_for_logging_masks_nested_secret_fields(): "database_url": nested_db_url, "database_extra_connection_params": {"password": nested_extra_pw}, "alert_to_webhook_url": {"budget_alerts": nested_webhook}, - "pass_through_endpoints": [ - {"path": "/up", "headers": {"Authorization": nested_bearer}} - ], + "pass_through_endpoints": [{"path": "/up", "headers": {"Authorization": nested_bearer}}], } } } @@ -451,16 +450,13 @@ def test_load_from_azure_key_vault_disabled_no_side_effect(monkeypatch): import litellm sentinel_secret_mgr = object() - monkeypatch.setattr( - litellm, "secret_manager_client", sentinel_secret_mgr, raising=False - ) + monkeypatch.setattr(litellm, "secret_manager_client", sentinel_secret_mgr, raising=False) result = load_from_azure_key_vault(use_azure_key_vault=False) observed = { "return_value": result, - "secret_manager_unchanged": litellm.secret_manager_client - is sentinel_secret_mgr, + "secret_manager_unchanged": litellm.secret_manager_client is sentinel_secret_mgr, "called_with": False, } assert normalize(observed) == { @@ -614,9 +610,7 @@ def test_get_litellm_model_info_uses_base_model_for_lookup(monkeypatch): observed = { "called_arg": ( - fake_get.call_args.args[0] - if fake_get.call_args.args - else fake_get.call_args.kwargs.get("model") + fake_get.call_args.args[0] if fake_get.call_args.args else fake_get.call_args.kwargs.get("model") ), "returned_max_tokens": result.get("max_tokens"), "returned_cost": result.get("input_cost_per_token"), @@ -663,9 +657,7 @@ def test_run_ollama_serve_invokes_subprocess_popen(monkeypatch): def test_run_ollama_serve_popen_failure_is_swallowed(monkeypatch): """Popen raising OSError must NOT propagate — function logs and returns.""" - monkeypatch.setattr( - ps.subprocess, "Popen", MagicMock(side_effect=OSError("no ollama binary")) - ) + monkeypatch.setattr(ps.subprocess, "Popen", MagicMock(side_effect=OSError("no ollama binary"))) result = run_ollama_serve() assert result is None @@ -685,8 +677,7 @@ async def test_proxy_startup_event_is_async_context_manager_with_expected_signat observed = { "param_count": len(sig.parameters), "has_app_param": "app" in sig.parameters, - "wrapped_is_async": inspect.iscoroutinefunction(wrapped) - or inspect.isasyncgenfunction(wrapped), + "wrapped_is_async": inspect.iscoroutinefunction(wrapped) or inspect.isasyncgenfunction(wrapped), "has_asynccontextmanager_wrapper": wrapped is not None, } assert normalize(observed) == { @@ -777,3 +768,164 @@ def test_proxy_startup_event_warns_for_global_budget_without_database(): assert budget_check_pos < warn_pos < next_startup_section_pos, ( "DB-less budget warning must run after Prisma setup and the DB-backed budget block" ) + + +# --------------------------------------------------------------------------- +# _initialize_slack_alerting_jobs — spend-report pod locking (issue #14809) +# --------------------------------------------------------------------------- + +SlackAlertingJobs = dict[str, Callable[[], Awaitable[None]]] + + +def _make_slack_alerting_proxy_logging(acquire_lock_result: bool | None) -> MagicMock: + proxy_logging_obj = MagicMock() + proxy_logging_obj.slack_alerting_instance.alerting = ["slack"] + proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report = AsyncMock() + proxy_logging_obj.slack_alerting_instance.send_monthly_spend_report = AsyncMock() + proxy_logging_obj.slack_alerting_instance.send_fallback_stats_from_prometheus = AsyncMock() + pod_lock_manager = proxy_logging_obj.db_spend_update_writer.pod_lock_manager + pod_lock_manager.acquire_lock = AsyncMock(return_value=acquire_lock_result) + pod_lock_manager.release_lock = AsyncMock() + return proxy_logging_obj + + +async def _init_slack_alerting_jobs( + acquire_lock_result: bool | None, + spend_report_frequency: str = "7d", +) -> tuple[SlackAlertingJobs, MagicMock]: + scheduler = MagicMock() + proxy_logging_obj = _make_slack_alerting_proxy_logging(acquire_lock_result) + + await ProxyStartupEvent._initialize_slack_alerting_jobs( + scheduler=scheduler, + general_settings={"spend_report_frequency": spend_report_frequency}, + proxy_logging_obj=proxy_logging_obj, + prisma_client=MagicMock(), + ) + + jobs = {call.kwargs["id"]: call.args[0] for call in scheduler.add_job.call_args_list} + return jobs, proxy_logging_obj + + +@pytest.mark.parametrize("spend_report_frequency", ["0d", "-1d", "7h"]) +@pytest.mark.asyncio +async def test_initialize_slack_alerting_jobs_invalid_frequency_raises(spend_report_frequency: str): + """A non-positive window used to become an every-second APScheduler interval, and now also + computes a negative lock TTL that expires instantly and suppresses the report for good. + match= is load-bearing: drop the guard and "-1d" still raises, but from duration_in_seconds.""" + with pytest.raises(ValueError, match="positive number of days"): + await _init_slack_alerting_jobs( + acquire_lock_result=True, + spend_report_frequency=spend_report_frequency, + ) + + +@pytest.mark.asyncio +async def test_weekly_spend_report_skipped_when_another_pod_holds_the_lock(): + """regression: issue #14809 - every pod ran its own weekly spend report job.""" + jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=False) + + await jobs["weekly_spend_report_job"]() + + proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report.assert_not_awaited() + proxy_logging_obj.db_spend_update_writer.pod_lock_manager.acquire_lock.assert_awaited_once_with( + cronjob_id="weekly_spend_report_job", + ttl=7 * 86400 - 3600, + allow_reentrant=False, + ) + + +@pytest.mark.parametrize("acquire_lock_result", [True, None]) +@pytest.mark.asyncio +async def test_weekly_spend_report_sent_when_the_lock_is_free_or_absent(acquire_lock_result): + """None means redis isn't configured; a single-pod deploy must still report.""" + jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=acquire_lock_result) + + await jobs["weekly_spend_report_job"]() + + proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report.assert_awaited_once_with("7d") + + +@pytest.mark.asyncio +async def test_weekly_spend_report_lock_ttl_tracks_the_configured_window(): + """TTL is the window less an hour: long enough that no second pod re-sends inside the + window, short enough that the lock is gone before the next one opens. A fixed TTL would + break one end or the other as soon as spend_report_frequency changes.""" + jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=True, spend_report_frequency="1d") + + await jobs["weekly_spend_report_job"]() + + proxy_logging_obj.db_spend_update_writer.pod_lock_manager.acquire_lock.assert_awaited_once_with( + cronjob_id="weekly_spend_report_job", + ttl=86400 - 3600, + allow_reentrant=False, + ) + proxy_logging_obj.slack_alerting_instance.send_weekly_spend_report.assert_awaited_once_with("1d") + + +@pytest.mark.asyncio +async def test_monthly_spend_report_skipped_when_another_pod_holds_the_lock(): + jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=False) + + await jobs["monthly_spend_report_job"]() + + proxy_logging_obj.slack_alerting_instance.send_monthly_spend_report.assert_not_awaited() + proxy_logging_obj.db_spend_update_writer.pod_lock_manager.acquire_lock.assert_awaited_once_with( + cronjob_id="monthly_spend_report_job", + ttl=3600, + allow_reentrant=False, + ) + + +@pytest.mark.parametrize("acquire_lock_result", [True, None]) +@pytest.mark.asyncio +async def test_monthly_spend_report_sent_when_the_lock_is_free_or_absent(acquire_lock_result): + jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=acquire_lock_result) + + await jobs["monthly_spend_report_job"]() + + proxy_logging_obj.slack_alerting_instance.send_monthly_spend_report.assert_awaited_once_with() + + +@pytest.mark.asyncio +async def test_spend_report_locks_are_never_released(): + """The lock is a per-window marker, not a mutex: releasing it lets the next pod re-send.""" + jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=True) + + await jobs["weekly_spend_report_job"]() + await jobs["monthly_spend_report_job"]() + + proxy_logging_obj.db_spend_update_writer.pod_lock_manager.release_lock.assert_not_awaited() + + +@pytest.mark.asyncio +async def test_prometheus_fallback_stats_job_skipped_when_another_pod_holds_the_lock(monkeypatch): + """The boot-time send goes through the same gate, so a losing pod sends nothing at all: + startup and the scheduled job both stay at zero.""" + monkeypatch.setenv("PROMETHEUS_URL", "http://prometheus.invalid") + jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=False) + send_fallback_stats = proxy_logging_obj.slack_alerting_instance.send_fallback_stats_from_prometheus + assert send_fallback_stats.await_count == 0 + + await jobs["prometheus_fallback_stats_job"]() + + assert send_fallback_stats.await_count == 0 + proxy_logging_obj.db_spend_update_writer.pod_lock_manager.acquire_lock.assert_awaited_with( + cronjob_id="prometheus_fallback_stats_job", + ttl=3600, + allow_reentrant=False, + ) + assert proxy_logging_obj.db_spend_update_writer.pod_lock_manager.acquire_lock.await_count == 2 + + +@pytest.mark.parametrize("acquire_lock_result", [True, None]) +@pytest.mark.asyncio +async def test_prometheus_fallback_stats_job_runs_when_the_lock_is_free_or_absent(monkeypatch, acquire_lock_result): + monkeypatch.setenv("PROMETHEUS_URL", "http://prometheus.invalid") + jobs, proxy_logging_obj = await _init_slack_alerting_jobs(acquire_lock_result=acquire_lock_result) + send_fallback_stats = proxy_logging_obj.slack_alerting_instance.send_fallback_stats_from_prometheus + assert send_fallback_stats.await_count == 1 + + await jobs["prometheus_fallback_stats_job"]() + + assert send_fallback_stats.await_count == 2 diff --git a/tests/test_litellm/proxy/test_proxy_server.py b/tests/test_litellm/proxy/test_proxy_server.py index d9ac45c0531..3acb9fcafd3 100644 --- a/tests/test_litellm/proxy/test_proxy_server.py +++ b/tests/test_litellm/proxy/test_proxy_server.py @@ -19,9 +19,7 @@ from fastapi import FastAPI from fastapi.staticfiles import StaticFiles from fastapi.testclient import TestClient -sys.path.insert( - 0, os.path.abspath("../../..") -) # Adds the parent directory to the system-path +sys.path.insert(0, os.path.abspath("../../..")) # Adds the parent directory to the system-path import litellm import litellm.proxy.proxy_server as proxy_server_module @@ -179,9 +177,7 @@ def test_login_v2_returns_json_on_http_exception(monkeypatch): from fastapi import HTTPException mock_prisma_client = MagicMock() - mock_authenticate_user = AsyncMock( - side_effect=HTTPException(status_code=401, detail="Unauthorized") - ) + mock_authenticate_user = AsyncMock(side_effect=HTTPException(status_code=401, detail="Unauthorized")) monkeypatch.setattr( "litellm.proxy.auth.login_utils.authenticate_user", @@ -477,9 +473,7 @@ def test_fallback_login_has_no_deprecation_banner(client_no_auth): "relative/path/logo.png", ], ) -def test_get_logo_url_does_not_disclose_local_paths( - client_no_auth, monkeypatch, ui_logo_path -): +def test_get_logo_url_does_not_disclose_local_paths(client_no_auth, monkeypatch, ui_logo_path): # ``/get_logo_url`` is unauthenticated. Returning a local filesystem # path verbatim discloses admin-only config to any caller. Only # browser-loadable HTTP(S) URLs should be returned; for local paths @@ -579,9 +573,7 @@ def test_restructure_ui_html_files_handles_nested_routes(tmp_path): assert not (ui_root / "home.html").exists() assert (ui_root / "home" / "index.html").read_text() == "home" assert not (ui_root / "mcp" / "oauth" / "callback.html").exists() - assert ( - ui_root / "mcp" / "oauth" / "callback" / "index.html" - ).read_text() == "callback" + assert (ui_root / "mcp" / "oauth" / "callback" / "index.html").read_text() == "callback" assert (ui_root / "existing" / "index.html").read_text() == "keep" assert (ui_root / "_next" / "ignore.html").read_text() == "asset" assert (ui_root / "litellm-asset-prefix" / "ignore.html").read_text() == "asset" @@ -626,9 +618,7 @@ def test_admin_ui_export_serves_nested_extensionless_routes(): and "_next" not in path.parts and "litellm-asset-prefix" not in path.parts ] - assert not nested_html_offenders, ( - "Nested routes must be named index.html. Offenders: " f"{nested_html_offenders}" - ) + assert not nested_html_offenders, f"Nested routes must be named index.html. Offenders: {nested_html_offenders}" callback_index = out_dir / "mcp" / "oauth" / "callback" / "index.html" assert callback_index.is_file(), ( @@ -645,9 +635,7 @@ def test_admin_ui_export_serves_nested_extensionless_routes(): follow_redirects=False, ) assert redirect.status_code == 307 - assert redirect.headers["location"].endswith( - "/ui/mcp/oauth/callback/?code=abc&state=xyz" - ) + assert redirect.headers["location"].endswith("/ui/mcp/oauth/callback/?code=abc&state=xyz") landed = client.get("/ui/mcp/oauth/callback?code=abc&state=xyz") assert landed.status_code == 200 @@ -712,6 +700,7 @@ async def test_initialize_scheduled_jobs_credentials(monkeypatch): mock_prisma_client = MagicMock() mock_proxy_logging = MagicMock(spec=ProxyLogging) mock_proxy_logging.slack_alerting_instance = MagicMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() mock_proxy_config = AsyncMock() with ( @@ -750,9 +739,7 @@ async def test_initialize_scheduled_jobs_credentials(monkeypatch): assert mock_proxy_config.get_credentials.call_count == 1 # Direct call # Verify a scheduled job was added for get_credentials - mock_scheduler_calls = [ - call[0] for call in mock_proxy_config.get_credentials.mock_calls - ] + mock_scheduler_calls = [call[0] for call in mock_proxy_config.get_credentials.mock_calls] assert len(mock_scheduler_calls) > 0 @@ -773,6 +760,7 @@ async def test_periodic_reload_job_scheduled_without_store_model_in_db(monkeypat mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=None) mock_proxy_logging = MagicMock(spec=ProxyLogging) mock_proxy_logging.slack_alerting_instance = MagicMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() mock_proxy_config = AsyncMock() scheduler = AsyncIOScheduler() @@ -813,6 +801,7 @@ async def test_initialize_scheduled_jobs_uses_configured_config_reload_interval( mock_prisma_client = MagicMock() mock_proxy_logging = MagicMock(spec=ProxyLogging) mock_proxy_logging.slack_alerting_instance = MagicMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() mock_proxy_config = AsyncMock() mock_scheduler = MagicMock() @@ -861,6 +850,7 @@ async def test_initialize_scheduled_jobs_rejects_non_positive_config_reload_inte mock_prisma_client = MagicMock() mock_proxy_logging = MagicMock(spec=ProxyLogging) mock_proxy_logging.slack_alerting_instance = MagicMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() mock_proxy_config = AsyncMock() mock_scheduler = MagicMock() @@ -907,6 +897,7 @@ async def test_initialize_scheduled_jobs_hydrates_mcp_when_store_model_in_db_fal mock_prisma_client = MagicMock() mock_proxy_logging = MagicMock(spec=ProxyLogging) mock_proxy_logging.slack_alerting_instance = MagicMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() mock_proxy_config = AsyncMock() with ( @@ -1051,9 +1042,7 @@ def test_get_config_custom_callback_api_env_vars(monkeypatch): assert response.status_code == 200 callbacks = response.json()["callbacks"] - custom_cb = next( - (cb for cb in callbacks if cb["name"] == "custom_callback_api"), None - ) + custom_cb = next((cb for cb in callbacks if cb["name"] == "custom_callback_api"), None) assert custom_cb is not None assert custom_cb["variables"] == { @@ -1101,9 +1090,7 @@ def test_get_config_callbacks_fall_back_to_process_env(mock_env_vars, monkeypatc app.dependency_overrides = original_overrides assert response.status_code == 200 - langfuse_cb = next( - (cb for cb in response.json()["callbacks"] if cb["name"] == "langfuse"), None - ) + langfuse_cb = next((cb for cb in response.json()["callbacks"] if cb["name"] == "langfuse"), None) assert langfuse_cb is not None assert langfuse_cb["variables"] == { "LANGFUSE_PUBLIC_KEY": "pk-env-only", @@ -1150,9 +1137,7 @@ def test_get_config_callback_env_secrets_redacted_for_non_admin(mock_env_vars, m app.dependency_overrides = original_overrides assert response.status_code == 200 - langfuse_cb = next( - (cb for cb in response.json()["callbacks"] if cb["name"] == "langfuse"), None - ) + langfuse_cb = next((cb for cb in response.json()["callbacks"] if cb["name"] == "langfuse"), None) assert langfuse_cb is not None assert langfuse_cb["variables"]["LANGFUSE_SECRET_KEY"] == "REDACTED" assert langfuse_cb["variables"]["LANGFUSE_HOST"] == "https://cloud.langfuse.com" @@ -1202,9 +1187,7 @@ def test_get_config_returns_email_settings(monkeypatch): app.dependency_overrides = original_overrides assert response.status_code == 200 - email_alert = next( - (a for a in response.json()["alerts"] if a["name"] == "email"), None - ) + email_alert = next((a for a in response.json()["alerts"] if a["name"] == "email"), None) assert email_alert is not None variables = email_alert["variables"] @@ -1349,9 +1332,7 @@ def test_get_config_returns_slack_webhook(monkeypatch): mock_logging = MagicMock() mock_logging.slack_alerting_instance.alert_types = ["budget_alerts"] - mock_logging.slack_alerting_instance._all_possible_alert_types.return_value = [ - "budget_alerts" - ] + mock_logging.slack_alerting_instance._all_possible_alert_types.return_value = ["budget_alerts"] mock_logging.slack_alerting_instance.alert_to_webhook_url = {} monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_logging) monkeypatch.setattr(proxy_config, "get_config", AsyncMock(return_value=config_data)) @@ -1368,9 +1349,7 @@ def test_get_config_returns_slack_webhook(monkeypatch): app.dependency_overrides = original_overrides assert response.status_code == 200 - slack_alert = next( - (a for a in response.json()["alerts"] if a["name"] == "slack"), None - ) + slack_alert = next((a for a in response.json()["alerts"] if a["name"] == "slack"), None) assert slack_alert is not None masked_url = slack_alert["variables"]["SLACK_WEBHOOK_URL"] @@ -1390,9 +1369,7 @@ def test_get_config_cleared_slack_webhook_not_overridden_by_os_env(monkeypatch): """ from litellm.proxy.proxy_server import app, proxy_config, user_api_key_auth - monkeypatch.setenv( - "SLACK_WEBHOOK_URL", "https://hooks.slack.com/services/STALE/OS/ENVVALUE" - ) + monkeypatch.setenv("SLACK_WEBHOOK_URL", "https://hooks.slack.com/services/STALE/OS/ENVVALUE") config_data = { "litellm_settings": {}, "general_settings": {"alerting": ["slack"]}, @@ -1405,9 +1382,7 @@ def test_get_config_cleared_slack_webhook_not_overridden_by_os_env(monkeypatch): mock_logging = MagicMock() mock_logging.slack_alerting_instance.alert_types = ["budget_alerts"] - mock_logging.slack_alerting_instance._all_possible_alert_types.return_value = [ - "budget_alerts" - ] + mock_logging.slack_alerting_instance._all_possible_alert_types.return_value = ["budget_alerts"] mock_logging.slack_alerting_instance.alert_to_webhook_url = {} monkeypatch.setattr("litellm.proxy.proxy_server.proxy_logging_obj", mock_logging) monkeypatch.setattr(proxy_config, "get_config", AsyncMock(return_value=config_data)) @@ -1424,9 +1399,7 @@ def test_get_config_cleared_slack_webhook_not_overridden_by_os_env(monkeypatch): app.dependency_overrides = original_overrides assert response.status_code == 200 - slack_alert = next( - (a for a in response.json()["alerts"] if a["name"] == "slack"), None - ) + slack_alert = next((a for a in response.json()["alerts"] if a["name"] == "slack"), None) assert slack_alert is not None assert slack_alert["variables"]["SLACK_WEBHOOK_URL"] == "" @@ -1505,9 +1478,7 @@ async def test_aaaproxy_startup_master_key(mock_prisma, monkeypatch, tmp_path): # Test Case 3: Master key with os.environ prefix test_resolved_key = "sk-resolved-key" - test_config_with_prefix = { - "general_settings": {"master_key": "os.environ/CUSTOM_MASTER_KEY"} - } + test_config_with_prefix = {"general_settings": {"master_key": "os.environ/CUSTOM_MASTER_KEY"}} # Create config with os.environ prefix with open(config_path, "w") as f: @@ -1659,9 +1630,7 @@ async def test_get_all_team_models(): ) # Verify find_many was called with where clause for specific teams - mock_litellm_teamtable.find_many.assert_called_with( - where={"team_id": {"in": ["team1"]}} - ) + mock_litellm_teamtable.find_many.assert_called_with(where={"team_id": {"in": ["team1"]}}) # Verify router.get_model_list was called only for team1 models expected_calls = [ @@ -1856,14 +1825,10 @@ async def test_apply_search_filter_scopes_byok_to_caller_teams(): prisma_client = MagicMock() prisma_client.db.litellm_proxymodeltable.count = AsyncMock(return_value=2) - prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock( - return_value=[db_caller_row, db_other_row] - ) + prisma_client.db.litellm_proxymodeltable.find_many = AsyncMock(return_value=[db_caller_row, db_other_row]) caller_user_row = MagicMock() caller_user_row.teams = ["team-mine"] - prisma_client.db.litellm_usertable.find_unique = AsyncMock( - return_value=caller_user_row - ) + prisma_client.db.litellm_usertable.find_unique = AsyncMock(return_value=caller_user_row) proxy_config = MagicMock() proxy_config.decrypt_model_list_from_db = lambda rows: [ @@ -1893,12 +1858,10 @@ async def test_apply_search_filter_scopes_byok_to_caller_teams(): assert "byok-db-mine" in filtered_ids assert "public-id" in filtered_ids assert "byok-other" not in filtered_ids, ( - "router-side BYOK from another team must be dropped from search " - "when caller doesn't belong to that team" + "router-side BYOK from another team must be dropped from search when caller doesn't belong to that team" ) assert "byok-db-other" not in filtered_ids, ( - "DB-only BYOK from another team must be dropped from search when " - "caller doesn't belong to that team" + "DB-only BYOK from another team must be dropped from search when caller doesn't belong to that team" ) # total_count is router_models_count (2: caller_team_byok + public_model, # other_team_byok dropped router-side) + DB count (2 from the mocked @@ -2049,9 +2012,7 @@ async def test_filter_models_by_team_id_excludes_viewer_direct_access(): assert "byok-team-111" in visible_ids, "team-111's own BYOK must always be visible" assert "byok-team-222" not in visible_ids, "must not leak other teams' BYOK" - assert ( - "public-id" not in visible_ids - ), "viewer's direct_access must not widen the team's visible set" + assert "public-id" not in visible_ids, "viewer's direct_access must not widen the team's visible set" @pytest.mark.asyncio @@ -2234,9 +2195,7 @@ async def test_add_access_group_models_to_team_models(): mock_ag_row.access_model_names = ["claude-3", "gemini"] mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_accessgrouptable.find_many = AsyncMock( - return_value=[mock_ag_row] - ) + mock_prisma_client.db.litellm_accessgrouptable.find_many = AsyncMock(return_value=[mock_ag_row]) result = await _add_access_group_models_to_team_models( team_db_objects_typed=[ @@ -2312,9 +2271,7 @@ async def test_add_access_group_models_multiple_teams_shared_group(): mock_extra_row.access_model_names = ["gemini"] mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_accessgrouptable.find_many = AsyncMock( - return_value=[mock_shared_row, mock_extra_row] - ) + mock_prisma_client.db.litellm_accessgrouptable.find_many = AsyncMock(return_value=[mock_shared_row, mock_extra_row]) result = await _add_access_group_models_to_team_models( team_db_objects_typed=[team_a, team_b], @@ -2507,24 +2464,14 @@ async def test_delete_deployment_type_mismatch(): # The two SHA-hash models have no corresponding entry in combined_id_list # and must be evicted. assert len(deleted_ids) == 2, f"Expected 2 deletions (SHA-hash models), got {deleted_ids}" - assert ( - "a96e12e76b36a57cfae57a41288eb41567629cac89b4828c6f7074afc3534695" - in deleted_ids - ) - assert ( - "a40186dd0fdb9b7282380277d7f57044d29de95bfbfcd7f4322b3493702d5cd3" - in deleted_ids - ) + assert "a96e12e76b36a57cfae57a41288eb41567629cac89b4828c6f7074afc3534695" in deleted_ids + assert "a40186dd0fdb9b7282380277d7f57044d29de95bfbfcd7f4322b3493702d5cd3" in deleted_ids # Models 12345678 and 12345679 exist in the config (as integers); str() # conversion in _delete_deployment makes them match the router's string IDs, # so they must NOT be evicted. - assert ( - "12345678" not in deleted_ids - ), f"Model 12345678 should NOT be deleted. Deleted IDs: {deleted_ids}" - assert ( - "12345679" not in deleted_ids - ), f"Model 12345679 should NOT be deleted. Deleted IDs: {deleted_ids}" + assert "12345678" not in deleted_ids, f"Model 12345678 should NOT be deleted. Deleted IDs: {deleted_ids}" + assert "12345679" not in deleted_ids, f"Model 12345679 should NOT be deleted. Deleted IDs: {deleted_ids}" assert still_desired is not None assert {"12345678", "12345679"} <= still_desired, ( @@ -2597,9 +2544,7 @@ async def test_get_config_from_file(tmp_path, monkeypatch): await proxy_config._get_config_from_file(str(empty_file)) # Test Case 5: Using global user_config_file_path when no config_file_path provided - monkeypatch.setattr( - "litellm.proxy.proxy_server.user_config_file_path", str(config_file) - ) + monkeypatch.setattr("litellm.proxy.proxy_server.user_config_file_path", str(config_file)) result = await proxy_config._get_config_from_file(None) assert result == test_config @@ -2718,9 +2663,7 @@ async def test_add_proxy_budget_to_db_only_creates_user_no_keys(): ) # Patch generate_key_helper_fn in proxy_server where it's being called from - with patch( - "litellm.proxy.proxy_server.generate_key_helper_fn", mock_generate_key_helper - ): + with patch("litellm.proxy.proxy_server.generate_key_helper_fn", mock_generate_key_helper): # Call the function under test ProxyStartupEvent._add_proxy_budget_to_db() @@ -2846,9 +2789,7 @@ async def test_custom_ui_sso_sign_in_handler_config_loading(): proxy_config = ProxyConfig() # Create a mock router since load_config requires it mock_router = MagicMock() - await proxy_config.load_config( - router=mock_router, config_file_path=config_file_path - ) + await proxy_config.load_config(router=mock_router, config_file_path=config_file_path) # Verify get_instance_fn was called with correct parameters mock_get_instance.assert_called_with( @@ -2888,9 +2829,7 @@ async def test_load_config_max_budget_env_var_coerced_to_float(tmp_path, monkeyp original_max_budget = litellm.max_budget try: proxy_config = ProxyConfig() - await proxy_config.load_config( - router=MagicMock(), config_file_path=str(config_file) - ) + await proxy_config.load_config(router=MagicMock(), config_file_path=str(config_file)) assert isinstance(litellm.max_budget, float) assert litellm.max_budget == 10.0 assert litellm.max_budget > 0 @@ -2925,9 +2864,7 @@ async def test_load_config_max_ui_session_budget_applied_and_coerced(tmp_path, m original_budget = litellm.max_ui_session_budget try: proxy_config = ProxyConfig() - await proxy_config.load_config( - router=MagicMock(), config_file_path=str(config_file) - ) + await proxy_config.load_config(router=MagicMock(), config_file_path=str(config_file)) assert isinstance(litellm.max_ui_session_budget, float) assert litellm.max_ui_session_budget == 2.5 finally: @@ -2953,9 +2890,7 @@ async def test_load_config_max_ui_session_budget_none_disables_cap(tmp_path): original_budget = litellm.max_ui_session_budget try: proxy_config = ProxyConfig() - await proxy_config.load_config( - router=MagicMock(), config_file_path=str(config_file) - ) + await proxy_config.load_config(router=MagicMock(), config_file_path=str(config_file)) assert litellm.max_ui_session_budget is None finally: litellm.max_ui_session_budget = original_budget @@ -3010,10 +2945,7 @@ async def test_load_config_default_internal_user_params_without_max_budget(tmp_p absent_config_file = tmp_path / "absent_config.yaml" absent_config_file.write_text( - "model_list: []\n" - "litellm_settings:\n" - " default_internal_user_params:\n" - " user_role: internal_user\n" + "model_list: []\nlitellm_settings:\n default_internal_user_params:\n user_role: internal_user\n" ) null_config_file = tmp_path / "null_config.yaml" @@ -3060,9 +2992,7 @@ async def test_load_config_user_url_validation_handles_null_and_string_false(tmp ) ) - await ProxyConfig().load_config( - router=MagicMock(), config_file_path=str(null_config_file) - ) + await ProxyConfig().load_config(router=MagicMock(), config_file_path=str(null_config_file)) assert litellm.user_url_validation is True assert litellm.user_url_allowed_hosts is None assert litellm.provider_url_destination_allowed_hosts is None @@ -3077,9 +3007,7 @@ async def test_load_config_user_url_validation_handles_null_and_string_false(tmp ) ) - await ProxyConfig().load_config( - router=MagicMock(), config_file_path=str(false_config_file) - ) + await ProxyConfig().load_config(router=MagicMock(), config_file_path=str(false_config_file)) assert litellm.user_url_validation is False @@ -3107,12 +3035,8 @@ async def test_load_environment_variables_direct_and_os_environ(): # Mock get_secret_str to return a resolved value mock_secret_value = "resolved_secret_value" - with patch( - "litellm.proxy.proxy_server.get_secret_str", return_value=mock_secret_value - ) as mock_get_secret: - with patch.dict( - os.environ, {}, clear=False - ): # Don't clear existing env vars, just track changes + with patch("litellm.proxy.proxy_server.get_secret_str", return_value=mock_secret_value) as mock_get_secret: + with patch.dict(os.environ, {}, clear=False): # Don't clear existing env vars, just track changes # Call the method under test proxy_config._load_environment_variables(test_config) @@ -3125,9 +3049,7 @@ async def test_load_environment_variables_direct_and_os_environ(): assert os.environ["SECRET_VAR"] == mock_secret_value # Verify get_secret_str was called with the correct value - mock_get_secret.assert_called_once_with( - secret_name="os.environ/ACTUAL_SECRET_VAR" - ) + mock_get_secret.assert_called_once_with(secret_name="os.environ/ACTUAL_SECRET_VAR") @pytest.mark.asyncio @@ -3180,9 +3102,7 @@ async def test_load_environment_variables_litellm_license_and_edge_cases(): assert result is None # Method returns None # Test Case 4: os.environ/ prefix but get_secret_str returns None - test_config_secret_none = { - "environment_variables": {"FAILED_SECRET": "os.environ/NONEXISTENT_SECRET"} - } + test_config_secret_none = {"environment_variables": {"FAILED_SECRET": "os.environ/NONEXISTENT_SECRET"}} with patch("litellm.proxy.proxy_server.get_secret_str", return_value=None): with patch.dict(os.environ, {}, clear=False): @@ -3221,9 +3141,7 @@ async def test_load_environment_variables_blocks_dangerous_keys(): # Blocked keys should not be set to the attacker value assert os.environ.get("PATH") != "/tmp/evil" - assert ( - "LD_PRELOAD" not in os.environ or os.environ["LD_PRELOAD"] != "/tmp/evil.so" - ) + assert "LD_PRELOAD" not in os.environ or os.environ["LD_PRELOAD"] != "/tmp/evil.so" assert os.environ.get("PYTHONPATH") != "/tmp/evil" # Safe keys should still be set @@ -3297,15 +3215,11 @@ async def test_write_config_to_file(monkeypatch): # Mock general_settings mock_general_settings = {"store_model_in_db": True} - monkeypatch.setattr( - "litellm.proxy.proxy_server.general_settings", mock_general_settings - ) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", mock_general_settings) # Mock user_config_file_path test_config_path = "/tmp/test_config.yaml" - monkeypatch.setattr( - "litellm.proxy.proxy_server.user_config_file_path", test_config_path - ) + monkeypatch.setattr("litellm.proxy.proxy_server.user_config_file_path", test_config_path) proxy_config = ProxyConfig() @@ -3326,9 +3240,7 @@ async def test_write_config_to_file(monkeypatch): # Verify the config passed to DB has model_list removed call_args = mock_prisma_client.insert_data.call_args - assert call_args.kwargs["data"] == { - "key": "value" - } # model_list should be popped + assert call_args.kwargs["data"] == {"key": "value"} # model_list should be popped assert call_args.kwargs["table_name"] == "config" @@ -3349,15 +3261,11 @@ async def test_write_config_to_file_when_store_model_in_db_false(monkeypatch): # Mock general_settings mock_general_settings = {"store_model_in_db": False} - monkeypatch.setattr( - "litellm.proxy.proxy_server.general_settings", mock_general_settings - ) + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", mock_general_settings) # Mock user_config_file_path test_config_path = "/tmp/test_config.yaml" - monkeypatch.setattr( - "litellm.proxy.proxy_server.user_config_file_path", test_config_path - ) + monkeypatch.setattr("litellm.proxy.proxy_server.user_config_file_path", test_config_path) proxy_config = ProxyConfig() @@ -3412,22 +3320,20 @@ async def test_async_data_generator_midstream_error(): for chunk in mock_chunks: yield chunk - mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = ( - mock_streaming_iterator - ) + mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = mock_streaming_iterator # Mock async_post_call_streaming_hook to return error on third chunk def mock_streaming_hook(*args, **kwargs): chunk = kwargs.get("response") # Return error message for the third chunk (simulating guardrail trigger) if chunk == mock_chunks[2]: - return 'data: {"error": {"error": "Azure Content Safety Guardrail: Hate crossed severity 2, Got severity: 2"}}' + return ( + 'data: {"error": {"error": "Azure Content Safety Guardrail: Hate crossed severity 2, Got severity: 2"}}' + ) # Return normal chunks for first two return chunk - mock_proxy_logging_obj.async_post_call_streaming_hook = AsyncMock( - side_effect=mock_streaming_hook - ) + mock_proxy_logging_obj.async_post_call_streaming_hook = AsyncMock(side_effect=mock_streaming_hook) mock_proxy_logging_obj.post_call_failure_hook = AsyncMock() # Mock the global proxy_logging_obj @@ -3438,26 +3344,18 @@ async def test_async_data_generator_midstream_error(): # Collect all yielded data from the generator yielded_data = [] try: - async for data in async_data_generator( - mock_response, mock_user_api_key_dict, mock_request_data - ): + async for data in async_data_generator(mock_response, mock_user_api_key_dict, mock_request_data): yielded_data.append(data) except Exception as e: # If there's an exception, that's also part of what we want to test pass # Verify the results - assert ( - len(yielded_data) >= 3 - ), f"Expected at least 3 chunks, got {len(yielded_data)}: {yielded_data}" + assert len(yielded_data) >= 3, f"Expected at least 3 chunks, got {len(yielded_data)}: {yielded_data}" # First two chunks should be normal data - assert yielded_data[0].startswith( - "data: " - ), f"First chunk should start with 'data: ', got: {yielded_data[0]}" - assert yielded_data[1].startswith( - "data: " - ), f"Second chunk should start with 'data: ', got: {yielded_data[1]}" + assert yielded_data[0].startswith("data: "), f"First chunk should start with 'data: ', got: {yielded_data[0]}" + assert yielded_data[1].startswith("data: "), f"Second chunk should start with 'data: ', got: {yielded_data[1]}" # The error message should be yielded error_found = False @@ -3469,15 +3367,11 @@ async def test_async_data_generator_midstream_error(): if "data: [DONE]" in data: done_found = True - assert ( - error_found - ), f"Error message should be found in yielded data. Got: {yielded_data}" + assert error_found, f"Error message should be found in yielded data. Got: {yielded_data}" assert done_found, f"[DONE] message should be found at the end. Got: {yielded_data}" # Verify that the streaming hook was called for each chunk - assert mock_proxy_logging_obj.async_post_call_streaming_hook.call_count == len( - mock_chunks - ) + assert mock_proxy_logging_obj.async_post_call_streaming_hook.call_count == len(mock_chunks) # Verify that post_call_failure_hook was NOT called (since this is not an exception case) mock_proxy_logging_obj.post_call_failure_hook.assert_not_called() @@ -3564,15 +3458,11 @@ async def test_chat_completion_result_no_nested_none_values(): # Verify the mock has None values before serialization raw_dict = mock_model_response.model_dump() none_paths_before = _has_nested_none_values(raw_dict) - assert ( - len(none_paths_before) > 0 - ), "Mock should have None values before exclude_none=True" + assert len(none_paths_before) > 0, "Mock should have None values before exclude_none=True" # Mock the request processing to return our mock response mock_base_processor = MagicMock() - mock_base_processor.base_process_llm_request = AsyncMock( - return_value=mock_model_response - ) + mock_base_processor.base_process_llm_request = AsyncMock(return_value=mock_model_response) # Mock other dependencies mock_request = MagicMock(spec=Request) @@ -3601,9 +3491,9 @@ async def test_chat_completion_result_no_nested_none_values(): # Check that there are no nested None values in the result none_paths_after = _has_nested_none_values(result) - assert ( - len(none_paths_after) == 0 - ), f"Result should not contain nested None values. Found None at: {none_paths_after}" + assert len(none_paths_after) == 0, ( + f"Result should not contain nested None values. Found None at: {none_paths_after}" + ) # Verify essential fields are present assert "id" in result @@ -3629,9 +3519,7 @@ async def test_chat_completion_result_no_nested_none_values(): "annotations", ] for field in excluded_fields: - assert ( - field not in message - ), f"Field '{field}' should be excluded when it's None" + assert field not in message, f"Field '{field}' should be excluded when it's None" # ============================================================================ @@ -3686,9 +3574,7 @@ class TestPriceDataReloadAPI: with patch( "litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new=AsyncMock( - return_value=ModelCostMapReloaded( - model_cost_map={"gpt-3.5-turbo": {"input_cost_per_token": 0.001}} - ) + return_value=ModelCostMapReloaded(model_cost_map={"gpt-3.5-turbo": {"input_cost_per_token": 0.001}}) ), ): # Mock the database connection @@ -3706,10 +3592,7 @@ class TestPriceDataReloadAPI: assert "timestamp" in data assert "models_count" in data # The new implementation immediately reloads and returns the count - assert ( - "Price data reloaded successfully! 1 models updated." - in data["message"] - ) + assert "Price data reloaded successfully! 1 models updated." in data["message"] assert data["models_count"] == 1 finally: # Restore the full model cost map so subsequent tests are not affected @@ -3732,9 +3615,7 @@ class TestPriceDataReloadAPI: def test_get_model_cost_map_public_access(self, client_no_auth): """Test that the model cost map endpoint is publicly accessible""" - with patch( - "litellm.model_cost", {"gpt-3.5-turbo": {"input_cost_per_token": 0.001}} - ): + with patch("litellm.model_cost", {"gpt-3.5-turbo": {"input_cost_per_token": 0.001}}): response = client_no_auth.get("/public/litellm_model_cost_map") assert response.status_code == 200 @@ -3756,9 +3637,7 @@ class TestPriceDataReloadAPI: response = client_with_auth.post("/reload/model_cost_map") - assert ( - response.status_code == 500 - ) # An unexpected exception still maps to 500 + assert response.status_code == 500 # An unexpected exception still maps to 500 data = response.json() assert "Failed to reload model cost map" in data["detail"] @@ -3966,9 +3845,7 @@ class TestPriceDataReloadIntegration: try: with patch( "litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", - new=AsyncMock( - return_value=ModelCostMapReloaded(model_cost_map=mock_cost_map) - ), + new=AsyncMock(return_value=ModelCostMapReloaded(model_cost_map=mock_cost_map)), ): # Mock the database connection with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: @@ -4036,10 +3913,14 @@ class TestPriceDataReloadIntegration: original_model_cost = litellm.model_cost.copy() try: with ( - patch("litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock) as mock_get_map, + patch( + "litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock + ) as mock_get_map, patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), ): - mock_get_map.return_value = ModelCostMapReloaded(model_cost_map={"gpt-3.5-turbo": {"input_cost_per_token": 0.001}}) + mock_get_map.return_value = ModelCostMapReloaded( + model_cost_map={"gpt-3.5-turbo": {"input_cost_per_token": 0.001}} + ) asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) @@ -4080,7 +3961,9 @@ class TestPriceDataReloadIntegration: original_model_cost = litellm.model_cost.copy() try: with ( - patch("litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock) as mock_get_map, + patch( + "litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock + ) as mock_get_map, patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), ): asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) @@ -4115,10 +3998,14 @@ class TestPriceDataReloadIntegration: original_model_cost = litellm.model_cost.copy() try: with ( - patch("litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock) as mock_get_map, + patch( + "litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock + ) as mock_get_map, patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), ): - mock_get_map.return_value = ModelCostMapReloaded(model_cost_map={"gpt-4-test": {"input_cost_per_token": 0.5}}) + mock_get_map.return_value = ModelCostMapReloaded( + model_cost_map={"gpt-4-test": {"input_cost_per_token": 0.5}} + ) asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) @@ -4156,10 +4043,14 @@ class TestPriceDataReloadIntegration: original_model_cost = litellm.model_cost.copy() try: with ( - patch("litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock) as mock_get_map, + patch( + "litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock + ) as mock_get_map, patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), ): - mock_get_map.return_value = ModelCostMapReloaded(model_cost_map={"gpt-4": {"input_cost_per_token": 0.001}}) + mock_get_map.return_value = ModelCostMapReloaded( + model_cost_map={"gpt-4": {"input_cost_per_token": 0.001}} + ) for _ in range(3): for pod in pods: @@ -4196,10 +4087,14 @@ class TestPriceDataReloadIntegration: original_model_cost = litellm.model_cost.copy() try: with ( - patch("litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock) as mock_get_map, + patch( + "litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock + ) as mock_get_map, patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), ): - mock_get_map.return_value = ModelCostMapReloaded(model_cost_map={"gpt-4": {"input_cost_per_token": 0.001}}) + mock_get_map.return_value = ModelCostMapReloaded( + model_cost_map={"gpt-4": {"input_cost_per_token": 0.001}} + ) asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) @@ -4228,10 +4123,14 @@ class TestPriceDataReloadIntegration: original_model_cost = litellm.model_cost.copy() try: with ( - patch("litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock) as mock_get_map, + patch( + "litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock + ) as mock_get_map, patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), ): - mock_get_map.return_value = ModelCostMapReloaded(model_cost_map={"gpt-4": {"input_cost_per_token": 0.001}}) + mock_get_map.return_value = ModelCostMapReloaded( + model_cost_map={"gpt-4": {"input_cost_per_token": 0.001}} + ) asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) @@ -4271,7 +4170,9 @@ class TestPriceDataReloadIntegration: ) as mock_get_map, patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), ): - mock_get_map.return_value = ModelCostMapReloaded(model_cost_map={"gpt-4": {"input_cost_per_token": 0.1}}) + mock_get_map.return_value = ModelCostMapReloaded( + model_cost_map={"gpt-4": {"input_cost_per_token": 0.1}} + ) asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) @@ -4317,8 +4218,7 @@ class TestPriceDataReloadIntegration: asyncio.run(proxy_config._check_and_reload_model_cost_map(mock_prisma)) assert litellm.model_cost is original_model_cost, ( - "a failed reload must keep the currently loaded cost map, " - "not swap in the packaged backup" + "a failed reload must keep the currently loaded cost map, not swap in the packaged backup" ) assert proxy_config.model_cost_map_loaded_at == pod_data_loaded_at, ( "a failed reload must not stamp the pod's data age, otherwise the retry waits a full interval" @@ -4428,11 +4328,15 @@ class TestPriceDataReloadIntegration: original_model_cost = litellm.model_cost.copy() try: with ( - patch("litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock) as mock_get_map, + patch( + "litellm.litellm_core_utils.get_model_cost_map.refetch_model_cost_map", new_callable=AsyncMock + ) as mock_get_map, patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma, patch("litellm.proxy.proxy_server.utc_now", return_value=frozen_now), ): - mock_get_map.return_value = ModelCostMapReloaded(model_cost_map={"gpt-4": {"input_cost_per_token": 0.001}}) + mock_get_map.return_value = ModelCostMapReloaded( + model_cost_map={"gpt-4": {"input_cost_per_token": 0.001}} + ) mock_prisma.db.litellm_config.upsert = AsyncMock( return_value=_reload_schedule_row({}, reload_revision=9) ) @@ -4480,14 +4384,10 @@ class TestPriceDataReloadIntegration: mock_prisma.get_generic_data = AsyncMock(return_value=mock_config) mock_prisma.db.litellm_config.upsert = AsyncMock(return_value=_reload_schedule_row({}, reload_revision=1)) - with patch( - "litellm.anthropic_beta_headers_manager.reload_beta_headers_config" - ) as mock_reload: + with patch("litellm.anthropic_beta_headers_manager.reload_beta_headers_config") as mock_reload: mock_reload.return_value = {"anthropic": {"beta_header": "test-value"}} - asyncio.run( - proxy_config._check_and_reload_anthropic_beta_headers(mock_prisma) - ) + asyncio.run(proxy_config._check_and_reload_anthropic_beta_headers(mock_prisma)) # Verify the upsert update branch preserves interval_hours mock_prisma.db.litellm_config.upsert.assert_called() @@ -4519,9 +4419,7 @@ class TestPriceDataReloadIntegration: app.dependency_overrides[user_api_key_auth] = lambda: mock_auth client = TestClient(app) - with patch( - "litellm.anthropic_beta_headers_manager.reload_beta_headers_config" - ) as mock_reload: + with patch("litellm.anthropic_beta_headers_manager.reload_beta_headers_config") as mock_reload: mock_reload.return_value = {"anthropic": {"beta_header": "test-value"}} with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: @@ -4619,9 +4517,7 @@ async def test_add_router_settings_from_db_config_merge_logic(): # Mock prisma client mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_config.find_first = AsyncMock( - return_value=mock_db_config - ) + mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=mock_db_config) # Call the method under test await proxy_config._add_router_settings_from_db_config( @@ -4631,9 +4527,7 @@ async def test_add_router_settings_from_db_config_merge_logic(): ) # Verify find_first was called with correct parameters - mock_prisma_client.db.litellm_config.find_first.assert_called_once_with( - where={"param_name": "router_settings"} - ) + mock_prisma_client.db.litellm_config.find_first.assert_called_once_with(where={"param_name": "router_settings"}) # Verify update_settings was called mock_router.update_settings.assert_called_once() @@ -4713,9 +4607,7 @@ async def test_add_router_settings_from_db_config_edge_cases(): # Test Case 4: Config has no router_settings mock_db_config = MagicMock() mock_db_config.param_value = {"db_setting": "db_value"} - mock_prisma_client.db.litellm_config.find_first = AsyncMock( - return_value=mock_db_config - ) + mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=mock_db_config) await proxy_config._add_router_settings_from_db_config( config_data={}, # No router_settings in config @@ -4740,9 +4632,7 @@ async def test_add_router_settings_from_db_config_edge_cases(): # Test Case 6: DB config exists but param_value is not a dict mock_db_config_invalid = MagicMock() mock_db_config_invalid.param_value = "not_a_dict" - mock_prisma_client.db.litellm_config.find_first = AsyncMock( - return_value=mock_db_config_invalid - ) + mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=mock_db_config_invalid) config_data = {"router_settings": {"config_setting": "config_value"}} @@ -4794,9 +4684,7 @@ async def test_add_router_settings_shallow_merge_behavior(): } mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_config.find_first = AsyncMock( - return_value=mock_db_config - ) + mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=mock_db_config) await proxy_config._add_router_settings_from_db_config( config_data=config_data, @@ -4873,9 +4761,7 @@ async def test_model_info_v1_oci_secrets_not_leaked(): patch("litellm.proxy.proxy_server.user_model", None), ): # Call the model_info_v1 endpoint - result = await model_info_v1( - user_api_key_dict=mock_user_api_key_dict, litellm_model_id=None - ) + result = await model_info_v1(user_api_key_dict=mock_user_api_key_dict, litellm_model_id=None) # Verify the result structure assert "data" in result @@ -4886,40 +4772,24 @@ async def test_model_info_v1_oci_secrets_not_leaked(): # Verify that sensitive OCI fields are masked assert "****" in litellm_params["oci_key"], "oci_key should be masked" - assert ( - "****" in litellm_params["oci_fingerprint"] - ), "oci_fingerprint should be masked" + assert "****" in litellm_params["oci_fingerprint"], "oci_fingerprint should be masked" assert "****" in litellm_params["oci_tenancy"], "oci_tenancy should be masked" assert "****" in litellm_params["oci_key_file"], "oci_key_file should be masked" # Verify that non-sensitive fields are NOT masked - assert ( - litellm_params["model"] == "oci/xai.grok-4" - ), "model field should not be masked" - assert ( - litellm_params["oci_region"] == "us-phoenix-1" - ), "oci_region should not be masked" + assert litellm_params["model"] == "oci/xai.grok-4", "model field should not be masked" + assert litellm_params["oci_region"] == "us-phoenix-1", "oci_region should not be masked" assert litellm_params["drop_params"] is True, "drop_params should not be masked" # Verify the model field specifically is not masked (this was the original issue) - assert ( - "****" not in litellm_params["model"] - ), "model field should never be masked" - assert litellm_params["model"].startswith( - "oci/" - ), "model should retain its full value" + assert "****" not in litellm_params["model"], "model field should never be masked" + assert litellm_params["model"].startswith("oci/"), "model should retain its full value" # Verify that actual secret values are not present in the response result_str = str(result) - assert ( - "ocid1.api_key.oc1..aaaaaaaa7kbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbk" - not in result_str - ) + assert "ocid1.api_key.oc1..aaaaaaaa7kbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbk" not in result_str assert "aa:bb:cc:dd:ee:ff:11:22:33:44:55:66:77:88:99:00" not in result_str - assert ( - "ocid1.tenancy.oc1..aaaaaaaa7kbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbk" - not in result_str - ) + assert "ocid1.tenancy.oc1..aaaaaaaa7kbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbkbk" not in result_str assert "/path/to/oci_api_key.pem" not in result_str @@ -4949,9 +4819,7 @@ def test_add_callback_from_db_to_in_memory_litellm_callbacks(): event_types=["success"], existing_callbacks=mock_success_callbacks, ) - mock_callback_manager.add_litellm_success_callback.assert_called_once_with( - "prometheus" - ) + mock_callback_manager.add_litellm_success_callback.assert_called_once_with("prometheus") mock_callback_manager.reset_mock() # Test Case 2: Add failure callback @@ -4961,9 +4829,7 @@ def test_add_callback_from_db_to_in_memory_litellm_callbacks(): event_types=["failure"], existing_callbacks=mock_failure_callbacks, ) - mock_callback_manager.add_litellm_failure_callback.assert_called_once_with( - "langfuse" - ) + mock_callback_manager.add_litellm_failure_callback.assert_called_once_with("langfuse") mock_callback_manager.reset_mock() # Test Case 3: Add callback for both success and failure @@ -5064,10 +4930,7 @@ def test_should_load_db_object_with_supported_db_objects(): assert proxy_config._should_load_db_object(object_type="mcp") is True assert proxy_config._should_load_db_object(object_type="guardrails") is True assert proxy_config._should_load_db_object(object_type="vector_stores") is True - assert ( - proxy_config._should_load_db_object(object_type="pass_through_endpoints") - is True - ) + assert proxy_config._should_load_db_object(object_type="pass_through_endpoints") is True assert proxy_config._should_load_db_object(object_type="prompts") is True assert proxy_config._should_load_db_object(object_type="model_cost_map") is True @@ -5093,12 +4956,8 @@ async def test_tag_cache_update_called(): "spend": 10.0, } - with patch.object( - cache, "async_get_cache", new=AsyncMock(return_value=mock_tag_obj) - ) as mock_get_cache: - with patch.object( - cache, "async_set_cache_pipeline", new=AsyncMock() - ) as mock_set_cache: + with patch.object(cache, "async_get_cache", new=AsyncMock(return_value=mock_tag_obj)) as mock_get_cache: + with patch.object(cache, "async_set_cache_pipeline", new=AsyncMock()) as mock_set_cache: await litellm.proxy.proxy_server.update_cache( token=None, user_id=None, @@ -5152,9 +5011,7 @@ async def test_tag_cache_update_multiple_tags(): with patch.object( cache, "async_get_cache", new=AsyncMock(side_effect=mock_get_cache_side_effect) ) as mock_get_cache: - with patch.object( - cache, "async_set_cache_pipeline", new=AsyncMock() - ) as mock_set_cache: + with patch.object(cache, "async_set_cache_pipeline", new=AsyncMock()) as mock_set_cache: await litellm.proxy.proxy_server.update_cache( token=None, user_id=None, @@ -5175,9 +5032,7 @@ async def test_tag_cache_update_multiple_tags(): assert len(cache_list) == 2 - tag_updates = { - cache_key: cache_value for cache_key, cache_value in cache_list - } + tag_updates = {cache_key: cache_value for cache_key, cache_value in cache_list} assert "tag:tag1" in tag_updates assert "tag:tag2" in tag_updates assert tag_updates["tag:tag1"]["spend"] == 15.0 @@ -5203,9 +5058,7 @@ async def test_update_cache_pipeline_honors_user_api_key_cache_ttl(): "async_get_cache", new=AsyncMock(return_value={"tag_name": "active-tag", "spend": 1.0}), ): - with patch.object( - cache, "async_set_cache_pipeline", new=AsyncMock() - ) as mock_set_cache: + with patch.object(cache, "async_set_cache_pipeline", new=AsyncMock()) as mock_set_cache: await litellm.proxy.proxy_server.update_cache( token=None, user_id=None, @@ -5248,9 +5101,7 @@ async def test_spend_tracking_never_writes_the_auth_object_back(): model_type=UserAPIKeyAuth, ) with ( - patch.object( - cache, "async_set_cache_pipeline", new=AsyncMock() - ) as mock_pipeline, + patch.object(cache, "async_set_cache_pipeline", new=AsyncMock()) as mock_pipeline, patch.object(cache, "async_set_cache", new=AsyncMock()) as mock_set, ): await litellm.proxy.proxy_server.update_cache( @@ -5261,9 +5112,7 @@ async def test_spend_tracking_never_writes_the_auth_object_back(): response_cost=5.0, parent_otel_span=None, ) - pending = [ - t for t in asyncio.all_tasks() if t is not asyncio.current_task() - ] + pending = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()] if pending: await asyncio.wait(pending, timeout=5) @@ -5305,12 +5154,8 @@ async def test_update_cache_global_proxy_spend_scalar_stays_shared(): cache = DualCache(default_in_memory_ttl=300) setattr(litellm.proxy.proxy_server, "user_api_key_cache", cache) try: - with patch.object( - cache, "async_get_cache", new=AsyncMock(side_effect=fake_get) - ): - with patch.object( - cache, "async_set_cache_pipeline", new=AsyncMock() - ) as mock_set_cache: + with patch.object(cache, "async_get_cache", new=AsyncMock(side_effect=fake_get)): + with patch.object(cache, "async_set_cache_pipeline", new=AsyncMock()) as mock_set_cache: await litellm.proxy.proxy_server.update_cache( token=None, user_id="user-lit", @@ -5320,24 +5165,14 @@ async def test_update_cache_global_proxy_spend_scalar_stays_shared(): parent_otel_span=None, ) - pending = [ - t for t in asyncio.all_tasks() if t is not asyncio.current_task() - ] + pending = [t for t in asyncio.all_tasks() if t is not asyncio.current_task()] if pending: await asyncio.wait(pending, timeout=5) calls = mock_set_cache.await_args_list - local_keys = [ - k - for c in calls - if c.kwargs.get("local_only") is True - for k, _ in c.kwargs["cache_list"] - ] + local_keys = [k for c in calls if c.kwargs.get("local_only") is True for k, _ in c.kwargs["cache_list"]] shared_keys = [ - k - for c in calls - if c.kwargs.get("local_only") is not True - for k, _ in c.kwargs["cache_list"] + k for c in calls if c.kwargs.get("local_only") is not True for k, _ in c.kwargs["cache_list"] ] assert "user-lit" in local_keys assert global_key not in local_keys @@ -5368,20 +5203,14 @@ async def test_init_sso_settings_in_db(): } mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock( - return_value=mock_sso_config - ) + mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=mock_sso_config) # Mock _decrypt_and_set_db_env_variables - with patch.object( - proxy_config, "_decrypt_and_set_db_env_variables" - ) as mock_decrypt_and_set: + with patch.object(proxy_config, "_decrypt_and_set_db_env_variables") as mock_decrypt_and_set: await proxy_config._init_sso_settings_in_db(prisma_client=mock_prisma_client) # Verify find_unique was called with correct parameters - mock_prisma_client.db.litellm_ssoconfig.find_unique.assert_awaited_once_with( - where={"id": "sso_config"} - ) + mock_prisma_client.db.litellm_ssoconfig.find_unique.assert_awaited_once_with(where={"id": "sso_config"}) # Verify _decrypt_and_set_db_env_variables was called with uppercased keys mock_decrypt_and_set.assert_called_once() @@ -5421,15 +5250,11 @@ async def test_init_sso_settings_in_db_no_settings(): mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=None) # Mock _decrypt_and_set_db_env_variables - with patch.object( - proxy_config, "_decrypt_and_set_db_env_variables" - ) as mock_decrypt_and_set: + with patch.object(proxy_config, "_decrypt_and_set_db_env_variables") as mock_decrypt_and_set: await proxy_config._init_sso_settings_in_db(prisma_client=mock_prisma_client) # Verify find_unique was called - mock_prisma_client.db.litellm_ssoconfig.find_unique.assert_awaited_once_with( - where={"id": "sso_config"} - ) + mock_prisma_client.db.litellm_ssoconfig.find_unique.assert_awaited_once_with(where={"id": "sso_config"}) # Verify _decrypt_and_set_db_env_variables was NOT called when no settings exist mock_decrypt_and_set.assert_not_called() @@ -5448,9 +5273,7 @@ async def test_init_sso_settings_in_db_error_handling(): # Mock prisma client to raise an exception mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock( - side_effect=Exception("Database connection error") - ) + mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock(side_effect=Exception("Database connection error")) # The method should not raise an exception, it should log it instead try: @@ -5459,9 +5282,7 @@ async def test_init_sso_settings_in_db_error_handling(): assert True except Exception as e: # The exception should be caught and logged, not propagated - pytest.fail( - f"Exception should have been caught and logged, but was raised: {e}" - ) + pytest.fail(f"Exception should have been caught and logged, but was raised: {e}") @pytest.mark.asyncio @@ -5480,20 +5301,14 @@ async def test_init_sso_settings_in_db_empty_settings(): mock_sso_config.sso_settings = {} mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock( - return_value=mock_sso_config - ) + mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock(return_value=mock_sso_config) # Mock _decrypt_and_set_db_env_variables - with patch.object( - proxy_config, "_decrypt_and_set_db_env_variables" - ) as mock_decrypt_and_set: + with patch.object(proxy_config, "_decrypt_and_set_db_env_variables") as mock_decrypt_and_set: await proxy_config._init_sso_settings_in_db(prisma_client=mock_prisma_client) # Verify find_unique was called - mock_prisma_client.db.litellm_ssoconfig.find_unique.assert_awaited_once_with( - where={"id": "sso_config"} - ) + mock_prisma_client.db.litellm_ssoconfig.find_unique.assert_awaited_once_with(where={"id": "sso_config"}) # Verify _decrypt_and_set_db_env_variables was called with empty dict mock_decrypt_and_set.assert_called_once() @@ -5526,16 +5341,12 @@ async def test_init_sso_settings_in_db_retries_on_transport_error(): return mock_sso_config mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock( - side_effect=_flaky_find_unique - ) + mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock(side_effect=_flaky_find_unique) mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 - with patch.object( - proxy_config, "_decrypt_and_set_db_env_variables" - ) as mock_decrypt: + with patch.object(proxy_config, "_decrypt_and_set_db_env_variables") as mock_decrypt: await proxy_config._init_sso_settings_in_db(prisma_client=mock_prisma_client) assert len(invocations) == 2 @@ -5556,9 +5367,7 @@ async def test_init_sso_settings_in_db_propagates_when_reconnect_fails(): proxy_config = ProxyConfig() mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock( - side_effect=prisma.errors.ClientNotConnectedError() - ) + mock_prisma_client.db.litellm_ssoconfig.find_unique = AsyncMock(side_effect=prisma.errors.ClientNotConnectedError()) mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=False) mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 @@ -5589,24 +5398,17 @@ async def test_init_hashicorp_vault_config_override_retries_on_transport_error() return None # No config in DB → function returns early after retry. mock_prisma_client = MagicMock() - mock_prisma_client.db.litellm_configoverrides.find_unique = AsyncMock( - side_effect=_flaky_find_unique - ) + mock_prisma_client.db.litellm_configoverrides.find_unique = AsyncMock(side_effect=_flaky_find_unique) mock_prisma_client.attempt_db_reconnect = AsyncMock(return_value=True) mock_prisma_client._db_auth_reconnect_timeout_seconds = 2.0 mock_prisma_client._db_auth_reconnect_lock_timeout_seconds = 0.1 - await proxy_config._init_hashicorp_vault_config_override( - prisma_client=mock_prisma_client - ) + await proxy_config._init_hashicorp_vault_config_override(prisma_client=mock_prisma_client) assert len(invocations) == 2 mock_prisma_client.attempt_db_reconnect.assert_awaited_once() reconnect_kwargs = mock_prisma_client.attempt_db_reconnect.await_args.kwargs - assert ( - reconnect_kwargs["reason"] - == "init_hashicorp_vault_config_override_lookup_failure" - ) + assert reconnect_kwargs["reason"] == "init_hashicorp_vault_config_override_lookup_failure" def test_update_config_fields_uppercases_env_vars(monkeypatch): @@ -5656,37 +5458,20 @@ def test_encrypt_env_variables_for_db_is_idempotent(monkeypatch): plaintext = "pk-langfuse-secret-value" # First write: plaintext in -> single-encrypted out. - enc1 = proxy_config._encrypt_env_variables_for_db( - {"LANGFUSE_PUBLIC_KEY": plaintext} - ) + enc1 = proxy_config._encrypt_env_variables_for_db({"LANGFUSE_PUBLIC_KEY": plaintext}) assert enc1["LANGFUSE_PUBLIC_KEY"] != plaintext - assert ( - decrypt_value_helper( - value=enc1["LANGFUSE_PUBLIC_KEY"], key="LANGFUSE_PUBLIC_KEY" - ) - == plaintext - ) + assert decrypt_value_helper(value=enc1["LANGFUSE_PUBLIC_KEY"], key="LANGFUSE_PUBLIC_KEY") == plaintext # UI round-trip: feed the ciphertext back in. Must NOT double-encrypt. enc2 = proxy_config._encrypt_env_variables_for_db(enc1) - assert ( - decrypt_value_helper( - value=enc2["LANGFUSE_PUBLIC_KEY"], key="LANGFUSE_PUBLIC_KEY" - ) - == plaintext - ) + assert decrypt_value_helper(value=enc2["LANGFUSE_PUBLIC_KEY"], key="LANGFUSE_PUBLIC_KEY") == plaintext # And again, ×3 total ciphertext re-feeds — still exactly one layer, # never stacked, no matter how many times the UI re-saves. enc3 = proxy_config._encrypt_env_variables_for_db(enc2) enc4 = proxy_config._encrypt_env_variables_for_db(enc3) for stacked in (enc3, enc4): - assert ( - decrypt_value_helper( - value=stacked["LANGFUSE_PUBLIC_KEY"], key="LANGFUSE_PUBLIC_KEY" - ) - == plaintext - ) + assert decrypt_value_helper(value=stacked["LANGFUSE_PUBLIC_KEY"], key="LANGFUSE_PUBLIC_KEY") == plaintext # Write path must not leak the value into the process environment. assert os.environ.get("LANGFUSE_PUBLIC_KEY") is None @@ -5728,15 +5513,11 @@ def test_get_prompt_spec_for_db_prompt_with_versions(): } # Test version 1 - prompt_spec_v1 = proxy_config._get_prompt_spec_for_db_prompt( - db_prompt=mock_prompt_v1 - ) + prompt_spec_v1 = proxy_config._get_prompt_spec_for_db_prompt(db_prompt=mock_prompt_v1) assert prompt_spec_v1.prompt_id == "chat_prompt.v1" # Test version 2 - prompt_spec_v2 = proxy_config._get_prompt_spec_for_db_prompt( - db_prompt=mock_prompt_v2 - ) + prompt_spec_v2 = proxy_config._get_prompt_spec_for_db_prompt(db_prompt=mock_prompt_v2) assert prompt_spec_v2.prompt_id == "chat_prompt.v2" @@ -5804,9 +5585,7 @@ async def test_get_image_non_root_uses_var_lib_assets_dir(monkeypatch): with ( patch("litellm.proxy.proxy_server.os.makedirs") as mock_makedirs, - patch( - "litellm.proxy.proxy_server.os.path.exists", side_effect=exists_side_effect - ), + patch("litellm.proxy.proxy_server.os.path.exists", side_effect=exists_side_effect), patch("litellm.proxy.proxy_server.os.access", return_value=True), patch("litellm.proxy.proxy_server.os.getenv") as mock_getenv, patch("litellm.proxy.proxy_server.FileResponse") as mock_file_response, @@ -5856,9 +5635,7 @@ async def test_get_image_non_root_fallback_to_default_logo(monkeypatch): # Mock os.path operations with ( patch("litellm.proxy.proxy_server.os.makedirs") as mock_makedirs, - patch( - "litellm.proxy.proxy_server.os.path.exists", side_effect=exists_side_effect - ), + patch("litellm.proxy.proxy_server.os.path.exists", side_effect=exists_side_effect), patch("litellm.proxy.proxy_server.os.access", return_value=True), patch("litellm.proxy.proxy_server.os.getenv") as mock_getenv, patch("litellm.proxy.proxy_server.FileResponse") as mock_file_response, @@ -5881,9 +5658,7 @@ async def test_get_image_non_root_fallback_to_default_logo(monkeypatch): # Verify that exists was called to check /var/lib/litellm/assets/logo.jpg assets_logo_path = "/var/lib/litellm/assets/logo.jpg" - assert any( - assets_logo_path in str(call) for call in exists_calls - ), f"Should check if {assets_logo_path} exists" + assert any(assets_logo_path in str(call) for call in exists_calls), f"Should check if {assets_logo_path} exists" # Verify FileResponse was called (with fallback logo) assert mock_file_response.called, "FileResponse should be called" @@ -5923,14 +5698,8 @@ async def test_get_image_root_case_uses_current_dir(monkeypatch): await get_image() # Verify makedirs was NOT called with /var/lib/litellm/assets (should not create it for root case) - var_lib_assets_calls = [ - call - for call in mock_makedirs.call_args_list - if "/var/lib/litellm/assets" in str(call) - ] - assert ( - len(var_lib_assets_calls) == 0 - ), "Should not create /var/lib/litellm/assets for root case" + var_lib_assets_calls = [call for call in mock_makedirs.call_args_list if "/var/lib/litellm/assets" in str(call)] + assert len(var_lib_assets_calls) == 0, "Should not create /var/lib/litellm/assets for root case" # Verify FileResponse was called assert mock_file_response.called, "FileResponse should be called" @@ -5961,15 +5730,11 @@ async def test_get_image_custom_local_logo_bypasses_cache(monkeypatch, tmp_path) return MagicMock() with ( - patch( - "litellm.proxy.proxy_server.FileResponse", side_effect=fake_file_response - ), + patch("litellm.proxy.proxy_server.FileResponse", side_effect=fake_file_response), ): await get_image() - assert ( - len(calls_to_file_response) == 1 - ), "FileResponse should be called exactly once" + assert len(calls_to_file_response) == 1, "FileResponse should be called exactly once" assert calls_to_file_response[0] == str(custom_logo.resolve()), ( f"Expected custom logo path, got {calls_to_file_response[0]}. " "A stale cached_logo.jpg may have been returned instead." @@ -5999,24 +5764,18 @@ async def test_get_image_default_logo_ignores_stale_cache(monkeypatch, tmp_path) return MagicMock() with ( - patch( - "litellm.proxy.proxy_server.FileResponse", side_effect=fake_file_response - ), + patch("litellm.proxy.proxy_server.FileResponse", side_effect=fake_file_response), ): await get_image() - assert ( - len(calls_to_file_response) == 1 - ), "FileResponse should be called exactly once" + assert len(calls_to_file_response) == 1, "FileResponse should be called exactly once" served_path = calls_to_file_response[0] assert served_path != str(cache_path.resolve()) assert served_path.endswith("logo.jpg") @pytest.mark.asyncio -async def test_get_image_custom_logo_missing_falls_through_to_default( - monkeypatch, tmp_path -): +async def test_get_image_custom_logo_missing_falls_through_to_default(monkeypatch, tmp_path): """ Test that when UI_LOGO_PATH points to a non-existent local file, get_image falls through to the default logo instead of failing. @@ -6037,26 +5796,18 @@ async def test_get_image_custom_logo_missing_falls_through_to_default( return MagicMock() with ( - patch( - "litellm.proxy.proxy_server.FileResponse", side_effect=fake_file_response - ), + patch("litellm.proxy.proxy_server.FileResponse", side_effect=fake_file_response), ): await get_image() - assert ( - len(calls_to_file_response) == 1 - ), "FileResponse should be called exactly once" + assert len(calls_to_file_response) == 1, "FileResponse should be called exactly once" served_path = calls_to_file_response[0] - assert served_path != str( - custom_logo_path - ), "Should not attempt to serve a non-existent custom logo" + assert served_path != str(custom_logo_path), "Should not attempt to serve a non-existent custom logo" assert served_path.endswith("logo.jpg") @pytest.mark.asyncio -async def test_get_image_custom_logo_missing_no_cache_serves_default( - monkeypatch, tmp_path -): +async def test_get_image_custom_logo_missing_no_cache_serves_default(monkeypatch, tmp_path): """ Test that when UI_LOGO_PATH points to a non-existent file AND there is no cached_logo.jpg, get_image serves the default logo instead of the non-existent @@ -6078,22 +5829,14 @@ async def test_get_image_custom_logo_missing_no_cache_serves_default( return MagicMock() with ( - patch( - "litellm.proxy.proxy_server.FileResponse", side_effect=fake_file_response - ), + patch("litellm.proxy.proxy_server.FileResponse", side_effect=fake_file_response), ): await get_image() - assert ( - len(calls_to_file_response) == 1 - ), "FileResponse should be called exactly once" + assert len(calls_to_file_response) == 1, "FileResponse should be called exactly once" served_path = calls_to_file_response[0] - assert served_path != str( - custom_logo_path - ), "Should not attempt to serve a non-existent custom logo" - assert served_path.endswith( - "logo.jpg" - ), f"Expected fallback to default logo.jpg, got {served_path}" + assert served_path != str(custom_logo_path), "Should not attempt to serve a non-existent custom logo" + assert served_path.endswith("logo.jpg"), f"Expected fallback to default logo.jpg, got {served_path}" def test_get_config_normalizes_string_callbacks(monkeypatch): @@ -6133,9 +5876,7 @@ def test_get_config_normalizes_string_callbacks(monkeypatch): success_callbacks = [cb["name"] for cb in callbacks if cb.get("type") == "success"] failure_callbacks = [cb["name"] for cb in callbacks if cb.get("type") == "failure"] - success_and_failure_callbacks = [ - cb["name"] for cb in callbacks if cb.get("type") == "success_and_failure" - ] + success_and_failure_callbacks = [cb["name"] for cb in callbacks if cb.get("type") == "success_and_failure"] assert "langfuse" in success_callbacks assert len(failure_callbacks) == 0 @@ -6172,9 +5913,7 @@ def test_deep_merge_dicts_skips_none_and_empty_lists(monkeypatch): }, } - result = proxy_config._update_config_fields( - current_config, "general_settings", db_param_value - ) + result = proxy_config._update_config_fields(current_config, "general_settings", db_param_value) assert result["general_settings"]["max_parallel_requests"] == 10 assert result["general_settings"]["allowed_models"] == ["gpt-3.5-turbo", "gpt-4"] @@ -6241,9 +5980,7 @@ class TestInvitationEndpoints: ), ], ) - def test_invitation_endpoints_proxy_admin_success( - self, client_with_auth, endpoint, payload, mock_return - ): + def test_invitation_endpoints_proxy_admin_success(self, client_with_auth, endpoint, payload, mock_return): """Proxy admin can successfully create and delete invitations.""" with patch("litellm.proxy.proxy_server.prisma_client") as mock_prisma: mock_prisma.db.litellm_invitationlink = MagicMock() @@ -6258,9 +5995,7 @@ class TestInvitationEndpoints: mock_prisma.db.litellm_invitationlink.find_unique = AsyncMock( return_value={**mock_return, "created_by": "admin-user-id"} ) - mock_prisma.db.litellm_invitationlink.delete = AsyncMock( - return_value=mock_return - ) + mock_prisma.db.litellm_invitationlink.delete = AsyncMock(return_value=mock_return) response = client_with_auth.post(endpoint, json=payload) assert response.status_code == 200 @@ -6275,9 +6010,7 @@ class TestInvitationEndpoints: ("/invitation/delete", {"invitation_id": "inv-456"}), ], ) - def test_invitation_endpoints_non_admin_denied( - self, client_with_auth, endpoint, payload - ): + def test_invitation_endpoints_non_admin_denied(self, client_with_auth, endpoint, payload): """Non-admin users cannot access invitation endpoints.""" from litellm.proxy._types import LitellmUserRoles @@ -6332,9 +6065,7 @@ async def test_async_data_generator_cleanup_on_early_exit(): for chunk in mock_chunks: yield chunk - mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = ( - mock_streaming_iterator - ) + mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = mock_streaming_iterator mock_proxy_logging_obj.async_post_call_streaming_hook = AsyncMock( side_effect=lambda **kwargs: kwargs.get("response") ) @@ -6346,9 +6077,7 @@ async def test_async_data_generator_cleanup_on_early_exit(): with patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj): # Consume only the first chunk then abandon the generator (simulates client disconnect) - gen = async_data_generator( - mock_response, mock_user_api_key_dict, mock_request_data - ) + gen = async_data_generator(mock_response, mock_user_api_key_dict, mock_request_data) first_chunk = await gen.__anext__() assert first_chunk.startswith("data: ") @@ -6401,19 +6130,12 @@ async def test_async_data_generator_uses_direct_stream_fast_path_without_callbac mock_proxy_logging_obj.post_call_failure_hook = AsyncMock() with patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj): - with patch.object( - ProxyLogging, "_fire_deferred_stream_logging" - ) as mock_deferred_logging: + with patch.object(ProxyLogging, "_fire_deferred_stream_logging") as mock_deferred_logging: yielded_data = [] - async for data in async_data_generator( - mock_response, mock_user_api_key_dict, mock_request_data - ): + async for data in async_data_generator(mock_response, mock_user_api_key_dict, mock_request_data): yielded_data.append(data) - yielded_text = [ - chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk - for chunk in yielded_data - ] + yielded_text = [chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk for chunk in yielded_data] assert len([chunk for chunk in yielded_text if chunk.startswith("data: {")]) == 2 assert yielded_text[-1] == "data: [DONE]\n\n" mock_proxy_logging_obj.async_post_call_streaming_iterator_hook.assert_not_called() @@ -6466,18 +6188,13 @@ async def test_async_data_generator_preserves_non_raw_sse_like_bytes(): with patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj): with patch.object(ProxyLogging, "_fire_deferred_stream_logging"): yielded_data = [] - async for data in async_data_generator( - mock_response, mock_user_api_key_dict, mock_request_data - ): + async for data in async_data_generator(mock_response, mock_user_api_key_dict, mock_request_data): yielded_data.append(data) - yielded_text = [ - chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk - for chunk in yielded_data - ] + yielded_text = [chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk for chunk in yielded_data] assert yielded_text[0] == gemini_event.decode("utf-8") assert yielded_text[1] == gemini_event_without_terminator.decode("utf-8") + "\n\n" - assert yielded_text[2] == f'data: {raw_payload.decode("utf-8")}\n\n' + assert yielded_text[2] == f"data: {raw_payload.decode('utf-8')}\n\n" assert "b'data:" not in "".join(yielded_text) assert yielded_text[-1] == "data: [DONE]\n\n" @@ -6500,12 +6217,8 @@ async def test_async_data_generator_buffers_split_google_native_sse_json_frame() ) raw_chunks = [ payload[:2].encode("utf-8"), - payload[ - 2 : payload.index("thoughtSignature") + len('thoughtSignature": "abc') - ].encode("utf-8"), - payload[ - payload.index("thoughtSignature") + len('thoughtSignature": "abc') : - ].encode("utf-8"), + payload[2 : payload.index("thoughtSignature") + len('thoughtSignature": "abc')].encode("utf-8"), + payload[payload.index("thoughtSignature") + len('thoughtSignature": "abc') :].encode("utf-8"), ] class MockStream: @@ -6532,15 +6245,10 @@ async def test_async_data_generator_buffers_split_google_native_sse_json_frame() with patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj): with patch.object(ProxyLogging, "_fire_deferred_stream_logging"): yielded_data = [] - async for data in async_data_generator( - mock_response, mock_user_api_key_dict, mock_request_data - ): + async for data in async_data_generator(mock_response, mock_user_api_key_dict, mock_request_data): yielded_data.append(data) - yielded_text = [ - chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk - for chunk in yielded_data - ] + yielded_text = [chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk for chunk in yielded_data] assert yielded_text == [payload] for chunk in yielded_text: @@ -6586,15 +6294,10 @@ async def test_async_data_generator_flushes_raw_sse_stream_without_trailing_deli patch.object(ProxyLogging, "_fire_deferred_stream_logging"), ): yielded_data = [] - async for data in async_data_generator( - mock_response, mock_user_api_key_dict, mock_request_data - ): + async for data in async_data_generator(mock_response, mock_user_api_key_dict, mock_request_data): yielded_data.append(data) - yielded_text = [ - chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk - for chunk in yielded_data - ] + yielded_text = [chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk for chunk in yielded_data] assert len(yielded_text) == 1 assert yielded_text[0] == 'data: {"candidates": [{"content": "unterminated"}]\n\n' assert "[DONE]" not in yielded_text[0] @@ -6641,15 +6344,10 @@ async def test_async_data_generator_errors_when_raw_sse_frame_exceeds_buffer_lim patch.object(ProxyLogging, "_fire_deferred_stream_logging"), ): yielded_data = [] - async for data in async_data_generator( - mock_response, mock_user_api_key_dict, mock_request_data - ): + async for data in async_data_generator(mock_response, mock_user_api_key_dict, mock_request_data): yielded_data.append(data) - yielded_text = [ - chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk - for chunk in yielded_data - ] + yielded_text = [chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk for chunk in yielded_data] assert len(yielded_text) == 1 assert "maximum buffered size" in yielded_text[0] assert "[DONE]" not in yielded_text[0] @@ -6702,15 +6400,10 @@ async def test_async_data_generator_checks_raw_sse_buffer_limit_after_complete_f patch.object(ProxyLogging, "_fire_deferred_stream_logging"), ): yielded_data = [] - async for data in async_data_generator( - mock_response, mock_user_api_key_dict, mock_request_data - ): + async for data in async_data_generator(mock_response, mock_user_api_key_dict, mock_request_data): yielded_data.append(data) - yielded_text = [ - chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk - for chunk in yielded_data - ] + yielded_text = [chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk for chunk in yielded_data] assert yielded_text[0] == complete_frame assert yielded_text[1] == partial_frame + "\n\n" assert "[DONE]" not in "".join(yielded_text) @@ -6731,9 +6424,7 @@ async def test_async_data_generator_google_genai_stream_omits_openai_done(): "model": "gemini-2.0-flash", "_litellm_skip_openai_stream_done": True, } - gemini_event = ( - b'data: {"candidates": [{"content": {"parts": [{"text": "Hi"}]}}]}\n\n' - ) + gemini_event = b'data: {"candidates": [{"content": {"parts": [{"text": "Hi"}]}}]}\n\n' class MockStream: def __aiter__(self): @@ -6758,15 +6449,10 @@ async def test_async_data_generator_google_genai_stream_omits_openai_done(): with patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj): with patch.object(ProxyLogging, "_fire_deferred_stream_logging"): yielded_data = [] - async for data in async_data_generator( - mock_response, mock_user_api_key_dict, mock_request_data - ): + async for data in async_data_generator(mock_response, mock_user_api_key_dict, mock_request_data): yielded_data.append(data) - yielded_text = [ - chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk - for chunk in yielded_data - ] + yielded_text = [chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk for chunk in yielded_data] assert yielded_text == [gemini_event.decode("utf-8")] assert "[DONE]" not in "".join(yielded_text) @@ -6855,15 +6541,10 @@ async def test_async_data_generator_google_genai_stream_forwards_error_without_d with patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj): with patch.object(ProxyLogging, "_fire_deferred_stream_logging"): yielded_data = [] - async for data in async_data_generator( - mock_response, mock_user_api_key_dict, mock_request_data - ): + async for data in async_data_generator(mock_response, mock_user_api_key_dict, mock_request_data): yielded_data.append(data) - yielded_text = [ - chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk - for chunk in yielded_data - ] + yielded_text = [chunk.decode("utf-8") if isinstance(chunk, bytes) else chunk for chunk in yielded_data] assert yielded_text == [error_sse] assert "[DONE]" not in "".join(yielded_text) @@ -6893,9 +6574,7 @@ async def test_async_data_generator_cleanup_on_normal_completion(): for chunk in mock_chunks: yield chunk - mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = ( - mock_streaming_iterator - ) + mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = mock_streaming_iterator mock_proxy_logging_obj.async_post_call_streaming_hook = AsyncMock( side_effect=lambda **kwargs: kwargs.get("response") ) @@ -6906,9 +6585,7 @@ async def test_async_data_generator_cleanup_on_normal_completion(): with patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj): yielded_data = [] - async for data in async_data_generator( - mock_response, mock_user_api_key_dict, mock_request_data - ): + async for data in async_data_generator(mock_response, mock_user_api_key_dict, mock_request_data): yielded_data.append(data) # Should have completed normally with [DONE] @@ -6939,9 +6616,7 @@ async def test_async_data_generator_cleanup_on_midstream_error(): yield {"choices": [{"delta": {"content": "Hello"}}]} raise RuntimeError("upstream connection reset") - mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = ( - mock_streaming_iterator_with_error - ) + mock_proxy_logging_obj.async_post_call_streaming_iterator_hook = mock_streaming_iterator_with_error mock_proxy_logging_obj.async_post_call_streaming_hook = AsyncMock( side_effect=lambda **kwargs: kwargs.get("response") ) @@ -6952,9 +6627,7 @@ async def test_async_data_generator_cleanup_on_midstream_error(): with patch("litellm.proxy.proxy_server.proxy_logging_obj", mock_proxy_logging_obj): yielded_data = [] - async for data in async_data_generator( - mock_response, mock_user_api_key_dict, mock_request_data - ): + async for data in async_data_generator(mock_response, mock_user_api_key_dict, mock_request_data): yielded_data.append(data) # Should have yielded data chunk and then an error chunk @@ -7009,9 +6682,7 @@ async def test_update_general_settings_store_model_in_db_true(): patch("litellm.proxy.proxy_server.store_model_in_db", False) as mock_store, patch("litellm.proxy.proxy_server.general_settings", {}) as mock_gs, ): - await proxy_config._update_general_settings( - db_general_settings={"store_model_in_db": True} - ) + await proxy_config._update_general_settings(db_general_settings={"store_model_in_db": True}) import litellm.proxy.proxy_server as ps @@ -7033,9 +6704,7 @@ async def test_update_general_settings_store_model_in_db_false(): patch("litellm.proxy.proxy_server.store_model_in_db", True), patch("litellm.proxy.proxy_server.general_settings", {}), ): - await proxy_config._update_general_settings( - db_general_settings={"store_model_in_db": False} - ) + await proxy_config._update_general_settings(db_general_settings={"store_model_in_db": False}) import litellm.proxy.proxy_server as ps @@ -7116,9 +6785,7 @@ async def test_update_general_settings_store_model_in_db_string_normalization(): patch("litellm.proxy.proxy_server.store_model_in_db", False), patch("litellm.proxy.proxy_server.general_settings", {}), ): - await proxy_config._update_general_settings( - db_general_settings={"store_model_in_db": "true"} - ) + await proxy_config._update_general_settings(db_general_settings={"store_model_in_db": "true"}) import litellm.proxy.proxy_server as ps assert ps.store_model_in_db is True @@ -7128,9 +6795,7 @@ async def test_update_general_settings_store_model_in_db_string_normalization(): patch("litellm.proxy.proxy_server.store_model_in_db", False), patch("litellm.proxy.proxy_server.general_settings", {}), ): - await proxy_config._update_general_settings( - db_general_settings={"store_model_in_db": "True"} - ) + await proxy_config._update_general_settings(db_general_settings={"store_model_in_db": "True"}) import litellm.proxy.proxy_server as ps assert ps.store_model_in_db is True @@ -7140,9 +6805,7 @@ async def test_update_general_settings_store_model_in_db_string_normalization(): patch("litellm.proxy.proxy_server.store_model_in_db", True), patch("litellm.proxy.proxy_server.general_settings", {}), ): - await proxy_config._update_general_settings( - db_general_settings={"store_model_in_db": "false"} - ) + await proxy_config._update_general_settings(db_general_settings={"store_model_in_db": "false"}) import litellm.proxy.proxy_server as ps assert ps.store_model_in_db is False @@ -7163,9 +6826,7 @@ async def test_update_general_settings_store_model_in_db_none_keeps_current(): patch("litellm.proxy.proxy_server.store_model_in_db", True), patch("litellm.proxy.proxy_server.general_settings", {}), ): - await proxy_config._update_general_settings( - db_general_settings={"store_model_in_db": None} - ) + await proxy_config._update_general_settings(db_general_settings={"store_model_in_db": None}) import litellm.proxy.proxy_server as ps assert ps.store_model_in_db is True @@ -7175,9 +6836,7 @@ async def test_update_general_settings_store_model_in_db_none_keeps_current(): patch("litellm.proxy.proxy_server.store_model_in_db", False), patch("litellm.proxy.proxy_server.general_settings", {}), ): - await proxy_config._update_general_settings( - db_general_settings={"store_model_in_db": None} - ) + await proxy_config._update_general_settings(db_general_settings={"store_model_in_db": None}) import litellm.proxy.proxy_server as ps assert ps.store_model_in_db is False @@ -7197,12 +6856,11 @@ async def test_store_model_in_db_db_override_when_config_false(): # Mock DB returning store_model_in_db=True in general_settings mock_db_record = MagicMock() mock_db_record.param_value = {"store_model_in_db": True} - mock_prisma_client.db.litellm_config.find_first = AsyncMock( - return_value=mock_db_record - ) + mock_prisma_client.db.litellm_config.find_first = AsyncMock(return_value=mock_db_record) mock_proxy_logging = MagicMock(spec=ProxyLogging) mock_proxy_logging.slack_alerting_instance = MagicMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() mock_proxy_config = AsyncMock() with ( @@ -7245,6 +6903,7 @@ async def test_store_model_in_db_db_check_skipped_when_already_true(monkeypatch) mock_proxy_logging = MagicMock(spec=ProxyLogging) mock_proxy_logging.slack_alerting_instance = MagicMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() mock_proxy_config = AsyncMock() with ( @@ -7283,12 +6942,11 @@ async def test_store_model_in_db_db_failure_graceful(monkeypatch): mock_prisma_client = MagicMock() # Simulate DB failure - mock_prisma_client.db.litellm_config.find_first = AsyncMock( - side_effect=Exception("DB connection error") - ) + mock_prisma_client.db.litellm_config.find_first = AsyncMock(side_effect=Exception("DB connection error")) mock_proxy_logging = MagicMock(spec=ProxyLogging) mock_proxy_logging.slack_alerting_instance = MagicMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() mock_proxy_config = AsyncMock() with ( @@ -7423,9 +7081,7 @@ async def test_increment_spend_counters_initializes_and_increments(): ) # Counter should be: base(5.0) + increment(0.50) = 5.50 - counter = counter_cache.in_memory_cache.get_cache( - key=f"spend:key:{hashed_token}" - ) + counter = counter_cache.in_memory_cache.get_cache(key=f"spend:key:{hashed_token}") assert counter == 5.50 # Second increment — counter already exists, just increment @@ -7436,9 +7092,7 @@ async def test_increment_spend_counters_initializes_and_increments(): response_cost=0.25, ) - counter = counter_cache.in_memory_cache.get_cache( - key=f"spend:key:{hashed_token}" - ) + counter = counter_cache.in_memory_cache.get_cache(key=f"spend:key:{hashed_token}") assert counter == 5.75 finally: ps.user_api_key_cache = original_key_cache @@ -7484,9 +7138,7 @@ async def test_increment_spend_counters_team_and_member(): team_counter = counter_cache.in_memory_cache.get_cache(key="spend:team:team-1") assert team_counter == 2.30 - member_counter = counter_cache.in_memory_cache.get_cache( - key="spend:team_member:user-1:team-1" - ) + member_counter = counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") assert member_counter == 1.30 finally: ps.user_api_key_cache = original_key_cache @@ -7544,14 +7196,10 @@ async def test_init_and_increment_spend_counter_reseeds_from_db_on_counter_miss( increment=1.5, ) - fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with( - where={"team_id": "team-9"} - ) + fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with(where={"team_id": "team-9"}) # Seed uses SET NX with db_spend (42) — cross-pod safe, no INCR of 42. # Only the per-request delta (1.5) goes through INCRBYFLOAT. - fake_redis.async_set_cache.assert_awaited_once_with( - key="spend:team:team-9", value=42.0, nx=True - ) + fake_redis.async_set_cache.assert_awaited_once_with(key="spend:team:team-9", value=42.0, nx=True) writes = [(c["key"], c["value"]) for c in recorded_increments] assert writes == [("spend:team:team-9", 1.5)] finally: @@ -7620,9 +7268,7 @@ async def test_primary_spend_counter_redis_concurrent_seed_does_not_double_seed( return row fake_prisma = MagicMock() - fake_prisma.db.litellm_teamtable.find_unique = AsyncMock( - side_effect=slow_find_unique - ) + fake_prisma.db.litellm_teamtable.find_unique = AsyncMock(side_effect=slow_find_unique) pod_a = DualCache() pod_a.redis_cache = fake_redis @@ -7655,11 +7301,7 @@ async def test_primary_spend_counter_redis_concurrent_seed_does_not_double_seed( # (winner) and one was rejected (loser). assert db_read_count == 2 assert fake_redis.async_set_cache.await_count == 2 - nx_writes = [ - call - for call in fake_redis.async_set_cache.await_args_list - if call.kwargs.get("nx") is True - ] + nx_writes = [call for call in fake_redis.async_set_cache.await_args_list if call.kwargs.get("nx") is True] assert len(nx_writes) == 2 assert sorted(set_results) == [ False, @@ -7668,9 +7310,7 @@ async def test_primary_spend_counter_redis_concurrent_seed_does_not_double_seed( # Loser path executed: after the winner's SET NX returned True, the # losing coalesced() call falls back to async_get_cache to read the # winner's value rather than re-seeding. - assert ( - get_after_set_count >= 1 - ), "loser branch (else: read back winner's value) was never exercised" + assert get_after_set_count >= 1, "loser branch (else: read back winner's value) was never exercised" @pytest.mark.asyncio @@ -7692,14 +7332,10 @@ async def test_reseed_spend_from_db_user_and_org_prefixes(): fake_prisma.db.litellm_usertable.find_unique = AsyncMock(return_value=user_row) fake_prisma.db.litellm_endusertable.find_unique = AsyncMock() fake_prisma.db.litellm_tagtable.find_unique = AsyncMock() - fake_prisma.db.litellm_organizationtable.find_unique = AsyncMock( - return_value=org_row - ) + fake_prisma.db.litellm_organizationtable.find_unique = AsyncMock(return_value=org_row) assert await SpendCounterReseed.from_db(fake_prisma, "spend:user:alice") == 17.0 - fake_prisma.db.litellm_usertable.find_unique.assert_awaited_once_with( - where={"user_id": "alice"} - ) + fake_prisma.db.litellm_usertable.find_unique.assert_awaited_once_with(where={"user_id": "alice"}) assert ( await SpendCounterReseed.from_db( @@ -7714,9 +7350,7 @@ async def test_reseed_spend_from_db_user_and_org_prefixes(): fake_prisma.db.litellm_tagtable.find_unique.assert_not_awaited() assert await SpendCounterReseed.from_db(fake_prisma, "spend:org:acme") == 305.0 - fake_prisma.db.litellm_organizationtable.find_unique.assert_awaited_once_with( - where={"organization_id": "acme"} - ) + fake_prisma.db.litellm_organizationtable.find_unique.assert_awaited_once_with(where={"organization_id": "acme"}) @pytest.mark.asyncio @@ -7730,14 +7364,8 @@ async def test_reseed_spend_from_db_skips_window_variant_keys(): fake_prisma.db.litellm_verificationtoken.find_unique = AsyncMock() fake_prisma.db.litellm_teamtable.find_unique = AsyncMock() - assert ( - await SpendCounterReseed.from_db(fake_prisma, "spend:key:sk-abc:window:1h") - is None - ) - assert ( - await SpendCounterReseed.from_db(fake_prisma, "spend:team:team-1:window:1d") - is None - ) + assert await SpendCounterReseed.from_db(fake_prisma, "spend:key:sk-abc:window:1h") is None + assert await SpendCounterReseed.from_db(fake_prisma, "spend:team:team-1:window:1d") is None fake_prisma.db.litellm_verificationtoken.find_unique.assert_not_awaited() fake_prisma.db.litellm_teamtable.find_unique.assert_not_awaited() @@ -7773,9 +7401,7 @@ async def test_window_spend_counter_reseeds_from_spend_logs_on_counter_miss(): where={"api_key": "key-window", "startTime": {"gte": window_start}}, sum={"spend": True}, ) - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-window:window:1h" - ) == pytest.approx(2.75) + assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-window:window:1h") == pytest.approx(2.75) finally: ps.spend_counter_cache = orig_counter ps.prisma_client = orig_prisma @@ -7830,14 +7456,10 @@ async def test_init_spend_counter_redis_clean_miss_skips_stale_in_memory(): increment=1.5, ) - fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with( - where={"team_id": "team-stale-local"} - ) + fake_prisma.db.litellm_teamtable.find_unique.assert_awaited_once_with(where={"team_id": "team-stale-local"}) # Seed via SET NX (42) + delta via INCRBYFLOAT (1.5) = 43.5. assert redis_store[counter_key] == pytest.approx(43.5) - assert counter_cache.in_memory_cache.get_cache( - key=counter_key - ) == pytest.approx(43.5) + assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(43.5) finally: ps.spend_counter_cache = orig_counter ps.prisma_client = orig_prisma @@ -7900,9 +7522,7 @@ async def test_window_spend_counter_redis_clean_miss_skips_stale_in_memory(): sum={"spend": True}, ) assert redis_store[counter_key] == pytest.approx(2.75) - assert counter_cache.in_memory_cache.get_cache( - key=counter_key - ) == pytest.approx(2.75) + assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(2.75) finally: ps.spend_counter_cache = orig_counter ps.prisma_client = orig_prisma @@ -7938,9 +7558,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() fake_prisma = MagicMock() fake_prisma.db.litellm_spendlogs.group_by = AsyncMock( - return_value=[ - {"api_key": "key-window-concurrent-seed", "_sum": {"spend": 2.25}} - ] + return_value=[{"api_key": "key-window-concurrent-seed", "_sum": {"spend": 2.25}}] ) import litellm.proxy.proxy_server as ps @@ -7963,9 +7581,7 @@ async def test_window_spend_counter_redis_concurrent_seed_does_not_double_seed() nx=True, ) assert redis_store[counter_key] == pytest.approx(3.25) - assert counter_cache.in_memory_cache.get_cache( - key=counter_key - ) == pytest.approx(3.25) + assert counter_cache.in_memory_cache.get_cache(key=counter_key) == pytest.approx(3.25) finally: ps.spend_counter_cache = orig_counter ps.prisma_client = orig_prisma @@ -7991,12 +7607,7 @@ async def test_window_spend_counter_skips_invalid_window_start(): increment=0.5, ) - assert ( - counter_cache.in_memory_cache.get_cache( - key="spend:key:key-invalid-window:window:not-a-duration" - ) - is None - ) + assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-invalid-window:window:not-a-duration") is None finally: ps.spend_counter_cache = orig_counter @@ -8078,9 +7689,9 @@ async def test_increment_spend_counters_finalizes_after_unreserved_increments(): assert incremented_counters == ["spend:team:team-finalize-after-increments"] assert budget_reservation["finalized"] is True - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-finalize-after-increments" - ) == pytest.approx(0.25) + assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-finalize-after-increments") == pytest.approx( + 0.25 + ) finally: ps.spend_counter_cache = orig_counter ps.user_api_key_cache = orig_user @@ -8124,9 +7735,7 @@ async def test_increment_spend_counters_finalizes_none_cost_reservation(): ) assert budget_reservation["finalized"] is True - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-finalize-none-cost" - ) == pytest.approx(0.0) + assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-finalize-none-cost") == pytest.approx(0.0) finally: ps.spend_counter_cache = orig_counter @@ -8176,9 +7785,7 @@ async def test_increment_spend_counters_reseeds_from_db_on_bad_reserved_counter( assert budget_reservation["finalized"] is True # counter reseeded to the authoritative DB value, not deleted/left None # and not double-counted via a direct increment - assert counter_cache.in_memory_cache.get_cache( - key="spend:key:key-bad-reserved-counter" - ) == pytest.approx(0.6) + assert counter_cache.in_memory_cache.get_cache(key="spend:key:key-bad-reserved-counter") == pytest.approx(0.6) finally: ps.spend_counter_cache = orig_counter ps.prisma_client = orig_prisma @@ -8207,12 +7814,8 @@ async def test_increment_spend_counter_invalidates_stale_cache_on_redis_failure( increment=0.5, ) - assert ( - counter_cache.in_memory_cache.get_cache(key="spend:team:redis-fail") is None - ) - fake_redis.async_delete_cache.assert_awaited_once_with( - key="spend:team:redis-fail" - ) + assert counter_cache.in_memory_cache.get_cache(key="spend:team:redis-fail") is None + fake_redis.async_delete_cache.assert_awaited_once_with(key="spend:team:redis-fail") finally: ps.spend_counter_cache = orig_counter @@ -8258,16 +7861,13 @@ async def test_get_current_spend_reseeds_from_db_when_counter_missing(): fallback_spend=30.0, ) assert spend == 362.0, ( - f"expected DB reseed to return 362.0, got {spend} " - f"(fallback would have returned 30.0 and caused bypass)" + f"expected DB reseed to return 362.0, got {spend} (fallback would have returned 30.0 and caused bypass)" ) # Counter warmed via SET NX so subsequent reads are fast. assert ("spend:team_member:user-1:team-1", 362.0, True) in [ (s["key"], s["value"], s["nx"]) for s in recorded_seeds ] - assert counter_cache.in_memory_cache.get_cache( - key="spend:team_member:user-1:team-1" - ) == pytest.approx(362.0) + assert counter_cache.in_memory_cache.get_cache(key="spend:team_member:user-1:team-1") == pytest.approx(362.0) finally: ps.spend_counter_cache = orig_counter ps.prisma_client = orig_prisma @@ -8352,9 +7952,7 @@ async def test_get_current_spend_coalesces_concurrent_reseeds(): counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() - fake_prisma.db.litellm_teammembership.find_unique = AsyncMock( - side_effect=slow_find_unique - ) + fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(side_effect=slow_find_unique) import litellm.proxy.proxy_server as ps @@ -8363,15 +7961,10 @@ async def test_get_current_spend_coalesces_concurrent_reseeds(): ps.prisma_client = fake_prisma try: results = await _asyncio.gather( - *[ - get_current_spend(counter_key=counter_key, fallback_spend=0.0) - for _ in range(5) - ] + *[get_current_spend(counter_key=counter_key, fallback_spend=0.0) for _ in range(5)] ) assert results == [100.0] * 5, f"all callers should see DB value, got {results}" - assert ( - db_call_count == 1 - ), f"expected exactly 1 DB query for 5 concurrent reseeds, got {db_call_count}" + assert db_call_count == 1, f"expected exactly 1 DB query for 5 concurrent reseeds, got {db_call_count}" finally: ps.spend_counter_cache = orig_counter ps.prisma_client = orig_prisma @@ -8408,9 +8001,7 @@ async def test_get_current_spend_uses_db_zero_over_stale_fallback(): counter_key="spend:team_member:user-1:team-after-reset", fallback_spend=42.0, ) - assert ( - spend == 0.0 - ), f"DB authoritative 0 must override stale fallback 42, got {spend}" + assert spend == 0.0, f"DB authoritative 0 must override stale fallback 42, got {spend}" finally: ps.spend_counter_cache = orig_counter ps.prisma_client = orig_prisma @@ -8468,9 +8059,7 @@ async def test_concurrent_read_and_write_paths_share_one_db_query(): counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() - fake_prisma.db.litellm_teammembership.find_unique = AsyncMock( - side_effect=slow_find_unique - ) + fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(side_effect=slow_find_unique) import litellm.proxy.proxy_server as ps @@ -8492,9 +8081,7 @@ async def test_concurrent_read_and_write_paths_share_one_db_query(): ), get_current_spend(counter_key=counter_key, fallback_spend=0.0), ) - assert ( - db_call_count == 1 - ), f"expected 1 DB query for concurrent read+write+read, got {db_call_count}" + assert db_call_count == 1, f"expected 1 DB query for concurrent read+write+read, got {db_call_count}" # Read-path callers see the warmed counter; the write path's # increment may or may not have landed by then, so accept either # the seeded value or seeded+increment. @@ -8530,9 +8117,7 @@ async def test_reseed_locks_dict_is_bounded(): try: for i in range(7): await SpendCounterReseed._get_lock(f"spend:key:test-key-{i}") - assert ( - len(SpendCounterReseed._locks) == 5 - ), f"got {len(SpendCounterReseed._locks)}" + assert len(SpendCounterReseed._locks) == 5, f"got {len(SpendCounterReseed._locks)}" # Oldest two evicted assert "spend:key:test-key-0" not in SpendCounterReseed._locks assert "spend:key:test-key-1" not in SpendCounterReseed._locks @@ -8589,9 +8174,7 @@ async def test_reseed_warms_cache_even_on_zero_db_spend(): return row fake_prisma = MagicMock() - fake_prisma.db.litellm_teammembership.find_unique = AsyncMock( - side_effect=find_unique - ) + fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(side_effect=find_unique) import litellm.proxy.proxy_server as ps @@ -8604,9 +8187,7 @@ async def test_reseed_warms_cache_even_on_zero_db_spend(): # Second call: cache should be warmed at 0, no second DB query. spend2 = await get_current_spend(counter_key=counter_key, fallback_spend=0.0) assert spend1 == 0.0 and spend2 == 0.0 - assert ( - db_call_count == 1 - ), f"second read should hit warmed cache, got {db_call_count} DB queries" + assert db_call_count == 1, f"second read should hit warmed cache, got {db_call_count} DB queries" assert redis_store.get(counter_key) == 0.0, "cache must be warmed at 0" finally: ps.spend_counter_cache = orig_counter @@ -8669,9 +8250,7 @@ def _update_config_setup(monkeypatch): def _install(initial_rows=None, store_model_in_db=True): prisma = _FakePrismaClient(initial_rows=initial_rows) monkeypatch.setattr("litellm.proxy.proxy_server.prisma_client", prisma) - monkeypatch.setattr( - "litellm.proxy.proxy_server.store_model_in_db", store_model_in_db - ) + monkeypatch.setattr("litellm.proxy.proxy_server.store_model_in_db", store_model_in_db) monkeypatch.setattr( "litellm.proxy.proxy_server.encrypt_value_helper", lambda value, **_: f"enc:{value}", @@ -8682,9 +8261,7 @@ def _update_config_setup(monkeypatch): ) from litellm.proxy.proxy_server import proxy_config as real_proxy_config - monkeypatch.setattr( - real_proxy_config, "add_deployment", AsyncMock(return_value=None) - ) + monkeypatch.setattr(real_proxy_config, "add_deployment", AsyncMock(return_value=None)) original_overrides = app.dependency_overrides.copy() app.dependency_overrides[auth_dep] = lambda: UserAPIKeyAuth( @@ -8719,19 +8296,13 @@ def test_update_config_writes_only_sent_section(_update_config_setup): assert resp.status_code == 200 written = {name for name, _ in prisma.db.litellm_config.upsert_calls} assert written == {"general_settings"} - assert prisma.db.litellm_config.rows["litellm_settings"] == { - "drop_params": True - } - assert prisma.db.litellm_config.rows["environment_variables"] == { - "FOO": "enc:bar" - } + assert prisma.db.litellm_config.rows["litellm_settings"] == {"drop_params": True} + assert prisma.db.litellm_config.rows["environment_variables"] == {"FOO": "enc:bar"} finally: restore() -def test_update_config_env_var_round_trip_not_double_encrypted( - _update_config_setup, monkeypatch -): +def test_update_config_env_var_round_trip_not_double_encrypted(_update_config_setup, monkeypatch): """Endpoint-level regression for the /config/update double-encryption bug. The Admin UI reads config back via /get/config/callbacks (which returns @@ -8744,16 +8315,12 @@ def test_update_config_env_var_round_trip_not_double_encrypted( code this stored "enc:enc:..."; the assertions below would fail there. """ - def _fake_decrypt( - value, key=None, exception_type="error", return_original_value=False - ): + def _fake_decrypt(value, key=None, exception_type="error", return_original_value=False): if isinstance(value, str) and value.startswith("enc:"): return value[len("enc:") :] return value if return_original_value else None - monkeypatch.setattr( - "litellm.proxy.proxy_server.decrypt_value_helper", _fake_decrypt - ) + monkeypatch.setattr("litellm.proxy.proxy_server.decrypt_value_helper", _fake_decrypt) client, prisma, restore = _update_config_setup( initial_rows={"environment_variables": {"PREEXISTING_KEY": "enc:keepme"}} @@ -8771,21 +8338,14 @@ def test_update_config_env_var_round_trip_not_double_encrypted( # UI round-trip: re-POST the stored ciphertext (no field change). resp = client.post( "/config/update", - json={ - "environment_variables": { - "LANGFUSE_SECRET_KEY": stored["LANGFUSE_SECRET_KEY"] - } - }, + json={"environment_variables": {"LANGFUSE_SECRET_KEY": stored["LANGFUSE_SECRET_KEY"]}}, ) assert resp.status_code == 200 stored = prisma.db.litellm_config.rows["environment_variables"] # The bug: this would be "enc:enc:sk-secret". The fix keeps it single. assert stored["LANGFUSE_SECRET_KEY"] == "enc:sk-secret" - assert ( - _fake_decrypt(stored["LANGFUSE_SECRET_KEY"], return_original_value=True) - == "sk-secret" - ) + assert _fake_decrypt(stored["LANGFUSE_SECRET_KEY"], return_original_value=True) == "sk-secret" # Untouched key preserved byte-for-byte (only sent keys rewritten). assert stored["PREEXISTING_KEY"] == "enc:keepme" @@ -8800,14 +8360,9 @@ def test_update_config_can_flip_store_model_in_db_when_currently_false( False, blocking the very request that would flip it to True.""" client, prisma, restore = _update_config_setup(store_model_in_db=False) try: - resp = client.post( - "/config/update", json={"general_settings": {"store_model_in_db": True}} - ) + resp = client.post("/config/update", json={"general_settings": {"store_model_in_db": True}}) assert resp.status_code == 200 - assert ( - prisma.db.litellm_config.rows["general_settings"]["store_model_in_db"] - is True - ) + assert prisma.db.litellm_config.rows["general_settings"]["store_model_in_db"] is True finally: restore() @@ -8840,9 +8395,7 @@ def test_update_config_litellm_settings_request_wins_for_non_callback_keys( } ) try: - resp = client.post( - "/config/update", json={"litellm_settings": {"drop_params": False}} - ) + resp = client.post("/config/update", json={"litellm_settings": {"drop_params": False}}) assert resp.status_code == 200 stored = prisma.db.litellm_config.rows["litellm_settings"] assert stored["drop_params"] is False @@ -8938,9 +8491,7 @@ class TestLazyFeaturesNotImportedAtStartup: from litellm.proxy._lazy_features import LAZY_FEATURES - proxy_server_src = ( - Path(__file__).resolve().parents[3] / "litellm/proxy/proxy_server.py" - ).read_text() + proxy_server_src = (Path(__file__).resolve().parents[3] / "litellm/proxy/proxy_server.py").read_text() leaks = [] for feat in LAZY_FEATURES: @@ -9045,9 +8596,7 @@ class TestLazyFeatureMiddleware: ("/api/v1", "/api/v1/unrelated", False, "unrelated path under root"), ], ) - async def test_root_path_handling( - self, monkeypatch, server_root_path, request_path, should_load, case - ): + async def test_root_path_handling(self, monkeypatch, server_root_path, request_path, should_load, case): """ The middleware must strip SERVER_ROOT_PATH before prefix-matching so lazy features load under deployments that set a server root path, @@ -9157,9 +8706,7 @@ class TestLazyFeatureMiddleware: ) await asyncio.gather(hit(), hit(), hit(), hit(), hit()) - assert loads == [ - "json" - ], f"expected one registration despite concurrent first hits, got {loads}" + assert loads == ["json"], f"expected one registration despite concurrent first hits, got {loads}" @pytest.mark.asyncio async def test_failing_import_does_not_loop(self): @@ -9209,9 +8756,9 @@ class TestLazyFeatureMiddleware: receive, send, ) - assert attempts == [ - "called" - ], f"failing register_fn should be invoked once, not on every request; got {attempts}" + assert attempts == ["called"], ( + f"failing register_fn should be invoked once, not on every request; got {attempts}" + ) @pytest.mark.asyncio @@ -9279,9 +8826,7 @@ async def test_get_current_spend_redis_error_falls_back_to_in_memory(): counter_cache.redis_cache = fake_redis fake_prisma = MagicMock() - fake_prisma.db.litellm_teammembership.find_unique = AsyncMock( - return_value=MagicMock(spend=999.0) - ) + fake_prisma.db.litellm_teammembership.find_unique = AsyncMock(return_value=MagicMock(spend=999.0)) import litellm.proxy.proxy_server as ps @@ -9291,8 +8836,7 @@ async def test_get_current_spend_redis_error_falls_back_to_in_memory(): try: spend = await get_current_spend(counter_key=counter_key, fallback_spend=0.0) assert spend == 42.0, ( - f"expected in-memory fallback 42.0 on Redis error, got {spend} " - f"(should not have hit DB when Redis errored)" + f"expected in-memory fallback 42.0 on Redis error, got {spend} (should not have hit DB when Redis errored)" ) # DB query should NOT have fired - in-memory short-circuits. fake_prisma.db.litellm_teammembership.find_unique.assert_not_awaited() @@ -9315,9 +8859,7 @@ def test_realtime_websocket_route_aliases_registered(): from litellm.proxy.proxy_server import app from litellm.types.utils import API_ROUTE_TO_CALL_TYPES, CallTypes - websocket_paths = { - route.path for route in app.routes if isinstance(route, WebSocketRoute) - } + websocket_paths = {route.path for route in app.routes if isinstance(route, WebSocketRoute)} openai_routes = LiteLLMRoutes.openai_routes.value for expected in ("/openai/v1/realtime", "/v1/realtime", "/realtime"): @@ -9329,9 +8871,7 @@ def test_realtime_websocket_route_aliases_registered(): f"{expected!r} missing from LiteLLMRoutes.openai_routes; " f"non-admin / team / key-scoped users will get 403 on this path." ) - assert tuple(API_ROUTE_TO_CALL_TYPES.get(expected) or ()) == ( - CallTypes.arealtime, - ), ( + assert tuple(API_ROUTE_TO_CALL_TYPES.get(expected) or ()) == (CallTypes.arealtime,), ( f"{expected!r} missing from API_ROUTE_TO_CALL_TYPES; call-type " f"resolution will return None and break call-type-aware features." ) @@ -9381,8 +8921,7 @@ class TestTransformRequestBannedParams: }, ) assert response.status_code == 400, ( - f"Expected 400 for banned param '{banned}', " - f"got {response.status_code}: {response.json()}" + f"Expected 400 for banned param '{banned}', got {response.status_code}: {response.json()}" ) @@ -9408,13 +8947,8 @@ class TestSortModelsByDisplayName: {"model_name": "gpt-4o", "model_info": {}}, ] - sorted_models = _sort_models( - all_models=models, sort_by="model_name", sort_order="asc" - ) - displayed_order = [ - m["model_info"].get("team_public_model_name") or m["model_name"] - for m in sorted_models - ] + sorted_models = _sort_models(all_models=models, sort_by="model_name", sort_order="asc") + displayed_order = [m["model_info"].get("team_public_model_name") or m["model_name"] for m in sorted_models] assert displayed_order == [ "anthropic/claude", "claude-haiku-4-5", @@ -9433,13 +8967,8 @@ class TestSortModelsByDisplayName: {"model_name": "gpt-4o", "model_info": {}}, ] - sorted_models = _sort_models( - all_models=models, sort_by="model_name", sort_order="desc" - ) - displayed_order = [ - m["model_info"].get("team_public_model_name") or m["model_name"] - for m in sorted_models - ] + sorted_models = _sort_models(all_models=models, sort_by="model_name", sort_order="desc") + displayed_order = [m["model_info"].get("team_public_model_name") or m["model_name"] for m in sorted_models] assert displayed_order == [ "zeta/model", "gpt-4o", @@ -9457,9 +8986,7 @@ class TestSortModelsByDisplayName: {"model_name": "beta", "model_info": {}}, ] - sorted_models = _sort_models( - all_models=models, sort_by="model_name", sort_order="asc" - ) + sorted_models = _sort_models(all_models=models, sort_by="model_name", sort_order="asc") assert [m["model_name"] for m in sorted_models] == ["alpha", "beta"] @@ -9481,9 +9008,7 @@ class TestDeleteDeploymentSync: mock_router.delete_deployment.return_value = MagicMock() with patch("litellm.proxy.proxy_server.llm_router", mock_router): - with patch.object( - proxy_config, "get_config", AsyncMock(return_value={"model_list": []}) - ): + with patch.object(proxy_config, "get_config", AsyncMock(return_value={"model_list": []})): still_desired = await proxy_config._delete_deployment(db_models=[]) mock_router.delete_deployment.assert_called_once_with(id="model-id-to-evict") @@ -9507,9 +9032,7 @@ class TestDeleteDeploymentSync: with patch("litellm.proxy.proxy_server.llm_router", mock_router): with patch.object(proxy_config, "get_config", AsyncMock(return_value={})): - await proxy_config._update_llm_router( - new_models=None, proxy_logging_obj=MagicMock() - ) + await proxy_config._update_llm_router(new_models=None, proxy_logging_obj=MagicMock()) mock_router.delete_deployment.assert_not_called() mock_router.upsert_deployment.assert_not_called() @@ -9526,15 +9049,11 @@ class TestDeleteDeploymentSync: proxy_config = ProxyConfig() mock_prisma = MagicMock() - mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock( - side_effect=Exception("DB connection lost") - ) + mock_prisma.db.litellm_proxymodeltable.find_many = AsyncMock(side_effect=Exception("DB connection lost")) result = await proxy_config._get_models_from_db(prisma_client=mock_prisma) - assert ( - result is None - ), f"Expected None on DB failure to signal fetch error, got {result!r}" + assert result is None, f"Expected None on DB failure to signal fetch error, got {result!r}" def test_get_config_list_includes_cancel_on_disconnect(monkeypatch): @@ -9816,9 +9335,18 @@ def test_general_settings_ui_defaults_unchanged_for_existing_fields(): _general_settings_ui_litellm_default, ) - assert _general_settings_ui_litellm_default(_GENERAL_SETTINGS_UI_LITELLM_FIELDS["budget_exceeded_throttle_percentage"]) is None - assert _general_settings_ui_litellm_default(_GENERAL_SETTINGS_UI_LITELLM_FIELDS["enable_anthropic_prompt_caching"]) is False - assert _general_settings_ui_litellm_default(_GENERAL_SETTINGS_UI_LITELLM_FIELDS["anthropic_prompt_caching_ttl"]) is None + assert ( + _general_settings_ui_litellm_default(_GENERAL_SETTINGS_UI_LITELLM_FIELDS["budget_exceeded_throttle_percentage"]) + is None + ) + assert ( + _general_settings_ui_litellm_default(_GENERAL_SETTINGS_UI_LITELLM_FIELDS["enable_anthropic_prompt_caching"]) + is False + ) + assert ( + _general_settings_ui_litellm_default(_GENERAL_SETTINGS_UI_LITELLM_FIELDS["anthropic_prompt_caching_ttl"]) + is None + ) @pytest.mark.parametrize( @@ -10084,16 +9612,10 @@ def test_preserve_redacted_plugin_keys_keeps_stored_credential(): existing = [{"name": "p1", "url": "https://p1", "plugin_key": "sk-real-1"}] - redacted = _preserve_redacted_plugin_keys( - [{"name": "p1", "url": "https://p1-new", "plugin_key": "***"}], existing - ) - assert redacted == [ - {"name": "p1", "url": "https://p1-new", "plugin_key": "sk-real-1"} - ] + redacted = _preserve_redacted_plugin_keys([{"name": "p1", "url": "https://p1-new", "plugin_key": "***"}], existing) + assert redacted == [{"name": "p1", "url": "https://p1-new", "plugin_key": "sk-real-1"}] - blanked = _preserve_redacted_plugin_keys( - [{"name": "p1", "url": "https://p1", "plugin_key": ""}], existing - ) + blanked = _preserve_redacted_plugin_keys([{"name": "p1", "url": "https://p1", "plugin_key": ""}], existing) assert blanked[0]["plugin_key"] == "sk-real-1" @@ -10103,14 +9625,10 @@ def test_preserve_redacted_plugin_keys_sets_new_and_drops_orphan_placeholder(): existing = [{"name": "p1", "url": "https://p1", "plugin_key": "sk-real-1"}] - rotated = _preserve_redacted_plugin_keys( - [{"name": "p1", "url": "https://p1", "plugin_key": "sk-new"}], existing - ) + rotated = _preserve_redacted_plugin_keys([{"name": "p1", "url": "https://p1", "plugin_key": "sk-new"}], existing) assert rotated[0]["plugin_key"] == "sk-new" - new_plugin = _preserve_redacted_plugin_keys( - [{"name": "p2", "url": "https://p2", "plugin_key": "***"}], existing - ) + new_plugin = _preserve_redacted_plugin_keys([{"name": "p2", "url": "https://p2", "plugin_key": "***"}], existing) assert "plugin_key" not in new_plugin[0] @@ -10143,9 +9661,7 @@ def _config_field_info_client(monkeypatch, user_role): mock_prisma = MagicMock() mock_prisma.db = types.SimpleNamespace(litellm_config=mock_config_table) monkeypatch.setattr(ps, "prisma_client", mock_prisma) - app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth( - user_id="u", user_role=user_role - ) + app.dependency_overrides[ps.user_api_key_auth] = lambda: UserAPIKeyAuth(user_id="u", user_role=user_role) return TestClient(app) @@ -10156,9 +9672,7 @@ def test_config_field_info_redacts_secrets_for_view_only_admin(monkeypatch): is not a FULL PROXY_ADMIN, while non-secret fields stay readable.""" from litellm.proxy._types import LitellmUserRoles - client = _config_field_info_client( - monkeypatch, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY - ) + client = _config_field_info_client(monkeypatch, LitellmUserRoles.PROXY_ADMIN_VIEW_ONLY) try: for secret_field in ("master_key", "database_url", "pass_through_endpoints"): resp = client.get("/config/field/info", params={"field_name": secret_field}) @@ -10168,9 +9682,7 @@ def test_config_field_info_redacts_secrets_for_view_only_admin(monkeypatch): assert "secret" not in str(body["field_value"]) assert "p4ssw0rd" not in str(body["field_value"]) - resp = client.get( - "/config/field/info", params={"field_name": "max_parallel_requests"} - ) + resp = client.get("/config/field/info", params={"field_name": "max_parallel_requests"}) assert resp.status_code == 200, resp.text assert resp.json()["field_value"] == 100 finally: @@ -10188,14 +9700,9 @@ def test_config_field_info_returns_raw_secrets_for_full_admin(monkeypatch): assert resp.status_code == 200, resp.text assert resp.json()["field_value"] == "sk-super-secret-master" - resp = client.get( - "/config/field/info", params={"field_name": "pass_through_endpoints"} - ) + resp = client.get("/config/field/info", params={"field_name": "pass_through_endpoints"}) assert resp.status_code == 200, resp.text - assert ( - resp.json()["field_value"][0]["headers"]["Authorization"] - == "Bearer sk-upstream-secret" - ) + assert resp.json()["field_value"][0]["headers"]["Authorization"] == "Bearer sk-upstream-secret" finally: app.dependency_overrides.clear() @@ -10437,9 +9944,7 @@ async def test_delete_config_general_settings_emits_deleted_audit_log(monkeypatc user_role=LitellmUserRoles.PROXY_ADMIN, ) await delete_config_general_settings( - data=ConfigFieldDelete( - field_name="max_parallel_requests", config_type="general_settings" - ), + data=ConfigFieldDelete(field_name="max_parallel_requests", config_type="general_settings"), user_api_key_dict=admin, ) # Audit is scheduled via asyncio.create_task; yield so it runs. @@ -10462,9 +9967,7 @@ def test_update_config_audits_every_written_section(_update_config_setup, monkey is the row that holds default_internal_user_params ("default user settings").""" import litellm.proxy.proxy_server as proxy_server_module - client, prisma, restore = _update_config_setup( - initial_rows={"litellm_settings": {"drop_params": True}} - ) + client, prisma, restore = _update_config_setup(initial_rows={"litellm_settings": {"drop_params": True}}) audit_create = AsyncMock() prisma.db.litellm_auditlog.create = audit_create monkeypatch.setattr(proxy_server_module, "premium_user", True) @@ -10475,17 +9978,14 @@ def test_update_config_audits_every_written_section(_update_config_setup, monkey json={ "general_settings": {"store_prompts_in_spend_logs": True}, "environment_variables": {"FOO": "bar"}, - "litellm_settings": { - "default_internal_user_params": {"max_budget": 10} - }, + "litellm_settings": {"default_internal_user_params": {"max_budget": 10}}, "router_settings": {"routing_strategy": "latency-based-routing"}, }, ) assert resp.status_code == 200, resp.text audited = { - call.kwargs["data"]["object_id"]: call.kwargs["data"]["action"] - for call in audit_create.await_args_list + call.kwargs["data"]["object_id"]: call.kwargs["data"]["action"] for call in audit_create.await_args_list } assert audited == { "general_settings": "updated", @@ -10497,20 +9997,14 @@ def test_update_config_audits_every_written_section(_update_config_setup, monkey assert call.kwargs["data"]["table_name"] == "LiteLLM_Config" assert call.kwargs["data"]["changed_by"] == "test_admin" - ls_call = next( - c - for c in audit_create.await_args_list - if c.kwargs["data"]["object_id"] == "litellm_settings" - ) + ls_call = next(c for c in audit_create.await_args_list if c.kwargs["data"]["object_id"] == "litellm_settings") after = json.loads(ls_call.kwargs["data"]["updated_values"]) assert after["default_internal_user_params"] == {"max_budget": 10} finally: restore() -def test_delete_callback_audits_litellm_settings_deletion( - _update_config_setup, monkeypatch -): +def test_delete_callback_audits_litellm_settings_deletion(_update_config_setup, monkeypatch): """/config/callback/delete must emit a deleted audit row for litellm_settings capturing the success_callback list before and after removal.""" import litellm.proxy.proxy_server as proxy_server_module @@ -10526,19 +10020,11 @@ def test_delete_callback_audits_litellm_settings_deletion( monkeypatch.setattr( real_proxy_config, "get_config", - AsyncMock( - return_value={ - "litellm_settings": {"success_callback": ["langfuse", "datadog"]} - } - ), - ) - monkeypatch.setattr( - real_proxy_config, "save_config", AsyncMock(return_value=None) + AsyncMock(return_value={"litellm_settings": {"success_callback": ["langfuse", "datadog"]}}), ) + monkeypatch.setattr(real_proxy_config, "save_config", AsyncMock(return_value=None)) try: - resp = client.post( - "/config/callback/delete", json={"callback_name": "datadog"} - ) + resp = client.post("/config/callback/delete", json={"callback_name": "datadog"}) assert resp.status_code == 200, resp.text audit_create.assert_awaited_once() @@ -10567,24 +10053,16 @@ def test_delete_callback_audits_before_reload_failure(_update_config_setup, monk monkeypatch.setattr( real_proxy_config, "get_config", - AsyncMock( - return_value={ - "litellm_settings": {"success_callback": ["langfuse", "datadog"]} - } - ), - ) - monkeypatch.setattr( - real_proxy_config, "save_config", AsyncMock(return_value=None) + AsyncMock(return_value={"litellm_settings": {"success_callback": ["langfuse", "datadog"]}}), ) + monkeypatch.setattr(real_proxy_config, "save_config", AsyncMock(return_value=None)) monkeypatch.setattr( real_proxy_config, "add_deployment", AsyncMock(side_effect=RuntimeError("reload failed")), ) try: - resp = client.post( - "/config/callback/delete", json={"callback_name": "datadog"} - ) + resp = client.post("/config/callback/delete", json={"callback_name": "datadog"}) assert resp.status_code == 500, resp.text audit_create.assert_awaited_once() @@ -10595,9 +10073,7 @@ def test_delete_callback_audits_before_reload_failure(_update_config_setup, monk restore() -def test_update_config_redacts_all_environment_variable_values( - _update_config_setup, monkeypatch -): +def test_update_config_redacts_all_environment_variable_values(_update_config_setup, monkeypatch): """environment_variables hold credentials under arbitrary uppercase keys (DATABASE_URL) that key-name secret matching misses, so every value in the section must be redacted before the audit row is written; a plaintext @@ -10607,11 +10083,7 @@ def test_update_config_redacts_all_environment_variable_values( # DATABASE_URL is the bug class: an uppercase env key that key-name secret # matching does NOT flag, so only whole-section value redaction protects it. client, prisma, restore = _update_config_setup( - initial_rows={ - "environment_variables": { - "DATABASE_URL": "enc:postgresql://OLDsecret@old.host:5432/db" - } - } + initial_rows={"environment_variables": {"DATABASE_URL": "enc:postgresql://OLDsecret@old.host:5432/db"}} ) audit_create = AsyncMock() prisma.db.litellm_auditlog.create = audit_create @@ -10630,9 +10102,7 @@ def test_update_config_redacts_all_environment_variable_values( assert resp.status_code == 200, resp.text env_call = next( - c - for c in audit_create.await_args_list - if c.kwargs["data"]["object_id"] == "environment_variables" + c for c in audit_create.await_args_list if c.kwargs["data"]["object_id"] == "environment_variables" ) data = env_call.kwargs["data"] @@ -10796,11 +10266,7 @@ def test_init_coordination_redis_startup_nodes_builds_cluster_client(): """A coordination_redis block with startup_nodes must construct a cluster client, so cluster-aware consumers (v3 rate limiter) take the cluster path.""" usage_cache, _, _ = _run_init_coordination_redis( - config={ - "general_settings": { - "coordination_redis": {"startup_nodes": [{"host": "node-1", "port": 7000}]} - } - }, + config={"general_settings": {"coordination_redis": {"startup_nodes": [{"host": "node-1", "port": 7000}]}}}, ) assert isinstance(usage_cache, _EnvBuiltClusterCache) @@ -11050,17 +10516,13 @@ async def _collect_async_data_generator_frames(request_data: dict) -> list: with patch.object(proxy_server_module.ProxyLogging, "_fire_deferred_stream_logging"): return [ frame.decode("utf-8") if isinstance(frame, bytes) else frame - async for frame in async_data_generator( - MockStream(), MagicMock(spec=UserAPIKeyAuth), request_data - ) + async for frame in async_data_generator(MockStream(), MagicMock(spec=UserAPIKeyAuth), request_data) ] @pytest.mark.asyncio async def test_async_data_generator_strips_injected_usage_chunk(): - frames = await _collect_async_data_generator_frames( - {"model": "gpt-5.4-nano", "_litellm_strip_stream_usage": True} - ) + frames = await _collect_async_data_generator_frames({"model": "gpt-5.4-nano", "_litellm_strip_stream_usage": True}) data_frames = [frame for frame in frames if frame.startswith("data: {")] assert len(data_frames) == 2 @@ -11138,9 +10600,7 @@ def test_startup_warns_when_mock_testing_params_enabled(caplog): ) with caplog.at_level(logging.WARNING, logger="LiteLLM Proxy"): - ProxyStartupEvent._warn_if_mock_testing_params_enabled( - general_settings={MOCK_TESTING_CONFIG_KEY: True} - ) + ProxyStartupEvent._warn_if_mock_testing_params_enabled(general_settings={MOCK_TESTING_CONFIG_KEY: True}) assert MOCK_TESTING_CONFIG_KEY in caplog.text for param_name in GATED_MOCK_PARAM_NAMES: @@ -11201,9 +10661,7 @@ async def test_setup_prisma_client_retains_connected_client_when_startup_health_ {"allow_requests_on_db_unavailable": True}, ) - mock_client = _mock_startup_prisma_client( - health_check_error=httpx.ReadTimeout("startup health check timed out") - ) + mock_client = _mock_startup_prisma_client(health_check_error=httpx.ReadTimeout("startup health check timed out")) result = await _run_setup_prisma_client(mock_client) assert mock_client.connect.await_count == 1 @@ -11227,9 +10685,7 @@ async def test_setup_prisma_client_arms_health_watchdog_before_startup_health_ch {"allow_requests_on_db_unavailable": True}, ) - mock_client = _mock_startup_prisma_client( - health_check_error=httpx.ReadTimeout("startup health check timed out") - ) + mock_client = _mock_startup_prisma_client(health_check_error=httpx.ReadTimeout("startup health check timed out")) call_order = MagicMock() call_order.attach_mock(mock_client.start_db_health_watchdog_task, "watchdog") call_order.attach_mock(mock_client.health_check, "health_check") @@ -11253,9 +10709,7 @@ async def test_setup_prisma_client_raises_when_db_unavailable_is_not_allowed(mon {"allow_requests_on_db_unavailable": False}, ) - mock_client = _mock_startup_prisma_client( - health_check_error=httpx.ReadTimeout("startup health check timed out") - ) + mock_client = _mock_startup_prisma_client(health_check_error=httpx.ReadTimeout("startup health check timed out")) with pytest.raises(httpx.ReadTimeout): await _run_setup_prisma_client(mock_client) @@ -11289,6 +10743,7 @@ async def _run_scheduled_background_jobs(): mock_proxy_logging = MagicMock(spec=ProxyLogging) mock_proxy_logging.slack_alerting_instance = MagicMock() + mock_proxy_logging.db_spend_update_writer = MagicMock() mock_proxy_config = AsyncMock() with ( diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py b/tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py index cf906259246..db842802435 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_lifecycle.py @@ -8,6 +8,7 @@ because they are direct dependents on the lifecycle state. from __future__ import annotations +import asyncio from typing import Any, Dict, List from unittest.mock import AsyncMock, MagicMock, patch @@ -138,6 +139,41 @@ def test_startup_event_propagates_init_callbacks_failure_raises(proxy_logging): proxy_logging.startup_event(llm_router=None, redis_usage_cache=None) +@pytest.mark.asyncio +async def test_startup_event_hands_the_daily_report_this_pods_lock_manager(proxy_logging): + """regression: issue #14809 - the daily report's dedupe lock only works if startup_event + passes the writer's pod_lock_manager down; dropping the argument silently restores the + every-pod-reports behavior.""" + proxy_logging.slack_alerting_instance = MagicMock() + proxy_logging.slack_alerting_instance.alert_types = ["daily_reports"] + proxy_logging.slack_alerting_instance._run_scheduled_daily_report = AsyncMock() + proxy_logging._init_litellm_callbacks = MagicMock() + proxy_logging.update_values = MagicMock() + llm_router = MagicMock() + + proxy_logging.startup_event(llm_router=llm_router, redis_usage_cache=None) + await asyncio.sleep(0) + + call = proxy_logging.slack_alerting_instance._run_scheduled_daily_report.call_args + assert proxy_logging.slack_alerting_instance._run_scheduled_daily_report.call_count == 1 + assert call.kwargs["pod_lock_manager"] is proxy_logging.db_spend_update_writer.pod_lock_manager + assert call.kwargs["llm_router"] is llm_router + + +@pytest.mark.asyncio +async def test_startup_event_skips_the_daily_report_when_it_is_not_an_alert_type(proxy_logging): + proxy_logging.slack_alerting_instance = MagicMock() + proxy_logging.slack_alerting_instance.alert_types = [] + proxy_logging.slack_alerting_instance._run_scheduled_daily_report = AsyncMock() + proxy_logging._init_litellm_callbacks = MagicMock() + proxy_logging.update_values = MagicMock() + + proxy_logging.startup_event(llm_router=None, redis_usage_cache=None) + await asyncio.sleep(0) + + proxy_logging.slack_alerting_instance._run_scheduled_daily_report.assert_not_called() + + # --------------------------------------------------------------------------- # _add_proxy_hooks # ---------------------------------------------------------------------------