From 703159eaabaf6a2c52e334041a9b144a45c3ff31 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?R=C4=81na=28Bass=20Ver=2E=29?= <1759138827@qq.com> Date: Mon, 27 Jul 2026 12:43:26 +0800 Subject: [PATCH 01/49] fix(proxy): allow unblocking customers via /customer/update update_end_user filtered out non-default values with v not in ([], {}, 0). Since False == 0 in Python, blocked: False was stripped from the update payload. Treat bools as explicit values while preserving the existing skips for empty containers and numeric zero Fixes #34379 --- .../customer_endpoints.py | 6 +----- .../test_customer_endpoints.py | 20 +++++++++++++++++++ 2 files changed, 21 insertions(+), 5 deletions(-) diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index a46481d5bb7..388888d960d 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -553,11 +553,7 @@ async def update_end_user( # get non default values for key non_default_values = {} for k, v in data_json.items(): - if v is not None and v not in ( - [], - {}, - 0, - ): # models default to [], spend defaults to 0, we should not reset these values + if v is not None and (isinstance(v, bool) or v not in ([], {}, 0)): non_default_values[k] = v ## Get end user table data ## diff --git a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py index 98e93eea5f9..c1479d4539c 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py @@ -85,6 +85,26 @@ def test_update_customer_success(mock_prisma_client, mock_user_api_key_auth): assert response.json()["alias"] == "Updated Test User" +def test_update_customer_unblock(mock_prisma_client, mock_user_api_key_auth): + mock_end_user = LiteLLM_EndUserTable(user_id="test-user-1", blocked=True) + updated_mock_end_user = LiteLLM_EndUserTable(user_id="test-user-1", blocked=False) + + mock_prisma_client.db.litellm_endusertable.find_first = AsyncMock(return_value=mock_end_user) + mock_prisma_client.db.litellm_endusertable.update = AsyncMock(return_value=updated_mock_end_user) + + response = client.post( + "/customer/update", + json={"user_id": "test-user-1", "blocked": False}, + headers={"Authorization": "Bearer test-key"}, + ) + + assert response.status_code == 200 + assert response.json()["blocked"] is False + update_mock = mock_prisma_client.db.litellm_endusertable.update + update_mock.assert_called_once() + assert update_mock.call_args.kwargs["data"]["blocked"] is False + + def test_update_customer_not_found(mock_prisma_client, mock_user_api_key_auth): """ Test that update_end_user raises a 404 ProxyException when user_id does not exist. From 463e9cd7ff3612070aaab7657ce31d98373e5f5b Mon Sep 17 00:00:00 2001 From: cat0825 Date: Tue, 4 Aug 2026 11:27:09 +0800 Subject: [PATCH 02/49] fix(proxy): only apply blocked when explicitly supplied The model default blocked=False was being written on every customer update that omitted the field, silently unblocking blocked customers when admins changed unrelated fields like alias or budget. Only accept bool values for fields the caller explicitly supplied (data.fields_set()), keeping the isinstance(v, bool) semantics for explicit updates like blocked=True/False. Adds a regression test: updating a blocked customer without the blocked field must not reset the block. --- .../customer_endpoints.py | 5 +++- .../test_customer_endpoints.py | 24 +++++++++++++++++++ 2 files changed, 28 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index 388888d960d..4127bc973e4 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -553,7 +553,10 @@ async def update_end_user( # get non default values for key non_default_values = {} for k, v in data_json.items(): - if v is not None and (isinstance(v, bool) or v not in ([], {}, 0)): + if v is not None and ( + (isinstance(v, bool) and k in data.fields_set()) + or v not in ([], {}, 0) + ): non_default_values[k] = v ## Get end user table data ## diff --git a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py index c1479d4539c..6eb03256b67 100644 --- a/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py +++ b/tests/test_litellm/proxy/management_endpoints/test_customer_endpoints.py @@ -105,6 +105,30 @@ def test_update_customer_unblock(mock_prisma_client, mock_user_api_key_auth): assert update_mock.call_args.kwargs["data"]["blocked"] is False +def test_update_customer_keeps_blocked_when_omitted(mock_prisma_client, mock_user_api_key_auth): + """ + Regression test: updating a blocked customer without supplying `blocked` + must NOT reset it to unblocked. `blocked=False` is the model default and + should only be applied when explicitly provided by the caller. + """ + mock_end_user = LiteLLM_EndUserTable(user_id="test-user-1", blocked=True) + updated_mock_end_user = LiteLLM_EndUserTable(user_id="test-user-1", blocked=True) + + mock_prisma_client.db.litellm_endusertable.find_first = AsyncMock(return_value=mock_end_user) + mock_prisma_client.db.litellm_endusertable.update = AsyncMock(return_value=updated_mock_end_user) + + response = client.post( + "/customer/update", + json={"user_id": "test-user-1", "alias": "Updated Test User"}, + headers={"Authorization": "Bearer test-key"}, + ) + + assert response.status_code == 200 + update_mock = mock_prisma_client.db.litellm_endusertable.update + update_mock.assert_called_once() + assert "blocked" not in update_mock.call_args.kwargs["data"] + + def test_update_customer_not_found(mock_prisma_client, mock_user_api_key_auth): """ Test that update_end_user raises a 404 ProxyException when user_id does not exist. From c12b2e82f74d49f2377e463addbf9248bf796d94 Mon Sep 17 00:00:00 2001 From: cat0825 Date: Tue, 4 Aug 2026 12:18:23 +0800 Subject: [PATCH 03/49] style: ruff format customer_endpoints.py --- litellm/proxy/management_endpoints/customer_endpoints.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/litellm/proxy/management_endpoints/customer_endpoints.py b/litellm/proxy/management_endpoints/customer_endpoints.py index 4127bc973e4..6efb3365cd9 100644 --- a/litellm/proxy/management_endpoints/customer_endpoints.py +++ b/litellm/proxy/management_endpoints/customer_endpoints.py @@ -553,10 +553,7 @@ async def update_end_user( # get non default values for key non_default_values = {} for k, v in data_json.items(): - if v is not None and ( - (isinstance(v, bool) and k in data.fields_set()) - or v not in ([], {}, 0) - ): + if v is not None and ((isinstance(v, bool) and k in data.fields_set()) or v not in ([], {}, 0)): non_default_values[k] = v ## Get end user table data ## From cb65bf08b8c937196270f14468916de3a627388b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Wed, 12 Aug 2026 19:41:30 -0700 Subject: [PATCH 04/49] chore(typing): clear 1.2k basedpyright Any errors across 16 hotspot files Replace Any-typed seams with real types in the files carrying the highest remaining reportAny/reportExplicitAny density: Literal-keyed structural Protocols for deployment dicts in tag-based routing, typed Prisma table wrappers and row protocols in the key and internal-user management endpoints, TypedDict views for websearch interception kwargs, typed streaming state in the responses iterator and background polling, and concrete request/response types in the google_genai, vertex_ai files, runwayml, rubrik, anthropic context-management, and guardrail translation modules. Mutable annotations introduced along the way were rewritten as read-only views (Mapping/Sequence/tuple) built functionally. No casts, no type: ignore, no noqa, no suppression comments, no new Any annotations, no behavior changes. Whole-tree basedpyright: reportAny 15,496 -> 14,523, reportExplicitAny 5,356 -> 5,102, all rules 145,547 -> 143,989, with no rule increased repo-wide or per-file. Budgets ratcheted: basedpyright -1,545, ruff-strict -73, type-discipline -237. --- basedpyright-code-budget.json | 22 +- .../google_genai/adapters/transformation.py | 186 ++++++++--- litellm/integrations/rubrik.py | 228 +++++++++---- .../websearch_interception/handler.py | 157 +++++++-- .../chat/guardrail_translation/handler.py | 202 +++++++----- .../context_management/editors/compact.py | 213 ++++++++---- .../llms/runwayml/videos/transformation.py | 138 ++++---- .../llms/vertex_ai/files/transformation.py | 114 +++++-- litellm/proxy/db/tool_registry_writer.py | 197 ++++++++---- .../internal_user_endpoints.py | 302 +++++++++++++----- .../key_management_endpoints.py | 241 ++++++++++---- .../tool_management_endpoints.py | 213 ++++++++++-- litellm/proxy/prompts/prompt_endpoints.py | 161 ++++++---- .../response_polling/background_streaming.py | 117 +++++-- .../responses/file_search/emulated_handler.py | 286 +++++++++-------- litellm/responses/streaming_iterator.py | 227 +++++++++---- litellm/router_strategy/tag_based_routing.py | 174 ++++++---- ruff-strict-budget.json | 16 +- type-discipline-budget.json | 10 +- 19 files changed, 2261 insertions(+), 943 deletions(-) diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index 6e3cbdff9d0..c71ef7a0020 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -1,9 +1,9 @@ { "reportAny": { - "limit": 23919 + "limit": 21974 }, "reportArgumentType": { - "limit": 2580 + "limit": 2575 }, "reportAssignmentType": { "limit": 323 @@ -24,7 +24,7 @@ "limit": 19 }, "reportExplicitAny": { - "limit": 7573 + "limit": 7068 }, "reportFunctionMemberAccess": { "limit": 7 @@ -54,10 +54,10 @@ "limit": 0 }, "reportMissingParameterType": { - "limit": 5719 + "limit": 5697 }, "reportMissingTypeArgument": { - "limit": 15657 + "limit": 15627 }, "reportMissingTypeStubs": { "limit": 40 @@ -99,19 +99,19 @@ "limit": 0 }, "reportUnknownArgumentType": { - "limit": 44832 + "limit": 44549 }, "reportUnknownLambdaType": { - "limit": 113 + "limit": 112 }, "reportUnknownMemberType": { - "limit": 39269 + "limit": 39156 }, "reportUnknownParameterType": { - "limit": 19988 + "limit": 19951 }, "reportUnknownVariableType": { - "limit": 30923 + "limit": 30798 }, "reportUnnecessaryCast": { "limit": 118 @@ -123,7 +123,7 @@ "limit": 5 }, "reportUnnecessaryIsInstance": { - "limit": 853 + "limit": 852 }, "reportUntypedBaseClass": { "limit": 0 diff --git a/litellm/google_genai/adapters/transformation.py b/litellm/google_genai/adapters/transformation.py index 4f127f476c3..4c8e77d9feb 100644 --- a/litellm/google_genai/adapters/transformation.py +++ b/litellm/google_genai/adapters/transformation.py @@ -1,6 +1,9 @@ import json -from collections.abc import AsyncIterator, Iterator -from typing import Any, Final, cast +from collections.abc import AsyncIterator, Callable, Iterator, Mapping, Sequence +from types import MappingProxyType +from typing import Any, Final, TypeAlias, cast + +from typing_extensions import TypedDict from litellm import verbose_logger from litellm.litellm_core_utils.json_validation_rule import normalize_tool_schema @@ -9,7 +12,6 @@ from litellm.types.llms.openai import ( ChatCompletionAssistantMessage, ChatCompletionAssistantToolCall, ChatCompletionImageObject, - ChatCompletionRequest, ChatCompletionSystemMessage, ChatCompletionTextObject, ChatCompletionToolCallFunctionChunk, @@ -21,12 +23,79 @@ from litellm.types.llms.openai import ( from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import ( AdapterCompletionStreamWrapper, + ChatCompletionDeltaCustomToolCall, + ChatCompletionDeltaToolCall, + ChatCompletionMessageCustomToolCall, + ChatCompletionMessageToolCall, Choices, + Delta, + Function, + Message, ModelResponse, ModelResponseStream, StreamingChoices, ) +_JsonDict: TypeAlias = dict[str, object] +_JsonDictList: TypeAlias = list[_JsonDict] + + +class _ToolCallAccumulator(TypedDict): + name: str + arguments: str + + +class _GenAIFunctionCall(TypedDict): + name: str + args: Mapping[str, object] + + +class _GenAIPart(TypedDict, total=False): + text: str + functionCall: _GenAIFunctionCall + + +class _GenAIFunctionResponse(TypedDict, total=False): + name: str + response: object + + +class _GenAIRequestFunctionCall(TypedDict, total=False): + name: str + args: Mapping[str, object] + + +class _GenAIContentPart(TypedDict, total=False): + text: str + inline_data: Mapping[str, str] + functionResponse: _GenAIFunctionResponse + functionCall: _GenAIRequestFunctionCall + + +class _GenAIFunctionDeclaration(TypedDict, total=False): + name: str + description: str + parametersJsonSchema: object + + +class _GenAITool(TypedDict, total=False): + functionDeclarations: Sequence[_GenAIFunctionDeclaration] + + +class _GenAIFunctionCallingConfig(TypedDict, total=False): + mode: str + + +class _GenAIToolConfig(TypedDict, total=False): + functionCallingConfig: _GenAIFunctionCallingConfig + + +class _GenAISystemInstruction(TypedDict, total=False): + parts: Sequence[Mapping[str, str]] + + +_EMPTY_STR_MAPPING: Final[Mapping[str, str]] = MappingProxyType({}) + class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): """ @@ -35,12 +104,12 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): """ sent_first_chunk: bool = False - # State tracking for accumulating partial tool calls - accumulated_tool_calls: dict[str, dict[str, Any]] + _parse_accumulated_args: Callable[[str], Mapping[str, object]] = staticmethod(json.loads) - def __init__(self, completion_stream: Any): + def __init__(self, completion_stream: object): self.sent_first_chunk = False - self.accumulated_tool_calls = {} + # State tracking for accumulating partial tool calls + self.accumulated_tool_calls = dict[int, _ToolCallAccumulator]() self._returned_response = False super().__init__(completion_stream) @@ -85,7 +154,7 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): # After the stream is exhausted, check for any remaining accumulated tool calls if self.accumulated_tool_calls: try: - parts: Final = [] + parts: Final = list[_GenAIPart]() for ( tool_call_index, tool_call_data, @@ -93,8 +162,10 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): try: # For tool calls with no arguments, accumulated_args will be "", which is not valid JSON. # We default to an empty JSON object in this case. - parsed_args = json.loads(tool_call_data["arguments"] or "{}") - function_call_part = { + parsed_args: Mapping[str, object] = self._parse_accumulated_args( + tool_call_data["arguments"] or "{}" + ) + function_call_part: _GenAIPart = { "functionCall": { "name": tool_call_data["name"] or "undefined_tool_name", "args": parsed_args, @@ -172,14 +243,16 @@ class GoogleGenAIStreamWrapper(AdapterCompletionStreamWrapper): class GoogleGenAIAdapter: """Adapter for transforming Google GenAI generate_content requests to/from litellm.completion format""" + _parse_tool_call_args: Callable[[str], Mapping[str, object]] = staticmethod(json.loads) + def __init__(self) -> None: pass def translate_generate_content_to_completion( self, model: str, - contents: list[dict[str, Any]] | dict[str, Any], - config: dict[str, Any] | None = None, + contents: _JsonDictList | _JsonDict, + config: Mapping[str, object] | None = None, litellm_params: GenericLiteLLMParams | None = None, **kwargs, ) -> dict[str, Any]: @@ -211,7 +284,7 @@ class GoogleGenAIAdapter: messages: Final = self._transform_contents_to_messages(contents_list, system_instruction=system_instruction) # Create base request as dict (which is compatible with ChatCompletionRequest) - completion_request: Final[ChatCompletionRequest] = { + completion_request: Final[_JsonDict] = { "model": model, "messages": messages, } @@ -273,9 +346,9 @@ class GoogleGenAIAdapter: def _add_generic_litellm_params_to_request( self, - completion_request_dict: dict[str, Any], + completion_request_dict: _JsonDict, litellm_params: GenericLiteLLMParams | None = None, - ) -> dict: + ) -> _JsonDict: """Add generic litellm params to request. e.g add api_base, api_key, api_version, etc. Args: @@ -287,7 +360,7 @@ class GoogleGenAIAdapter: """ allowed_fields: Final = GenericLiteLLMParams.model_fields.keys() if litellm_params: - litellm_dict: Final = litellm_params.model_dump(exclude_none=True) + litellm_dict: Final[_JsonDict] = litellm_params.model_dump(exclude_none=True) for key, value in litellm_dict.items(): if key in allowed_fields: completion_request_dict[key] = value @@ -295,7 +368,7 @@ class GoogleGenAIAdapter: def translate_completion_output_params_streaming( self, - completion_stream: Any, + completion_stream: object, ) -> AsyncIterator[bytes] | None: """Transform streaming completion output to Google GenAI format""" google_genai_wrapper: Final = GoogleGenAIStreamWrapper(completion_stream=completion_stream) @@ -304,15 +377,15 @@ class GoogleGenAIAdapter: def _transform_google_genai_tools_to_openai( self, - tools: list[dict[str, Any]], + tools: Sequence[_GenAITool], ) -> list[ChatCompletionToolParam]: """Transform Google GenAI tools to OpenAI tools format""" - openai_tools: Final[list[dict[str, Any]]] = [] + openai_tools: Final = list[_JsonDict]() for tool in tools: if "functionDeclarations" in tool: for func_decl in tool["functionDeclarations"]: - function_chunk: dict[str, Any] = { + function_chunk: _JsonDict = { "name": func_decl.get("name", ""), } @@ -321,7 +394,7 @@ class GoogleGenAIAdapter: if "parametersJsonSchema" in func_decl: function_chunk["parameters"] = func_decl["parametersJsonSchema"] - openai_tool = {"type": "function", "function": function_chunk} + openai_tool: _JsonDict = {"type": "function", "function": function_chunk} openai_tools.append(openai_tool) # normalize the tool schemas @@ -331,7 +404,7 @@ class GoogleGenAIAdapter: def _transform_google_genai_tool_config_to_openai( self, - tool_config: dict[str, Any], + tool_config: _GenAIToolConfig, ) -> ChatCompletionToolChoiceValues | None: """Transform Google GenAI tool_config to OpenAI tool_choice""" function_calling_config: Final = tool_config.get("functionCallingConfig", {}) @@ -345,20 +418,20 @@ class GoogleGenAIAdapter: def _transform_contents_to_messages( self, contents: list[dict[str, Any]], - system_instruction: dict[str, Any] | None = None, + system_instruction: _GenAISystemInstruction | None = None, ) -> list[AllMessageValues]: """Transform Google GenAI contents to OpenAI messages format""" messages: Final[list[AllMessageValues]] = [] # Handle system instruction if system_instruction: - system_parts: Final = system_instruction.get("parts", []) + system_parts: Final[Sequence[Mapping[str, str]]] = system_instruction.get("parts", []) if system_parts and "text" in system_parts[0]: messages.append(ChatCompletionSystemMessage(role="system", content=system_parts[0]["text"])) for content in contents: role = content.get("role", "user") - parts = content.get("parts", []) + parts: Sequence[_GenAIContentPart | str | None] = content.get("parts", []) if role == "user": # Handle user messages with potential function responses @@ -461,7 +534,7 @@ class GoogleGenAIAdapter: def translate_completion_to_generate_content( self, response: ModelResponse, - ) -> dict[str, Any]: + ) -> _JsonDict: """ Transform litellm completion response to Google GenAI generate_content format @@ -484,13 +557,13 @@ class GoogleGenAIAdapter: parts = self._transform_openai_message_to_google_genai_parts(choice.message) else: # Fallback for generic choice objects - message_content = getattr(choice, "message", {}).get("content", "") or getattr(choice, "delta", {}).get( - "content", "" - ) + message_content: str = getattr(choice, "message", _EMPTY_STR_MAPPING).get("content", "") or getattr( + choice, "delta", _EMPTY_STR_MAPPING + ).get("content", "") parts = [{"text": message_content}] if message_content else [] # Create Google GenAI format response - generate_content_response: Final[dict[str, Any]] = { + generate_content_response: Final[_JsonDict] = { "candidates": [ { "content": {"parts": parts, "role": "model"}, @@ -524,7 +597,7 @@ class GoogleGenAIAdapter: self, response: ModelResponse | ModelResponseStream, wrapper: GoogleGenAIStreamWrapper, - ) -> dict[str, Any] | None: + ) -> Mapping[str, object] | None: """ Transform streaming litellm completion chunk to Google GenAI generate_content format @@ -548,10 +621,10 @@ class GoogleGenAIAdapter: parts = self._transform_openai_delta_to_google_genai_parts_with_accumulation(choice.delta, wrapper) else: parts = [] - finish_reason = getattr(choice, "finish_reason", None) + finish_reason: str | None = getattr(choice, "finish_reason", None) else: # Fallback for generic choice objects - message_content: Final = getattr(choice, "delta", {}).get("content", "") + message_content: Final[str] = getattr(choice, "delta", _EMPTY_STR_MAPPING).get("content", "") parts = [{"text": message_content}] if message_content else [] finish_reason = getattr(choice, "finish_reason", None) @@ -560,7 +633,7 @@ class GoogleGenAIAdapter: return None # Create Google GenAI streaming format response - streaming_chunk: Final[dict[str, Any]] = { + streaming_chunk: Final[_JsonDict] = { "candidates": [ { "content": {"parts": parts, "role": "model"}, @@ -596,10 +669,10 @@ class GoogleGenAIAdapter: def _transform_openai_message_to_google_genai_parts( self, - message: Any, - ) -> list[dict[str, Any]]: + message: Message, + ) -> Sequence[_GenAIPart]: """Transform OpenAI message to Google GenAI parts format""" - parts: Final[list[dict[str, Any]]] = [] + parts: Final = list[_GenAIPart]() # Add text content if present if hasattr(message, "content") and message.content: @@ -607,16 +680,22 @@ class GoogleGenAIAdapter: # Add tool calls if present if hasattr(message, "tool_calls") and message.tool_calls: - for tool_call in message.tool_calls: - if hasattr(tool_call, "function") and tool_call.function: + tool_calls: Final[Sequence[ChatCompletionMessageToolCall | ChatCompletionMessageCustomToolCall]] = ( + message.tool_calls + ) + for tool_call in tool_calls: + function: Function | None = getattr(tool_call, "function", None) + if function: try: - args = json.loads(tool_call.function.arguments) if tool_call.function.arguments else {} + args: Mapping[str, object] = ( + self._parse_tool_call_args(function.arguments) if function.arguments else {} + ) except json.JSONDecodeError: args = {} - function_call_part = { + function_call_part: _GenAIPart = { "functionCall": { - "name": tool_call.function.name or "undefined_tool_name", + "name": function.name or "undefined_tool_name", "args": args, } } @@ -625,28 +704,30 @@ class GoogleGenAIAdapter: return parts if parts else [{"text": ""}] def _transform_openai_delta_to_google_genai_parts_with_accumulation( - self, delta: Any, wrapper: GoogleGenAIStreamWrapper - ) -> list[dict[str, Any]]: + self, delta: Delta, wrapper: GoogleGenAIStreamWrapper + ) -> Sequence[_GenAIPart]: """Transforms OpenAI delta to Google GenAI parts, accumulating streaming tool calls.""" # 1. Initialize wrapper state if it doesn't exist if not hasattr(wrapper, "accumulated_tool_calls"): wrapper.accumulated_tool_calls = {} - parts: Final[list[dict[str, Any]]] = [] + parts: Final = list[_GenAIPart]() if hasattr(delta, "content") and delta.content: parts.append({"text": delta.content}) # 2. Ensure tool_calls is iterable - tool_calls: Final = delta.tool_calls or [] + tool_calls: Final[Sequence[ChatCompletionDeltaToolCall | ChatCompletionDeltaCustomToolCall]] = ( + delta.tool_calls or [] + ) for tool_call in tool_calls: if not hasattr(tool_call, "function"): continue # 3. Use `index` as the primary key for accumulation - tool_call_index = getattr(tool_call, "index", None) + tool_call_index: int | None = getattr(tool_call, "index", None) if tool_call_index is None: continue # Index is essential for tracking streaming tool calls @@ -658,8 +739,9 @@ class GoogleGenAIAdapter: } # Accumulate name and arguments - function_name = getattr(tool_call.function, "name", None) - args_chunk = getattr(tool_call.function, "arguments", None) + delta_function: Function | None = getattr(tool_call, "function", None) + function_name: str | None = getattr(delta_function, "name", None) + args_chunk: str | None = getattr(delta_function, "arguments", None) # Optimization: Skip chunks that have no new data if not function_name and not args_chunk: @@ -680,13 +762,13 @@ class GoogleGenAIAdapter: # 5. Attempt to parse arguments even if name hasn't arrived. try: # Attempt to parse the accumulated arguments string - parsed_args = json.loads(accumulated_args) + parsed_args: Mapping[str, object] = self._parse_tool_call_args(accumulated_args) # If parsing succeeds, but we don't have a name yet, wait. # The part will be created by a later chunk that brings the name. if accumulated_name: # If successful, create the part and clean up - function_call_part = {"functionCall": {"name": accumulated_name, "args": parsed_args}} + function_call_part: _GenAIPart = {"functionCall": {"name": accumulated_name, "args": parsed_args}} parts.append(function_call_part) # Remove the completed tool call from the accumulator @@ -714,7 +796,7 @@ class GoogleGenAIAdapter: return mapping.get(finish_reason, "STOP") - def _map_usage(self, usage: Any) -> dict[str, int]: + def _map_usage(self, usage: object) -> Mapping[str, int]: """Map OpenAI usage to Google GenAI usage format""" return { "promptTokenCount": getattr(usage, "prompt_tokens", 0) or 0, diff --git a/litellm/integrations/rubrik.py b/litellm/integrations/rubrik.py index 97e831f5822..c206849c86f 100644 --- a/litellm/integrations/rubrik.py +++ b/litellm/integrations/rubrik.py @@ -6,12 +6,14 @@ import random import time import uuid from collections import Counter -from collections.abc import Mapping, Sequence +from collections.abc import Awaitable, Mapping, Sequence from dataclasses import dataclass +from datetime import datetime from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Optional +from typing import TYPE_CHECKING, Final, Literal, Optional, Protocol, TypedDict, overload import httpx +from typing_extensions import Never, Required from litellm._logging import verbose_logger from litellm.integrations.custom_batch_logger import CustomBatchLogger @@ -29,6 +31,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.types.guardrails import GuardrailEventHooks +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolCallChunk from litellm.types.utils import ( ChatCompletionMessageToolCall, Function, @@ -48,7 +51,105 @@ _WEBHOOK_PATH_PROMPT_MODERATION: Final = "/v1/before_prompt/openai/v1" _WEBHOOK_PATH_LOGGING_BATCH: Final = "/v1/litellm/batch" _MAX_QUEUE_SIZE: Final = 10_000 _DROP_WARNING_INTERVAL_SECONDS: Final = 60.0 -_EMPTY_MAPPING: Final[Mapping[str, Any]] = MappingProxyType({}) +_EMPTY_MAPPING: Final[Mapping[str, Never]] = MappingProxyType({}) + + +class _ModerationToolCall(TypedDict, total=False): + id: Required[str] + + +class _ModerationMessage(TypedDict, total=False): + content: str | None + tool_calls: Sequence[_ModerationToolCall] | None + + +class _ModerationChoice(TypedDict, total=False): + message: _ModerationMessage | None + + +class _ModerationResponse(TypedDict, total=False): + choices: Sequence[_ModerationChoice] + + +class _LogEventKwargs(TypedDict, total=False): + standard_logging_object: Required[StandardLoggingPayload] + litellm_call_id: str + + +class _HasCallId(Protocol): + def get(self, key: Literal["litellm_call_id"], /) -> str | None: ... + + +class _HasModelAttr(Protocol): + model: str | None + + +class _ResponseSource(Protocol): + def get(self, key: Literal["response"], /) -> "_HasModelAttr | None": ... + + +class _ModelSource(Protocol): + def get(self, key: Literal["model"], default: str, /) -> str: ... + + +class _FallbackSource(Protocol): + @overload + def get(self, key: Literal["start_time"], /) -> datetime | None: ... + @overload + def get(self, key: str, /) -> object | None: ... + + +class _RequestContextSource(Protocol): + @overload + def get(self, key: Literal["optional_params"], /) -> Mapping[str, object] | None: ... + @overload + def get(self, key: str, /) -> object | None: ... + def __contains__(self, key: object, /) -> bool: ... + def __getitem__(self, key: str, /) -> object: ... + + +class _ToolCallLike(Protocol): + id: str | None + type: str | None + function: Function + + +class _ModerationSourceToolCall(TypedDict, total=False): + function: Mapping[str, object] | None + + +class _ModerationSourceMessage(TypedDict, total=False): + role: str + function_call: Mapping[str, object] | None + tool_calls: Sequence[_ModerationSourceToolCall | None] | None + + +class _FlattenedModerationMessage(TypedDict): + role: str | None + content: str + + +class _CorrelatablePayload(TypedDict): + id: str + + +class _SystemPromptCarrier(TypedDict, total=False): + messages: object + + +class _BlockFailurePayload(TypedDict, total=False): + id: object + model: object + model_group: object + model_id: str + model_parameters: object + startTime: float | None + endTime: float | None + completionStartTime: float | None + messages: object + metadata: StandardLoggingUserAPIKeyMetadata + response: str + status: str class _MalformedToolBlockingResponseError(Exception): @@ -143,7 +244,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): else {"Content-Type": "application/json"} ) - self._periodic_flush_task: asyncio.Task[Any] | None = self._start_periodic_flush_task() + self._periodic_flush_task: asyncio.Task[None] | None = self._start_periodic_flush_task() @classmethod def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: @@ -191,7 +292,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): params={"timeout": httpx.Timeout(5.0, connect=2.0)}, ) - def _start_periodic_flush_task(self) -> asyncio.Task[Any] | None: + def _start_periodic_flush_task(self) -> asyncio.Task[None] | None: """Start the periodic flush task only when an event loop is already running.""" try: loop: Final = asyncio.get_running_loop() @@ -212,7 +313,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): Closing them here would close the shared connection pool for every other logger instance; let LiteLLM manage their lifecycle instead. """ - task: Final = getattr(self, "_periodic_flush_task", None) + task: Final[asyncio.Task[None] | None] = getattr(self, "_periodic_flush_task", None) if task is not None: task.cancel() @@ -253,7 +354,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): @staticmethod async def _guarded( - coro: Any, + coro: Awaitable[GenericGuardrailAPIInputs], inputs: GenericGuardrailAPIInputs, label: str, ) -> GenericGuardrailAPIInputs: @@ -371,7 +472,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): @staticmethod def _stash_block_context( logging_obj: Optional["LiteLLMLoggingObj"], - request_data: dict, + request_data: dict[str, object], ) -> None: """Stash signals so the deferred success-event skips this request and ``async_post_call_failure_hook`` can build the failure payload. @@ -400,12 +501,16 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): request_data["_rubrik_logging_obj"] = logging_obj @staticmethod - def _normalize_tool_calls(tool_calls: Any) -> tuple[ChatCompletionMessageToolCall, ...]: + def _normalize_tool_calls( + tool_calls: Sequence[ChatCompletionToolCallChunk | ChatCompletionMessageToolCall | _ToolCallLike], + ) -> tuple[ChatCompletionMessageToolCall, ...]: """Convert tool_calls from inputs to ChatCompletionMessageToolCall objects.""" return tuple(RubrikLogger._normalize_tool_call(tc) for tc in tool_calls) @staticmethod - def _normalize_tool_call(tc: Any) -> ChatCompletionMessageToolCall: + def _normalize_tool_call( + tc: ChatCompletionToolCallChunk | ChatCompletionMessageToolCall | _ToolCallLike, + ) -> ChatCompletionMessageToolCall: if isinstance(tc, ChatCompletionMessageToolCall): return tc if isinstance(tc, dict): @@ -427,7 +532,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): raise TypeError(f"Cannot normalize tool_call of type {type(tc).__name__}: {tc!r}") @staticmethod - def _join_texts(texts: Any) -> str: + def _join_texts(texts: Sequence[str] | None) -> str: """Join response text segments into the single content string the webhook evaluates. Empty when there is no assistant text.""" if not texts: @@ -439,19 +544,22 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): tool_calls: Sequence[ChatCompletionMessageToolCall], content: str, request_id: str | None, - ) -> Mapping[str, Any]: + ) -> Mapping[str, object]: """Build an OpenAI ChatCompletion-format dict (assistant text + tool calls) for the after_completion webhook. ``content`` is sent so the webhook can moderate the response text; ``None`` when the assistant produced no text (tool-call-only response). """ - message: Final[dict[str, Any]] = { + message: Final[Mapping[str, object]] = { "role": "assistant", "content": content or None, + **( + {"tool_calls": tuple(tc.model_dump(exclude_none=True) for tc in tool_calls)} + if tool_calls + else _EMPTY_MAPPING + ), } - if tool_calls: - message["tool_calls"] = tuple(tc.model_dump(exclude_none=True) for tc in tool_calls) return { "id": request_id or f"chatcmpl-{uuid.uuid4()}", "object": "chat.completion", @@ -467,7 +575,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): } @staticmethod - def _flatten_messages_for_moderation(messages: Any) -> tuple[Mapping[str, Any], ...]: + def _flatten_messages_for_moderation( + messages: Sequence[AllMessageValues | None] | None, + ) -> tuple[_FlattenedModerationMessage, ...]: """Collapse each message's content to a plain string for the webhook. litellm normalizes Anthropic ``/v1/messages`` requests to OpenAI shape, @@ -488,7 +598,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): ) @staticmethod - def _moderation_text_parts(message: Mapping[str, Any]) -> tuple[str, ...]: + def _moderation_text_parts(message: _ModerationSourceMessage) -> tuple[str, ...]: """Every attacker-controlled text segment of a message: its content plus the arguments of any tool call or deprecated function call.""" fc: Final = message.get("function_call") @@ -506,8 +616,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): @staticmethod def _build_prompt_moderation_payload( inputs: GenericGuardrailAPIInputs, - request_data: Mapping[str, Any], - ) -> Mapping[str, Any]: + request_data: Mapping[str, object], + ) -> Mapping[str, object]: """Build the bare OpenAI request the before_prompt webhook consumes. Unlike the after_completion envelope, this endpoint takes a raw OpenAI @@ -516,16 +626,8 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): ``/v1/messages`` requests too. Optional fields are sent only when present so the payload stays clean. """ - payload: Final[dict[str, Any]] = { - "model": inputs.get("model") or request_data.get("model") or "", - "messages": RubrikLogger._flatten_messages_for_moderation(inputs.get("structured_messages")), - } tools: Final = inputs.get("tools") - if tools is not None: - payload["tools"] = tools user: Final = request_data.get("user") - if user: - payload["user"] = user # Fall back to litellm_call_id, the stable cross-provider join key the # response/tool path uses (see _correlation_id). LiteLLM does not # populate request_data["correlation_key"]; it carries litellm_call_id. @@ -533,15 +635,19 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): # when correlation_key is empty, so without this the block fires but no # log is ever written. An explicit correlation_key still wins. correlation_key: Final = request_data.get("correlation_key") or request_data.get("litellm_call_id") - if correlation_key: - payload["correlation_key"] = correlation_key - return payload + return { + "model": inputs.get("model") or request_data.get("model") or "", + "messages": RubrikLogger._flatten_messages_for_moderation(inputs.get("structured_messages")), + **({"tools": tools} if tools is not None else _EMPTY_MAPPING), + **({"user": user} if user else _EMPTY_MAPPING), + **({"correlation_key": correlation_key} if correlation_key else _EMPTY_MAPPING), + } @staticmethod def _extract_request_data( - call_details: Mapping[str, Any], - request_data: Mapping[str, Any] | None, - ) -> Mapping[str, Any]: + call_details: _RequestContextSource, + request_data: _RequestContextSource | None, + ) -> Mapping[str, object]: """Extract original request data from model_call_details for the response moderation service envelope. @@ -576,7 +682,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): } @staticmethod - def _sanitize_proxy_server_request(proxy_server_request: Any) -> Any: + def _sanitize_proxy_server_request(proxy_server_request: Mapping[str, object] | str | None) -> object: """Allowlist only routing fields (``url``, ``method``) when forwarding ``proxy_server_request`` to an external webhook, dropping inbound ``headers`` (Authorization, Cookie, x-api-key, ...) and the raw @@ -586,7 +692,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): return {key: proxy_server_request[key] for key in ("url", "method") if key in proxy_server_request} @staticmethod - def _resolve_model(request_data: Mapping[str, Any], call_details: Mapping[str, Any]) -> str: + def _resolve_model(request_data: _ResponseSource, call_details: _ModelSource) -> str: """Get the model name for the ModifyResponseException.""" response: Final = request_data.get("response") if response and hasattr(response, "model"): @@ -596,7 +702,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): # -- Logging hooks --------------------------------------------------------- @staticmethod - def _correlation_id(call_details: Mapping[str, Any], request_data: Mapping[str, Any] | None = None) -> str | None: + def _correlation_id( + call_details: _HasCallId | _LogEventKwargs, request_data: _HasCallId | None = None + ) -> str | None: """The id that joins a blocked request's two S3 logs by filename: the moderation (``_blocking``) log and the failure (response) log. @@ -610,7 +718,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): return call_details.get("litellm_call_id") or (request_data or _EMPTY_MAPPING).get("litellm_call_id") @classmethod - def _apply_correlation_id(cls, payload: dict[str, Any], source: Mapping[str, Any]) -> None: + def _apply_correlation_id(cls, payload: _CorrelatablePayload, source: _HasCallId | _LogEventKwargs) -> None: """Pin ``payload["id"]`` to ``litellm_call_id`` in place so this log shares its S3 filename id with the moderation (``_blocking``) and failure logs for the same request -- for every provider. @@ -630,7 +738,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): payload["id"] = correlated @staticmethod - def _prepend_system_prompt(payload: dict[str, Any], source: Mapping[str, Any]) -> None: + def _prepend_system_prompt(payload: _SystemPromptCarrier, source: Mapping[str, object]) -> None: """Prepend ``source["system"]`` onto ``payload["messages"]``. Builds a NEW messages list rather than mutating ``payload["messages"]`` @@ -658,7 +766,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): exc_info=True, ) - async def _prepare_log_payload(self, kwargs: Mapping[str, Any], event_type: str) -> StandardLoggingPayload | None: + async def _prepare_log_payload(self, kwargs: _LogEventKwargs, event_type: str) -> StandardLoggingPayload | None: """Shared logic for success logging (sampled).""" if random.random() > self.sampling_rate: verbose_logger.debug("Skipping Rubrik %s logging (sampling_rate=%s)", event_type, self.sampling_rate) @@ -667,12 +775,12 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): # Deep-copy so mutations don't affect other callbacks sharing this object standard_logging_payload: Final[StandardLoggingPayload] = safe_deep_copy(kwargs["standard_logging_object"]) - self._apply_correlation_id(standard_logging_payload, kwargs) # pyright: ignore[reportArgumentType] # StandardLoggingPayload is dict[str,Any] at runtime + self._apply_correlation_id(standard_logging_payload, kwargs) self._prepend_system_prompt(standard_logging_payload, kwargs) # pyright: ignore[reportArgumentType] # StandardLoggingPayload is dict[str,Any] at runtime return standard_logging_payload - async def _append_and_maybe_flush(self, payload) -> None: + async def _append_and_maybe_flush(self, payload: Mapping[str, object]) -> None: self._ensure_periodic_flush_task() self.log_queue.append(payload) self._enforce_max_queue_size() @@ -697,7 +805,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): self._dropped_since_warning = 0 self._last_drop_warning_time = now - async def _enqueue_log_event(self, kwargs: Mapping[str, Any], event_type: str): + async def _enqueue_log_event(self, kwargs: _LogEventKwargs, event_type: str): try: payload: Final = await self._prepare_log_payload(kwargs, event_type) if payload is None: @@ -818,7 +926,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): logging_obj: "LiteLLMLoggingObj", exception: "ModifyResponseException", user_api_key_dict: "UserAPIKeyAuth", - ) -> StandardLoggingPayload: + ) -> _BlockFailurePayload: """Build a failure-style payload using the exception text as response. Blocked-tool events are security-relevant and **bypass sampling**: @@ -860,9 +968,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): call_details: Final = logging_obj.model_call_details exception_text: Final = f"{type(exception).__name__}: {exception.message}" - base: Final = call_details.get("standard_logging_object") + base: Final[StandardLoggingPayload | None] = call_details.get("standard_logging_object") if base is not None: - payload: dict = safe_deep_copy(base) + payload: _BlockFailurePayload = self._copy_block_payload_base(base) else: verbose_logger.debug( "Rubrik: standard_logging_object not yet on model_call_details " @@ -884,6 +992,10 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): return payload + @staticmethod + def _copy_block_payload_base(base: StandardLoggingPayload) -> _BlockFailurePayload: + return safe_deep_copy(base) + @staticmethod def _caller_metadata(user_api_key_dict: "UserAPIKeyAuth") -> StandardLoggingUserAPIKeyMetadata: """Identify the caller whose request was blocked. @@ -906,9 +1018,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): @classmethod def _build_fallback_payload( cls, - call_details: Mapping[str, Any], + call_details: _FallbackSource, user_api_key_dict: "UserAPIKeyAuth", - ) -> dict[str, Any]: + ) -> _BlockFailurePayload: # Convert datetime to a Unix float so json.dumps can serialize it. # httpx's json= parameter uses stdlib json.dumps with no custom encoder. _raw_start: Final = call_details.get("start_time") @@ -942,7 +1054,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): response: Final = await self.async_httpx_client.post( url=self.logging_endpoint, json=data, - headers=self._headers, + headers=dict(self._headers), ) response.raise_for_status() except httpx.HTTPStatusError as e: @@ -996,7 +1108,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): # -- Webhook services ------------------------------------------------------ - async def _post_json(self, endpoint: str, payload: Mapping[str, Any], service_name: str) -> Mapping[str, Any]: + async def _post_json(self, endpoint: str, payload: Mapping[str, object], service_name: str) -> _ModerationResponse: """POST ``payload`` to a Rubrik webhook and return its dict response. Raises: @@ -1006,11 +1118,11 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): verbose_logger.debug("Sending request to %s: %s", service_name, endpoint) http_response: Final = await self.moderation_client.post( endpoint, - json=payload, - headers=self._headers, + json=dict(payload), + headers=dict(self._headers), ) http_response.raise_for_status() - result: Final = http_response.json() + result: Final[_ModerationResponse | None] = http_response.json() if not isinstance(result, dict): raise TypeError( f"{service_name} returned non-dict JSON " @@ -1021,9 +1133,9 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): async def _post_to_response_moderation_endpoint( self, - response_data: Mapping[str, Any], - request_data: Mapping[str, Any], - ) -> Mapping[str, Any]: + response_data: Mapping[str, object], + request_data: Mapping[str, object], + ) -> _ModerationResponse: """Post the ``{request, response}`` envelope to the after_completion webhook and return its (possibly rewritten) response. @@ -1039,7 +1151,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): "Response moderation service", ) - async def _post_to_prompt_moderation_endpoint(self, payload: Mapping[str, Any]) -> Mapping[str, Any]: + async def _post_to_prompt_moderation_endpoint(self, payload: Mapping[str, object]) -> _ModerationResponse: """Post a bare OpenAI request to the before_prompt webhook. Returns ``{}`` (passthrough) or a synthetic chat.completion (block). @@ -1047,7 +1159,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): return await self._post_json(self.prompt_moderation_endpoint, payload, "Prompt moderation service") @staticmethod - def _extract_prompt_refusal(service_response: Mapping[str, Any]) -> str | None: + def _extract_prompt_refusal(service_response: _ModerationResponse) -> str | None: """Return the refusal text when the prompt was blocked, else None. The before_prompt webhook returns ``{}`` (passthrough) or a synthetic @@ -1063,7 +1175,7 @@ class RubrikLogger(CustomGuardrail, CustomBatchLogger): @staticmethod def _extract_response_block( - service_response: Mapping[str, Any], + service_response: _ModerationResponse, all_tool_calls: Sequence[ChatCompletionMessageToolCall], sent_content: str, ) -> BlockedResponseResult | None: diff --git a/litellm/integrations/websearch_interception/handler.py b/litellm/integrations/websearch_interception/handler.py index 972ae1d9856..edd3fdb8c61 100644 --- a/litellm/integrations/websearch_interception/handler.py +++ b/litellm/integrations/websearch_interception/handler.py @@ -10,7 +10,7 @@ import asyncio import math import uuid from collections.abc import AsyncIterator, Mapping, Sequence -from typing import TYPE_CHECKING, Any, Final, TypedDict, cast +from typing import TYPE_CHECKING, Any, Final, Literal, Never, TypedDict, TypeVar, cast import litellm from litellm._logging import verbose_logger @@ -41,7 +41,13 @@ from litellm.types.integrations.websearch_interception import ( AnthropicServerToolUseBlock, WebSearchInterceptionConfig, ) -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.anthropic import AnthropicThinkingParam +from litellm.types.llms.openai import ( + AllMessageValues, + ChatCompletionAudioParam, + ChatCompletionPredictionContentParam, + OpenAIWebSearchOptions, +) from litellm.types.utils import ( AgenticLoopParams, CallTypes, @@ -51,6 +57,8 @@ from litellm.types.utils import ( from litellm.utils import ProviderConfigManager if TYPE_CHECKING: + from aiohttp import ClientSession + from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.base_llm.anthropic_messages.transformation import ( BaseAnthropicMessagesConfig, @@ -72,6 +80,8 @@ WEBSEARCH_EMIT_NATIVE_BLOCKS_KEY: Final = "_websearch_interception_emit_native_b # ``web_search_tool_result`` blocks to inject into the final response. WEBSEARCH_NATIVE_BLOCKS_METADATA_KEY: Final = "websearch_native_blocks" +_ResponseT = TypeVar("_ResponseT") + class _PlanMetadataView(TypedDict): websearch_native_blocks: Sequence[Mapping[str, object]] | None @@ -85,9 +95,96 @@ class _WebSearchSettingsView(TypedDict): websearch_interception_params: WebSearchInterceptionConfig +class _SearchToolLitellmParams(TypedDict, total=False): + search_provider: str | None + + class _SearchToolConfig(TypedDict, total=False): search_tool_name: str - litellm_params: Mapping[str, object] | None + litellm_params: _SearchToolLitellmParams | None + + +class _LitellmParamsProviderView(TypedDict, total=False): + custom_llm_provider: str + + +class _DeploymentCallKwargsView(TypedDict): + custom_llm_provider: str + litellm_params: _LitellmParamsProviderView + model: str + + +class _AcreateNamedParams(TypedDict, total=False): + metadata: Never + stop_sequences: Never + stream: bool | None + system: str | None + temperature: float | None + thinking: Never + tool_choice: Never + tools: Never + top_k: int | None + top_p: float | None + container: Never + + +class _AsearchNamedParams(TypedDict, total=False): + max_results: int | None + search_domain_filter: Never + max_tokens_per_page: int | None + country: str | None + api_key: str | None + api_base: str | None + timeout: float | None + extra_headers: Never + + +class _AcompletionNamedParams(TypedDict, total=False): + functions: Never + function_call: str | None + timeout: float | None + temperature: float | None + top_p: float | None + n: int | None + stream: bool | None + stream_options: Never + stop: Never + max_tokens: int | None + max_completion_tokens: int | None + modalities: Never + prediction: ChatCompletionPredictionContentParam | None + audio: ChatCompletionAudioParam | None + presence_penalty: float | None + frequency_penalty: float | None + logit_bias: Never + user: str | None + response_format: Never + seed: int | None + tools: Never + tool_choice: Never + parallel_tool_calls: bool | None + logprobs: bool | None + top_logprobs: int | None + deployment_id: str | None + reasoning_effort: Literal["none", "minimal", "low", "medium", "high", "xhigh", "default"] | None + verbosity: Literal["low", "medium", "high"] | None + safety_identifier: str | None + service_tier: str | None + base_url: str | None + api_version: str | None + api_key: str | None + model_list: Never + extra_headers: Never + thinking: AnthropicThinkingParam | None + web_search_options: OpenAIWebSearchOptions | None + include_server_side_tool_invocations: bool | None + shared_session: "ClientSession | None" + enable_json_schema_validation: bool | None + + +_NO_ACREATE_NAMED: Final[_AcreateNamedParams] = {} +_NO_ASEARCH_NAMED: Final[_AsearchNamedParams] = {} +_NO_ACOMPLETION_NAMED: Final[_AcompletionNamedParams] = {} class WebSearchInterceptionLogger(CustomLogger): @@ -275,12 +372,17 @@ class WebSearchInterceptionLogger(CustomLogger): """ # Check if this is for an enabled provider # Try top-level kwargs first, then nested litellm_params, then derive from model name - custom_llm_provider = kwargs.get("custom_llm_provider", "") or kwargs.get("litellm_params", {}).get( + call_kwargs_view: Final[_DeploymentCallKwargsView] = { + "custom_llm_provider": kwargs.get("custom_llm_provider", ""), + "litellm_params": kwargs.get("litellm_params", {}), + "model": kwargs.get("model", ""), + } + custom_llm_provider = call_kwargs_view["custom_llm_provider"] or call_kwargs_view["litellm_params"].get( "custom_llm_provider", "" ) if not custom_llm_provider: try: - _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=kwargs.get("model", "")) + _, custom_llm_provider, _, _ = litellm.get_llm_provider(model=call_kwargs_view["model"]) except Exception: custom_llm_provider = "" if custom_llm_provider not in self.enabled_providers: @@ -903,17 +1005,18 @@ class WebSearchInterceptionLogger(CustomLogger): ) @staticmethod - def _inject_native_blocks(response: Any, native_blocks: Sequence[Mapping[str, object]]) -> Any: + def _inject_native_blocks(response: _ResponseT, native_blocks: Sequence[Mapping[str, object]]) -> _ResponseT: """Prepend native blocks to response content, dict or object form.""" if not native_blocks: return response if isinstance(response, dict): - existing = response.get("content") or [] + existing: Sequence[object] = response.get("content") or [] response["content"] = list(native_blocks) + list(existing) return response existing = getattr(response, "content", None) or [] + content_attribute: Final = "content" try: - response.content = list(native_blocks) + list(existing) + setattr(response, content_attribute, list(native_blocks) + list(existing)) except (AttributeError, TypeError): # Object refused write — fall through and leave the response # untouched rather than crash the request. @@ -1169,10 +1272,10 @@ class WebSearchInterceptionLogger(CustomLogger): messages: list[dict], tool_calls: list[dict], thinking_blocks: list[dict], - anthropic_messages_optional_request_params: dict, + anthropic_messages_optional_request_params: Mapping[str, object], logging_obj: "LiteLLMLoggingObj | None", stream: bool, - kwargs: dict, + kwargs: Mapping[str, object], ) -> "AnthropicMessagesResponse | AsyncIterator[object]": """Legacy path: execute search + build patch + run follow-up call.""" request_patch, structured_results = await self._build_anthropic_request_patch( @@ -1180,9 +1283,9 @@ class WebSearchInterceptionLogger(CustomLogger): messages=messages, tool_calls=tool_calls, thinking_blocks=thinking_blocks, - anthropic_messages_optional_request_params=anthropic_messages_optional_request_params, + anthropic_messages_optional_request_params=dict[str, object](anthropic_messages_optional_request_params), logging_obj=logging_obj, - kwargs=kwargs, + kwargs=dict[str, object](kwargs), ) if request_patch.messages is None: raise ValueError("WebSearchInterception: missing follow-up messages") @@ -1197,12 +1300,14 @@ class WebSearchInterceptionLogger(CustomLogger): if max_tokens is None: max_tokens = cast(int, kwargs.get("max_tokens", 1024)) + patch_kwargs: Final = dict[str, object](request_patch.kwargs) response: AnthropicMessagesResponse | AsyncIterator[object] = await anthropic_messages.acreate( max_tokens=max_tokens, messages=request_patch.messages, model=request_patch.model or model, + **_NO_ACREATE_NAMED, **optional_params, - **request_patch.kwargs, + **patch_kwargs, ) # Legacy path: the new path goes through the typed plan + core @@ -1344,12 +1449,13 @@ class WebSearchInterceptionLogger(CustomLogger): search_tool: Final = self._select_search_tool_from_router(llm_router=llm_router) search_provider: str | None = None - search_litellm_params: dict[str, Any] = {} + search_litellm_params: Mapping[str, object] = {} search_tool_name: Final = self._selected_search_tool_name(search_tool=search_tool) if search_tool is not None: await self._authorize_search_tool(search_tool=search_tool, kwargs=kwargs) - search_litellm_params = dict(search_tool.get("litellm_params", {}) or {}) - search_provider = search_litellm_params.get("search_provider") + tool_params: Final[_SearchToolLitellmParams] = search_tool.get("litellm_params", {}) or {} + search_litellm_params = dict[str, object](tool_params) + search_provider = tool_params.get("search_provider") # Fallback to perplexity if no router or no search tools configured if not search_provider: @@ -1377,12 +1483,15 @@ class WebSearchInterceptionLogger(CustomLogger): if key != "search_provider" and value is not None } result: Final = ( - await litellm.asearch(query=query, search_provider=search_provider, **search_kwargs) + await litellm.asearch( + query=query, search_provider=search_provider, **_NO_ASEARCH_NAMED, **search_kwargs + ) if search_metadata is None else await litellm.asearch( query=query, search_provider=search_provider, litellm_metadata=search_metadata, + **_NO_ASEARCH_NAMED, **search_kwargs, ) ) @@ -1422,7 +1531,7 @@ class WebSearchInterceptionLogger(CustomLogger): valid_token=user_api_key_auth, ) - team_id: Final = getattr(user_api_key_auth, "team_id", None) + team_id: Final[str | None] = getattr(user_api_key_auth, "team_id", None) if team_id: from litellm.proxy.proxy_server import ( prisma_client, @@ -1537,10 +1646,10 @@ class WebSearchInterceptionLogger(CustomLogger): model: str, messages: list[dict], tool_calls: list[dict], - optional_params: dict, + optional_params: Mapping[str, object], logging_obj: "LiteLLMLoggingObj | None", stream: bool, - kwargs: dict, + kwargs: Mapping[str, object], response_format: str = "openai", ) -> "ModelResponse | CustomStreamWrapper": """Legacy path: execute search + build patch + run follow-up call.""" @@ -1548,8 +1657,8 @@ class WebSearchInterceptionLogger(CustomLogger): model=model, messages=messages, tool_calls=tool_calls, - optional_params=optional_params, - kwargs=kwargs, + optional_params=dict[str, object](optional_params), + kwargs=dict[str, object](kwargs), response_format=response_format, ) if request_patch.messages is None: @@ -1557,11 +1666,13 @@ class WebSearchInterceptionLogger(CustomLogger): params: Final = dict(optional_params) params.update(request_patch.optional_params) params.pop("tool_choice", None) + patch_kwargs: Final = dict[str, object](request_patch.kwargs) return await litellm.acompletion( model=request_patch.model or model, messages=request_patch.messages, + **_NO_ACOMPLETION_NAMED, **params, - **request_patch.kwargs, + **patch_kwargs, ) async def _build_chat_completion_request_patch( diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index e4a4d23b438..47dfe8de292 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -13,12 +13,12 @@ Pattern Overview: """ import json -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from copy import deepcopy from dataclasses import dataclass -from typing import TYPE_CHECKING, Any, Final, cast +from typing import TYPE_CHECKING, Any, Final, Protocol, cast, overload, runtime_checkable -from typing_extensions import assert_never +from typing_extensions import TypedDict, assert_never from litellm._logging import verbose_proxy_logger from litellm.llms.anthropic.chat.transformation import AnthropicConfig @@ -61,6 +61,7 @@ if TYPE_CHECKING: ModifyResponseException, ) from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.proxy._types import UserAPIKeyAuth from litellm.types.llms.anthropic_messages.anthropic_response import ( AnthropicMessagesResponse, ) @@ -95,6 +96,48 @@ InputWriteBackTarget = ( ) +class _SSEDelta(TypedDict, total=False): + type: str + text: str + stop_reason: str | None + + +class _SSEEventData(TypedDict, total=False): + delta: _SSEDelta + + +def _as_str_mapping(value: Mapping[str, object]) -> Mapping[str, object]: + return value + + +def _content_block_at(blocks: Sequence[object], index: int) -> object: + return blocks[index] + + +@runtime_checkable +class _ModelDumpBlock(Protocol): + def model_dump(self) -> Mapping[str, object]: ... + + +@runtime_checkable +class _TextAttrBlock(Protocol): + text: str + + +class _WritableMessage(Protocol): + @overload + def get(self, key: str, /) -> object | None: ... + + @overload + def get(self, key: str, default: object, /) -> object: ... + + def __setitem__(self, key: str, value: object, /) -> None: ... + + +def _as_writable(value: _WritableMessage) -> _WritableMessage: + return value + + @dataclass(frozen=True, slots=True) class ScannedText: text: str @@ -123,7 +166,7 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _build_streaming_usage_response( - responses_so_far: list[Any], + responses_so_far: Sequence[object], request_data: dict | None, ) -> ModelResponse | None: chunks: Final = tuple(response for response in responses_so_far if isinstance(response, (str, bytes))) @@ -141,7 +184,7 @@ class AnthropicMessagesHandler(BaseTranslation): self, exc: "ModifyResponseException", stream_started: bool = False, - responses_so_far: list[Any] | None = None, + responses_so_far: Sequence[object] | None = None, ) -> list[bytes]: """ Build an Anthropic SSE sequence delivering the guardrail block message @@ -159,7 +202,7 @@ class AnthropicMessagesHandler(BaseTranslation): would make Anthropic clients reject the stream. """ if stream_started: - return self._block_continuation_chunks(exc, responses_so_far or []) + return list(self._block_continuation_chunks(exc, responses_so_far or [])) return self._standalone_block_chunks(exc) def _standalone_block_chunks(self, exc: "ModifyResponseException") -> list[bytes]: @@ -184,7 +227,9 @@ class AnthropicMessagesHandler(BaseTranslation): ) return list(FakeAnthropicMessagesStreamIterator(response=block_response)) - def _block_continuation_chunks(self, exc: "ModifyResponseException", responses_so_far: list[Any]) -> list[bytes]: + def _block_continuation_chunks( + self, exc: "ModifyResponseException", responses_so_far: Sequence[object] + ) -> Sequence[bytes]: """Continue an already-started message: close the open content block, append the block message as a new text block, then end the message -- without a second message_start.""" @@ -234,7 +279,7 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _content_block_state( - responses_so_far: list[Any], + responses_so_far: Sequence[object], ) -> tuple[int | None, int | None]: """From the SSE chunks already sent to the client, return (open content-block index or None, highest content-block index seen or None). @@ -260,7 +305,20 @@ class AnthropicMessagesHandler(BaseTranslation): return open_index, max_index @staticmethod - def _iter_sse_events(item: Any) -> list[dict]: + def _parse_sse_data_line(raw_line: str) -> tuple[Mapping[str, object], ...]: + line: Final = raw_line.strip() + if not line.startswith("data:"): + return () + try: + parsed: Final[object] = json.loads(line[len("data:") :].strip()) + except json.JSONDecodeError: + return () + if not isinstance(parsed, dict): + return () + return (_as_str_mapping(parsed),) + + @staticmethod + def _iter_sse_events(item: object) -> Sequence[Mapping[str, object]]: """Yield the event-data dicts in one stream chunk. Handles both formats this stream can carry (see @@ -268,22 +326,15 @@ class AnthropicMessagesHandler(BaseTranslation): several events separated by a blank line -- and an already-parsed event ``dict``.""" if isinstance(item, dict): - return [item] + return (_as_str_mapping(item),) if not isinstance(item, (bytes, bytearray)): - return [] - events: Final[list[dict]] = [] - for block in item.decode("utf-8", errors="replace").split("\n\n"): - for line in block.split("\n"): - line = line.strip() - if not line.startswith("data:"): - continue - try: - parsed = json.loads(line[len("data:") :].strip()) - except json.JSONDecodeError: - continue - if isinstance(parsed, dict): - events.append(parsed) - return events + return () + return tuple( + event + for block in item.decode("utf-8", errors="replace").split("\n\n") + for line in block.split("\n") + for event in AnthropicMessagesHandler._parse_sse_data_line(line) + ) def _translate_to_openai(self, data: dict) -> ChatCompletionRequest: """Translate Anthropic request to OpenAI chat completion format.""" @@ -315,8 +366,8 @@ class AnthropicMessagesHandler(BaseTranslation): self, data: dict, guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, - ) -> Any: + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, + ) -> Mapping[str, object]: """ Process input messages by applying guardrails to text content. """ @@ -467,8 +518,8 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _openai_system_message_to_anthropic( - message: dict[str, Any], - ) -> dict[str, Any] | None: # mutable-ok: API message payload + message: Mapping[str, object], + ) -> dict[str, object] | None: # mutable-ok: API message payload """Convert an OpenAI system message to the client's Anthropic-shaped entry.""" content: Final = message.get("content") if isinstance(content, str): @@ -477,14 +528,14 @@ class AnthropicMessagesHandler(BaseTranslation): ) # mutable-ok: API message payload if not isinstance(content, list): return None - blocks: Final[list[dict[str, Any]]] = [] # mutable-ok: API message payload + blocks: Final[list[dict[str, object]]] = [] # mutable-ok: API message payload for block in content: if not isinstance(block, dict) or block.get("type") != "text": continue text = block.get("text") if not isinstance(text, str) or not text: continue - anthropic_block: dict[str, Any] = { # mutable-ok: API message payload + anthropic_block: dict[str, object] = { # mutable-ok: API message payload "type": "text", "text": text, } # mutable-ok: API message payload @@ -514,7 +565,7 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _defer_systems_inside_tool_exchanges( - structured_messages: list, # mutable-ok: API message payload + structured_messages: Sequence[Mapping[str, object]], ) -> list: """Hold a system row until the tool exchange around it completes so the call/result pair converts together.""" from litellm.litellm_core_utils.prompt_templates.factory import group_tool_exchanges @@ -602,7 +653,7 @@ class AnthropicMessagesHandler(BaseTranslation): @staticmethod def _extract_midturn_system_text( - message: dict[str, Any], # mutable-ok: API message payload + message: Mapping[str, object], msg_idx: int, ) -> ExtractedInput: """Match the adapter's filtering so positional guardrail write-back stays aligned.""" @@ -636,7 +687,7 @@ class AnthropicMessagesHandler(BaseTranslation): @classmethod def _extract_input_text_and_images( cls, - message: dict[str, Any], + message: Mapping[str, object], msg_idx: int, skip_system_message: bool = False, skip_tool_message: bool = False, @@ -696,7 +747,7 @@ class AnthropicMessagesHandler(BaseTranslation): if scan_only_tool_results: return EMPTY_EXTRACTED_INPUT - text_str: Final = content_item.get("text", None) + text_str: Final[str | None] = content_item.get("text") return ExtractedInput( scanned=( () if text_str is None else (ScannedText(text_str, ContentBlockTextTarget(msg_idx, content_idx)),) @@ -707,7 +758,7 @@ class AnthropicMessagesHandler(BaseTranslation): @classmethod def _extract_tool_result( cls, - content_item: Mapping[str, Any], + content_item: Mapping[str, object], msg_idx: int, content_idx: int, ) -> ExtractedInput: @@ -736,7 +787,7 @@ class AnthropicMessagesHandler(BaseTranslation): ) @staticmethod - def _image_sources(block: Mapping[str, Any]) -> tuple[str, ...]: + def _image_sources(block: Mapping[str, object]) -> tuple[str, ...]: source: Final = block.get("source") if not isinstance(source, Mapping): return () @@ -746,7 +797,7 @@ class AnthropicMessagesHandler(BaseTranslation): async def _apply_guardrail_responses_to_input( self, - messages: list[dict[str, Any]], + messages: Sequence[_WritableMessage], responses: list[str], scanned: tuple[ScannedText, ...], ) -> None: @@ -788,10 +839,10 @@ class AnthropicMessagesHandler(BaseTranslation): self, response: "AnthropicMessagesResponse", guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, - user_api_key_dict: Any | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, + user_api_key_dict: "UserAPIKeyAuth | None" = None, request_data: dict | None = None, - ) -> Any: + ) -> "AnthropicMessagesResponse": """ Process output response by applying guardrails to text content and tool calls. @@ -869,10 +920,10 @@ class AnthropicMessagesHandler(BaseTranslation): self, responses_so_far: list[Any], guardrail_to_apply: "CustomGuardrail", - litellm_logging_obj: Any | None = None, - user_api_key_dict: Any | None = None, + litellm_logging_obj: "LiteLLMLoggingObj | None" = None, + user_api_key_dict: "UserAPIKeyAuth | None" = None, request_data: dict | None = None, - ) -> list[Any]: + ) -> Sequence[object]: """ Process output streaming response by applying guardrails to text content. @@ -950,8 +1001,8 @@ class AnthropicMessagesHandler(BaseTranslation): def _prepare_request_data( self, request_data: dict | None, - response: Any, - user_api_key_dict: Any | None, + response: object, + user_api_key_dict: "UserAPIKeyAuth | None", key: str, ) -> dict: """Ensure request_data has the response/responses_so_far key and metadata.""" @@ -968,7 +1019,7 @@ class AnthropicMessagesHandler(BaseTranslation): return request_data @staticmethod - def _get_response_content(response: Any) -> list[Any]: + def _get_response_content(response: object) -> Sequence[object]: """Extract content list from a dict or object response.""" if isinstance(response, dict): return response.get("content", []) or [] @@ -978,7 +1029,7 @@ class AnthropicMessagesHandler(BaseTranslation): def _extract_from_content_blocks( self, - response_content: list[Any], + response_content: Sequence[object], texts_to_check: list[str], images_to_check: list[str], task_mappings: list[tuple[int, int | None]], @@ -986,21 +1037,10 @@ class AnthropicMessagesHandler(BaseTranslation): ) -> None: """Extract text, images, and tool calls from content blocks.""" for content_idx, content_block in enumerate(response_content): - block_dict: dict[str, Any] = {} - if isinstance(content_block, dict): - block_type = content_block.get("type") - block_dict = cast(dict[str, Any], content_block) - elif hasattr(content_block, "type"): - block_type = getattr(content_block, "type", None) - if hasattr(content_block, "model_dump"): - block_dict = content_block.model_dump() - else: - block_dict = { - "type": block_type, - "text": getattr(content_block, "text", None), - } - else: + fields = self._output_block_fields(content_block) + if fields is None: continue + block_type, block_dict = fields if block_type in ["text", "tool_use"]: self._extract_output_text_and_images( @@ -1012,12 +1052,27 @@ class AnthropicMessagesHandler(BaseTranslation): tool_calls_to_check=tool_calls_to_check, ) + @staticmethod + def _output_block_fields(content_block: object) -> "tuple[object, Mapping[str, object]] | None": + if isinstance(content_block, dict): + block_dict: Final = _as_str_mapping(content_block) + return block_dict.get("type"), block_dict + if not hasattr(content_block, "type"): + return None + block_type: Final = getattr(content_block, "type", None) + if isinstance(content_block, _ModelDumpBlock): + return block_type, content_block.model_dump() + return block_type, { + "type": block_type, + "text": getattr(content_block, "text", None), + } + @staticmethod def _build_guardrail_inputs( texts_to_check: list[str], images_to_check: list[str], tool_calls_to_check: list["ChatCompletionToolCallChunk"], - response: Any, + response: object, ) -> "GenericGuardrailAPIInputs": """Build GenericGuardrailAPIInputs with optional images, tool calls, model.""" inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) @@ -1034,7 +1089,7 @@ class AnthropicMessagesHandler(BaseTranslation): inputs["model"] = response_model return inputs - def get_streaming_string_so_far(self, responses_so_far: list[Any]) -> str: + def get_streaming_string_so_far(self, responses_so_far: Sequence[object]) -> str: """ Parse streaming responses and extract accumulated text content. @@ -1105,7 +1160,7 @@ class AnthropicMessagesHandler(BaseTranslation): # Only process content_block_delta events if event_type == "content_block_delta" and data_line: try: - data = json.loads(data_line) + data: _SSEEventData = json.loads(data_line) delta = data.get("delta", {}) if delta.get("type") == "text_delta": text += delta.get("text", "") @@ -1117,7 +1172,7 @@ class AnthropicMessagesHandler(BaseTranslation): return text - def _check_streaming_has_ended(self, responses_so_far: list[Any]) -> bool: + def _check_streaming_has_ended(self, responses_so_far: Sequence[object]) -> bool: """ Check if streaming response has ended by looking for non-null stop_reason. @@ -1168,7 +1223,7 @@ class AnthropicMessagesHandler(BaseTranslation): # Check for message_delta event with stop_reason if event_type == "message_delta" and data_line: try: - data = json.loads(data_line) + data: _SSEEventData = json.loads(data_line) delta = data.get("delta", {}) stop_reason = delta.get("stop_reason") if stop_reason is not None: @@ -1212,7 +1267,7 @@ class AnthropicMessagesHandler(BaseTranslation): def _extract_output_text_and_images( self, - content_block: dict[str, Any], + content_block: Mapping[str, object], content_idx: int, texts_to_check: list[str], images_to_check: list[str], @@ -1235,7 +1290,7 @@ class AnthropicMessagesHandler(BaseTranslation): task_mappings.append((content_idx, None)) # Extract tool calls - elif content_type == "tool_use": + elif content_type == "tool_use" and isinstance(content_block, dict): tool_call: Final = AnthropicConfig.convert_tool_use_to_openai_format( anthropic_tool_content=content_block, index=content_idx, @@ -1260,7 +1315,7 @@ class AnthropicMessagesHandler(BaseTranslation): content_idx = cast(int, mapping[0]) # Handle both dict and object responses - response_content: list[Any] = [] + response_content: Sequence[object] = [] if isinstance(response, dict): response_content = response.get("content", []) or [] elif hasattr(response, "content"): @@ -1276,14 +1331,15 @@ class AnthropicMessagesHandler(BaseTranslation): if content_idx >= len(response_content): continue - content_block = response_content[content_idx] + content_block = _content_block_at(response_content, content_idx) # Verify it's a text block and update the text field # Handle both dict and Pydantic object content blocks if isinstance(content_block, dict): - if content_block.get("type") == "text": - cast(dict[str, Any], content_block)["text"] = guardrail_response + block = _as_writable(content_block) + if block.get("type") == "text": + block["text"] = guardrail_response elif hasattr(content_block, "type") and getattr(content_block, "type", None) == "text": # Update Pydantic object's text attribute - if hasattr(content_block, "text"): + if isinstance(content_block, _TextAttrBlock): content_block.text = guardrail_response diff --git a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py index dbeac453791..d230b438086 100644 --- a/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py +++ b/litellm/llms/anthropic/experimental_pass_through/context_management/editors/compact.py @@ -13,8 +13,10 @@ Mirrors Anthropic's native ``compact_20260112`` for non-Anthropic providers: """ import re -from collections.abc import Mapping -from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Union, cast +from collections.abc import Awaitable, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, TypeVar, Union, cast + +from typing_extensions import NotRequired, TypedDict, Unpack import litellm from litellm._logging import verbose_logger @@ -27,11 +29,11 @@ from litellm.types.llms.anthropic import ( if TYPE_CHECKING: from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.hooks.parallel_request_limiter_v3 import RateLimitDescriptor, RateLimitResponse from litellm.router import Router from litellm.types.llms.anthropic import ( + AllAnthropicPassThroughMessageValues, AllAnthropicToolsValues, - AnthopicMessagesAssistantMessageParam, - AnthropicMessagesUserMessageParam, ) from litellm.types.llms.openai import ChatCompletionToolParam from litellm.types.utils import ModelResponse @@ -82,6 +84,69 @@ _PROPAGATED_METADATA_KEYS: Final = ( _SUMMARY_TAG_RE: Final = re.compile(r"(.*?)", re.IGNORECASE | re.DOTALL) +_MsgT: Final = TypeVar("_MsgT", bound=Mapping[str, object]) + + +def _as_object(value: object) -> object: + return value + + +def _is_tool_result_block(block: object) -> bool: + return isinstance(block, dict) and block.get("type") in ("tool_result",) + + +class _SummaryCallKwargs(TypedDict): + model: str + max_tokens: int + timeout: float + litellm_metadata: Mapping[str, object] + user: NotRequired[str] + allowed_model_region: NotRequired[str] + + +class _SummaryAcompletion(Protocol): + def __call__( + self, *, messages: Sequence[Mapping[str, object]], **kwargs: Unpack[_SummaryCallKwargs] + ) -> "Awaitable[ModelResponse | CustomStreamWrapper]": ... + + +class _CreateRateLimitDescriptors(Protocol): + def __call__( + self, + *, + user_api_key_dict: "UserAPIKeyAuth", + data: Mapping[str, str], + rpm_limit_type: object, + tpm_limit_type: object, + model_has_failures: bool, + ) -> "Sequence[RateLimitDescriptor]": ... + + +class _AddModelRateLimitDescriptor(Protocol): + def __call__( + self, + *, + user_api_key_dict: "UserAPIKeyAuth", + requested_model: str, + descriptors: "Sequence[RateLimitDescriptor]", + ) -> None: ... + + +class _CreateOrgRateLimitDescriptors(Protocol): + def __call__( + self, user_api_key_dict: "UserAPIKeyAuth", requested_model: str | None = None + ) -> "Sequence[RateLimitDescriptor]": ... + + +class _ShouldRateLimit(Protocol): + def __call__( + self, + *, + descriptors: "Sequence[RateLimitDescriptor]", + parent_otel_span: object, + read_only: bool, + ) -> "Awaitable[RateLimitResponse]": ... + def _read_summary_model_setting() -> str | None: """Look up the configured summarization model from proxy general_settings.""" @@ -157,11 +222,11 @@ async def _check_summary_model_access( return True key_models: Final = list(getattr(user_api_key_auth, "models", None) or []) - team_id: Final = getattr(user_api_key_auth, "team_id", None) + team_id: Final[str | None] = getattr(user_api_key_auth, "team_id", None) team_model_aliases: Final = getattr(user_api_key_auth, "team_model_aliases", None) team_models: Final = list(getattr(user_api_key_auth, "team_models", None) or []) - user_id: Final = getattr(user_api_key_auth, "user_id", None) - project_id: Final = getattr(user_api_key_auth, "project_id", None) + user_id: Final[str | None] = getattr(user_api_key_auth, "user_id", None) + project_id: Final[str | None] = getattr(user_api_key_auth, "project_id", None) checks: Final[tuple[tuple[Literal["key", "team"], list[str]], ...]] = ( ("key", key_models), @@ -347,7 +412,7 @@ async def _check_summary_model_budget( return False end_user_model_max_budget: Final = getattr(user_api_key_auth, "end_user_model_max_budget", None) - end_user_id: Final = getattr(user_api_key_auth, "end_user_id", None) + end_user_id: Final[str | None] = getattr(user_api_key_auth, "end_user_id", None) if isinstance(end_user_model_max_budget, dict) and end_user_model_max_budget and end_user_id is not None: try: await model_max_budget_limiter.is_end_user_within_model_budget( @@ -399,40 +464,57 @@ async def _check_summary_model_rate_limit( except Exception: return True - limiter: Final = getattr(proxy_logging_obj, "max_parallel_request_limiter", None) + limiter: Final[object] = getattr(proxy_logging_obj, "max_parallel_request_limiter", None) + should_rate_limit_check: Final[_ShouldRateLimit | None] = getattr(limiter, "should_rate_limit", None) + create_descriptors: Final[_CreateRateLimitDescriptors | None] = getattr( + limiter, "_create_rate_limit_descriptors", None + ) + add_team_descriptor: Final[_AddModelRateLimitDescriptor | None] = getattr( + limiter, "_add_team_model_rate_limit_descriptor_from_metadata", None + ) + add_project_descriptor: Final[_AddModelRateLimitDescriptor | None] = getattr( + limiter, "_add_project_model_rate_limit_descriptor_from_metadata", None + ) + create_org_descriptors: Final[_CreateOrgRateLimitDescriptors | None] = getattr( + limiter, "create_organization_rate_limit_descriptor", None + ) if ( limiter is None - or not hasattr(limiter, "should_rate_limit") - or not hasattr(limiter, "_create_rate_limit_descriptors") + or should_rate_limit_check is None + or create_descriptors is None + or add_team_descriptor is None + or add_project_descriptor is None + or create_org_descriptors is None ): return True try: - metadata: Final = getattr(user_api_key_auth, "metadata", None) or {} + metadata: Final[Mapping[str, object]] = getattr(user_api_key_auth, "metadata", None) or {} data: Final = {"model": summary_model} - descriptors: Final = limiter._create_rate_limit_descriptors( + base_descriptors: Final = create_descriptors( user_api_key_dict=user_api_key_auth, data=data, rpm_limit_type=metadata.get("rpm_limit_type"), tpm_limit_type=metadata.get("tpm_limit_type"), model_has_failures=False, ) - limiter._add_team_model_rate_limit_descriptor_from_metadata( + add_team_descriptor( user_api_key_dict=user_api_key_auth, requested_model=summary_model, - descriptors=descriptors, + descriptors=base_descriptors, ) - limiter._add_project_model_rate_limit_descriptor_from_metadata( + add_project_descriptor( user_api_key_dict=user_api_key_auth, requested_model=summary_model, - descriptors=descriptors, + descriptors=base_descriptors, ) - descriptors.extend(limiter.create_organization_rate_limit_descriptor(user_api_key_auth, summary_model)) + descriptors: Final = (*base_descriptors, *create_org_descriptors(user_api_key_auth, summary_model)) if not descriptors: return True - response: Final = await limiter.should_rate_limit( + parent_otel_span: Final[object] = getattr(user_api_key_auth, "parent_otel_span", None) + response: Final[RateLimitResponse] = await should_rate_limit_check( descriptors=descriptors, - parent_otel_span=getattr(user_api_key_auth, "parent_otel_span", None), + parent_otel_span=parent_otel_span, read_only=True, ) except Exception as e: @@ -446,7 +528,7 @@ async def _check_summary_model_rate_limit( def _find_latest_compaction_index( - messages: list[dict[str, object]], + messages: Sequence[Mapping[str, object]], ) -> tuple[int | None, int | None]: """Return (message_index, block_index) of the most recent compaction block. @@ -465,8 +547,8 @@ def _find_latest_compaction_index( def _slice_around_compaction_block( - messages: list[dict[str, Any]], -) -> tuple[list[dict[str, object]], dict[str, object] | None]: + messages: Sequence[_MsgT], +) -> tuple[Sequence[_MsgT | dict[str, object]], dict[str, object] | None]: """Apply Anthropic's "drop everything before the compaction block" rule. Returns ``(sliced_messages_with_compaction_block, compaction_block_dict)`` @@ -481,19 +563,21 @@ def _slice_around_compaction_block( original_msg: Final = messages[msg_idx] original_content: Final = original_msg["content"] - compaction_block: Final = cast(dict[str, object], original_content[blk_idx]) + if not isinstance(original_content, list): + return messages, None + original_blocks: Final = cast("Sequence[dict[str, object]]", original_content) + compaction_block: Final = original_blocks[blk_idx] # Per Anthropic's contract everything before the compaction block is # dropped, including earlier blocks within the same assistant message. - sliced_content: Final = list(original_content[blk_idx:]) + sliced_content: Final = list(original_blocks[blk_idx:]) - sliced_messages: Final[list[dict[str, object]]] = [{**original_msg, "content": sliced_content}] - sliced_messages.extend(messages[msg_idx + 1 :]) + sliced_messages: Final = [{**original_msg, "content": sliced_content}, *messages[msg_idx + 1 :]] return sliced_messages, compaction_block def _strip_compaction_blocks( - messages: list[dict[str, object]], + messages: Sequence[dict[str, object]], ) -> list[dict[str, object]]: """Drop any ``compaction`` content blocks from messages. @@ -600,7 +684,7 @@ def _propagate_metadata( def _count_effective_tokens( model: str, - effective_messages: list[dict[str, object]], + effective_messages: Sequence[dict[str, object]], compaction_block: CompactionBlock | None, tools: list[dict[str, object]] | None, system: str | list[dict[str, object]] | None = None, @@ -623,7 +707,7 @@ def _count_effective_tokens( try: openai_shape = adapter.translate_anthropic_messages_to_openai( messages=cast( - "list[AnthropicMessagesUserMessageParam | AnthopicMessagesAssistantMessageParam]", + "list[AllAnthropicPassThroughMessageValues]", messages_without_compaction, ) ) @@ -679,17 +763,18 @@ def _system_to_text( return "" if isinstance(system, str): return system - parts: Final[list[str]] = [] - for block in system: - if isinstance(block, dict) and block.get("type") == "text": - text = block.get("text") - if isinstance(text, str) and text: - parts.append(text) - return "\n".join(parts) + return "\n".join( + text + for block in system + if isinstance(block, dict) + and block.get("type") == "text" + and isinstance(text := block.get("text"), str) + and text + ) def _select_last_user_question( - messages: list[dict[str, object]], + messages: Sequence[dict[str, object]], ) -> list[dict[str, object]]: """Pick the most recent ``user`` turn that is a real question. @@ -704,16 +789,18 @@ def _select_last_user_question( turns, or contained no user turns at all). The downstream call always needs a non-empty user message. """ + blocks: Sequence[object] for msg in reversed(messages): if msg.get("role") != "user": continue content = msg.get("content") if isinstance(content, list): - filtered = [blk for blk in content if not (isinstance(blk, dict) and blk.get("type") == "tool_result")] + blocks = [*map(_as_object, content)] + filtered = [blk for blk in blocks if not _is_tool_result_block(blk)] if not filtered: # Purely tool_result — skip and look for an earlier turn. continue - if len(filtered) < len(content): + if len(filtered) < len(blocks): return [{**msg, "content": filtered}] return [msg] return [ @@ -736,7 +823,7 @@ def _extract_summary_text(raw: str | None) -> str | None: def _system_to_openai_message( system: str | list[dict[str, Any]] | None, -) -> dict[str, Any] | None: +) -> Mapping[str, object] | None: """Translate Anthropic-shaped ``system`` to an OpenAI system message. Accepts a bare string or a list of Anthropic content blocks; returns @@ -747,17 +834,19 @@ def _system_to_openai_message( if isinstance(system, str): return {"role": "system", "content": system} if system else None if isinstance(system, list): - parts = [block.get("text", "") for block in system if isinstance(block, dict) and block.get("type") == "text"] + parts: Final[tuple[str, ...]] = tuple( + block.get("text", "") for block in system if isinstance(block, dict) and block.get("type") == "text" + ) joined: Final = "\n\n".join(part for part in parts if part) return {"role": "system", "content": joined} if joined else None return None def _build_summary_messages( - effective_messages: list[dict[str, object]], + effective_messages: Sequence[dict[str, object]], prompt: str, system: str | list[dict[str, object]] | None = None, -) -> list[dict[str, object]]: +) -> Sequence[Mapping[str, object]]: """Build the OpenAI-shape message list for the summary call. The caller's ``system`` prompt is prepended (the default summarization @@ -773,7 +862,7 @@ def _build_summary_messages( try: openai_messages = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( messages=cast( - "list[AnthropicMessagesUserMessageParam | AnthopicMessagesAssistantMessageParam]", + "list[AllAnthropicPassThroughMessageValues]", stripped, ) ) @@ -785,7 +874,7 @@ def _build_summary_messages( ) openai_messages = stripped - summary_messages: Final[list[dict[str, object]]] = [] + summary_messages: Final[list[Mapping[str, object]]] = [] system_message: Final = _system_to_openai_message(system) if system_message is not None: summary_messages.append(system_message) @@ -809,7 +898,7 @@ def _is_user_message(msg: object) -> bool: return isinstance(msg, dict) and msg.get("role") == "user" -def _append_text_to_content(content: Any, extra_text: str) -> Any: +def _append_text_to_content(content: object, extra_text: str) -> object: """Append ``extra_text`` to an OpenAI-shape message ``content`` field. Handles the two common shapes: ``str`` and ``list`` of content parts. @@ -820,16 +909,17 @@ def _append_text_to_content(content: Any, extra_text: str) -> Any: if isinstance(content, str): return f"{content}\n\n{extra_text}" if isinstance(content, list): - return [*content, {"type": "text", "text": extra_text}] + appended: Final[Sequence[object]] = [*map(_as_object, content), {"type": "text", "text": extra_text}] + return appended return [content, {"type": "text", "text": extra_text}] async def _call_summary_model( *, summary_model: str, - summary_messages: list[dict[str, object]], + summary_messages: Sequence[Mapping[str, object]], metadata: Mapping[str, object], - llm_router: Any, + llm_router: object, allowed_model_region: str | None = None, max_tokens: int = COMPACT_SUMMARY_MAX_TOKENS, ) -> Union["ModelResponse", "CustomStreamWrapper"]: @@ -860,9 +950,8 @@ async def _call_summary_model( # the parent ``/v1/messages`` request. On timeout the caller catches the # exception and surfaces ``applied_edits[0].error = "summary_call_failed"``, # forwarding the request without compaction rather than hanging. - call_kwargs: Final[dict[str, Any]] = { + call_kwargs: Final[_SummaryCallKwargs] = { "model": summary_model, - "messages": summary_messages, "max_tokens": max_tokens, "timeout": COMPACT_SUMMARY_TIMEOUT_SECONDS, "litellm_metadata": metadata, @@ -872,19 +961,23 @@ async def _call_summary_model( # than from ``litellm_metadata``, so without it the summary tokens would not # debit the caller's end-user counters. end_user_id: Final = metadata.get("user_api_key_end_user_id") - if end_user_id: + if isinstance(end_user_id, str) and end_user_id: call_kwargs["user"] = end_user_id if allowed_model_region is not None: call_kwargs["allowed_model_region"] = allowed_model_region - if llm_router is not None and hasattr(llm_router, "acompletion"): - return await llm_router.acompletion(**call_kwargs) - return await litellm.acompletion(**call_kwargs) + router_acompletion: Final[_SummaryAcompletion | None] = getattr(llm_router, "acompletion", None) + if llm_router is not None and router_acompletion is not None: + return await router_acompletion(messages=summary_messages, **call_kwargs) + return await litellm.acompletion(messages=[*summary_messages], **call_kwargs) -def _extract_response_text(response: Any) -> str | None: +def _extract_response_text(response: object) -> str | None: try: - choice: Final = response.choices[0] - message: Final = choice.message + choices: Final[Sequence[object] | None] = getattr(response, "choices", None) + if choices is None: + return None + choice: Final = choices[0] + message: Final = getattr(choice, "message", None) content: Final = getattr(message, "content", None) if isinstance(content, str): return content @@ -900,7 +993,7 @@ def _extract_response_text(response: Any) -> str | None: def _extract_usage(response: object) -> tuple[int, int]: - usage: Final = getattr(response, "usage", None) + usage: Final[object] = getattr(response, "usage", None) if usage is None: return 0, 0 return ( diff --git a/litellm/llms/runwayml/videos/transformation.py b/litellm/llms/runwayml/videos/transformation.py index 2e0ae30a192..6e720c058fb 100644 --- a/litellm/llms/runwayml/videos/transformation.py +++ b/litellm/llms/runwayml/videos/transformation.py @@ -1,8 +1,10 @@ +from collections.abc import Mapping, Sequence from datetime import datetime -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Any, Final, Literal, TypedDict import httpx -from httpx._types import RequestFiles +from httpx._types import FileTypes, RequestFiles +from typing_extensions import NotRequired import litellm from litellm.constants import RUNWAYML_DEFAULT_API_VERSION @@ -31,6 +33,31 @@ else: LiteLLMLoggingObj = Any +class _RunwayTaskResponse(TypedDict, total=False): + id: str + status: str + createdAt: str + completedAt: str + output: Sequence[str] | str + progress: int + failureCode: str + failure: str + + +class _RunwayVideoData(TypedDict): + id: str + object: Literal["video"] + status: str + created_at: int + output_url: NotRequired[str] + completed_at: NotRequired[int] + progress: NotRequired[int] + error: NotRequired[Mapping[str, str]] + model: NotRequired[str] + size: NotRequired[str] + seconds: NotRequired[str] + + class RunwayMLVideoConfig(BaseVideoConfig): """ Configuration class for RunwayML video generation. @@ -44,6 +71,10 @@ class RunwayMLVideoConfig(BaseVideoConfig): def __init__(self): super().__init__() + @staticmethod + def _parse_task_response(raw_response: httpx.Response) -> _RunwayTaskResponse: + return raw_response.json() + def get_supported_openai_params(self, model: str) -> list: """ Get the list of supported OpenAI parameters for video generation. @@ -68,7 +99,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): video_create_optional_params: VideoCreateOptionalRequestParams, model: str, drop_params: bool, - ) -> dict: + ) -> dict[str, object]: """ Map OpenAI parameters to RunwayML format. @@ -78,37 +109,44 @@ class RunwayMLVideoConfig(BaseVideoConfig): - size -> ratio (convert "WIDTHxHEIGHT" to "WIDTH:HEIGHT") - seconds -> duration (convert to integer) """ - mapped_params: Final[dict[str, Any]] = {} + supported_openai_params: Final = self.get_supported_openai_params(model) + return { + **self._prompt_image_param(video_create_optional_params), + **self._ratio_param(video_create_optional_params), + **self._duration_param(video_create_optional_params), + # Pass through other parameters that aren't OpenAI-specific + **{key: value for key, value in video_create_optional_params.items() if key not in supported_openai_params}, + } + @staticmethod + def _prompt_image_param(video_create_optional_params: VideoCreateOptionalRequestParams) -> Mapping[str, object]: # Handle input_reference parameter - map to promptImage + # RunwayML supports URLs and data URIs directly if "input_reference" in video_create_optional_params: - input_reference: Final = video_create_optional_params["input_reference"] - # RunwayML supports URLs and data URIs directly - mapped_params["promptImage"] = input_reference + return {"promptImage": video_create_optional_params["input_reference"]} + return {} + @staticmethod + def _ratio_param(video_create_optional_params: VideoCreateOptionalRequestParams) -> Mapping[str, str]: # Handle size parameter - convert "1280x720" to "1280:720" if "size" in video_create_optional_params: size: Final = video_create_optional_params["size"] if isinstance(size, str) and "x" in size: - mapped_params["ratio"] = size.replace("x", ":") + return {"ratio": size.replace("x", ":")} + return {} + @staticmethod + def _duration_param(video_create_optional_params: VideoCreateOptionalRequestParams) -> Mapping[str, int]: # Handle seconds parameter - convert to integer if "seconds" in video_create_optional_params: seconds: Final = video_create_optional_params["seconds"] if seconds is not None: try: - mapped_params["duration"] = int(float(seconds)) if isinstance(seconds, str) else int(seconds) + return {"duration": int(float(seconds)) if isinstance(seconds, str) else int(seconds)} except (ValueError, TypeError): # If conversion fails, use default duration pass - - # Pass through other parameters that aren't OpenAI-specific - supported_openai_params: Final = self.get_supported_openai_params(model) - for key, value in video_create_optional_params.items(): - if key not in supported_openai_params: - mapped_params[key] = value - - return mapped_params + return {} def validate_environment( self, @@ -163,7 +201,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): model: str, prompt: str, api_base: str, - video_create_optional_request_params: dict, + video_create_optional_request_params: dict[str, object], litellm_params: GenericLiteLLMParams, headers: dict, ) -> tuple[dict, RequestFiles, str]: @@ -179,17 +217,15 @@ class RunwayMLVideoConfig(BaseVideoConfig): "duration": 5 } """ - # Build the request data - request_data: Final[dict[str, Any]] = { + # Build the request data with the mapped parameters merged in + request_data: Final = { "model": model, "promptText": prompt, + **video_create_optional_request_params, } - # Add mapped parameters - request_data.update(video_create_optional_request_params) - # RunwayML uses JSON body, no files multipart - files_list: Final[list[tuple[str, Any]]] = [] + files_list: Final[Sequence[tuple[str, FileTypes]]] = [] # Append the specific endpoint for video generation full_api_base: Final = f"{api_base}/image_to_video" @@ -216,10 +252,10 @@ class RunwayMLVideoConfig(BaseVideoConfig): We map this to OpenAI VideoObject format. """ - response_data: Final = raw_response.json() + response_data: Final[_RunwayTaskResponse] = self._parse_task_response(raw_response) # Map RunwayML task response to VideoObject format - video_data: Final[dict[str, Any]] = { + video_data: Final[_RunwayVideoData] = { "id": response_data.get("id", ""), "object": "video", "status": self._map_runway_status(response_data.get("status", "pending")), @@ -229,9 +265,8 @@ class RunwayMLVideoConfig(BaseVideoConfig): # Add optional fields if present if "output" in response_data and response_data["output"]: # RunwayML returns output as array of URLs when task succeeds - video_data["output_url"] = ( - response_data["output"][0] if isinstance(response_data["output"], list) else response_data["output"] - ) + output: Final = response_data["output"] + video_data["output_url"] = output if isinstance(output, str) else output[0] if "completedAt" in response_data: video_data["completed_at"] = self._parse_runway_timestamp(response_data.get("completedAt")) @@ -254,7 +289,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): if "duration" in request_data: video_data["seconds"] = str(request_data["duration"]) - video_obj: Final = VideoObject(**video_data) + video_obj: Final = VideoObject.model_validate(video_data) if custom_llm_provider and video_obj.id: video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, model) @@ -326,20 +361,18 @@ class RunwayMLVideoConfig(BaseVideoConfig): # Get task status to retrieve video URL url: Final = f"{api_base}/tasks/{encoded_video_id}" - params: Final[dict[str, Any]] = {} + return url, dict[str, str]() - return url, params - - def _extract_video_url_from_response(self, response_data: dict[str, Any]) -> str: + def _extract_video_url_from_response(self, response_data: _RunwayTaskResponse) -> str: """ Helper method to extract video URL from RunwayML response. Shared between sync and async transforms. """ # Extract video URL from the output field video_url = None - if "output" in response_data and response_data["output"]: - output: Final = response_data["output"] - video_url = output[0] if isinstance(output, list) else output + raw_output: Final = response_data.get("output") + if raw_output: + video_url = raw_output if isinstance(raw_output, str) else raw_output[0] if not video_url: # Check if the video generation failed or is still processing @@ -373,7 +406,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): "output":["https://dnznrvs05pmza.cloudfront.net/.../video.mp4?_jwt=..."] } """ - response_data: Final = raw_response.json() + response_data: Final[_RunwayTaskResponse] = self._parse_task_response(raw_response) video_url: Final = self._extract_video_url_from_response(response_data) # Download the video from the CloudFront URL synchronously @@ -402,7 +435,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): "output":["https://dnznrvs05pmza.cloudfront.net/.../video.mp4?_jwt=..."] } """ - response_data: Final = raw_response.json() + response_data: Final[_RunwayTaskResponse] = self._parse_task_response(raw_response) video_url: Final = self._extract_video_url_from_response(response_data) # Download the video from the CloudFront URL asynchronously @@ -421,7 +454,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): api_base: str, litellm_params: GenericLiteLLMParams, headers: dict, - extra_body: dict[str, Any] | None = None, + extra_body: Mapping[str, object] | None = None, ) -> tuple[str, dict]: """ Transform the video remix request for RunwayML API. @@ -448,7 +481,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): after: str | None = None, limit: int | None = None, order: str | None = None, - extra_query: dict[str, Any] | None = None, + extra_query: Mapping[str, object] | None = None, ) -> tuple[str, dict]: """ Transform the video list request for RunwayML API. @@ -484,9 +517,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): # Construct the URL for task cancellation url: Final = f"{api_base}/tasks/{encoded_video_id}/cancel" - data: Final[dict[str, Any]] = {} - - return url, data + return url, dict[str, str]() def transform_video_delete_response( self, @@ -494,7 +525,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): logging_obj: LiteLLMLoggingObj, ) -> VideoObject: """Transform the RunwayML video delete/cancel response.""" - response_data: Final = raw_response.json() + response_data: Final[_RunwayTaskResponse] = self._parse_task_response(raw_response) video_obj: Final = VideoObject( id=response_data.get("id", ""), @@ -524,9 +555,7 @@ class RunwayMLVideoConfig(BaseVideoConfig): url: Final = f"{api_base}/tasks/{encoded_video_id}" # Empty dict for GET request (no body) - data: Final[dict[str, Any]] = {} - - return url, data + return url, dict[str, str]() def transform_video_status_retrieve_response( self, @@ -537,10 +566,10 @@ class RunwayMLVideoConfig(BaseVideoConfig): """ Transform the RunwayML video status retrieve response. """ - response_data: Final = raw_response.json() + response_data: Final[_RunwayTaskResponse] = self._parse_task_response(raw_response) # Map RunwayML task response to VideoObject format - video_data: Final[dict[str, Any]] = { + video_data: Final[_RunwayVideoData] = { "id": response_data.get("id", ""), "object": "video", "status": self._map_runway_status(response_data.get("status", "pending")), @@ -549,9 +578,8 @@ class RunwayMLVideoConfig(BaseVideoConfig): # Add optional fields if present if "output" in response_data and response_data["output"]: - video_data["output_url"] = ( - response_data["output"][0] if isinstance(response_data["output"], list) else response_data["output"] - ) + output: Final = response_data["output"] + video_data["output_url"] = output if isinstance(output, str) else output[0] if "completedAt" in response_data: video_data["completed_at"] = self._parse_runway_timestamp(response_data.get("completedAt")) @@ -565,14 +593,14 @@ class RunwayMLVideoConfig(BaseVideoConfig): "message": response_data.get("failure", "Video generation failed"), } - video_obj: Final = VideoObject(**video_data) + video_obj: Final = VideoObject.model_validate(video_data) if custom_llm_provider and video_obj.id: video_obj.id = encode_video_id_with_provider(video_obj.id, custom_llm_provider, None) return video_obj - def transform_video_create_character_request(self, name, video, api_base, litellm_params, headers): + def transform_video_create_character_request(self, name, video: object, api_base, litellm_params, headers): raise NotImplementedError("video create character is not supported for RunwayML") def transform_video_create_character_response(self, raw_response, logging_obj): diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 3538fc5b1a7..131ee41ea8b 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -5,12 +5,13 @@ import json import os import re import time -from collections.abc import Callable, Iterable, Iterator -from typing import Any, Final +from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence +from typing import Final, TypedDict import httpx from httpx import Headers, Response from openai.types.file_deleted import FileDeleted +from typing_extensions import Required import litellm from litellm._uuid import uuid @@ -50,9 +51,10 @@ from litellm.types.llms.openai import ( HttpxBinaryResponseContent, OpenAICreateFileRequestOptionalParams, OpenAIFileObject, + OpenAIFilesPurpose, PathLike, ) -from litellm.types.llms.vertex_ai import GcsBucketResponse +from litellm.types.llms.vertex_ai import GcsBucketResponse, GenerateContentResponseBody from litellm.types.utils import LlmProviders, ModelResponse from ..common_utils import VertexAIError @@ -62,6 +64,47 @@ _GCP_LABEL_VALUE_MAX_LEN: Final = 63 _CUSTOM_ID_RAW_LABEL_PREFIX: Final = "b32_" +class _OpenAIBatchRequestBody(TypedDict, total=False): + model: str + messages: Sequence[AllMessageValues] + + +class _OpenAIBatchJsonlEntry(TypedDict, total=False): + custom_id: Required[object] + body: _OpenAIBatchRequestBody + + +class _VertexBatchOutputRequest(TypedDict, total=False): + labels: Mapping[str, str] + + +class _VertexBatchResponse(GenerateContentResponseBody, total=False): + modelVersion: str + + +class _VertexBatchOutputRow(TypedDict, total=False): + request: _VertexBatchOutputRequest + status: str + processed_time: str + response: _VertexBatchResponse + + +class _GcsObjectMetadata(TypedDict, total=False): + purpose: OpenAIFilesPurpose + + +class _GcsObjectResponse(GcsBucketResponse, total=False): + metadata: _GcsObjectMetadata + + +def _parse_gcs_object_response(raw_response: Response) -> _GcsObjectResponse: + return raw_response.json() + + +def _parse_vertex_batch_output_row(line: str) -> _VertexBatchOutputRow: + return json.loads(line) + + def _sanitize_gcp_label_value(value: str) -> str: """ Sanitize a string to meet GCP label value constraints. @@ -106,7 +149,7 @@ def _decode_gcp_label_value_chunks(values: list[str]) -> str | None: return None -def _set_litellm_batch_custom_id_labels(labels: dict[str, str], custom_id: Any) -> None: +def _litellm_batch_custom_id_labels(custom_id: object) -> Mapping[str, str]: """ Store OpenAI batch custom_id for Vertex batch correlation. @@ -115,15 +158,19 @@ def _set_litellm_batch_custom_id_labels(labels: dict[str, str], custom_id: Any) round-trip correlation in batch output transforms. """ custom_id_str: Final = str(custom_id) - labels["litellm_custom_id"] = _sanitize_gcp_label_value(custom_id_str) raw_label_chunks: Final = _encode_gcp_label_value_chunks(custom_id_str) - labels["litellm_custom_id_raw"] = raw_label_chunks[0] - for index, raw_label_chunk in enumerate(raw_label_chunks[1:], start=1): - labels[f"litellm_custom_id_raw_{index}"] = raw_label_chunk + return { + "litellm_custom_id": _sanitize_gcp_label_value(custom_id_str), + "litellm_custom_id_raw": raw_label_chunks[0], + **{ + f"litellm_custom_id_raw_{index}": raw_label_chunk + for index, raw_label_chunk in enumerate(raw_label_chunks[1:], start=1) + }, + } -def _get_litellm_batch_custom_id_from_labels(labels: dict[str, Any]) -> str: - """Prefer encoded custom_id when present (see _set_litellm_batch_custom_id_labels).""" +def _get_litellm_batch_custom_id_from_labels(labels: Mapping[str, str]) -> str: + """Prefer encoded custom_id when present (see _litellm_batch_custom_id_labels).""" raw: Final = labels.get("litellm_custom_id_raw") if raw: raw_chunks: Final = [str(raw)] @@ -141,9 +188,9 @@ def _get_litellm_batch_custom_id_from_labels(labels: dict[str, Any]) -> str: def _openai_batch_jsonl_entry_to_vertex_wrapped_request( - openai_entry: dict[str, Any], - map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, Any]], -) -> dict[str, Any]: + openai_entry: _OpenAIBatchJsonlEntry, + map_openai_to_vertex_params: Callable[[_OpenAIBatchRequestBody], Mapping[str, object]], +) -> Mapping[str, object]: """ Transforms a single OpenAI JSONL batch entry into its Vertex wrapped request. @@ -151,11 +198,11 @@ def _openai_batch_jsonl_entry_to_vertex_wrapped_request( Example Vertex jsonl {"request":{"contents": [{"role": "user", "parts": [{"text": "What is the relation between the following video and image samples?"}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/video/animals.mp4", "mimeType": "video/mp4"}}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/image/cricket.jpeg", "mimeType": "image/jpeg"}}]}]}} """ - openai_request_body: Final = openai_entry.get("body") or {} + openai_request_body: Final[_OpenAIBatchRequestBody] = openai_entry.get("body") or {} vertex_request_body: Final = _transform_request_body( - messages=openai_request_body.get("messages", []), + messages=[*openai_request_body.get("messages", [])], model=openai_request_body.get("model", ""), - optional_params=map_openai_to_vertex_params(openai_request_body), + optional_params=dict(map_openai_to_vertex_params(openai_request_body)), custom_llm_provider="vertex_ai", litellm_params={}, cached_content=None, @@ -163,9 +210,10 @@ def _openai_batch_jsonl_entry_to_vertex_wrapped_request( custom_id: Final = openai_entry.get("custom_id") if custom_id is not None: - if "labels" not in vertex_request_body: - vertex_request_body["labels"] = {} - _set_litellm_batch_custom_id_labels(vertex_request_body["labels"], custom_id) + vertex_request_body["labels"] = { + **vertex_request_body.get("labels", {}), + **_litellm_batch_custom_id_labels(custom_id), + } return {"request": vertex_request_body} @@ -186,7 +234,7 @@ def _iter_openai_jsonl_lines(openai_file_content: FileTypes) -> Iterator[str]: ``str.splitlines()`` + ``line.strip()`` for ``\\n`` / ``\\r\\n`` delimited JSONL. """ - content: Any = openai_file_content + content: FileTypes | str = openai_file_content if isinstance(content, tuple): content = content[1] @@ -241,7 +289,7 @@ def _iter_openai_jsonl_lines(openai_file_content: FileTypes) -> Iterator[str]: def _iter_openai_jsonl_entries( openai_file_content: FileTypes, -) -> Iterator[dict[str, Any]]: +) -> Iterator[_OpenAIBatchJsonlEntry]: for line in _iter_openai_jsonl_lines(openai_file_content): yield json.loads(line) @@ -257,7 +305,7 @@ class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): def __init__( self, openai_file_content: FileTypes, - map_openai_to_vertex_params: Callable[[dict[str, Any]], dict[str, Any]], + map_openai_to_vertex_params: Callable[[_OpenAIBatchRequestBody], Mapping[str, object]], ) -> None: self._openai_file_content = openai_file_content self._map_openai_to_vertex_params = map_openai_to_vertex_params @@ -308,7 +356,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _get_gcs_object_name_from_batch_jsonl( self, - openai_jsonl_content: list[dict[str, Any]], + openai_jsonl_content: Sequence[_OpenAIBatchJsonlEntry], ) -> str: """ Gets a unique GCS object name for the VertexAI batch prediction job @@ -396,8 +444,8 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _map_openai_to_vertex_params( self, - openai_request_body: dict[str, Any], - ) -> dict[str, Any]: + openai_request_body: _OpenAIBatchRequestBody, + ) -> Mapping[str, object]: """ wrapper to call VertexGeminiConfig.map_openai_params """ @@ -409,7 +457,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): _model: Final = openai_request_body.get("model", "") vertex_params: Final = config.map_openai_params( model=_model, - non_default_params=openai_request_body, + non_default_params=dict(openai_request_body), optional_params={}, drop_params=False, ) @@ -463,10 +511,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): """ Transform VertexAI File upload response into OpenAI-style FileObject """ - response_json: Final = raw_response.json() + response_json: Final = _parse_gcs_object_response(raw_response) try: - response_object: Final = GcsBucketResponse(**response_json) + response_object: Final = _GcsObjectResponse(**response_json) except Exception as e: raise VertexAIError( status_code=raw_response.status_code, @@ -523,7 +571,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): logging_obj: LiteLLMLoggingObj, litellm_params: dict, ) -> OpenAIFileObject: - response_json: Final = raw_response.json() + response_json: Final = _parse_gcs_object_response(raw_response) gcs_id = response_json.get("id", "") gcs_id = "/".join(gcs_id.split("/")[:-1]) if gcs_id else "" return OpenAIFileObject( @@ -682,7 +730,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): # discriminating fields. Anything else (e.g. a binary file whose # first line is not valid UTF-8/JSON) raises and falls through to the # passthrough below, leaving the content untouched. - first_row: Final = json.loads(first_line) + first_row: Final = _parse_vertex_batch_output_row(first_line) is_vertex_batch_output: Final = ( "request" in first_row and "response" in first_row @@ -723,7 +771,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): for line in itertools.chain([first_line], lines): try: openai_output = self._transform_single_vertex_batch_output_to_openai( - vertex_output=json.loads(line), + vertex_output=_parse_vertex_batch_output_row(line), vertex_gemini_config=vertex_gemini_config, logging_obj=batch_transform_logging_obj, mock_httpx_response=mock_httpx_response, @@ -742,11 +790,11 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def _transform_single_vertex_batch_output_to_openai( self, - vertex_output: dict[str, Any], + vertex_output: _VertexBatchOutputRow, vertex_gemini_config: VertexGeminiConfig, logging_obj: Logging, mock_httpx_response: httpx.Response, - ) -> dict[str, Any]: + ) -> Mapping[str, object]: """ Transform a single Vertex AI batch output line to OpenAI format. Uses the existing VertexGeminiConfig transformation for the response. diff --git a/litellm/proxy/db/tool_registry_writer.py b/litellm/proxy/db/tool_registry_writer.py index b2fa538b1cc..e6fcff548d6 100644 --- a/litellm/proxy/db/tool_registry_writer.py +++ b/litellm/proxy/db/tool_registry_writer.py @@ -6,8 +6,11 @@ Admins use the management endpoints to read and update input_policy / output_pol """ import uuid +from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from typing import TYPE_CHECKING, Any, Final +from typing import TYPE_CHECKING, Final, Protocol + +from pydantic import TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ToolDiscoveryQueueItem @@ -23,33 +26,109 @@ if TYPE_CHECKING: from litellm.proxy.utils import PrismaClient -def _row_to_model(row: dict | Any) -> LiteLLM_ToolTableRow: +class _ToolTableRecord(Protocol): + tool_name: str + input_policy: str | None + output_policy: str | None + + +class _TokenRelationRecord(Protocol): + token: str | None + key_alias: str | None + + +class _TeamRelationRecord(Protocol): + team_id: str | None + team_alias: str | None + + +class _ObjectPermissionRecord(Protocol): + object_permission_id: str + blocked_tools: Sequence[str] | None + verification_tokens: Sequence[_TokenRelationRecord] | None + teams: Sequence[_TeamRelationRecord] | None + + +class _ToolTable(Protocol): + async def find_many( + self, + *, + where: Mapping[str, object] | None = None, + order: Mapping[str, str] | None = None, + ) -> Sequence[_ToolTableRecord]: ... + + async def find_unique(self, *, where: Mapping[str, object]) -> _ToolTableRecord | None: ... + + async def upsert(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> object: ... + + +class _ObjectPermissionTable(Protocol): + async def find_many( + self, + *, + where: Mapping[str, object] | None = None, + include: Mapping[str, bool] | None = None, + ) -> Sequence[_ObjectPermissionRecord]: ... + + async def find_unique(self, *, where: Mapping[str, object]) -> _ObjectPermissionRecord | None: ... + + async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> object: ... + + +class _ModelDumpMethod(Protocol): + def __call__(self) -> Mapping: ... + + +_ROW_DICT: Final = TypeAdapter(dict) + + +class _ToolTableHolder(Protocol): + @property + def table(self) -> _ToolTable: ... + + +class _ObjectPermissionTableHolder(Protocol): + @property + def table(self) -> _ObjectPermissionTable: ... + + +def _tool_table(repo: _ToolTableHolder) -> _ToolTable: + return repo.table + + +def _object_permission_table(repo: _ObjectPermissionTableHolder) -> _ObjectPermissionTable: + return repo.table + + +def _row_to_model(row: object) -> LiteLLM_ToolTableRow: """Convert a Prisma model instance or dict to LiteLLM_ToolTableRow.""" - model_dump: Final = getattr(row, "model_dump", None) + model_dump: Final[_ModelDumpMethod | None] = getattr(row, "model_dump", None) if callable(model_dump): row = model_dump() elif not isinstance(row, dict): - row = { - k: getattr(row, k, None) - for k in ( - "tool_id", - "tool_name", - "origin", - "input_policy", - "output_policy", - "call_count", - "assignments", - "key_hash", - "team_id", - "key_alias", - "user_agent", - "last_used_at", - "created_at", - "updated_at", - "created_by", - "updated_by", - ) - } + row = _ROW_DICT.validate_python( + { + k: getattr(row, k, None) + for k in ( + "tool_id", + "tool_name", + "origin", + "input_policy", + "output_policy", + "call_count", + "assignments", + "key_hash", + "team_id", + "key_alias", + "user_agent", + "last_used_at", + "created_at", + "updated_at", + "created_by", + "updated_by", + ) + } + ) return LiteLLM_ToolTableRow( tool_id=row.get("tool_id", ""), tool_name=row.get("tool_name", ""), @@ -87,7 +166,7 @@ async def batch_upsert_tools( if not data: return now: Final = datetime.now(timezone.utc) - table: Final = ToolRepository(prisma_client).table + table: Final = _tool_table(ToolRepository(prisma_client)) for item in data: tool_name = item.get("tool_name", "") origin = item.get("origin") or "user_defined" @@ -132,8 +211,8 @@ async def list_tools( ) -> list[LiteLLM_ToolTableRow]: """Return all tools, optionally filtered by input_policy.""" try: - where: Final = {"input_policy": input_policy} if input_policy is not None else {} - rows: Final = await ToolRepository(prisma_client).table.find_many( + where: Final[Mapping[str, str]] = {"input_policy": input_policy} if input_policy is not None else {} + rows: Final = await _tool_table(ToolRepository(prisma_client)).find_many( where=where, order={"created_at": "desc"}, ) @@ -149,7 +228,7 @@ async def get_tool( ) -> LiteLLM_ToolTableRow | None: """Return a single tool row by tool_name.""" try: - row: Final = await ToolRepository(prisma_client).table.find_unique( + row: Final = await _tool_table(ToolRepository(prisma_client)).find_unique( where={"tool_name": tool_name}, ) if row is None: @@ -172,7 +251,7 @@ async def update_tool_policy( _updated_by: Final = updated_by or "system" now: Final = datetime.now(timezone.utc) - create_data: Final[dict] = { + create_data: Final[Mapping[str, str | datetime]] = { "tool_id": str(uuid.uuid4()), "tool_name": tool_name, "input_policy": input_policy or "untrusted", @@ -182,16 +261,18 @@ async def update_tool_policy( "created_at": now, "updated_at": now, } - update_data: Final[dict] = { - "updated_by": _updated_by, - "updated_at": now, + update_data: Final[Mapping[str, str | datetime]] = { + key: value + for key, value in ( + ("updated_by", _updated_by), + ("updated_at", now), + ("input_policy", input_policy), + ("output_policy", output_policy), + ) + if value is not None } - if input_policy is not None: - update_data["input_policy"] = input_policy - if output_policy is not None: - update_data["output_policy"] = output_policy - await ToolRepository(prisma_client).table.upsert( + await _tool_table(ToolRepository(prisma_client)).upsert( where={"tool_name": tool_name}, data={ "create": create_data, @@ -214,7 +295,7 @@ async def get_tools_by_names( if not tool_names: return {} try: - rows: Final = await ToolRepository(prisma_client).table.find_many( + rows: Final = await _tool_table(ToolRepository(prisma_client)).find_many( where={"tool_name": {"in": tool_names}}, ) return { @@ -239,7 +320,7 @@ async def list_overrides_for_tool( """ out: Final[list[ToolPolicyOverrideRow]] = [] try: - perms: Final = await ObjectPermissionRepository(prisma_client).table.find_many( + perms: Final = await _object_permission_table(ObjectPermissionRepository(prisma_client)).find_many( where={"blocked_tools": {"has": tool_name}}, include={ "verification_tokens": True, @@ -248,8 +329,8 @@ async def list_overrides_for_tool( ) for perm in perms: op_id = getattr(perm, "object_permission_id", None) or "" - tokens = getattr(perm, "verification_tokens", []) or [] - teams = getattr(perm, "teams", []) or [] + tokens: Sequence[_TokenRelationRecord] = getattr(perm, "verification_tokens", []) or [] + teams: Sequence[_TeamRelationRecord] = getattr(perm, "teams", []) or [] for t in tokens: out.append( ToolPolicyOverrideRow( @@ -302,7 +383,7 @@ class ToolPolicyRegistry: try: tools: Final = await call_with_db_reconnect_retry( prisma_client, - lambda: ToolRepository(prisma_client).table.find_many(), + lambda: _tool_table(ToolRepository(prisma_client)).find_many(), reason="sync_tool_policy_from_db_tools_lookup_failure", ) self._tool_input_policies = { @@ -314,13 +395,13 @@ class ToolPolicyRegistry: perms: Final = await call_with_db_reconnect_retry( prisma_client, - lambda: ObjectPermissionRepository(prisma_client).table.find_many(), + lambda: _object_permission_table(ObjectPermissionRepository(prisma_client)).find_many(), reason="sync_tool_policy_from_db_perms_lookup_failure", ) self._blocked_tools_by_op_id = {} for row in perms: op_id = getattr(row, "object_permission_id", None) - blocked = getattr(row, "blocked_tools", None) or [] + blocked: Sequence[str] = getattr(row, "blocked_tools", None) or [] if op_id: self._blocked_tools_by_op_id[op_id] = list(blocked) @@ -352,10 +433,12 @@ class ToolPolicyRegistry: """ if not tool_names: return {} - blocked: Final[set] = set() - for op_id in (object_permission_id, team_object_permission_id): - if op_id and op_id.strip(): - blocked.update(self._blocked_tools_by_op_id.get(op_id.strip(), [])) + blocked: Final[frozenset[str]] = frozenset( + tool + for op_id in (object_permission_id, team_object_permission_id) + if op_id and op_id.strip() + for tool in self._blocked_tools_by_op_id.get(op_id.strip(), []) + ) result: Final[dict[str, str]] = {} for name in tool_names: if name in blocked: @@ -385,18 +468,17 @@ async def add_tool_to_object_permission_blocked( if not object_permission_id or not tool_name: return False try: - row: Final = await ObjectPermissionRepository(prisma_client).table.find_unique( + row: Final = await _object_permission_table(ObjectPermissionRepository(prisma_client)).find_unique( where={"object_permission_id": object_permission_id}, ) if row is None: return False - current: Final = list(getattr(row, "blocked_tools", []) or []) + current: Final[Sequence[str]] = getattr(row, "blocked_tools", []) or [] if tool_name in current: return True - current.append(tool_name) - await ObjectPermissionRepository(prisma_client).table.update( + await _object_permission_table(ObjectPermissionRepository(prisma_client)).update( where={"object_permission_id": object_permission_id}, - data={"blocked_tools": current}, + data={"blocked_tools": [*current, tool_name]}, ) return True except Exception as e: @@ -413,18 +495,17 @@ async def remove_tool_from_object_permission_blocked( if not object_permission_id or not tool_name: return False try: - row: Final = await ObjectPermissionRepository(prisma_client).table.find_unique( + row: Final = await _object_permission_table(ObjectPermissionRepository(prisma_client)).find_unique( where={"object_permission_id": object_permission_id}, ) if row is None: return False - current = list(getattr(row, "blocked_tools", []) or []) + current: Final[Sequence[str]] = getattr(row, "blocked_tools", []) or [] if tool_name not in current: return False - current = [t for t in current if t != tool_name] - await ObjectPermissionRepository(prisma_client).table.update( + await _object_permission_table(ObjectPermissionRepository(prisma_client)).update( where={"object_permission_id": object_permission_id}, - data={"blocked_tools": current}, + data={"blocked_tools": [t for t in current if t != tool_name]}, ) return True except Exception as e: diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index a416a197ab8..7dfc80d2b2a 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -17,7 +17,7 @@ import json import traceback from collections.abc import Mapping, Sequence from datetime import datetime, timezone -from typing import Any, Final, Literal, cast +from typing import Any, Final, Literal, Protocol, TypeAlias, TypeVar, cast, overload import fastapi from fastapi import APIRouter, Depends, Header, HTTPException, Request, status @@ -85,14 +85,6 @@ from litellm.types.proxy.management_endpoints.scim_v2 import ( if TYPE_CHECKING: from prisma import models as prisma_models from prisma import types as prisma_types - from prisma.actions import ( - LiteLLM_InvitationLinkActions, - LiteLLM_OrganizationMembershipActions, - LiteLLM_TeamMembershipActions, - LiteLLM_TeamTableActions, - LiteLLM_UserTableActions, - LiteLLM_VerificationTokenActions, - ) from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache from litellm.proxy.proxy_server import PrismaClient @@ -100,55 +92,151 @@ if TYPE_CHECKING: router: Final = APIRouter() +_PrismaTableT = TypeVar("_PrismaTableT", covariant=True) + + +class _TableActions(Protocol[_PrismaTableT]): + async def find_unique( + self, + *, + where: Mapping[str, object], + include: Mapping[str, object] | None = None, + ) -> "_PrismaTableT | None": ... + + async def find_first(self, *, where: Mapping[str, object]) -> "_PrismaTableT | None": ... + + async def find_many( + self, + *, + where: Mapping[str, object] | None = None, + order: Mapping[str, object] | Sequence[Mapping[str, object]] | None = None, + skip: int | None = None, + take: int | None = None, + ) -> "Sequence[_PrismaTableT]": ... + + async def create(self, *, data: Mapping[str, object]) -> "_PrismaTableT": ... + + async def update(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> "_PrismaTableT | None": ... + + async def delete_many(self, *, where: Mapping[str, object] | None = None) -> int: ... + + +class _PrismaTableHolder(Protocol[_PrismaTableT]): + @property + def table(self) -> "_TableActions[_PrismaTableT]": ... + + +def _typed_table(holder: "_PrismaTableHolder[_PrismaTableT]") -> "_TableActions[_PrismaTableT]": + return holder.table + + +class _LenientTableActions(Protocol[_PrismaTableT]): + async def find_first(self, *, where: Mapping[str, object]) -> "_PrismaTableT | None": ... + + async def find_many( + self, + *, + where: Mapping[str, object] | None = None, + order: Mapping[str, object] | Sequence[Mapping[str, object]] | None = None, + skip: int | None = None, + take: int | None = None, + ) -> "Sequence[_PrismaTableT] | None": ... + + async def count(self, *, where: Mapping[str, object] | None = None) -> int: ... + + async def update_many(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... + + +class _LenientTableHolder(Protocol[_PrismaTableT]): + @property + def table(self) -> "_LenientTableActions[_PrismaTableT]": ... + + +def _lenient_table(holder: "_LenientTableHolder[_PrismaTableT]") -> "_LenientTableActions[_PrismaTableT]": + return holder.table + + +class _UserDeleteRow(Protocol): + user_id: str + user_email: str | None + + @property + def teams(self) -> Sequence[str]: ... + + def json(self, *, exclude_none: bool) -> str: ... + + +class _TeamCleanupRow(Protocol): + team_id: str + members_with_roles: str + + def model_dump(self) -> Mapping[str, object]: ... + def _user_table( prisma_client: "PrismaClient | None", -) -> "LiteLLM_UserTableActions[prisma_models.LiteLLM_UserTable]": - user_table: Final[LiteLLM_UserTableActions[prisma_models.LiteLLM_UserTable]] = UserRepository(prisma_client).table - return user_table +) -> "_TableActions[prisma_models.LiteLLM_UserTable]": + return _typed_table(UserRepository(prisma_client)) + + +def _user_table_lenient( + prisma_client: "PrismaClient | None", +) -> "_LenientTableActions[prisma_models.LiteLLM_UserTable]": + return _lenient_table(UserRepository(prisma_client)) + + +def _user_delete_table( + prisma_client: "PrismaClient | None", +) -> "_TableActions[_UserDeleteRow]": + return _typed_table(UserRepository(prisma_client)) def _team_table( prisma_client: "PrismaClient | None", -) -> "LiteLLM_TeamTableActions[prisma_models.LiteLLM_TeamTable]": - team_table: Final[LiteLLM_TeamTableActions[prisma_models.LiteLLM_TeamTable]] = TeamRepository(prisma_client).table - return team_table +) -> "_TableActions[prisma_models.LiteLLM_TeamTable]": + return _typed_table(TeamRepository(prisma_client)) + + +def _team_cleanup_table( + prisma_client: "PrismaClient | None", +) -> "_TableActions[_TeamCleanupRow]": + return _typed_table(TeamRepository(prisma_client)) def _verification_token_table( prisma_client: "PrismaClient | None", -) -> "LiteLLM_VerificationTokenActions[prisma_models.LiteLLM_VerificationToken]": - token_table: Final[LiteLLM_VerificationTokenActions[prisma_models.LiteLLM_VerificationToken]] = ( - VerificationTokenRepository(prisma_client).table - ) - return token_table +) -> "_TableActions[prisma_models.LiteLLM_VerificationToken]": + return _typed_table(VerificationTokenRepository(prisma_client)) + + +def _verification_token_table_lenient( + prisma_client: "PrismaClient | None", +) -> "_LenientTableActions[prisma_models.LiteLLM_VerificationToken]": + return _lenient_table(VerificationTokenRepository(prisma_client)) + + +def _organization_table( + prisma_client: "PrismaClient | None", +) -> "_TableActions[prisma_models.LiteLLM_OrganizationTable]": + return _typed_table(OrganizationRepository(prisma_client)) def _organization_membership_table( prisma_client: "PrismaClient | None", -) -> "LiteLLM_OrganizationMembershipActions[prisma_models.LiteLLM_OrganizationMembership]": - membership_table: Final[LiteLLM_OrganizationMembershipActions[prisma_models.LiteLLM_OrganizationMembership]] = ( - OrganizationMembershipRepository(prisma_client).table - ) - return membership_table +) -> "_TableActions[prisma_models.LiteLLM_OrganizationMembership]": + return _typed_table(OrganizationMembershipRepository(prisma_client)) def _invitation_link_table( prisma_client: "PrismaClient | None", -) -> "LiteLLM_InvitationLinkActions[prisma_models.LiteLLM_InvitationLink]": - invitation_table: LiteLLM_InvitationLinkActions[prisma_models.LiteLLM_InvitationLink] = InvitationLinkRepository( - prisma_client - ).table - return invitation_table +) -> "_TableActions[prisma_models.LiteLLM_InvitationLink]": + return _typed_table(InvitationLinkRepository(prisma_client)) def _team_membership_table( prisma_client: "PrismaClient | None", -) -> "LiteLLM_TeamMembershipActions[prisma_models.LiteLLM_TeamMembership]": - team_membership_table: Final[LiteLLM_TeamMembershipActions[prisma_models.LiteLLM_TeamMembership]] = ( - TeamMembershipRepository(prisma_client).table - ) - return team_membership_table +) -> "_TableActions[prisma_models.LiteLLM_TeamMembership]": + return _typed_table(TeamMembershipRepository(prisma_client)) def _hash_password_in_dict(data: dict) -> None: @@ -234,7 +322,7 @@ async def _check_duplicate_user_field( if case_insensitive: where_clause[field_name]["mode"] = "insensitive" - existing_user: Final = await UserRepository(prisma_client).table.find_first(where=where_clause) + existing_user: Final = await _user_table_lenient(prisma_client).find_first(where=where_clause) if existing_user is not None: existing_value: Final = getattr(existing_user, field_name, value) @@ -650,7 +738,7 @@ async def ui_get_available_role( def get_team_from_list( - team_list: list[LiteLLM_TeamTable] | list[TeamListResponseObject] | None, + team_list: Sequence[LiteLLM_TeamTable] | Sequence[TeamListResponseObject] | None, team_id: str, ) -> LiteLLM_TeamTable | LiteLLM_TeamMembership | None: if team_list is None: @@ -732,18 +820,59 @@ def _enforce_user_info_access(user_id: str | None, user_api_key_dict: UserAPIKey ) -async def _get_user_info_teams( - prisma_client: Any, +_TeamIdList: TypeAlias = list[str] + + +class _UserInfoDataClient(Protocol): + @overload + async def get_data(self, *, user_id: str) -> "prisma_models.LiteLLM_UserTable | None": ... + + @overload + async def get_data( + self, + *, + user_id: str | None, + table_name: Literal["key"], + query_type: Literal["find_all"], + ) -> "Sequence[LiteLLM_VerificationToken] | None": ... + + @overload + async def get_data( + self, + *, + team_id_list: _TeamIdList, + table_name: Literal["team"], + query_type: Literal["find_all"], + ) -> "Sequence[TeamListResponseObject] | None": ... + + +async def _get_user_info_row( + prisma_client: "_UserInfoDataClient", + user_id: str, +) -> "prisma_models.LiteLLM_UserTable | None": + return await prisma_client.get_data(user_id=user_id) + + +async def _get_user_info_keys( + prisma_client: "_UserInfoDataClient", user_id: str | None, - user_info: Any | None, +) -> "Sequence[LiteLLM_VerificationToken] | None": + return await prisma_client.get_data( + user_id=user_id, + table_name="key", + query_type="find_all", + ) + + +async def _get_user_info_teams( + prisma_client: "_UserInfoDataClient", + user_id: str | None, + user_info: "prisma_models.LiteLLM_UserTable", user_api_key_dict: UserAPIKeyAuth, -) -> tuple[list[Any], list[Any] | None]: +) -> tuple[Sequence[TeamListResponseObject], Sequence[TeamListResponseObject] | None]: """Fetch and merge teams from membership + user.teams field.""" from litellm.proxy.management_endpoints.team_endpoints import list_team - team_list: list[Any] = [] - team_id_list: list[str] = [] - teams_1: Final = await list_team( http_request=Request( scope={"type": "http", "path": "/user/info"}, @@ -752,11 +881,10 @@ async def _get_user_info_teams( user_api_key_dict=user_api_key_dict, ) - if teams_1 is not None and isinstance(teams_1, list): - team_list = teams_1 - team_id_list = [team.team_id for team in teams_1] + team_list: Final = teams_1 if teams_1 is not None and isinstance(teams_1, list) else list[TeamListResponseObject]() + team_id_list: Final = [team.team_id for team in team_list] - teams_2: list[Any] | None = None + teams_2: Sequence[TeamListResponseObject] | None = None target_team_ids: Final = getattr(user_info, "teams", None) if target_team_ids and isinstance(target_team_ids, list): @@ -767,7 +895,7 @@ async def _get_user_info_teams( ) elif user_api_key_dict.user_id is not None and user_id is None: caller_user_info: Final = await prisma_client.get_data(user_id=user_api_key_dict.user_id) - caller_team_ids: Final = getattr(caller_user_info, "teams", None) + caller_team_ids: Final = caller_user_info.teams if caller_user_info is not None else None if caller_team_ids: teams_2 = await prisma_client.get_data( team_id_list=caller_team_ids, @@ -804,9 +932,9 @@ def _redact_scim_enterprise_metadata( def _build_user_info_response( user_id: str | None, user_info: Any | None, - keys: list[LiteLLM_VerificationToken] | None, - team_list: list[Any], - teams_1: list[Any] | None, + keys: Sequence[LiteLLM_VerificationToken] | None, + team_list: Sequence[TeamListResponseObject], + teams_1: Sequence[TeamListResponseObject] | None, ) -> UserInfoResponse: """Create UserInfoResponse while filtering sensitive fields.""" if user_info is None and keys is not None: @@ -814,7 +942,7 @@ def _build_user_info_response( user_info = {"spend": spend} returned_keys: Final = _process_keys_for_user_info(keys=keys, all_teams=teams_1) - team_list.sort(key=lambda x: getattr(x, "team_alias", "") or "") + sorted_team_list: Final = sorted(team_list, key=lambda x: getattr(x, "team_alias", "") or "") _user_info: Final = user_info.model_dump() if isinstance(user_info, BaseModel) else user_info if isinstance(_user_info, dict): @@ -825,7 +953,7 @@ def _build_user_info_response( user_id=user_id, user_info=_user_info, keys=returned_keys, - teams=team_list, + teams=sorted_team_list, ) @@ -870,9 +998,9 @@ async def user_info( user_id = user_api_key_dict.user_id ## GET USER ROW ## - user_info = None + user_info: prisma_models.LiteLLM_UserTable | None = None if user_id is not None: - user_info = await prisma_client.get_data(user_id=user_id) + user_info = await _get_user_info_row(prisma_client, user_id) if user_info is None: raise HTTPException( @@ -888,11 +1016,7 @@ async def user_info( ) ## GET ALL KEYS ## - keys: Final = await prisma_client.get_data( - user_id=user_id, - table_name="key", - query_type="find_all", - ) + keys: Final = await _get_user_info_keys(prisma_client, user_id) response_data: Final = _build_user_info_response( user_id=user_id, @@ -1058,6 +1182,12 @@ async def user_info_v2( raise handle_exception_on_proxy(e) +async def _fetch_admin_teams_and_keys_rows( + prisma_client: "PrismaClient", sql_query: str +) -> Sequence[Mapping[str, Sequence[Mapping[str, object]] | None]]: + return await prisma_client.db.query_raw(sql_query) + + async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth): """ Admin UI Endpoint - Returns All Teams and Keys when Proxy Admin is querying @@ -1081,22 +1211,25 @@ async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth): "Database not connected. Connect a database to your proxy - https://docs.litellm.ai/docs/simple_proxy#managing-auth---virtual-keys" ) - results: Final = await prisma_client.db.query_raw(sql_query) + results: Final = await _fetch_admin_teams_and_keys_rows(prisma_client, sql_query) verbose_proxy_logger.debug("results_keys: %s", results) - _keys_in_db: Final[list] = results[0]["keys"] or [] + _keys_in_db: Final[Sequence[Mapping[str, object]]] = results[0]["keys"] or [] # cast all keys to LiteLLM_VerificationToken keys_in_db: Final = [] for key in _keys_in_db: - if key.get("models") is None: - key["models"] = [] - keys_in_db.append(LiteLLM_VerificationToken.model_validate(key)) + key_payload = dict[str, object](key) + if key_payload.get("models") is None: + key_payload["models"] = [] + keys_in_db.append(LiteLLM_VerificationToken.model_validate(key_payload)) # cast all teams to LiteLLM_TeamTable - _teams_in_db: list = results[0]["teams"] or [] - _teams_in_db = [LiteLLM_TeamTable.model_validate(team) for team in _teams_in_db] - _teams_in_db.sort(key=lambda x: getattr(x, "team_alias", "") or "") + _teams_rows: Final[Sequence[Mapping[str, object]]] = results[0]["teams"] or [] + _teams_in_db: Final = sorted( + (LiteLLM_TeamTable.model_validate(team) for team in _teams_rows), + key=lambda x: getattr(x, "team_alias", "") or "", + ) returned_keys: Final = _process_keys_for_user_info(keys=keys_in_db, all_teams=_teams_in_db) # Get admin's own user_id and user_info @@ -1121,8 +1254,8 @@ async def _get_user_info_for_proxy_admin(user_api_key_dict: UserAPIKeyAuth): def _process_keys_for_user_info( - keys: list[LiteLLM_VerificationToken] | None, - all_teams: list[LiteLLM_TeamTable] | list[TeamListResponseObject] | None, + keys: Sequence[LiteLLM_VerificationToken] | None, + all_teams: Sequence[LiteLLM_TeamTable] | Sequence[TeamListResponseObject] | None, ): from litellm.constants import UI_SESSION_TOKEN_TEAM_ID from litellm.proxy.proxy_server import general_settings, litellm_master_key_hash @@ -1212,7 +1345,7 @@ def _update_internal_user_params(data_json: dict, data: UpdateUserRequest | Upda async def _schedule_user_update_audit_log( - response: dict[str, Any], + response: Mapping[str, object], existing_user_row: BaseModel | None, litellm_changed_by: str | None, user_api_key_dict: UserAPIKeyAuth, @@ -1768,7 +1901,10 @@ async def bulk_user_update( # Apply update transformations (reuse existing logic) data_json: Final[dict] = data.user_updates.model_dump(exclude_unset=True) - non_default_values: Final = _update_internal_user_params(data_json=data_json, data=data.user_updates) + _raw_update_values: Final[Mapping[str, object]] = _update_internal_user_params( + data_json=data_json, data=data.user_updates + ) + non_default_values: Final = dict[str, object](_raw_update_values) # Remove user identification fields since we're updating by user_id non_default_values.pop("user_id", None) @@ -1780,7 +1916,7 @@ async def bulk_user_update( try: # Perform bulk database update - await UserRepository(prisma_client).table.update_many( + await _user_table_lenient(prisma_client).update_many( where={}, data=non_default_values, # Update all users ) @@ -1885,7 +2021,7 @@ async def get_user_key_counts( # Get count for each user_id individually for user_id in user_ids: - count = await VerificationTokenRepository(prisma_client).table.count( + count = await _verification_token_table_lenient(prisma_client).count( where={ "user_id": user_id, "OR": [ @@ -2122,7 +2258,7 @@ async def get_users( _validate_sort_params(sort_by, sort_order) if sort_by is not None and isinstance(sort_by, str) else None ) - users: Sequence[prisma_models.LiteLLM_UserTable] | None = await UserRepository(prisma_client).table.find_many( + users: Sequence[prisma_models.LiteLLM_UserTable] | None = await _user_table_lenient(prisma_client).find_many( where=where_conditions, skip=skip, take=page_size, @@ -2130,7 +2266,7 @@ async def get_users( ) # Get total count of user rows - total_count: Final[int] = await UserRepository(prisma_client).table.count(where=where_conditions) + total_count: Final[int] = await _user_table_lenient(prisma_client).count(where=where_conditions) # Get key count for each user if users is not None: @@ -2256,7 +2392,7 @@ async def delete_user( # check that all teams passed exist for user_id in data.user_ids: - user_row = await UserRepository(prisma_client).table.find_unique(where={"user_id": user_id}) + user_row = await _user_delete_table(prisma_client).find_unique(where={"user_id": user_id}) if user_row is None: raise HTTPException( @@ -2308,8 +2444,8 @@ async def delete_user( ) ## CLEANUP MEMBERS_WITH_ROLES - fetch_all_teams = await TeamRepository(prisma_client).table.find_many(where={"team_id": {"in": user_row.teams}}) - teams_to_update = [] + fetch_all_teams = await _team_cleanup_table(prisma_client).find_many(where={"team_id": {"in": user_row.teams}}) + teams_to_update = list[_TeamCleanupRow]() for team in fetch_all_teams: is_member_in_team, new_team_members = _cleanup_members_with_roles( existing_team_row=LiteLLM_TeamTable.model_validate(team.model_dump()), @@ -2327,7 +2463,7 @@ async def delete_user( ## update teams for team in teams_to_update: - await TeamRepository(prisma_client).table.update( + await _team_cleanup_table(prisma_client).update( where={"team_id": team.team_id}, data={"members_with_roles": team.members_with_roles}, ) @@ -2382,14 +2518,14 @@ async def add_internal_user_to_organization( try: # Check if organization_id exists - organization_row: Final = await OrganizationRepository(prisma_client).table.find_unique( + organization_row: Final = await _organization_table(prisma_client).find_unique( where={"organization_id": organization_id} ) if organization_row is None: raise Exception(f"Organization not found, passed organization_id={organization_id}") # Create a new organization membership entry - new_membership: Final = await OrganizationMembershipRepository(prisma_client).table.create( + new_membership: Final = await _organization_membership_table(prisma_client).create( data={ "user_id": user_id, "organization_id": organization_id, diff --git a/litellm/proxy/management_endpoints/key_management_endpoints.py b/litellm/proxy/management_endpoints/key_management_endpoints.py index 2a385c4c42a..8fe388d2fea 100644 --- a/litellm/proxy/management_endpoints/key_management_endpoints.py +++ b/litellm/proxy/management_endpoints/key_management_endpoints.py @@ -18,9 +18,10 @@ import os import re import secrets import traceback -from collections.abc import Awaitable, Callable, Mapping, Sequence +from collections.abc import Awaitable, Callable, Iterator, Mapping, Sequence +from contextlib import AbstractAsyncContextManager from datetime import datetime, timedelta, timezone -from typing import Any, Final, Literal, Optional, Protocol, TypeVar, cast +from typing import Any, Final, Literal, Optional, Protocol, TypeAlias, TypeVar, cast import fastapi import yaml @@ -88,6 +89,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 import object_permission_utils 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, @@ -144,6 +146,7 @@ from litellm.types.utils import ( ) _PrismaRowT = TypeVar("_PrismaRowT") +_PrismaRowCoT: Final = TypeVar("_PrismaRowCoT", covariant=True) _RepositoryModelT = TypeVar("_RepositoryModelT", bound=BaseModel) @@ -169,7 +172,7 @@ class _PrismaTableActions(Protocol[_PrismaRowT]): *, where: Mapping[str, object] | None = None, include: Mapping[str, object] | None = None, - order: Mapping[str, object] | None = None, + order: Mapping[str, object] | Sequence[Mapping[str, object]] | None = None, skip: int | None = None, take: int | None = None, ) -> list[_PrismaRowT]: ... @@ -189,17 +192,73 @@ class _PrismaTableActions(Protocol[_PrismaRowT]): data: Mapping[str, object], ) -> _PrismaRowT | None: ... + async def upsert( + self, + *, + where: Mapping[str, object], + data: Mapping[str, object], + ) -> _PrismaRowT: ... -class _UserRowLike(Protocol): - user_id: str | None - user_email: str | None - user_alias: str | None - def model_dump(self) -> Mapping[str, object]: ... +class _PrismaTableHolder(Protocol[_PrismaRowT]): + @property + def table(self) -> _PrismaTableActions[_PrismaRowT]: ... + +def _typed_table(holder: _PrismaTableHolder[_PrismaRowT]) -> _PrismaTableActions[_PrismaRowT]: + return holder.table + + +class _CustomKeyHooksModule(Protocol): + user_custom_key_generate: Callable[..., Awaitable[Mapping[str, object]]] | None + user_custom_key_update: Callable[..., Awaitable[Mapping[str, object]]] | None + + +def _custom_key_generate_hook( + hooks: _CustomKeyHooksModule, +) -> Callable[..., Awaitable[Mapping[str, object]]] | None: + return hooks.user_custom_key_generate + + +def _custom_key_update_hook( + hooks: _CustomKeyHooksModule, +) -> Callable[..., Awaitable[Mapping[str, object]]] | None: + return hooks.user_custom_key_update + + +class _LegacyDumpable(Protocol): def dict(self) -> Mapping[str, object]: ... +def _legacy_model_dict(row: _LegacyDumpable) -> Mapping[str, object]: + return row.dict() + + +class _PrismaTableLenient(Protocol[_PrismaRowCoT]): + async def find_unique(self, *, where: Mapping[str, object]) -> _PrismaRowCoT: ... + + async def find_many(self, *, where: Mapping[str, object] | None = None) -> Sequence[_PrismaRowCoT] | None: ... + + +class _PrismaTableLenientHolder(Protocol[_PrismaRowCoT]): + @property + def table(self) -> _PrismaTableLenient[_PrismaRowCoT]: ... + + +def _lenient_table(holder: _PrismaTableLenientHolder[_PrismaRowCoT]) -> _PrismaTableLenient[_PrismaRowCoT]: + return holder.table + + +def _prisma_table_lenient( + repository: BaseRepository[_RepositoryModelT], +) -> _PrismaTableLenient[_RepositoryModelT]: + return _lenient_table(repository) + + +def _jsonify_for_db(client: PrismaClient, data: Mapping[str, object]) -> Mapping[str, object]: + return client.jsonify_object(dict[str, object](data)) + + class _TxTables(Protocol): litellm_proxymodeltable: _PrismaTableActions[object] @@ -207,21 +266,82 @@ class _TxTables(Protocol): def _prisma_table( repository: BaseRepository[_RepositoryModelT], ) -> _PrismaTableActions[_RepositoryModelT]: - return repository.table + return _typed_table(repository) def _deleted_verification_token_table( prisma_client: PrismaClient, ) -> _PrismaTableActions[LiteLLM_DeletedVerificationToken]: - return DeletedVerificationTokenRepository(prisma_client).table + return _typed_table(DeletedVerificationTokenRepository(prisma_client)) def _credentials_table(prisma_client: PrismaClient) -> _PrismaTableActions[CredentialItem]: - return CredentialsRepository(prisma_client).table + return _typed_table(CredentialsRepository(prisma_client)) def _config_table(prisma_client: PrismaClient) -> _PrismaTableActions[ConfigParam]: - return ConfigRepository(prisma_client).table + return _typed_table(ConfigRepository(prisma_client)) + + +def _deprecated_verification_token_table(prisma_client: PrismaClient) -> _PrismaTableActions[object]: + return _typed_table(DeprecatedVerificationTokenRepository(prisma_client)) + + +_StringList: TypeAlias = list[str] + + +class _CreatedUserRow(Protocol): + models: _StringList + + +def _created_user_row(user_row: "_CreatedUserRow | None") -> "_CreatedUserRow | None": + return user_row + + +async def _query_raw_text_rows(prisma_client: PrismaClient, sql: str, *params: object) -> Sequence[Mapping[str, str]]: + return await prisma_client.db.query_raw(sql, *params) + + +def _as_object_dict(values: Mapping[str, object]) -> Mapping[str, object]: + return values + + +def _model_items(model: BaseModel) -> Iterator[tuple[str, object]]: + return iter(model) + + +class _SpendCache(Protocol): + async def async_get_cache(self, key: str) -> float | None: ... + + +def _spend_cache(cache: _SpendCache) -> _SpendCache: + return cache + + +class _ObjectPermissionUtils(Protocol): + @property + def attach_object_permission_to_dict( + self, + ) -> Callable[..., Awaitable[Mapping[str, object]]]: ... + + +def _object_permission_utils(module: _ObjectPermissionUtils) -> _ObjectPermissionUtils: + return module + + +class _EnvVarsParam(Protocol): + @property + def param_value(self) -> Mapping[str, str] | None: ... + + +def _env_vars_param_value(param: _EnvVarsParam) -> Mapping[str, str] | None: + return param.param_value + + +def _tx_tables_context( + open_tx: Callable[[], AbstractAsyncContextManager[_TxTables]], +) -> AbstractAsyncContextManager[_TxTables]: + return open_tx() async def _check_custom_key_allowed(custom_key_value: str | None) -> None: @@ -886,7 +1006,7 @@ async def _common_key_generation_helper( # check if user set default key/generate params on config.yaml if litellm.default_key_generate_params is not None: - for elem in data: + for elem in _model_items(data): key, value = elem if value is None and key in [ "max_budget", @@ -984,9 +1104,9 @@ async def _common_key_generation_helper( soft_budget=data.soft_budget, model_max_budget=data.model_max_budget or {}, ) - new_budget: Final = prisma_client.jsonify_object(budget_row.json(exclude_none=True)) + new_budget: Final = _jsonify_for_db(prisma_client, budget_row.json(exclude_none=True)) - _budget: Final = await BudgetRepository(prisma_client).table.create( + _budget: Final = await _prisma_table(BudgetRepository(prisma_client)).create( data={ **new_budget, "created_by": user_api_key_dict.user_id or litellm_proxy_admin_name, @@ -1655,11 +1775,11 @@ async def generate_key_fn( - user_id: (str) Unique user id - used for tracking spend across multiple keys for same user id. """ try: + from litellm.proxy import proxy_server from litellm.proxy._types import CommonProxyErrors from litellm.proxy.proxy_server import ( prisma_client, user_api_key_cache, - user_custom_key_generate, ) if prisma_client is None: @@ -1686,7 +1806,7 @@ async def generate_key_fn( ) custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = ( - user_custom_key_generate + _custom_key_generate_hook(proxy_server) ) if custom_key_generate_hook is not None: if inspect.iscoroutinefunction(custom_key_generate_hook): @@ -1855,11 +1975,11 @@ async def generate_service_account_key_fn( - user_id: (str) Unique user id - used for tracking spend across multiple keys for same user id. """ + from litellm.proxy import proxy_server from litellm.proxy._types import CommonProxyErrors from litellm.proxy.proxy_server import ( prisma_client, user_api_key_cache, - user_custom_key_generate, ) if prisma_client is None: @@ -1887,7 +2007,9 @@ async def generate_service_account_key_fn( verbose_proxy_logger.debug("entered /key/generate") - custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = user_custom_key_generate + custom_key_generate_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = _custom_key_generate_hook( + proxy_server + ) if custom_key_generate_hook is not None: if inspect.iscoroutinefunction(custom_key_generate_hook): result: Final = await custom_key_generate_hook(data) @@ -1959,7 +2081,7 @@ def prepare_metadata_fields(data: BaseModel, non_default_values: dict, existing_ ) casted_metadata[reserved_field] = existing_value - data_json: Final[Mapping[str, object]] = data.model_dump(exclude_unset=True, exclude_none=True) + data_json: Final = _as_object_dict(data.model_dump(exclude_unset=True, exclude_none=True)) try: for k, v in data_json.items(): @@ -2737,13 +2859,13 @@ async def update_key_fn( }' ``` """ + from litellm.proxy import proxy_server from litellm.proxy.proxy_server import ( llm_router, premium_user, prisma_client, proxy_logging_obj, user_api_key_cache, - user_custom_key_update, ) try: @@ -2774,7 +2896,9 @@ async def update_key_fn( ) # Custom key update hook - custom_key_update_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = user_custom_key_update + custom_key_update_hook: Final[Callable[..., Awaitable[Mapping[str, object]]] | None] = _custom_key_update_hook( + proxy_server + ) if custom_key_update_hook is not None: if inspect.iscoroutinefunction(custom_key_update_hook): result: Final = await custom_key_update_hook(data) @@ -2927,14 +3051,16 @@ async def bulk_update_keys( }' ``` """ + from litellm.proxy import proxy_server from litellm.proxy.proxy_server import ( llm_router, prisma_client, proxy_logging_obj, user_api_key_cache, - user_custom_key_update, ) + custom_key_update_hook: Final = _custom_key_update_hook(proxy_server) + if user_api_key_dict.user_role != LitellmUserRoles.PROXY_ADMIN.value: raise HTTPException( status_code=403, @@ -2980,7 +3106,7 @@ async def bulk_update_keys( user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, llm_router=llm_router, - user_custom_key_update=user_custom_key_update, + user_custom_key_update=custom_key_update_hook, ) successful_updates.append( @@ -3058,7 +3184,7 @@ def _build_failed_team_key_update( if hasattr(existing_key_row, "model_dump"): key_info = existing_key_row.model_dump() elif hasattr(existing_key_row, "dict"): - key_info = existing_key_row.dict() + key_info = dict[str, object](_legacy_model_dict(existing_key_row)) if key_info: key_info.pop("token", None) @@ -3089,14 +3215,16 @@ async def bulk_update_team_keys( Callable by proxy admins, or by team admins with `KEY_UPDATE` permission. """ + from litellm.proxy import proxy_server from litellm.proxy.proxy_server import ( llm_router, prisma_client, proxy_logging_obj, user_api_key_cache, - user_custom_key_update, ) + custom_key_update_hook: Final = _custom_key_update_hook(proxy_server) + if prisma_client is None: raise HTTPException( status_code=500, @@ -3223,7 +3351,7 @@ async def bulk_update_team_keys( user_api_key_cache=user_api_key_cache, proxy_logging_obj=proxy_logging_obj, llm_router=llm_router, - user_custom_key_update=user_custom_key_update, + user_custom_key_update=custom_key_update_hook, existing_key_row=existing_by_token[db_token], ) @@ -3435,7 +3563,7 @@ async def _get_model_max_budget_current_spend( virtual_key_model_spend_cache_key = ( f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:{api_key_hash}:{model}:{budget_config.budget_duration}" ) - current_spend: float | None = await user_api_key_cache.async_get_cache( + current_spend: float | None = await _spend_cache(user_api_key_cache).async_get_cache( key=virtual_key_model_spend_cache_key, ) if current_spend is None: @@ -3444,7 +3572,7 @@ async def _get_model_max_budget_current_spend( f"{VIRTUAL_KEY_SPEND_CACHE_KEY_PREFIX}:" f"{api_key_hash}:{model_without_prefix}:{budget_config.budget_duration}" ) - current_spend = await user_api_key_cache.async_get_cache( + current_spend = await _spend_cache(user_api_key_cache).async_get_cache( key=virtual_key_model_spend_cache_key, ) try: @@ -3635,7 +3763,7 @@ async def info_key_fn( hashed_key: str | None = key if key is not None: hashed_key = _hash_token_if_needed(token=key) - key_info = await VerificationTokenRepository(prisma_client).table.find_unique( + key_info = await _prisma_table(VerificationTokenRepository(prisma_client)).find_unique( where={"token": hashed_key}, include={"litellm_budget_table": True}, ) @@ -3945,7 +4073,7 @@ async def generate_key_helper_fn( if table_name is None or table_name == "user": # do not auto-create users for `/key/generate` ## CREATE USER (If necessary) if query_type == "insert_data": - user_row = await prisma_client.insert_data(data=user_data, table_name="user") + user_row = _created_user_row(await prisma_client.insert_data(data=user_data, table_name="user")) if user_row is None: raise Exception("Failed to create user") @@ -4232,7 +4360,7 @@ def _transform_verification_tokens_to_deleted_records( "litellm_changed_by": litellm_changed_by, } ) - record = deleted_record.model_dump() + record = dict[str, object](_as_object_dict(deleted_record.model_dump())) # Map org_id to organization_id (model uses org_id, but schema expects organization_id) org_id_value: object = record.pop("org_id", None) @@ -4355,13 +4483,12 @@ async def _rotate_master_key( should_create_model_in_db=False, ) if new_model: - _dumped = new_model.model_dump(exclude_none=True) + _dumped = dict[str, object](_as_object_dict(new_model.model_dump(exclude_none=True))) _dumped["litellm_params"] = prisma.Json(_dumped["litellm_params"]) _dumped["model_info"] = prisma.Json(_dumped["model_info"]) new_models.append(_dumped) verbose_proxy_logger.debug("Resetting proxy model table") - async with prisma_client.db.tx() as tx_ctx: - tx: Final[_TxTables] = tx_ctx + async with _tx_tables_context(prisma_client.db.tx) as tx: await tx.litellm_proxymodeltable.delete_many() verbose_proxy_logger.debug("Creating %s models", len(new_models)) await tx.litellm_proxymodeltable.create_many( @@ -4376,14 +4503,14 @@ async def _rotate_master_key( if config: """If environment_variables is found, decrypt it and encrypt it with the new master key""" - environment_variables_dict = {} + environment_variables_dict: Mapping[str, str] | None = {} for c in config: if c.param_name == "environment_variables": - environment_variables_dict = c.param_value + environment_variables_dict = _env_vars_param_value(c) if environment_variables_dict: decrypted_env_vars: Final = proxy_config._decrypt_and_set_db_env_variables( - environment_variables=environment_variables_dict + environment_variables=dict[str, str](environment_variables_dict) ) encrypted_env_vars: Final = proxy_config._encrypt_env_variables( environment_variables=decrypted_env_vars, @@ -4449,7 +4576,7 @@ async def _rotate_master_key( updated_patch=decrypted_cred, new_encryption_key=new_master_key, ) - _cred_data = encrypted_cred.model_dump(exclude_none=True) + _cred_data = dict[str, object](_as_object_dict(encrypted_cred.model_dump(exclude_none=True))) if "credential_values" in _cred_data: _cred_data["credential_values"] = prisma.Json(_cred_data["credential_values"]) if "credential_info" in _cred_data: @@ -4605,7 +4732,7 @@ async def _insert_deprecated_key( try: revoke_at: Final = datetime.now(timezone.utc) + timedelta(seconds=grace_seconds) - await DeprecatedVerificationTokenRepository(prisma_client).table.upsert( + await _deprecated_verification_token_table(prisma_client).upsert( where={"token": old_token_hash}, data={ "create": { @@ -4704,11 +4831,13 @@ async def _execute_virtual_key_regeneration( grace_period=data.grace_period if data else None, ) - updated_token: Final[Mapping[str, object] | None] = await VerificationTokenRepository(prisma_client).table.update( + updated_token: Final = await _prisma_table(VerificationTokenRepository(prisma_client)).update( where={"token": hashed_api_key}, 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: Final = ( + dict[str, object](_as_object_dict(dict(updated_token))) if updated_token is not None else dict[str, object]() + ) updated_token_dict["key"] = new_token updated_token_dict["token_id"] = updated_token_dict.pop("token") @@ -5247,7 +5376,7 @@ async def validate_key_list_check( if key_hash: try: - key_info: Final = await VerificationTokenRepository(prisma_client).table.find_unique( + key_info: Final = await _prisma_table_lenient(VerificationTokenRepository(prisma_client)).find_unique( where={"token": key_hash}, ) except Exception: @@ -5278,7 +5407,7 @@ async def _fetch_user_team_objects( if complete_user_info is None or not complete_user_info.teams: return [] - teams: Final[list[BaseModel] | None] = await TeamRepository(prisma_client).table.find_many( + teams: Final = await _prisma_table_lenient(TeamRepository(prisma_client)).find_many( where={"team_id": {"in": complete_user_info.teams}} ) if teams is None: @@ -5653,7 +5782,7 @@ async def key_aliases( where_sql: Final = " AND ".join(where_parts) count_sql: Final = f'SELECT COUNT(*) AS count FROM "LiteLLM_VerificationToken" WHERE {where_sql}' - count_rows: Final[Sequence[Mapping[str, int]]] = await prisma_client.db.query_raw(count_sql, *query_params) + count_rows: Final = await _query_raw_text_rows(prisma_client, count_sql, *query_params) total_count: Final = int(count_rows[0]["count"]) if count_rows else 0 aliases_params: Final = query_params + [size, (page - 1) * size] @@ -5666,7 +5795,7 @@ async def key_aliases( f" ORDER BY key_alias ASC" f" LIMIT ${limit_idx} OFFSET ${offset_idx}" ) - alias_rows: Final[Sequence[Mapping[str, str]]] = await prisma_client.db.query_raw(aliases_sql, *aliases_params) + alias_rows: Final = await _query_raw_text_rows(prisma_client, aliases_sql, *aliases_params) aliases: Final[list[str]] = [row["key_alias"] for row in alias_rows if row.get("key_alias")] total_pages: Final = -(-total_count // size) if total_count > 0 else 0 @@ -5952,7 +6081,7 @@ async def _list_key_helper( # Fetch keys with pagination if use_deleted_table: - keys = await DeletedVerificationTokenRepository(prisma_client).table.find_many( + keys = await _deleted_verification_token_table(prisma_client).find_many( where=where, skip=skip, take=size, @@ -5966,7 +6095,7 @@ async def _list_key_helper( ), ) else: - keys = await VerificationTokenRepository(prisma_client).table.find_many( + keys = await _prisma_table(VerificationTokenRepository(prisma_client)).find_many( where=where, skip=skip, take=size, @@ -5995,13 +6124,13 @@ async def _list_key_helper( total_pages: Final = -(-total_count // size) # Ceiling division # Fetch user information if expand includes "user" - user_map = {} + user_map: Mapping[str, LiteLLM_UserTable] = {} if expand and "user" in expand: user_ids: Final = [key.user_id for key in keys if key.user_id] 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[Sequence[_UserRowLike]] = await UserRepository(prisma_client).table.find_many( + users: Final = await _prisma_table(UserRepository(prisma_client)).find_many( where={"user_id": {"in": all_ids}} ) user_map = {user.user_id: user for user in users} @@ -6014,10 +6143,14 @@ async def _list_key_helper( key_dict = key.model_dump() except Exception: # Fallback for Pydantic v1 compatibility - key_dict = key.dict() + key_dict = dict[str, object](_legacy_model_dict(key)) # Attach object_permission if object_permission_id is set (only for non-deleted keys) if not use_deleted_table: - key_dict = await attach_object_permission_to_dict(key_dict, prisma_client) + key_dict = dict[str, object]( + await _object_permission_utils(object_permission_utils).attach_object_permission_to_dict( + key_dict, prisma_client + ) + ) # Include user information if expand includes "user" if expand and "user" in expand: @@ -6025,7 +6158,7 @@ async def _list_key_helper( try: key_dict["user"] = user_map[key.user_id].model_dump() except Exception: - key_dict["user"] = user_map[key.user_id].dict() + key_dict["user"] = _legacy_model_dict(user_map[key.user_id]) if key.created_by and key.created_by in user_map: created_by_user = user_map[key.created_by] key_dict["created_by_user"] = { @@ -6039,7 +6172,7 @@ async def _list_key_helper( # Use deleted key type to preserve deleted_at, deleted_by, etc. key_list.append(LiteLLM_DeletedVerificationToken.model_validate(key_dict)) else: - key_list.append(UserAPIKeyAuth(**key_dict)) # Return full key object + key_list.append(UserAPIKeyAuth.model_validate(key_dict)) # Return full key object else: _token = key_dict.get("token") key_list.append(cast(str, _token)) # Return only the token diff --git a/litellm/proxy/management_endpoints/tool_management_endpoints.py b/litellm/proxy/management_endpoints/tool_management_endpoints.py index 0fdedafb2bf..dfb1422308f 100644 --- a/litellm/proxy/management_endpoints/tool_management_endpoints.py +++ b/litellm/proxy/management_endpoints/tool_management_endpoints.py @@ -10,8 +10,9 @@ POST /v1/tool/policy - Update the input_policy / output_policy for a """ import uuid +from collections.abc import Mapping, Sequence from datetime import datetime, timedelta, timezone -from typing import TYPE_CHECKING, Annotated, Any, Final +from typing import TYPE_CHECKING, Annotated, Final, Protocol from fastapi import APIRouter, Depends, HTTPException, Query from pydantic import BaseModel, Field, TypeAdapter @@ -49,6 +50,144 @@ from litellm.types.tool_management import ( ToolUsageLogsResponse, ) + +class _DailyToolSpendRecord(Protocol): + date: str + tool_name: str + spend: float + request_count: int + + +class _SpendLogToolIndexRecord(Protocol): + request_id: str + + +class _SpendLogRecord(Protocol): + request_id: str + startTime: datetime + model: str | None + spend: float | None + total_tokens: int | None + messages: object + proxy_server_request: object + + +class _VerificationTokenRecord(Protocol): + object_permission_id: str | None + + +class _TeamRecord(Protocol): + object_permission_id: str | None + + +class _DailyToolSpendTable(Protocol): + async def group_by( + self, + *, + by: Sequence[str], + sum: Mapping[str, bool], + where: Mapping[str, object], + order: Mapping[str, object], + take: int, + ) -> Sequence[object] | None: ... + + async def find_many( + self, + *, + where: Mapping[str, object], + order: Sequence[Mapping[str, str]], + ) -> Sequence[_DailyToolSpendRecord]: ... + + +class _SpendLogToolIndexTable(Protocol): + async def count(self, *, where: Mapping[str, object]) -> int: ... + + async def find_many( + self, + *, + where: Mapping[str, object], + order: Mapping[str, str], + skip: int, + take: int, + ) -> Sequence[_SpendLogToolIndexRecord]: ... + + +class _SpendLogsTable(Protocol): + async def find_many(self, *, where: Mapping[str, object]) -> Sequence[_SpendLogRecord]: ... + + +class _VerificationTokenTable(Protocol): + async def find_unique(self, *, where: Mapping[str, object]) -> _VerificationTokenRecord | None: ... + + async def update_many(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... + + +class _TeamTable(Protocol): + async def find_unique(self, *, where: Mapping[str, object]) -> _TeamRecord | None: ... + + async def update_many(self, *, where: Mapping[str, object], data: Mapping[str, object]) -> int: ... + + +class _ObjectPermissionTable(Protocol): + async def create(self, *, data: Mapping[str, str | Sequence[str]]) -> object: ... + + async def delete(self, *, where: Mapping[str, object]) -> object: ... + + +class _DailyToolSpendTableHolder(Protocol): + @property + def table(self) -> _DailyToolSpendTable: ... + + +class _SpendLogToolIndexTableHolder(Protocol): + @property + def table(self) -> _SpendLogToolIndexTable: ... + + +class _SpendLogsTableHolder(Protocol): + @property + def table(self) -> _SpendLogsTable: ... + + +class _VerificationTokenTableHolder(Protocol): + @property + def table(self) -> _VerificationTokenTable: ... + + +class _TeamTableHolder(Protocol): + @property + def table(self) -> _TeamTable: ... + + +class _ObjectPermissionTableHolder(Protocol): + @property + def table(self) -> _ObjectPermissionTable: ... + + +def _daily_tool_spend_table(repo: _DailyToolSpendTableHolder) -> _DailyToolSpendTable: + return repo.table + + +def _spend_log_tool_index_table(repo: _SpendLogToolIndexTableHolder) -> _SpendLogToolIndexTable: + return repo.table + + +def _spend_logs_table(repo: _SpendLogsTableHolder) -> _SpendLogsTable: + return repo.table + + +def _verification_token_table(repo: _VerificationTokenTableHolder) -> _VerificationTokenTable: + return repo.table + + +def _team_table(repo: _TeamTableHolder) -> _TeamTable: + return repo.table + + +def _object_permission_table(repo: _ObjectPermissionTableHolder) -> _ObjectPermissionTable: + return repo.table + + router: Final = APIRouter() TOOL_POLICY_OPTIONS: Final = ToolPolicyOptionsResponse( @@ -154,6 +293,7 @@ class _TopToolRow(BaseModel): _TOP_TOOL_ROWS: Final = TypeAdapter(list[_TopToolRow]) +_PARSED_JSON: Final = TypeAdapter(object) @router.get( @@ -201,7 +341,7 @@ async def get_tool_spend( end_str: Final = end_day.strftime("%Y-%m-%d") date_window: Final = {"date": {"gte": start_str, "lte": end_str}} - table: Final = DailyToolSpendRepository(prisma_client).table + table: Final = _daily_tool_spend_table(DailyToolSpendRepository(prisma_client)) top_tools: Final = _TOP_TOOL_ROWS.validate_python( await table.group_by( by=["tool_name"], @@ -222,7 +362,7 @@ async def get_tool_spend( for row in top_tools ] - daily_rows: Final = ( + daily_rows: Final[Sequence[_DailyToolSpendRecord]] = ( await table.find_many( where={**date_window, "tool_name": {"in": [row.tool_name for row in top_tools]}}, order=[{"date": "asc"}, {"spend": "desc"}], @@ -270,23 +410,23 @@ async def get_tool_detail( raise HTTPException(status_code=500, detail=str(e)) -def _input_snippet_for_tool_log(sl: Any, max_len: int = 200) -> str | None: +def _input_snippet_for_tool_log(sl: _SpendLogRecord | None, max_len: int = 200) -> str | None: """Short snippet from messages or proxy_server_request for tool usage log row.""" if sl is None: return None - messages: Final = getattr(sl, "messages", None) + messages: Final[object] = getattr(sl, "messages", None) if messages is not None: s = _snippet_str(messages, max_len) if s: return s - psr = getattr(sl, "proxy_server_request", None) + psr: object = getattr(sl, "proxy_server_request", None) if not psr: return None if isinstance(psr, str): import json try: - psr = json.loads(psr) + psr = _PARSED_JSON.validate_python(json.loads(psr)) except Exception: return _snippet_str(psr, max_len) if isinstance(psr, dict): @@ -299,7 +439,7 @@ def _input_snippet_for_tool_log(sl: Any, max_len: int = 200) -> str | None: return _snippet_str(psr, max_len) -def _snippet_str(text: Any, max_len: int = 200) -> str | None: +def _snippet_str(text: object, max_len: int = 200) -> str | None: if text is None: return None if isinstance(text, str): @@ -344,10 +484,9 @@ async def get_tool_usage_logs( raise HTTPException(status_code=500, detail=CommonProxyErrors.db_not_connected_error.value) try: - where: Final[dict] = {"tool_name": tool_name} + start_time_filter: datetime | None = None + end_time_filter: datetime | None = None if start_date or end_date: - start_time_filter: datetime | None = None - end_time_filter: datetime | None = None if start_date: try: start_time_filter = datetime.strptime(start_date + "T00:00:00", "%Y-%m-%dT%H:%M:%S").replace( @@ -362,15 +501,15 @@ async def get_tool_usage_logs( ) except ValueError: pass - if start_time_filter is not None or end_time_filter is not None: - where["start_time"] = {} - if start_time_filter is not None: - where["start_time"]["gte"] = start_time_filter - if end_time_filter is not None: - where["start_time"]["lte"] = end_time_filter + start_time_range: Final[Mapping[str, datetime]] = { + key: value for key, value in (("gte", start_time_filter), ("lte", end_time_filter)) if value is not None + } + where: Final[Mapping[str, str | Mapping[str, datetime]]] = ( + {"tool_name": tool_name, "start_time": start_time_range} if start_time_range else {"tool_name": tool_name} + ) - total: Final = await SpendLogToolIndexRepository(prisma_client).table.count(where=where) - index_rows: Final = await SpendLogToolIndexRepository(prisma_client).table.find_many( + total: Final = await _spend_log_tool_index_table(SpendLogToolIndexRepository(prisma_client)).count(where=where) + index_rows: Final = await _spend_log_tool_index_table(SpendLogToolIndexRepository(prisma_client)).find_many( where=where, order={"start_time": "desc"}, skip=(page - 1) * page_size, @@ -380,7 +519,9 @@ async def get_tool_usage_logs( if not request_ids: return ToolUsageLogsResponse(logs=[], total=total, page=page, page_size=page_size) - spend_logs = await SpendLogsRepository(prisma_client).table.find_many(where={"request_id": {"in": request_ids}}) + spend_logs = await _spend_logs_table(SpendLogsRepository(prisma_client)).find_many( + where={"request_id": {"in": request_ids}} + ) log_by_id: Final = {s.request_id: s for s in spend_logs} logs_out: Final[list[ToolUsageLogEntry]] = [] @@ -449,23 +590,29 @@ async def _resolve_key_hash_to_object_permission_id( hashed: Final = key_hash if "sk-" not in (key_hash or "") else hash_token(key_hash) if not hashed: return None - row = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": hashed}) + row = await _verification_token_table(VerificationTokenRepository(prisma_client)).find_unique( + where={"token": hashed} + ) if row is None: return None - op_id: Final = getattr(row, "object_permission_id", None) + op_id: Final[str | None] = getattr(row, "object_permission_id", None) if op_id: return op_id new_id: Final = str(uuid.uuid4()) - await ObjectPermissionRepository(prisma_client).table.create( + await _object_permission_table(ObjectPermissionRepository(prisma_client)).create( data={"object_permission_id": new_id, "blocked_tools": []} ) - updated_count: Final = await VerificationTokenRepository(prisma_client).table.update_many( + updated_count: Final = await _verification_token_table(VerificationTokenRepository(prisma_client)).update_many( where={"token": hashed, "object_permission_id": None}, data={"object_permission_id": new_id}, ) if updated_count == 0: - await ObjectPermissionRepository(prisma_client).table.delete(where={"object_permission_id": new_id}) - row = await VerificationTokenRepository(prisma_client).table.find_unique(where={"token": hashed}) + await _object_permission_table(ObjectPermissionRepository(prisma_client)).delete( + where={"object_permission_id": new_id} + ) + row = await _verification_token_table(VerificationTokenRepository(prisma_client)).find_unique( + where={"token": hashed} + ) return getattr(row, "object_permission_id", None) if row else None return new_id @@ -478,23 +625,25 @@ async def _resolve_team_id_to_object_permission_id( if not team_id or not team_id.strip(): return None team_id_clean: Final = team_id.strip() - row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id_clean}) + row = await _team_table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id_clean}) if row is None: return None - op_id: Final = getattr(row, "object_permission_id", None) + op_id: Final[str | None] = getattr(row, "object_permission_id", None) if op_id: return op_id new_id: Final = str(uuid.uuid4()) - await ObjectPermissionRepository(prisma_client).table.create( + await _object_permission_table(ObjectPermissionRepository(prisma_client)).create( data={"object_permission_id": new_id, "blocked_tools": []} ) - updated_count: Final = await TeamRepository(prisma_client).table.update_many( + updated_count: Final = await _team_table(TeamRepository(prisma_client)).update_many( where={"team_id": team_id_clean, "object_permission_id": None}, data={"object_permission_id": new_id}, ) if updated_count == 0: - await ObjectPermissionRepository(prisma_client).table.delete(where={"object_permission_id": new_id}) - row = await TeamRepository(prisma_client).table.find_unique(where={"team_id": team_id_clean}) + await _object_permission_table(ObjectPermissionRepository(prisma_client)).delete( + where={"object_permission_id": new_id} + ) + row = await _team_table(TeamRepository(prisma_client)).find_unique(where={"team_id": team_id_clean}) return getattr(row, "object_permission_id", None) if row else None return new_id diff --git a/litellm/proxy/prompts/prompt_endpoints.py b/litellm/proxy/prompts/prompt_endpoints.py index d8e9f8dfaee..d1dc1482401 100644 --- a/litellm/proxy/prompts/prompt_endpoints.py +++ b/litellm/proxy/prompts/prompt_endpoints.py @@ -3,8 +3,9 @@ CRUD ENDPOINTS FOR PROMPTS """ import tempfile +from collections.abc import Mapping, Sequence from pathlib import Path -from typing import Any, Final, cast +from typing import TYPE_CHECKING, Final, Protocol, cast from fastapi import ( APIRouter, @@ -15,7 +16,7 @@ from fastapi import ( Response, UploadFile, ) -from pydantic import BaseModel +from pydantic import BaseModel, TypeAdapter from litellm._logging import verbose_proxy_logger from litellm.proxy._types import ( @@ -38,8 +39,47 @@ from litellm.types.prompts.init_prompts import ( ) from litellm.types.proxy.prompt_endpoints import TestPromptRequest +if TYPE_CHECKING: + from litellm.proxy.prompts.prompt_registry import InMemoryPromptRegistry + from litellm.proxy.utils import PrismaClient + + +class _PromptRecord(Protocol): + id: str + version: int + environment: str | None + + +class _PromptTable(Protocol): + async def find_many( + self, + *, + where: Mapping[str, object], + order: Mapping[str, str] | None = None, + take: int | None = None, + distinct: Sequence[str] | None = None, + ) -> Sequence[_PromptRecord]: ... + + async def create(self, *, data: Mapping[str, str | int | None]) -> _PromptRecord: ... + + async def update(self, *, where: Mapping[str, object], data: Mapping[str, str]) -> _PromptRecord: ... + + async def delete_many(self, *, where: Mapping[str, str]) -> object: ... + + +class _PromptTableHolder(Protocol): + @property + def table(self) -> _PromptTable: ... + + +def _prompt_table(repo: _PromptTableHolder) -> _PromptTable: + return repo.table + + router: Final = APIRouter() +_PARSED_VALUE: Final = TypeAdapter(object) + def get_base_prompt_id(prompt_id: str) -> str: """ @@ -132,7 +172,7 @@ def construct_versioned_prompt_id(prompt_id: str, version: int | None = None) -> return f"{base_id}.v{version}" -def get_latest_version_prompt_id(prompt_id: str, all_prompt_ids: dict[str, Any]) -> str: +def get_latest_version_prompt_id(prompt_id: str, all_prompt_ids: Mapping[str, object]) -> str: """ Find the latest version of a prompt from available prompt IDs. @@ -198,7 +238,9 @@ def get_latest_prompt_versions(prompts: list[PromptSpec]) -> list[PromptSpec]: return list(latest_prompts.values()) -async def get_next_version_for_prompt(prisma_client, prompt_id: str, environment: str = "development") -> int: +async def get_next_version_for_prompt( + prisma_client: "PrismaClient", prompt_id: str, environment: str = "development" +) -> int: """ Get the next version number for a prompt in a specific environment. @@ -210,7 +252,7 @@ async def get_next_version_for_prompt(prisma_client, prompt_id: str, environment Returns: Next version number (1 if no versions exist, max_version + 1 otherwise) """ - existing_prompts: Final = await PromptRepository(prisma_client).table.find_many( + existing_prompts: Final = await _prompt_table(PromptRepository(prisma_client)).find_many( where={"prompt_id": prompt_id, "environment": environment} ) @@ -431,10 +473,10 @@ async def get_prompt_versions( # Query DB for versions versioned_prompts: Final = [] if prisma_client is not None: - where_clause: Final[dict[str, Any]] = {"prompt_id": base_prompt_id} - if environment: - where_clause["environment"] = environment - db_prompts: Final = await PromptRepository(prisma_client).table.find_many( + where_clause: Final[Mapping[str, str]] = ( + {"prompt_id": base_prompt_id, "environment": environment} if environment else {"prompt_id": base_prompt_id} + ) + db_prompts: Final = await _prompt_table(PromptRepository(prisma_client)).find_many( where=where_clause, order={"version": "desc"}, ) @@ -590,7 +632,7 @@ async def get_prompt_info( # Query all environments this prompt exists in (lightweight: distinct on environment) all_environments: list[str] = [] if prisma_client is not None: - all_prompt_rows: Final = await PromptRepository(prisma_client).table.find_many( + all_prompt_rows: Final = await _prompt_table(PromptRepository(prisma_client)).find_many( where={"prompt_id": base_prompt_id}, distinct=["environment"], ) @@ -602,13 +644,16 @@ async def get_prompt_info( prompt_spec = None requested_version: Final = get_version_number(prompt_id=prompt_id) if prompt_id != base_prompt_id else None if environment and prisma_client is not None: - where_clause: Final[dict[str, Any]] = { - "prompt_id": base_prompt_id, - "environment": environment, + where_clause: Final[Mapping[str, str | int]] = { + key: value + for key, value in ( + ("prompt_id", base_prompt_id), + ("environment", environment), + ("version", requested_version), + ) + if value is not None } - if requested_version is not None: - where_clause["version"] = requested_version - env_prompts: Final = await PromptRepository(prisma_client).table.find_many( + env_prompts: Final = await _prompt_table(PromptRepository(prisma_client)).find_many( where=where_clause, order={"version": "desc"}, take=1, @@ -721,7 +766,7 @@ async def create_prompt( ) # Store prompt in db with version - prompt_db_entry: Final = await PromptRepository(prisma_client).table.create( + prompt_db_entry: Final = await _prompt_table(PromptRepository(prisma_client)).create( data={ "prompt_id": request.prompt_id, "version": new_version, @@ -811,7 +856,9 @@ async def update_prompt( ) # Check if any version of this prompt exists (in any environment) - existing_prompts = await PromptRepository(prisma_client).table.find_many(where={"prompt_id": base_prompt_id}) + existing_prompts = await _prompt_table(PromptRepository(prisma_client)).find_many( + where={"prompt_id": base_prompt_id} + ) if not existing_prompts: raise HTTPException( @@ -835,7 +882,7 @@ async def update_prompt( ) # Store new version in db - prompt_db_entry: Final = await PromptRepository(prisma_client).table.create( + prompt_db_entry: Final = await _prompt_table(PromptRepository(prisma_client)).create( data={ "prompt_id": base_prompt_id, "version": new_version, @@ -936,12 +983,12 @@ async def delete_prompt( base_prompt_id: Final = get_base_prompt_id(prompt_id=prompt_id) # Build delete filter; scope to environment if provided - delete_where: Final[dict[str, Any]] = {"prompt_id": base_prompt_id} - if environment: - delete_where["environment"] = environment + delete_where: Final[Mapping[str, str]] = ( + {"prompt_id": base_prompt_id, "environment": environment} if environment else {"prompt_id": base_prompt_id} + ) # Delete versions from the database (scoped to environment if provided) - await PromptRepository(prisma_client).table.delete_many(where=delete_where) + await _prompt_table(PromptRepository(prisma_client)).delete_many(where=delete_where) # Remove matching prompts from memory — scope to environment if provided if environment: @@ -967,7 +1014,9 @@ async def delete_prompt( raise HTTPException(status_code=500, detail=str(e)) -def _reload_prompt_in_registry(registry: Any, versioned_id: str, updated_prompt_spec: PromptSpec) -> PromptSpec: +def _reload_prompt_in_registry( + registry: "InMemoryPromptRegistry", versioned_id: str, updated_prompt_spec: PromptSpec +) -> PromptSpec: """Remove stale entry and re-initialize the prompt in the in-memory registry.""" if versioned_id in registry.IN_MEMORY_PROMPTS: del registry.IN_MEMORY_PROMPTS[versioned_id] @@ -1033,14 +1082,13 @@ async def patch_prompt( requested_version: Final = get_version_number(prompt_id=prompt_id) if prompt_id != base_prompt_id else None # Build query to find the exact row by composite unique key - find_where: Final[dict[str, Any]] = { - "prompt_id": base_prompt_id, - "environment": env, + find_where: Final[Mapping[str, str | int]] = { + key: value + for key, value in (("prompt_id", base_prompt_id), ("environment", env), ("version", requested_version)) + if value is not None } - if requested_version is not None: - find_where["version"] = requested_version - db_rows: Final = await PromptRepository(prisma_client).table.find_many( + db_rows: Final = await _prompt_table(PromptRepository(prisma_client)).find_many( where=find_where, order={"version": "desc"}, take=1, @@ -1084,15 +1132,18 @@ async def patch_prompt( raise HTTPException(status_code=400, detail="litellm_params cannot be None") # Build update data dict - update_data: Final[dict[str, Any]] = { - "litellm_params": updated_litellm_params.model_dump_json(), - "prompt_info": updated_prompt_info.model_dump_json(), + update_data: Final[Mapping[str, str]] = { + key: value + for key, value in ( + ("litellm_params", updated_litellm_params.model_dump_json()), + ("prompt_info", updated_prompt_info.model_dump_json()), + ("created_by", user_api_key_dict.user_id), + ) + if value } - if user_api_key_dict.user_id: - update_data["created_by"] = user_api_key_dict.user_id # Update by primary key (id) to target exactly one row - updated_prompt_db_entry: Final = await PromptRepository(prisma_client).table.update( + updated_prompt_db_entry: Final = await _prompt_table(PromptRepository(prisma_client)).update( where={"id": target_row.id}, data=update_data, ) @@ -1216,23 +1267,25 @@ async def test_prompt( # Use ProxyBaseLLMRequestProcessing to go through all proxy logic base_llm_response_processor: Final = ProxyBaseLLMRequestProcessing(data=data) - result: Final = await base_llm_response_processor.base_process_llm_request( - request=fastapi_request, - fastapi_response=fastapi_response, - user_api_key_dict=user_api_key_dict, - route_type="acompletion", - proxy_logging_obj=proxy_logging_obj, - llm_router=llm_router, - general_settings=general_settings, - proxy_config=proxy_config, - select_data_generator=select_data_generator, - model=None, - user_model=user_model, - user_temperature=user_temperature, - user_request_timeout=user_request_timeout, - user_max_tokens=user_max_tokens, - user_api_base=user_api_base, - version=version, + result: Final = _PARSED_VALUE.validate_python( + await base_llm_response_processor.base_process_llm_request( + request=fastapi_request, + fastapi_response=fastapi_response, + user_api_key_dict=user_api_key_dict, + route_type="acompletion", + proxy_logging_obj=proxy_logging_obj, + llm_router=llm_router, + general_settings=general_settings, + proxy_config=proxy_config, + select_data_generator=select_data_generator, + model=None, + user_model=user_model, + user_temperature=user_temperature, + user_request_timeout=user_request_timeout, + user_max_tokens=user_max_tokens, + user_api_base=user_api_base, + version=version, + ) ) if isinstance(result, BaseModel): @@ -1257,7 +1310,7 @@ async def test_prompt( async def convert_prompt_file_to_json( file: UploadFile = File(...), user_api_key_dict: UserAPIKeyAuth = Depends(user_api_key_auth), -) -> dict[str, Any]: +) -> Mapping[str, object]: """ Convert a .prompt file to JSON format. diff --git a/litellm/proxy/response_polling/background_streaming.py b/litellm/proxy/response_polling/background_streaming.py index 383ada5a1bc..a8ec8884d9b 100644 --- a/litellm/proxy/response_polling/background_streaming.py +++ b/litellm/proxy/response_polling/background_streaming.py @@ -10,9 +10,12 @@ https://platform.openai.com/docs/api-reference/responses-streaming import asyncio import json -from typing import Any, Final, cast +from collections.abc import Callable, Mapping, Sequence +from typing import TYPE_CHECKING, Final, TypeAlias from fastapi import Request, Response +from fastapi.responses import StreamingResponse +from typing_extensions import TypedDict from litellm._logging import verbose_proxy_logger from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth @@ -20,6 +23,56 @@ from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessin from litellm.proxy.response_polling.polling_handler import ResponsePollingHandler from litellm.types.llms.openai import ResponsesAPIStatus +if TYPE_CHECKING: + from litellm.proxy.proxy_server import ProxyConfig + from litellm.proxy.utils import ProxyLogging + from litellm.router import Router + + +_JsonDict: TypeAlias = dict[str, object] +_JsonList: TypeAlias = list[object] + + +class _OutputItem(TypedDict, total=False): + id: str + content: Sequence[object] + + +class _TerminalResponse(TypedDict, total=False): + status: ResponsesAPIStatus + error: _JsonDict + usage: _JsonDict + reasoning: _JsonDict + tool_choice: object + tools: _JsonList + model: str + instructions: str + temperature: float + top_p: float + max_output_tokens: int + previous_response_id: str + text: _JsonDict + truncation: str + parallel_tool_calls: bool + user: str + store: bool + incomplete_details: _JsonDict + output: Sequence[_OutputItem] + + +class _StreamEvent(TypedDict, total=False): + type: str + item: _OutputItem + item_id: str + content_index: int + delta: str + part: object + response: _TerminalResponse + + +class _StreamEventParser: + parse: Callable[[str], _StreamEvent] = staticmethod(json.loads) + async def background_streaming_task( polling_id: str, @@ -29,16 +82,16 @@ async def background_streaming_task( fastapi_response: Response, user_api_key_dict: UserAPIKeyAuth, general_settings: dict, - llm_router, - proxy_config, - proxy_logging_obj, + llm_router: "Router | None", + proxy_config: "ProxyConfig", + proxy_logging_obj: "ProxyLogging", select_data_generator, user_model, - user_temperature, - user_request_timeout, - user_max_tokens, - user_api_base, - version, + user_temperature: float | None, + user_request_timeout: float | None, + user_max_tokens: int | None, + user_api_base: str | None, + version: str | None, ): """ Background task to stream response and update cache @@ -69,7 +122,7 @@ async def background_streaming_task( # Make streaming request. # Pre-call checks (rate limits, guardrails, budget) were already run # before polling ID creation, so skip them here to avoid double-counting. - response: Final = await processor.base_process_llm_request( + response: Final[StreamingResponse] = await processor.base_process_llm_request( request=request, fastapi_response=fastapi_response, user_api_key_dict=user_api_key_dict, @@ -91,8 +144,10 @@ async def background_streaming_task( # Process streaming response following OpenAI events format # https://platform.openai.com/docs/api-reference/responses-streaming - output_items: Final[dict[str, dict[str, Any]]] = {} # Track output items by ID - accumulated_text: Final = {} # Track accumulated text deltas by (item_id, content_index) + output_items: Final = dict[str, _OutputItem]() # Track output items by ID + accumulated_text: Final = dict[ + tuple[str, int], str + ]() # Track accumulated text deltas by (item_id, content_index) # ResponsesAPIResponse fields to extract from response.completed usage_data = None @@ -121,7 +176,7 @@ async def background_streaming_task( None # Will be set by response.completed/failed/incomplete/cancelled ) terminal_error = None - _event_to_status: Final = { + _event_to_status: Final[Mapping[str, ResponsesAPIStatus]] = { "response.completed": "completed", "response.failed": "failed", "response.incomplete": "incomplete", @@ -162,7 +217,7 @@ async def background_streaming_task( break try: - event = json.loads(chunk_data) + event: _StreamEvent = _StreamEventParser.parse(chunk_data) event_type = event.get("type", "") # Process different event types based on OpenAI streaming spec @@ -181,9 +236,8 @@ async def background_streaming_task( if item_id and item_id in output_items: # Update the output item with new content - if "content" not in output_items[item_id]: - output_items[item_id]["content"] = [] - output_items[item_id]["content"].append(content_part) + added_item = output_items[item_id] + added_item["content"] = (*added_item.get("content", ()), content_part) state_dirty = True elif event_type == "response.output_text.delta": @@ -201,12 +255,14 @@ async def background_streaming_task( accumulated_text[key] += delta # Update the content in output_items - if "content" in output_items[item_id]: - content_list = output_items[item_id]["content"] + delta_item = output_items[item_id] + if "content" in delta_item: + content_list = delta_item["content"] if content_index < len(content_list): # Update existing content part with accumulated text - if isinstance(content_list[content_index], dict): - content_list[content_index]["text"] = accumulated_text[key] + content_entry = content_list[content_index] + if isinstance(content_entry, dict): + content_entry["text"] = accumulated_text[key] state_dirty = True elif event_type == "response.content_part.done": @@ -217,10 +273,14 @@ async def background_streaming_task( if item_id and item_id in output_items: # Update with final content from event - if "content" in output_items[item_id]: - content_list = output_items[item_id]["content"] + done_item = output_items[item_id] + if "content" in done_item: + content_list = done_item["content"] if content_index < len(content_list): - content_list[content_index] = content_part + done_item["content"] = tuple( + content_part if part_index == content_index else existing_part + for part_index, existing_part in enumerate(content_list) + ) state_dirty = True elif event_type == "response.output_item.done": @@ -248,12 +308,9 @@ async def background_streaming_task( # Terminal event - extract all ResponsesAPIResponse fields # https://platform.openai.com/docs/api-reference/responses-streaming response_data = event.get("response", {}) - terminal_status = cast( - ResponsesAPIStatus, - response_data.get( - "status", - _event_to_status.get(event_type, "completed"), - ), + terminal_status = response_data.get( + "status", + _event_to_status.get(event_type, "completed"), ) # Extract error for failed and incomplete responses diff --git a/litellm/responses/file_search/emulated_handler.py b/litellm/responses/file_search/emulated_handler.py index 7854b17a06f..205189a6043 100644 --- a/litellm/responses/file_search/emulated_handler.py +++ b/litellm/responses/file_search/emulated_handler.py @@ -14,16 +14,19 @@ Flow: import json import time import uuid -from collections.abc import Iterable -from typing import Any, Final, cast +from collections.abc import Iterable, Mapping, Sequence +from typing import TYPE_CHECKING, Any, Final, TypedDict, cast from litellm._internal_context import is_internal_call from litellm._logging import verbose_logger -from litellm.types.llms.openai import ResponseOutputItem, ResponsesAPIResponse +from litellm.types.llms.openai import ResponsesAPIResponse from litellm.types.vector_stores import VectorStoreSearchResult +if TYPE_CHECKING: + from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig + # Keep ToolParam broad so we stay compatible with both dict and Pydantic forms -ToolParam = Any +ToolParam = object FILE_SEARCH_FUNCTION_NAME: Final = "litellm_file_search" @@ -35,7 +38,7 @@ FILE_SEARCH_FUNCTION_NAME: Final = "litellm_file_search" def should_use_emulated_file_search( tools: Iterable[ToolParam] | None, - provider_config: Any, # BaseResponsesAPIConfig + provider_config: "BaseResponsesAPIConfig | None", ) -> bool: """Return True when there is a file_search tool and the provider can't handle it natively.""" if not tools: @@ -51,7 +54,7 @@ def should_use_emulated_file_search( # --------------------------------------------------------------------------- -def _build_function_tool(vector_store_ids: list[str]) -> dict[str, Any]: +def _build_function_tool(vector_store_ids: Sequence[str]) -> Mapping[str, object]: """ Create a Responses API function-tool definition that describes file search. The function accepts one or more natural-language queries (like OpenAI's native @@ -94,27 +97,26 @@ def _build_function_tool(vector_store_ids: list[str]) -> dict[str, Any]: } +def _file_search_tool_vector_store_ids(tool: object) -> Sequence[str] | None: + if not (isinstance(tool, dict) and tool.get("type") == "file_search"): + return None + return tool.get("vector_store_ids") or [] + + def _replace_file_search_tools( tools: Iterable[ToolParam] | None, -) -> tuple[list[dict[str, Any]], list[str]]: +) -> tuple[Sequence[object], Sequence[str]]: """ Replace all file_search tools with a single function tool. Returns: (new_tools_list, all_vector_store_ids) """ - non_file_search: Final[list[dict[str, Any]]] = [] - vector_store_ids: Final[list[str]] = [] - - for tool in tools or []: - if isinstance(tool, dict) and tool.get("type") == "file_search": - ids = tool.get("vector_store_ids") or [] - vector_store_ids.extend(ids) - else: - non_file_search.append(tool) + ids_and_tools: Final = tuple((_file_search_tool_vector_store_ids(tool), tool) for tool in tools or ()) # Deduplicate while preserving order - unique_ids: Final[list[str]] = list(dict.fromkeys(vector_store_ids)) + unique_ids: Final = list(dict.fromkeys(vs_id for ids, _ in ids_and_tools if ids is not None for vs_id in ids)) + non_file_search: Final = [tool for ids, tool in ids_and_tools if ids is None] if unique_ids: non_file_search.append(_build_function_tool(unique_ids)) @@ -127,9 +129,9 @@ def _replace_file_search_tools( async def _run_vector_searches( - queries: list[str], - vector_store_ids: list[str], -) -> tuple[list[str], list[VectorStoreSearchResult]]: + queries: Sequence[str], + vector_store_ids: Sequence[str], +) -> tuple[Sequence[str], Sequence[VectorStoreSearchResult]]: """ Run `asearch` against all vector stores for all queries and collect results. @@ -172,7 +174,7 @@ async def _run_vector_searches( # --------------------------------------------------------------------------- -def _get_field(result: Any, key: str, default: Any = None) -> Any: +def _get_field(result: object, key: str, default: object = None) -> object: """Read a field from either a dict/TypedDict or an attribute-based object.""" if isinstance(result, dict): return result.get(key, default) @@ -180,7 +182,7 @@ def _get_field(result: Any, key: str, default: Any = None) -> Any: def _format_search_results_as_tool_output( - results: list[VectorStoreSearchResult], + results: Sequence[VectorStoreSearchResult], ) -> str: """Serialize search results into a string to pass back as the tool's output.""" if not results: @@ -191,7 +193,8 @@ def _format_search_results_as_tool_output( score = _get_field(result, "score") file_id = _get_field(result, "file_id") filename = _get_field(result, "filename") - content_items = _get_field(result, "content") or [] + raw_content = _get_field(result, "content") + content_items = raw_content if isinstance(raw_content, list) else [] text_chunks = [c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") for c in content_items] text = " ".join(t for t in text_chunks if t) @@ -209,9 +212,24 @@ def _format_search_results_as_tool_output( return "\n\n".join(parts) +def _format_result_for_include(result: VectorStoreSearchResult) -> Mapping[str, object]: + file_id: Final = _get_field(result, "file_id") or "" + raw_content: Final = _get_field(result, "content") + content_items: Final = raw_content if isinstance(raw_content, list) else [] + text_chunks: Final = [c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") for c in content_items] + text: Final = " ".join(t for t in text_chunks if t) + return { + "file_id": file_id, + "filename": _get_field(result, "filename") or "", + "score": _get_field(result, "score"), + "text": text, + "attributes": _get_field(result, "attributes") or {}, + } + + def _build_search_results_for_include( - results: list[VectorStoreSearchResult], -) -> list[dict[str, Any]]: + results: Sequence[VectorStoreSearchResult], +) -> Sequence[Mapping[str, object]]: """ Convert VectorStoreSearchResult objects to the format expected in file_search_call.search_results (mirrors OpenAI's include= format). @@ -220,30 +238,15 @@ def _build_search_results_for_include( behaviour of OpenAI's native file_search which surfaces every relevant chunk even when multiple chunks originate from the same document. """ - formatted: Final[list[dict[str, Any]]] = [] - for result in results: - file_id = _get_field(result, "file_id") or "" - content_items = _get_field(result, "content") or [] - text_chunks = [c.get("text", "") if isinstance(c, dict) else getattr(c, "text", "") for c in content_items] - text = " ".join(t for t in text_chunks if t) - formatted.append( - { - "file_id": file_id, - "filename": _get_field(result, "filename") or "", - "score": _get_field(result, "score"), - "text": text, - "attributes": _get_field(result, "attributes") or {}, - } - ) - return formatted + return [_format_result_for_include(result) for result in results] def _build_file_search_call_output( call_id: str, - queries: list[str], - results: list[VectorStoreSearchResult] | None = None, + queries: Sequence[str], + results: Sequence[VectorStoreSearchResult] | None = None, include_search_results: bool = False, -) -> dict[str, Any]: +) -> Mapping[str, object]: """Build the file_search_call output item (mirrors OpenAI's format). Args: @@ -266,39 +269,34 @@ def _build_file_search_call_output( def _build_file_citation_annotations( - results: list[VectorStoreSearchResult], + results: Sequence[VectorStoreSearchResult], text: str, -) -> list[dict[str, Any]]: +) -> Sequence[Mapping[str, object]]: """ Build file_citation annotations for the text. Each result with a file_id gets a citation at the end of the text. """ - annotations: Final[list[dict[str, Any]]] = [] index: Final = len(text) # cite at end of text block - seen_file_ids: Final[set] = set() + id_filename_pairs: Final = tuple( + (_get_field(result, "file_id"), _get_field(result, "filename")) for result in results + ) + first_filename_by_id: Final = {file_id: filename for file_id, filename in reversed(id_filename_pairs) if file_id} - for result in results: - file_id = _get_field(result, "file_id") - filename = _get_field(result, "filename") - if not file_id or file_id in seen_file_ids: - continue - seen_file_ids.add(file_id) - annotations.append( - { - "type": "file_citation", - "index": index, - "file_id": file_id, - "filename": filename or "", - } - ) - - return annotations + return [ + { + "type": "file_citation", + "index": index, + "file_id": file_id, + "filename": first_filename_by_id[file_id] or "", + } + for file_id in dict.fromkeys(file_id for file_id, _ in id_filename_pairs if file_id) + ] def _build_message_output( response_text: str, - results: list[VectorStoreSearchResult], -) -> dict[str, Any]: + results: Sequence[VectorStoreSearchResult], +) -> Mapping[str, object]: """Build the message output item with optional file_citation annotations.""" annotations: Final = _build_file_citation_annotations(results, response_text) return { @@ -330,8 +328,8 @@ def _extract_text_from_responses_output(response: ResponsesAPIResponse) -> str: def _synthesize_responses_api_response( original_response: ResponsesAPIResponse, - file_search_call_output: dict[str, Any], - message_output: dict[str, Any], + file_search_call_output: Mapping[str, object], + message_output: Mapping[str, object], first_response: ResponsesAPIResponse | None = None, ) -> ResponsesAPIResponse: """ @@ -343,21 +341,20 @@ def _synthesize_responses_api_response( synthesized _hidden_params so that billing callbacks see the total cost of both provider calls that the emulated flow makes. """ - synthesized_output: Final[list[dict[str, Any]]] = [file_search_call_output, message_output] synthesized: Final = ResponsesAPIResponse( id=getattr(original_response, "id", f"resp_{uuid.uuid4().hex}"), object="response", created_at=getattr(original_response, "created_at", int(time.time())), status="completed", model=getattr(original_response, "model", ""), - output=cast(list[ResponseOutputItem | dict[str, Any]], synthesized_output), + output=[dict(file_search_call_output), dict(message_output)], usage=getattr(original_response, "usage", None), error=None, ) if hasattr(original_response, "_hidden_params"): hidden: Final = dict(getattr(original_response, "_hidden_params") or {}) if first_response is not None and hasattr(first_response, "_hidden_params"): - first_hidden: Final = getattr(first_response, "_hidden_params") or {} + first_hidden: Final[object] = getattr(first_response, "_hidden_params", None) or {} first_cost: Final = ( first_hidden.get("response_cost") if isinstance(first_hidden, dict) @@ -382,9 +379,10 @@ async def _call_aresponses(input, model, tools, **kwargs): # pragma: no cover def _prepare_emulated_file_search_call( - kwargs: dict[str, Any], -) -> tuple[bool, dict[str, Any]]: - include_items: Final[list[str]] = list(kwargs.get("include") or []) + kwargs: Mapping[str, object], +) -> tuple[bool, Mapping[str, object]]: + raw_include: Final = kwargs.get("include") + include_items: Final[Sequence[str]] = raw_include if isinstance(raw_include, list) else [] include_search_results: Final = "file_search_call.results" in include_items original_stream: Final = kwargs.get("stream") @@ -398,7 +396,7 @@ def _prepare_emulated_file_search_call( return include_search_results, updated_kwargs -def _extract_tool_call_fields(tool_call: Any, fallback_call_id: str) -> tuple[str, str]: +def _extract_tool_call_fields(tool_call: object, fallback_call_id: str) -> tuple[str, str]: """Extract (call_id, raw_arguments_string) from a dict or Pydantic tool_call item.""" if isinstance(tool_call, dict): call_id = str(tool_call.get("call_id") or tool_call.get("id") or fallback_call_id) @@ -410,7 +408,13 @@ def _extract_tool_call_fields(tool_call: Any, fallback_call_id: str) -> tuple[st return call_id, raw_args -def _resolve_queries_from_args(args: dict[str, Any], input: Any) -> list[str]: +class _FileSearchArguments(TypedDict, total=False): + queries: Sequence[str] + query: str + vector_store_id: str + + +def _resolve_queries_from_args(args: _FileSearchArguments, input: object) -> Sequence[str]: """Pull the queries list out of parsed tool-call arguments, with backward-compat fallbacks.""" queries_from_call: Final = args.get("queries") if not queries_from_call: @@ -422,76 +426,96 @@ def _resolve_queries_from_args(args: dict[str, Any], input: Any) -> list[str]: return queries_from_call -async def _execute_file_search_tool_calls( - file_search_calls: list[Any], - all_vs_ids: list[str], - input: Any, +def _parse_file_search_arguments(raw_args: str) -> _FileSearchArguments: + if not isinstance(raw_args, str): + return raw_args + try: + return json.loads(raw_args) + except json.JSONDecodeError: + return {} + + +async def _execute_single_file_search_call( + tool_call: object, + all_vs_ids: Sequence[str], + input: object, file_search_call_id: str, -) -> tuple[list[dict[str, Any]], list[str], list[VectorStoreSearchResult]]: +) -> tuple[Mapping[str, object], Sequence[str], Sequence[VectorStoreSearchResult]]: + call_id, raw_args = _extract_tool_call_fields(tool_call, fallback_call_id=file_search_call_id) + args: Final = _parse_file_search_arguments(raw_args) + queries_from_call: Final = _resolve_queries_from_args(args, input) + + vs_id_arg: Final = args.get("vector_store_id") + vs_ids_for_call: Final = [vs_id_arg] if vs_id_arg else all_vs_ids + + queries, results = await _run_vector_searches( + queries=queries_from_call, + vector_store_ids=vs_ids_for_call, + ) + + return ( + { + "type": "function_call_output", + "call_id": call_id, + "output": _format_search_results_as_tool_output(results), + }, + queries, + results, + ) + + +async def _execute_file_search_tool_calls( + file_search_calls: Sequence[object], + all_vs_ids: Sequence[str], + input: object, + file_search_call_id: str, +) -> tuple[Sequence[Mapping[str, object]], Sequence[str], Sequence[VectorStoreSearchResult]]: """Run the vector search for each file_search tool_call and collect results.""" - tool_results: Final[list[dict[str, Any]]] = [] - all_queries: Final[list[str]] = [] - all_results: Final[list[VectorStoreSearchResult]] = [] + per_call: Final = tuple( + [ + await _execute_single_file_search_call( + tool_call=tool_call, + all_vs_ids=all_vs_ids, + input=input, + file_search_call_id=file_search_call_id, + ) + for tool_call in file_search_calls + ] + ) - for tool_call in file_search_calls: - call_id, raw_args = _extract_tool_call_fields(tool_call, fallback_call_id=file_search_call_id) - - try: - args = json.loads(raw_args) if isinstance(raw_args, str) else raw_args - except json.JSONDecodeError: - args = {} - - queries_from_call = _resolve_queries_from_args(args, input) - - vs_id_arg = args.get("vector_store_id") - vs_ids_for_call = [vs_id_arg] if vs_id_arg else all_vs_ids - - queries, results = await _run_vector_searches( - queries=queries_from_call, - vector_store_ids=vs_ids_for_call, - ) - all_queries.extend(queries) - all_results.extend(results) - - tool_results.append( - { - "type": "function_call_output", - "call_id": call_id, - "output": _format_search_results_as_tool_output(results), - } - ) - - return tool_results, all_queries, all_results + return ( + [tool_result for tool_result, _, _ in per_call], + [query for _, queries, _ in per_call for query in queries], + [result for _, _, results in per_call for result in results], + ) def _build_follow_up_input( - input: Any, + input: object, first_response: ResponsesAPIResponse, - tool_results: list[dict[str, Any]], -) -> list[Any]: + tool_results: Sequence[Mapping[str, object]], +) -> Sequence[object]: """Assemble the follow-up call input: original messages + first-response output + tool results. Including all output items (text blocks, reasoning, non-file-search calls) ensures providers like Anthropic that emit text before the tool call have complete conversation context. Serializes Pydantic model instances to plain dicts so the transformation layer can call .get(). """ - original_input_items: Final = ( - list(input) if isinstance(input, (list, tuple)) else [{"role": "user", "content": str(input)}] + original_input_items: Final[tuple[object, ...]] = ( + tuple(input) if isinstance(input, (list, tuple)) else ({"role": "user", "content": str(input)},) + ) + first_response_output_items: Final[tuple[object, ...]] = tuple( + _item + if isinstance(_item, dict) + else (_item.model_dump(exclude_none=True) if hasattr(_item, "model_dump") else _item) + for _item in first_response.output ) - first_response_output_items: Final[list[Any]] = [] - for _item in first_response.output: - if isinstance(_item, dict): - first_response_output_items.append(_item) - elif hasattr(_item, "model_dump"): - first_response_output_items.append(_item.model_dump(exclude_none=True)) - else: - first_response_output_items.append(_item) - return original_input_items + first_response_output_items + tool_results + return [*original_input_items, *first_response_output_items, *tool_results] async def aresponses_with_emulated_file_search( - input: Any, + input: object, model: str, tools: Iterable[ToolParam] | None = None, # Pass-through params — forwarded as-is to the underlying aresponses call @@ -504,7 +528,7 @@ async def aresponses_with_emulated_file_search( runs vector search, and synthesizes an OpenAI-format response. """ # Determine whether caller wants search_results populated in the output. - _include_search_results, kwargs = _prepare_emulated_file_search_call(kwargs=kwargs) + _include_search_results, call_kwargs = _prepare_emulated_file_search_call(kwargs=kwargs) # 1. Replace file_search tools with function tool transformed_tools, all_vs_ids = _replace_file_search_tools(tools) @@ -521,7 +545,7 @@ async def aresponses_with_emulated_file_search( input=input, model=model, tools=transformed_tools or None, - **kwargs, + **call_kwargs, ), ) finally: @@ -585,7 +609,7 @@ async def aresponses_with_emulated_file_search( input=follow_up_input, model=model, tools=None, # no tools needed for the answer step - **kwargs, + **call_kwargs, ), ) finally: diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index d7f6ece5cd1..1385e329e93 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -5,11 +5,11 @@ import json import time import traceback import uuid -from collections.abc import Awaitable, Callable, Mapping +from collections.abc import Awaitable, Callable, Iterable, Mapping, Sequence from datetime import datetime from functools import lru_cache from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, runtime_checkable +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload, runtime_checkable import httpx from openai._streaming import SSEDecoder @@ -42,12 +42,14 @@ from litellm.types.utils import CallTypes from litellm.utils import async_post_call_success_deployment_hook if TYPE_CHECKING: + from litellm.caching.caching_handler import LLMCachingHandler from litellm.proxy._types import UserAPIKeyAuth from litellm.types.responses.streaming_websocket import ( PresidioGuardrailCallback, ResponsesBackendWebSocket, ResponsesClientWebSocket, ) + from litellm.types.router import LiteLLM_Params @lru_cache(maxsize=1) @@ -69,6 +71,79 @@ def _is_str_mapping(value: object) -> TypeIs[dict[str, str]]: # guard-ok: verif return _is_json_object(value) and all(isinstance(item, str) for item in value.values()) +class _MutableJsonObject(Protocol): + @overload + def get(self, key: str, /) -> object | None: ... + @overload + def get(self, key: str, default: object, /) -> object: ... + def __getitem__(self, key: str, /) -> object: ... + def __setitem__(self, key: str, value: object, /) -> None: ... + def __contains__(self, key: object, /) -> bool: ... + def items(self) -> Iterable[tuple[str, object]]: ... + + +class _LoadsJsonValue(Protocol): + def __call__(self, s: str | bytes, /) -> object: ... + + +class _LoadsJsonDict(Protocol): + def __call__(self, s: str | bytes, /) -> _MutableJsonObject: ... + + +class _GetsLitellmParams(Protocol): + def __call__(self, key: str, default: Mapping[str, object], /) -> LiteLLM_Params: ... + + +class _PopsOptionalStr(Protocol): + def __call__(self, key: str, default: None, /) -> str | None: ... + + +class _UnmasksPiiText(Protocol): + def __call__(self, text: str, pii_tokens: Mapping[str, str]) -> str: ... + + +class _ShouldStoreResultInCache(Protocol): + def __call__(self, *, original_function: Callable[..., object] | None, kwargs: Mapping[str, object]) -> bool: ... + + +class _PostStreamingDeploymentHook(Protocol): + def __call__( + self, + *, + request_data: Mapping[str, object], + response_chunk: ResponsesAPIStreamingResponse, + call_type: CallTypes | None, + ) -> Awaitable[ResponsesAPIStreamingResponse | None]: ... + + +@runtime_checkable +class _HasPostStreamingDeploymentHook(Protocol): + async_post_call_streaming_deployment_hook: _PostStreamingDeploymentHook + + +def _typed_loads_json_value(fn: _LoadsJsonValue) -> _LoadsJsonValue: + return fn + + +def _typed_loads_json_dict(fn: _LoadsJsonDict) -> _LoadsJsonDict: + return fn + + +def _typed_gets_litellm_params(fn: _GetsLitellmParams) -> _GetsLitellmParams: + return fn + + +def _typed_pops_optional_str(fn: _PopsOptionalStr) -> _PopsOptionalStr: + return fn + + +_LOADS_JSON_VALUE: Final = _typed_loads_json_value(json.loads) +_LOADS_JSON_DICT: Final = _typed_loads_json_dict(json.loads) + +_SHOULD_STORE_RESULT_IN_CACHE_ATTR: Final = "_should_store_result_in_cache" +_UNMASK_PII_TEXT_ATTR: Final = "_unmask_pii_text" + + def _model_id_from_metadata(litellm_metadata: dict[str, object] | None) -> str | None: model_info: Final = litellm_metadata.get("model_info") if litellm_metadata else None model_id: Final = model_info.get("id") if _is_json_object(model_info) else None @@ -185,7 +260,7 @@ class BaseResponsesAPIStreamingIterator: # This matches the stream wrapper in litellm/litellm_core_utils/streaming_handler.py _api_base: Final = get_api_base( model=model or "", - optional_params=self.logging_obj.model_call_details.get("litellm_params", {}), + optional_params=_typed_gets_litellm_params(self.logging_obj.model_call_details.get)("litellm_params", {}), ) self._hidden_params: dict[str, object] = { "model_id": _model_id_from_metadata(litellm_metadata), @@ -228,7 +303,7 @@ class BaseResponsesAPIStreamingIterator: try: # Parse the JSON chunk - parsed_chunk: Final = json.loads(chunk) + parsed_chunk: Final = _LOADS_JSON_VALUE(chunk) # Format as ResponsesAPIStreamingResponse if isinstance(parsed_chunk, dict): @@ -514,7 +589,7 @@ class BaseResponsesAPIStreamingIterator: if response_obj is None: return - caching_handler: Final = getattr(self.logging_obj, "_llm_caching_handler", None) + caching_handler: Final[LLMCachingHandler | None] = getattr(self.logging_obj, "_llm_caching_handler", None) if caching_handler is None: return @@ -532,8 +607,11 @@ class BaseResponsesAPIStreamingIterator: if preset_cache_key is not None: request_kwargs["cache_key"] = preset_cache_key - if not caching_handler._should_store_result_in_cache( - original_function=caching_handler.original_function, + should_store_result_in_cache: Final[_ShouldStoreResultInCache] = getattr( + caching_handler, _SHOULD_STORE_RESULT_IN_CACHE_ATTR + ) + if not should_store_result_in_cache( + original_function=getattr(caching_handler, "original_function", None), kwargs=request_kwargs, ): return @@ -586,12 +664,15 @@ class BaseResponsesAPIStreamingIterator: typed_call_type = None request_data: Final = self.request_data or getattr(self.logging_obj, "model_call_details", {}) - callbacks: Final = getattr(litellm, "callbacks", None) or [] + callbacks: Final[Sequence[object]] = getattr(litellm, "callbacks", None) or [] hooks_ran = False for callback in callbacks: - if hasattr(callback, "async_post_call_streaming_deployment_hook"): + if isinstance(callback, _HasPostStreamingDeploymentHook): hooks_ran = True - result = await callback.async_post_call_streaming_deployment_hook( + post_streaming_hook: _PostStreamingDeploymentHook = ( + callback.async_post_call_streaming_deployment_hook + ) + result = await post_streaming_hook( request_data=request_data, response_chunk=chunk, call_type=typed_call_type, @@ -1043,8 +1124,8 @@ 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"): +def _dump_response_object(obj: object) -> Mapping[str, object]: + if isinstance(obj, _HasModelDump): return obj.model_dump() if _is_json_object(obj): return obj @@ -1073,21 +1154,20 @@ def _build_content_part_done_event( item_id: str, output_index: int, content_index: int, - part_payload: dict[str, Any], + part_payload: Mapping[str, object], ) -> ResponsesAPIStreamingResponse | None: openai_types: Final = _get_openai_response_types() part_type: Final = part_payload.get("type") part: PART_UNION_TYPES if part_type == "output_text": - annotations: Final = [ - openai_types.BaseLiteLLMOpenAIResponseObject(**annotation) - for annotation in part_payload.get("annotations", []) or [] - ] - part = openai_types.ContentPartDonePartOutputText( - type="output_text", - text=str(part_payload.get("text") or ""), - annotations=annotations, - logprobs=part_payload.get("logprobs"), + raw_annotations: Final[object] = part_payload.get("annotations", []) or [] + part = openai_types.ContentPartDonePartOutputText.model_validate( + { + "type": "output_text", + "text": str(part_payload.get("text") or ""), + "annotations": raw_annotations, + "logprobs": part_payload.get("logprobs"), + } ) elif part_type == "refusal": part = openai_types.ContentPartDonePartRefusal( @@ -1117,7 +1197,7 @@ def _add_text_like_part_events( item_id: str, output_index: int, content_index: int, - part_payload: dict[str, Any], + part_payload: Mapping[str, object], chunk_size: int, ) -> None: openai_types: Final = _get_openai_response_types() @@ -1134,15 +1214,19 @@ def _add_text_like_part_events( delta=text[i : i + chunk_size], ) ) - for annotation_index, annotation in enumerate(part_payload.get("annotations", []) or []): + raw_annotation_items: Final = part_payload.get("annotations") + annotation_items: Final[Sequence[object]] = raw_annotation_items if _is_json_array(raw_annotation_items) else [] + for annotation_index, annotation in enumerate(annotation_items): events.append( - openai_types.OutputTextAnnotationAddedEvent( - type=openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED, - item_id=item_id, - output_index=output_index, - content_index=content_index, - annotation_index=annotation_index, - annotation=annotation, + openai_types.OutputTextAnnotationAddedEvent.model_validate( + { + "type": openai_types.ResponsesAPIStreamEvents.OUTPUT_TEXT_ANNOTATION_ADDED, + "item_id": item_id, + "output_index": output_index, + "content_index": content_index, + "annotation_index": annotation_index, + "annotation": annotation, + } ) ) events.append( @@ -1200,7 +1284,8 @@ def _build_synthetic_response_events( ] sequence_number = 0 - for output_index, output_item in enumerate(getattr(transformed, "output", []) or []): + output_items: Final[Sequence[object]] = getattr(transformed, "output", []) or [] + for output_index, output_item in enumerate(output_items): output_item_payload = _dump_response_object(output_item) item_id = str(output_item_payload.get("id") or transformed.id) item_type = output_item_payload.get("type") @@ -1214,7 +1299,9 @@ def _build_synthetic_response_events( ) if item_type == "message": - for content_index, part in enumerate(output_item_payload.get("content", []) or []): + raw_content_parts = output_item_payload.get("content") + content_parts: Sequence[object] = raw_content_parts if _is_json_array(raw_content_parts) else [] + for content_index, part in enumerate(content_parts): part_payload = _dump_response_object(part) events.append( openai_types.ContentPartAddedEvent( @@ -1261,7 +1348,9 @@ def _build_synthetic_response_events( ) ) elif item_type == "reasoning": - for summary_index, summary in enumerate(output_item_payload.get("summary", []) or []): + raw_summary_items = output_item_payload.get("summary") + summary_items: Sequence[object] = raw_summary_items if _is_json_array(raw_summary_items) else [] + for summary_index, summary in enumerate(summary_items): summary_payload = _dump_response_object(summary) summary_text = str(summary_payload.get("text") or "") for i in range(0, len(summary_text), chunk_size): @@ -1354,7 +1443,7 @@ class ResponsesWebSocketStreaming: user_api_key_dict: UserAPIKeyAuth | None = None, request_data: dict[str, object] | None = None, first_message: str | None = None, - guardrail_callbacks: list[Any] | None = None, + guardrail_callbacks: Sequence[PresidioGuardrailCallback] | None = None, output_guardrail_callbacks: list[PresidioGuardrailCallback] | None = None, authorized_model: str | None = None, ): @@ -1363,16 +1452,16 @@ class ResponsesWebSocketStreaming: self.logging_obj = logging_obj self.user_api_key_dict = user_api_key_dict self.request_data: dict[str, object] = request_data or {} - self.messages: list[dict[str, object]] = [] + self.messages: list[_MutableJsonObject] = [] self.input_messages: list[dict[str, object]] = [] self.first_message = first_message - self.guardrail_callbacks: list[Any] = guardrail_callbacks or [] + self.guardrail_callbacks: Sequence[PresidioGuardrailCallback] = guardrail_callbacks or [] self.output_guardrail_callbacks: list[PresidioGuardrailCallback] = output_guardrail_callbacks or [] # Model name authorized at connection time; enforced on every # response.create frame to prevent deployment-substitution attacks. self.authorized_model: str | None = authorized_model - def _should_store_event(self, event_obj: Mapping[str, object]) -> bool: + def _should_store_event(self, event_obj: _MutableJsonObject) -> bool: return event_obj.get("type") in RESPONSES_WS_LOGGED_EVENT_TYPES def _store_event(self, event: str | bytes | dict[str, object]) -> None: @@ -1380,7 +1469,7 @@ class ResponsesWebSocketStreaming: event = event.decode("utf-8") if isinstance(event, str): try: - event_obj = json.loads(event) + event_obj = _LOADS_JSON_DICT(event) except (json.JSONDecodeError, TypeError): return else: @@ -1393,7 +1482,7 @@ class ResponsesWebSocketStreaming: """Extract user input content from response.create for logging.""" try: if isinstance(message, str): - msg_obj = json.loads(message) + msg_obj = _LOADS_JSON_DICT(message) elif _is_json_object(message): msg_obj = message else: @@ -1463,7 +1552,7 @@ class ResponsesWebSocketStreaming: # masked response.completed. if self.output_guardrail_callbacks: try: - _evt_type = json.loads(response_str).get("type") + _evt_type = _LOADS_JSON_DICT(response_str).get("type") except (json.JSONDecodeError, TypeError): _evt_type = None if _evt_type in self._DELTA_EVENT_TYPES or _evt_type in self._OUTPUT_DONE_EVENT_TYPES: @@ -1485,7 +1574,7 @@ class ResponsesWebSocketStreaming: finally: await self._log_messages() - def _enforce_authorized_model(self, msg_obj: dict[str, object]) -> bool: + def _enforce_authorized_model(self, msg_obj: _MutableJsonObject) -> bool: """ Overwrite any ``model`` field in a ``response.create`` frame with the connection-authorized model to prevent deployment-substitution attacks. @@ -1527,7 +1616,7 @@ class ResponsesWebSocketStreaming: Non-``response.create`` messages are returned unchanged. """ try: - msg_obj: Final = json.loads(message) + msg_obj: Final = _LOADS_JSON_DICT(message) except (json.JSONDecodeError, TypeError): return message @@ -1553,7 +1642,7 @@ class ResponsesWebSocketStreaming: # forwarded unmasked regardless of where the client places it. nested_candidate = msg_obj.get("response") nested_response = nested_candidate if _is_json_object(nested_candidate) else None - text_containers: list[tuple[dict[str, object], str]] = [] + text_containers: list[tuple[_MutableJsonObject, str]] = [] for container in (msg_obj, nested_response): if container is None: continue @@ -1655,11 +1744,12 @@ class ResponsesWebSocketStreaming: return response_str try: - evt_obj: Final = json.loads(response_str) + evt_obj: Final = _LOADS_JSON_DICT(response_str) except (json.JSONDecodeError, TypeError): return response_str cb: Final = self.guardrail_callbacks[0] + unmask_pii_text: Final[_UnmasksPiiText] = getattr(cb, _UNMASK_PII_TEXT_ATTR) event_type: Final = evt_obj.get("type") if event_type == "response.completed": @@ -1679,7 +1769,7 @@ class ResponsesWebSocketStreaming: continue text = content_block.get("text") if isinstance(text, str): - unmasked = cb._unmask_pii_text(text, pii_tokens) + unmasked = unmask_pii_text(text, pii_tokens) if unmasked != text: content_block["text"] = unmasked modified = True @@ -1688,7 +1778,7 @@ class ResponsesWebSocketStreaming: if event_type in self._DELTA_EVENT_TYPES: delta: Final = evt_obj.get("delta") if isinstance(delta, str): - unmasked = cb._unmask_pii_text(delta, pii_tokens) + unmasked = unmask_pii_text(delta, pii_tokens) if unmasked != delta: evt_obj["delta"] = unmasked return json.dumps(evt_obj) @@ -1711,7 +1801,7 @@ class ResponsesWebSocketStreaming: return response_str try: - evt_obj: Final[Mapping[str, object]] = json.loads(response_str) + evt_obj: Final = _LOADS_JSON_DICT(response_str) except (json.JSONDecodeError, TypeError): return response_str @@ -1859,7 +1949,7 @@ class ManagedResponsesWebSocketHandler: model: str, logging_obj: LiteLLMLoggingObj, user_api_key_dict: UserAPIKeyAuth | None = None, - litellm_metadata: dict[str, Any] | None = None, + litellm_metadata: Mapping[str, object] | None = None, api_key: str | None = None, api_base: str | None = None, timeout: float | None = None, @@ -1871,10 +1961,11 @@ class ManagedResponsesWebSocketHandler: self.model = model self.logging_obj = logging_obj self.user_api_key_dict = user_api_key_dict - self.litellm_metadata: dict[str, Any] = litellm_metadata or {} - self.model_group: str | None = self.litellm_metadata.get("model_group") or self.litellm_metadata.get( + self.litellm_metadata: Mapping[str, object] = litellm_metadata or {} + raw_model_group: Final = self.litellm_metadata.get("model_group") or self.litellm_metadata.get( "deployment_model_name" ) + self.model_group: str | None = raw_model_group if isinstance(raw_model_group, str) else None self.api_key = api_key self.api_base = api_base self.timeout = timeout @@ -1894,7 +1985,7 @@ class ManagedResponsesWebSocketHandler: # ------------------------------------------------------------------ @staticmethod - def _serialize_chunk(chunk: Any) -> str | None: + def _serialize_chunk(chunk: object) -> str | None: """Serialize a streaming chunk to a JSON string for WebSocket transmission.""" try: if isinstance(chunk, _HasModelDumpJson): @@ -1937,7 +2028,7 @@ class ManagedResponsesWebSocketHandler: self._session_history[response_id] = messages @staticmethod - def _extract_response_id(completed_event: dict[str, object]) -> str | None: + def _extract_response_id(completed_event: _MutableJsonObject) -> str | None: """ Pull the raw (decoded) response ID out of a ``response.completed`` event. Returns *None* if the event doesn't contain a usable ID. @@ -1952,7 +2043,7 @@ class ManagedResponsesWebSocketHandler: @staticmethod def _extract_output_messages( - completed_event: dict[str, object], + completed_event: _MutableJsonObject, ) -> list[dict[str, object]]: """ Convert the output items in a ``response.completed`` event into @@ -2009,10 +2100,10 @@ class ManagedResponsesWebSocketHandler: # _process_response_create sub-methods # ------------------------------------------------------------------ - async def _parse_message(self, raw_message: str) -> dict[str, object] | None: + async def _parse_message(self, raw_message: str) -> _MutableJsonObject | None: """Parse raw WS text; return the message dict or None (JSON error / ignored type).""" try: - msg_obj: Final = json.loads(raw_message) + msg_obj: Final = _LOADS_JSON_DICT(raw_message) except json.JSONDecodeError: await self._send_error("Invalid JSON in response.create event", "invalid_request_error") return None @@ -2022,7 +2113,7 @@ class ManagedResponsesWebSocketHandler: return msg_obj @staticmethod - def _is_warmup_frame(msg_obj: dict[str, object]) -> bool: + def _is_warmup_frame(msg_obj: _MutableJsonObject) -> bool: """Return True for a response.create whose generate flag is false.""" nested: Final = msg_obj.get("response") source: Final = nested if _is_json_object(nested) and nested else msg_obj @@ -2038,13 +2129,13 @@ class ManagedResponsesWebSocketHandler: return str(raw_id).startswith(_WARMUP_RESPONSE_ID_PREFIX) @staticmethod - def _warmup_source_params(msg_obj: dict[str, object]) -> dict[str, object]: + def _warmup_source_params(msg_obj: _MutableJsonObject) -> dict[str, object]: nested: Final = msg_obj.get("response") if _is_json_object(nested) and nested: return nested return {k: v for k, v in msg_obj.items() if k != "type"} - def _build_warmup_response(self, msg_obj: dict[str, object]) -> dict[str, object]: + def _build_warmup_response(self, msg_obj: _MutableJsonObject) -> dict[str, object]: """Build a minimal completed Responses API object for a warmup ack.""" source: Final = self._warmup_source_params(msg_obj) wire_model: Final = source.get("model") or self.model_group or self.model @@ -2062,7 +2153,7 @@ class ManagedResponsesWebSocketHandler: }, } - async def _send_warmup_ack(self, msg_obj: dict[str, object]) -> None: + async def _send_warmup_ack(self, msg_obj: _MutableJsonObject) -> None: """ Acknowledge a generate=false prewarm without calling the provider. @@ -2085,7 +2176,7 @@ class ManagedResponsesWebSocketHandler: await self.websocket.send_text(serialized) @staticmethod - def _build_base_call_kwargs(msg_obj: dict[str, object]) -> dict[str, Any]: + def _build_base_call_kwargs(msg_obj: _MutableJsonObject) -> dict[str, Any]: """ Extract Responses API params from the event, handling both wire formats: Nested: {"type": "response.create", "response": {"input": [...], ...}} @@ -2194,7 +2285,7 @@ class ManagedResponsesWebSocketHandler: call_kwargs.setdefault("litellm_params", {}) call_kwargs["litellm_params"]["proxy_server_request"] = proxy_server_request - async def _stream_and_forward(self, model: str, call_kwargs: dict[str, Any]) -> dict[str, object] | None: + async def _stream_and_forward(self, model: str, call_kwargs: dict[str, Any]) -> _MutableJsonObject | None: """ Stream ``litellm.aresponses`` and forward every chunk over the WebSocket. @@ -2202,7 +2293,7 @@ class ManagedResponsesWebSocketHandler: directly (before serialization) to avoid a redundant JSON round-trip on every chunk. Returns the completed event dict, or ``None``. """ - completed_event: dict[str, object] | None = ( + completed_event: _MutableJsonObject | None = ( None # rebind-ok: captures the completed event once the stream yields it ) stream_response: Final = await litellm.aresponses(model=model, **call_kwargs) @@ -2216,7 +2307,7 @@ class ManagedResponsesWebSocketHandler: continue if chunk_type == "response.completed" and completed_event is None: try: - completed_event = json.loads(serialized) + completed_event = _LOADS_JSON_DICT(serialized) except Exception: pass try: @@ -2228,7 +2319,7 @@ class ManagedResponsesWebSocketHandler: def _save_turn_history( self, - completed_event: dict[str, object] | None, + completed_event: _MutableJsonObject | None, prior_history: list[dict[str, object]], current_messages: list[dict[str, object]], ) -> None: @@ -2293,13 +2384,15 @@ class ManagedResponsesWebSocketHandler: # reuse the router-resolved self.model; passing the alias raw to # litellm.aresponses fails in get_llm_provider. A genuinely different # provider-prefixed per-frame model is still honored. - requested_model: Final = call_kwargs.pop("model", None) + requested_model: Final = _typed_pops_optional_str(call_kwargs.pop)("model", None) if requested_model is None or requested_model == self.model_group: model = self.model else: model = requested_model - previous_response_id: Final[str | None] = call_kwargs.pop("previous_response_id", None) + previous_response_id: Final[str | None] = _typed_pops_optional_str(call_kwargs.pop)( + "previous_response_id", None + ) current_messages: Final = self._input_to_messages(call_kwargs.get("input")) # Fetch history once; reused in both _apply_history and _save_turn_history diff --git a/litellm/router_strategy/tag_based_routing.py b/litellm/router_strategy/tag_based_routing.py index bbe97613c57..209dea87cc5 100644 --- a/litellm/router_strategy/tag_based_routing.py +++ b/litellm/router_strategy/tag_based_routing.py @@ -10,10 +10,10 @@ Use this to route requests between Teams import re from collections.abc import Mapping, Sequence from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, Literal +from typing import TYPE_CHECKING, Any, Final, Literal, Protocol, overload from litellm._logging import verbose_logger -from litellm.types.router import RouterErrors +from litellm.types.router import DeploymentTypedDict, RouterErrors if TYPE_CHECKING: from litellm.router import Router as _Router @@ -23,9 +23,68 @@ else: LitellmRouter = Any +class _TagLitellmParamsLike(Protocol): + @overload + def get(self, key: Literal["tags"], /) -> Sequence[str] | None: ... + @overload + def get(self, key: Literal["tags"], default: Sequence[str], /) -> Sequence[str]: ... + @overload + def get(self, key: Literal["tag_regex"], /) -> Sequence[str] | None: ... + + +class _ModelInfoLike(Protocol): + @overload + def get(self, key: Literal["allow_fail_open"], /) -> bool | None: ... + @overload + def get(self, key: Literal["enable_tag_filtering"], /) -> bool | None: ... + + +class _DeploymentLike(Protocol): + @overload + def get(self, key: Literal["litellm_params"], default: Mapping[str, object], /) -> _TagLitellmParamsLike: ... + @overload + def get(self, key: Literal["model_info"], /) -> _ModelInfoLike | None: ... + @overload + def get(self, key: Literal["model_name"], /) -> object: ... + + +class _MetadataLike(Protocol): + @overload + def get(self, key: Literal["tags"], /) -> Sequence[str] | None: ... + @overload + def get(self, key: Literal["tags"], default: Sequence[str], /) -> Sequence[str] | None: ... + @overload + def get(self, key: Literal["user_agent"], default: str, /) -> str: ... + @overload + def get(self, key: Literal["inherited_tags"], /) -> object: ... + def __contains__(self, key: object, /) -> bool: ... + def __setitem__(self, key: Literal["tag_routing"], value: Mapping[str, object], /) -> None: ... + + +class _NestedLitellmParamsLike(Protocol): + def get( + self, key: Literal["metadata", "litellm_metadata"], default: Mapping[str, object], / + ) -> _MetadataLike | None: ... + + +class _RequestKwargsLike(Protocol): + @overload + def get(self, key: Literal["enable_tag_filtering"], /) -> bool | None: ... + @overload + def get(self, key: Literal["metadata", "litellm_metadata"], /) -> _MetadataLike | None: ... + def __contains__(self, key: object, /) -> bool: ... + @overload + def __getitem__(self, key: Literal["metadata", "litellm_metadata"], /) -> _MetadataLike: ... + @overload + def __getitem__(self, key: Literal["litellm_params"], /) -> _NestedLitellmParamsLike: ... + + +_DeploymentPool = Sequence[_DeploymentLike] | Mapping[_DeploymentLike, object] + + def _is_valid_deployment_tag_regex( - tag_regexes: list[str], - header_strings: list[str], + tag_regexes: Sequence[str], + header_strings: Sequence[str], ) -> str | None: """ Test compiled regex patterns against "Header-Name: value" strings. @@ -46,7 +105,9 @@ def _is_valid_deployment_tag_regex( return None -def is_valid_deployment_tag(deployment_tags: list[str], request_tags: list[str], match_any: bool = True) -> bool: +def is_valid_deployment_tag( + deployment_tags: Sequence[str], request_tags: Sequence[str], match_any: bool = True +) -> bool: """ Check if a tag is valid, the matching can be either any or all based on `match_any` flag """ @@ -73,7 +134,7 @@ def is_valid_deployment_tag(deployment_tags: list[str], request_tags: list[str], def _match_deployment( - deployment: Any, + deployment: _DeploymentLike, request_tags: list[str] | None, header_strings: list[str], match_any: bool, @@ -90,8 +151,8 @@ def _match_deployment( ran and failed, so the regex cannot override strict-tag policy. """ litellm_params: Final = deployment.get("litellm_params", {}) - deployment_tags: Final[list[str] | None] = litellm_params.get("tags") - deployment_tag_regex: Final[list[str] | None] = litellm_params.get("tag_regex") + deployment_tags: Final[Sequence[str] | None] = litellm_params.get("tags") + deployment_tag_regex: Final[Sequence[str] | None] = litellm_params.get("tag_regex") # 1. Exact tag match (existing behaviour). if deployment_tags and request_tags: @@ -162,38 +223,38 @@ def _split_tags(tags: Sequence[str]) -> tuple[tuple[str, ...], list[str], tuple[ def _exclude_deployments( - deployments: Sequence[Any] | Mapping[Any, Any], + deployments: _DeploymentPool, excluded_set: frozenset[str], -) -> list[Any]: +) -> Sequence[_DeploymentLike]: if not excluded_set: return list(deployments) return [d for d in deployments if not excluded_set.intersection(d.get("litellm_params", {}).get("tags") or [])] def _require_all_tags( - deployments: Sequence[Any] | Mapping[Any, Any], + deployments: _DeploymentPool, required_set: frozenset[str], -) -> tuple[Any, ...]: +) -> tuple[_DeploymentLike, ...]: 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, ...]: + deployments: _DeploymentPool, +) -> tuple[_DeploymentLike, ...]: 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]: +def _known_tag_values(deployments: _DeploymentPool) -> 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], + healthy_deployments: _DeploymentPool, excluded_set: frozenset[str], required_set: frozenset[str], routing_confirmed: frozenset[str], @@ -217,7 +278,7 @@ def _unknown_required_tag_hides_an_answer( def _chain_allows_fail_open( - healthy_deployments: Sequence[Any] | Mapping[Any, Any], + healthy_deployments: _DeploymentPool, excluded_set: frozenset[str], required_set: frozenset[str], routing_confirmed: frozenset[str], @@ -228,12 +289,12 @@ def _chain_allows_fail_open( def _trusted_only_pool( - healthy_deployments: Sequence[Any] | Mapping[Any, Any], + healthy_deployments: _DeploymentPool, excluded_set: frozenset[str], required_set: frozenset[str], inherited_excluded_set: frozenset[str] | None, inherited_required_set: frozenset[str] | None, -) -> tuple[Any, ...]: +) -> tuple[_DeploymentLike, ...]: # 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 @@ -260,8 +321,8 @@ def _trusted_only_pool( def _resolve_or_fail_open( - pool: Sequence[Any], - healthy_deployments: Sequence[Any] | Mapping[Any, Any], + pool: Sequence[_DeploymentLike], + healthy_deployments: _DeploymentPool, excluded_set: frozenset[str], required_set: frozenset[str], inherited_excluded_set: frozenset[str] | None, @@ -269,7 +330,7 @@ def _resolve_or_fail_open( routing_confirmed: frozenset[str], model: str, request_tags: object, -) -> tuple[Any, ...]: +) -> tuple[_DeploymentLike, ...]: if pool: return tuple(pool) if _chain_allows_fail_open(healthy_deployments, excluded_set, required_set, routing_confirmed): @@ -289,7 +350,7 @@ def _resolve_or_fail_open( def _resolve_constraint_only_pool( - healthy_deployments: Sequence[Any] | Mapping[Any, Any], + healthy_deployments: _DeploymentPool, excluded_set: frozenset[str], required_set: frozenset[str], inherited_excluded_set: frozenset[str] | None, @@ -297,7 +358,7 @@ def _resolve_constraint_only_pool( routing_confirmed: frozenset[str], model: str, request_tags: object, -) -> tuple[Any, ...]: +) -> tuple[_DeploymentLike, ...]: pool: Final = ( _require_all_tags(_exclude_deployments(healthy_deployments, excluded_set), required_set) if required_set @@ -319,8 +380,8 @@ def _resolve_constraint_only_pool( def _all_deployments_or_fallback( llm_router_instance: LitellmRouter, model: str, - fallback: Sequence[Any] | Mapping[Any, Any], -) -> Sequence[Any] | Mapping[Any, Any]: + fallback: _DeploymentPool, +) -> Sequence[_DeploymentLike | DeploymentTypedDict] | Mapping[_DeploymentLike, object]: 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 @@ -330,7 +391,7 @@ def _all_deployments_or_fallback( def _chain_tag_filtering_override( llm_router_instance: LitellmRouter, model: str, - healthy_deployments: Sequence[Any] | Mapping[Any, Any], + healthy_deployments: _DeploymentPool, ) -> bool | None: # Resolved from every deployment configured for this model group, not just the # ones that survived cooldown/health filtering (async_get_healthy_deployments @@ -392,10 +453,10 @@ def _tag_known_to_group( async def get_deployments_for_tag( llm_router_instance: LitellmRouter, model: str, # used to raise the correct error - healthy_deployments: list[Any] | dict[Any, Any], - request_kwargs: dict[Any, Any] | None = None, + healthy_deployments: _DeploymentPool, + request_kwargs: _RequestKwargsLike | None = None, metadata_variable_name: Literal["metadata", "litellm_metadata"] = "metadata", -): +) -> _DeploymentPool: """ Returns a list of deployments that match the requested model and tags in the request. @@ -473,25 +534,25 @@ async def get_deployments_for_tag( request_tags, ) - new_healthy_deployments: Final[list[Any]] = [] - default_deployments: Final[list[Any]] = [] - if has_positive_filter: verbose_logger.debug( "get_deployments_for_tag routing: request_tags=%s user_agent=%s", request_tags, user_agent, ) - for deployment in candidates: - deployment_tags = deployment.get("litellm_params", {}).get("tags") - - match_result = _match_deployment( - deployment=deployment, - request_tags=positive_tags, - header_strings=header_strings, - match_any=match_any, + deployment_matches: Final = tuple( + ( + deployment, + _match_deployment( + deployment=deployment, + request_tags=positive_tags, + header_strings=header_strings, + match_any=match_any, + ), ) - + for deployment in candidates + ) + for deployment, match_result in deployment_matches: if match_result is not None: verbose_logger.debug( "tag routing match: deployment=%s matched_via=%s matched_value=%s", @@ -507,10 +568,10 @@ async def get_deployments_for_tag( "request_tags": request_tags or [], "user_agent": user_agent, } - new_healthy_deployments.append(deployment) - - if deployment_tags and "default" in deployment_tags: - default_deployments.append(deployment) + new_healthy_deployments: Final = [d for d, result in deployment_matches if result is not None] + default_deployments: Final = [ + d for d, _ in deployment_matches if "default" in (d.get("litellm_params", {}).get("tags") or ()) + ] if len(new_healthy_deployments) == 0 and len(default_deployments) == 0: return _resolve_or_fail_open( @@ -545,10 +606,11 @@ async def get_deployments_for_tag( return new_healthy_deployments if len(new_healthy_deployments) > 0 else default_deployments # for Untagged requests use default deployments if set - _default_deployments_with_tags: Final = [] - for deployment in healthy_deployments: - if "default" in deployment.get("litellm_params", {}).get("tags", []): - _default_deployments_with_tags.append(deployment) + _default_deployments_with_tags: Final = [ + deployment + for deployment in healthy_deployments + if "default" in deployment.get("litellm_params", {}).get("tags", []) + ] if len(_default_deployments_with_tags) > 0: return _default_deployments_with_tags @@ -562,7 +624,7 @@ async def get_deployments_for_tag( def _get_tags_from_request_kwargs( - request_kwargs: dict[Any, Any] | None = None, + request_kwargs: _RequestKwargsLike | None = None, metadata_variable_name: Literal["metadata", "litellm_metadata"] = "metadata", ) -> list[str]: """ @@ -577,12 +639,12 @@ def _get_tags_from_request_kwargs( if request_kwargs is None: return [] if metadata_variable_name in request_kwargs: - metadata: Final = request_kwargs[metadata_variable_name] or {} + metadata: Final[_MetadataLike] = request_kwargs[metadata_variable_name] or {} tags = metadata.get("tags", []) - return tags if tags is not None else [] + return list(tags) if tags is not None else [] elif "litellm_params" in request_kwargs: - litellm_params: Final = request_kwargs["litellm_params"] or {} - _metadata: Final = litellm_params.get(metadata_variable_name, {}) or {} + litellm_params: Final[_NestedLitellmParamsLike] = request_kwargs["litellm_params"] or {} + _metadata: Final[_MetadataLike] = litellm_params.get(metadata_variable_name, {}) or {} tags = _metadata.get("tags", []) - return tags if tags is not None else [] + return list(tags) if tags is not None else [] return [] diff --git a/ruff-strict-budget.json b/ruff-strict-budget.json index dff010bfd30..90033af024b 100644 --- a/ruff-strict-budget.json +++ b/ruff-strict-budget.json @@ -1,6 +1,6 @@ { "ANN001": { - "limit": 3058 + "limit": 3036 }, "ANN002": { "limit": 71 @@ -9,7 +9,7 @@ "limit": 827 }, "ANN201": { - "limit": 2022 + "limit": 2020 }, "ANN202": { "limit": 855 @@ -24,7 +24,7 @@ "limit": 133 }, "ANN401": { - "limit": 1384 + "limit": 1286 }, "ASYNC230": { "limit": 11 @@ -39,7 +39,7 @@ "limit": 505 }, "B009": { - "limit": 64 + "limit": 63 }, "B010": { "limit": 190 @@ -123,7 +123,7 @@ "limit": 12 }, "PERF403": { - "limit": 34 + "limit": 33 }, "PIE804": { "limit": 18 @@ -180,7 +180,7 @@ "limit": 8 }, "RUF019": { - "limit": 38 + "limit": 36 }, "RUF046": { "limit": 4 @@ -201,7 +201,7 @@ "limit": 58 }, "SIM102": { - "limit": 321 + "limit": 318 }, "SIM103": { "limit": 119 @@ -234,7 +234,7 @@ "limit": 5 }, "TID251": { - "limit": 1224 + "limit": 1214 }, "TRY002": { "limit": 524 diff --git a/type-discipline-budget.json b/type-discipline-budget.json index fdacf375844..83def5fe2e3 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 23003 + "limit": 22780 }, "LIT002": { - "limit": 27146 + "limit": 27144 }, "LIT003": { "limit": 269 @@ -15,7 +15,7 @@ "limit": 0 }, "LIT006": { - "limit": 1077 + "limit": 1069 }, "LIT007": { "limit": 0 @@ -27,9 +27,9 @@ "limit": 0 }, "LIT010": { - "limit": 16731 + "limit": 16725 }, "LIT011": { - "limit": 5596 + "limit": 5590 } } From a90ad5fe5c6dc3d837a0f28f32b1cc9e2ce93c33 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Fri, 28 Aug 2026 17:35:15 -0700 Subject: [PATCH 05/49] feat(bedrock): honor streaming buffer/sampling config for unbuffered post_call scans --- .../guardrail_hooks/bedrock_guardrails.py | 46 ++++- .../guardrails/guardrail_initializers.py | 4 + litellm/types/guardrails.py | 32 ++++ .../test_bedrock_guardrails.py | 172 ++++++++++++++++++ 4 files changed, 253 insertions(+), 1 deletion(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index dd76a27c80f..d7f222db35d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -52,7 +52,12 @@ from litellm.proxy.guardrails.anthropic_sse import ( model_response_text, ) from litellm.secret_managers.main import get_secret_str -from litellm.types.guardrails import BedrockChecksConfigModel, GuardrailEventHooks +from litellm.types.guardrails import ( + BedrockChecksConfigModel, + BedrockGuardrailStreamingParams, + GuardrailEventHooks, + LitellmParams, +) from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage from litellm.types.proxy.guardrails.guardrail_hooks.bedrock_guardrails import ( BedrockChecksMessage, @@ -221,9 +226,21 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): prompt_attack_threshold: float | None = 0.5, pii_confidence_threshold: float | None = 0.5, chunk_budget_chars: int = BEDROCK_APPLY_GUARDRAIL_CHUNK_BUDGET_CHARS, + streaming_buffer_until_moderated: bool | None = None, + streaming_sampling_rate: int | None = None, + streaming_end_of_stream_only: bool | None = None, **kwargs, ): self.async_handler = get_async_httpx_client(llm_provider=httpxSpecialProvider.GuardrailCallback) + self._set_streaming_params( + BedrockGuardrailStreamingParams.from_extras( + { + "streaming_buffer_until_moderated": streaming_buffer_until_moderated, + "streaming_sampling_rate": streaming_sampling_rate, + "streaming_end_of_stream_only": streaming_end_of_stream_only, + } + ) + ) self.guardrailIdentifier = guardrailIdentifier self.guardrailVersion = guardrailVersion self.guardrail_provider = "bedrock" @@ -278,6 +295,18 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): list(self.checks.keys()) if self.checks else None, ) + def _set_streaming_params(self, streaming_params: BedrockGuardrailStreamingParams) -> None: + self.streaming_buffer_until_moderated = streaming_params.streaming_buffer_until_moderated + self.streaming_sampling_rate = streaming_params.streaming_sampling_rate + self.streaming_end_of_stream_only = streaming_params.streaming_end_of_stream_only + + def update_in_memory_litellm_params(self, litellm_params: LitellmParams) -> None: + super().update_in_memory_litellm_params(litellm_params) + self._set_streaming_params(BedrockGuardrailStreamingParams.from_extras(litellm_params.model_extra)) + + def _streams_incrementally(self) -> bool: + return not self.streaming_buffer_until_moderated and not self.mask_response_content + @classmethod def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: return [ @@ -2660,6 +2689,21 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): Collect content from the stream and run the bedrock OUTPUT scan (post_call only validates the response). """ + if self._streams_incrementally(): + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, + ) + + async for streamed_chunk in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=response, + request_data=request_data, + guardrail_to_apply=self, + buffer_until_moderated_default=False, + ): + yield streamed_chunk + return + # Import here to avoid circular imports from litellm.llms.base_llm.base_model_iterator import MockResponseIterator from litellm.main import stream_chunk_builder diff --git a/litellm/proxy/guardrails/guardrail_initializers.py b/litellm/proxy/guardrails/guardrail_initializers.py index 47aea62f4c2..76dea1b7784 100644 --- a/litellm/proxy/guardrails/guardrail_initializers.py +++ b/litellm/proxy/guardrails/guardrail_initializers.py @@ -11,6 +11,7 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail): BedrockGuardrail, ) + streaming_params: Final = BedrockGuardrailStreamingParams.from_extras(litellm_params.model_extra) _bedrock_callback: Final = BedrockGuardrail( guardrail_name=guardrail.get("guardrail_name", ""), event_hook=litellm_params.mode, @@ -38,6 +39,9 @@ def initialize_bedrock(litellm_params: LitellmParams, guardrail: Guardrail): aws_bedrock_runtime_endpoint=litellm_params.aws_bedrock_runtime_endpoint, experimental_use_latest_role_message_only=litellm_params.experimental_use_latest_role_message_only, only_scan_new_messages=litellm_params.only_scan_new_messages or False, + streaming_buffer_until_moderated=streaming_params.streaming_buffer_until_moderated, + streaming_sampling_rate=streaming_params.streaming_sampling_rate, + streaming_end_of_stream_only=streaming_params.streaming_end_of_stream_only, ) litellm.logging_callback_manager.add_litellm_callback(_bedrock_callback) return _bedrock_callback diff --git a/litellm/types/guardrails.py b/litellm/types/guardrails.py index 9be78757511..c5398160c69 100644 --- a/litellm/types/guardrails.py +++ b/litellm/types/guardrails.py @@ -1,3 +1,4 @@ +from collections.abc import Mapping from datetime import datetime from enum import Enum from typing import Any, Final, Literal @@ -550,6 +551,37 @@ class BedrockGuardrailConfigModel(BaseModel): ) +class BedrockGuardrailStreamingParams(BaseModel): + streaming_buffer_until_moderated: bool = Field( + default=True, + description="If True (default), withhold every streamed chunk until the end-of-stream " + "ApplyGuardrail scan passes, so no flagged content reaches the client before a block. " + "If False, chunks stream through unbuffered, so flagged content can reach the client " + "before the scan finishes; a flagged scan still ends the stream, with a block message " + "when disable_exception_on_block is true and an in-stream error frame otherwise.", + ) + streaming_sampling_rate: int = Field( + default=5, + ge=1, + description="When not buffering and not end-of-stream-only, scan the accumulated response " + "every Nth streamed chunk. Each sampled scan is a full ApplyGuardrail call that delays " + "that chunk, so lower values add latency and AWS text-unit cost.", + ) + streaming_end_of_stream_only: bool = Field( + default=False, + description="When not buffering, skip per-chunk sampling and run one ApplyGuardrail scan " + "on the assembled response at end of stream. Combined with " + "streaming_buffer_until_moderated=false the full response streams live before the scan " + "and the scan result lands in guardrail_information; a flagged response still ends the " + "stream with a block message (disable_exception_on_block=true) or an error frame.", + ) + + @classmethod + def from_extras(cls, extras: Mapping[str, object] | None) -> "BedrockGuardrailStreamingParams": + source: Final[Mapping[str, object]] = extras or {} + return cls.model_validate({name: source[name] for name in cls.model_fields if source.get(name) is not None}) + + class LakeraV2GuardrailConfigModel(BaseModel): """Configuration parameters for the Lakera AI v2 guardrail""" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index 36b356e34d0..ca4b0a65e0b 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5345,3 +5345,175 @@ def test_initialize_bedrock_forwards_aws_external_id(): assert guardrail.optional_params["aws_external_id"] == "external-id-123" finally: litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.callbacks, guardrail) + + +def _chat_chunk(content: str, finish_reason: str | None) -> litellm.ModelResponseStream: + return litellm.ModelResponseStream( + id="tid", + choices=[ + litellm.types.utils.StreamingChoices( + delta=litellm.types.utils.Delta(content=content, role="assistant"), + finish_reason=finish_reason, + index=0, + ) + ], + created=1, + model="gpt-4o-mini", + object="chat.completion.chunk", + ) + + +def _streaming_litellm_params(**extras): + from litellm.types.guardrails import LitellmParams + + return LitellmParams( + guardrail="bedrock", + mode="post_call", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + **extras, + ) + + +def test_initialize_bedrock_wires_streaming_flags(): + from litellm.proxy.guardrails.guardrail_initializers import initialize_bedrock + + configured = initialize_bedrock( + _streaming_litellm_params( + streaming_buffer_until_moderated=False, + streaming_sampling_rate=3, + streaming_end_of_stream_only=True, + ), + {"guardrail_name": "bedrock-streaming"}, + ) + defaulted = initialize_bedrock( + _streaming_litellm_params(), + {"guardrail_name": "bedrock-defaults"}, + ) + for registered in (configured, defaulted): + litellm.logging_callback_manager.remove_callback_from_list_by_object(litellm.callbacks, registered) + + assert configured.streaming_buffer_until_moderated is False + assert configured.streaming_sampling_rate == 3 + assert configured.streaming_end_of_stream_only is True + assert defaulted.streaming_buffer_until_moderated is True + assert defaulted.streaming_sampling_rate == 5 + assert defaulted.streaming_end_of_stream_only is False + + +def test_initialize_bedrock_rejects_non_positive_sampling_rate(): + from pydantic import ValidationError + + from litellm.proxy.guardrails.guardrail_initializers import initialize_bedrock + + with pytest.raises(ValidationError): + initialize_bedrock( + _streaming_litellm_params(streaming_sampling_rate=0), + {"guardrail_name": "bedrock-bad-rate"}, + ) + + +def test_update_in_memory_litellm_params_round_trips_streaming_flags(): + guardrail = BedrockGuardrail( + guardrail_name="bedrock-update", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + ) + + guardrail.update_in_memory_litellm_params( + _streaming_litellm_params( + streaming_buffer_until_moderated=False, + streaming_sampling_rate=7, + streaming_end_of_stream_only=True, + ) + ) + assert guardrail.streaming_buffer_until_moderated is False + assert guardrail.streaming_sampling_rate == 7 + assert guardrail.streaming_end_of_stream_only is True + + guardrail.update_in_memory_litellm_params(_streaming_litellm_params()) + assert guardrail.streaming_buffer_until_moderated is True + assert guardrail.streaming_sampling_rate == 5 + assert guardrail.streaming_end_of_stream_only is False + + +async def _run_streaming_hook_recording_order(guardrail: BedrockGuardrail) -> list: + events = [] + minimal = {"action": "NONE", "assessments": [], "outputs": []} + + async def record_scan(*args, **kwargs): + events.append("scan") + return minimal + + async def mock_stream(): + yield _chat_chunk("Hello", None) + yield _chat_chunk(" world", None) + yield _chat_chunk("", "stop") + + with patch.object(guardrail, "make_bedrock_api_request", AsyncMock(side_effect=record_scan)): + async for chunk in guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(), + response=mock_stream(), + request_data={"model": "gpt-4o-mini", "messages": [{"role": "user", "content": "hi"}]}, + ): + content = chunk.choices[0].delta.content if chunk.choices else None + events.append(("chunk", content)) + return events + + +@pytest.mark.asyncio +async def test_unbuffered_end_of_stream_hook_yields_chunks_before_scan(): + guardrail = BedrockGuardrail( + guardrail_name="bedrock-audit-mode", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + streaming_buffer_until_moderated=False, + streaming_end_of_stream_only=True, + ) + + events = await _run_streaming_hook_recording_order(guardrail) + + scan_index = events.index("scan") + chunk_events = [e for e in events if e != "scan"] + assert events.count("scan") == 1 + assert [e for e in events[:scan_index] if e != "scan"] == chunk_events[: scan_index] + assert ("chunk", "Hello") in events[:scan_index] + assert ("chunk", " world") in events[:scan_index] + assert len(chunk_events) == 3 + + +@pytest.mark.asyncio +async def test_buffered_default_hook_scans_before_any_chunk(): + guardrail = BedrockGuardrail( + guardrail_name="bedrock-buffered-default", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + ) + + events = await _run_streaming_hook_recording_order(guardrail) + + assert events[0] == "scan" + assert all(e == "scan" or e[0] == "chunk" for e in events) + assert len([e for e in events if e != "scan"]) >= 1 + + +@pytest.mark.asyncio +async def test_masking_keeps_buffered_path_even_when_unbuffered_configured(): + guardrail = BedrockGuardrail( + guardrail_name="bedrock-mask-buffered", + guardrailIdentifier="test-id", + guardrailVersion="DRAFT", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + mask_response_content=True, + streaming_buffer_until_moderated=False, + streaming_end_of_stream_only=True, + ) + + assert guardrail._streams_incrementally() is False + events = await _run_streaming_hook_recording_order(guardrail) + assert events[0] == "scan" From 64eec53fd844684df2f953558c405c66eb27ce2d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 01:25:08 -0700 Subject: [PATCH 06/49] fix(guardrails): surface post-flush stream blocks as in-stream error frames and keep guardrail_information in spend logs A guardrail block or failed scan that fires after SSE chunks have been flushed can no longer set an HTTP status, so raising HTTPException there silently truncated the stream. _emit_streaming_http_error now routes post-flush failures through the endpoint translation's build_stream_error_items, emitting the surface-correct error frame on chat completions (data: {error}), /v1/messages (event: error), and /v1/responses (ErrorEvent with the next sequence number). Pre-flush blocks still raise with a real HTTP status. Successful flags-on scans also logged metadata.guardrail_information as null: the chat handler planted litellm_metadata on a route whose bucket is metadata, flipping the bucket for every later write, and responses streams fired their spend log before the eos scan ran. The chat handler now merges user_api_key metadata through get_or_create_metadata_bucket, and deferred stream-complete logging is armed for aresponses like it already was for anthropic_messages. --- .../chat/guardrail_translation/handler.py | 15 ++ .../guardrail_translation/base_translation.py | 48 ++++ .../chat/guardrail_translation/handler.py | 32 +-- .../guardrail_translation/handler.py | 35 +++ litellm/proxy/common_request_processing.py | 14 +- .../guardrail_hooks/bedrock_guardrails.py | 4 +- .../unified_guardrail/unified_guardrail.py | 67 ++++- litellm/responses/streaming_iterator.py | 21 +- .../openai/test_moderations.py | 32 +-- .../test_openai_moderation_streaming.py | 32 ++- .../test_bedrock_guardrails.py | 73 ++++++ .../test_unified_guardrail.py | 235 +++++++++++++++++- .../test_deferred_guardrail_logging.py | 69 +++++ .../proxy/test_common_request_processing.py | 16 +- 14 files changed, 608 insertions(+), 85 deletions(-) diff --git a/litellm/llms/anthropic/chat/guardrail_translation/handler.py b/litellm/llms/anthropic/chat/guardrail_translation/handler.py index b9ca18c7843..6cb2e568d0f 100644 --- a/litellm/llms/anthropic/chat/guardrail_translation/handler.py +++ b/litellm/llms/anthropic/chat/guardrail_translation/handler.py @@ -58,6 +58,8 @@ from litellm.types.utils import ( ) if TYPE_CHECKING: + from fastapi import HTTPException + from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, @@ -165,6 +167,19 @@ class AnthropicMessagesHandler(BaseTranslation): return self._block_continuation_chunks(exc, responses_so_far or []) return self._standalone_block_chunks(exc) + def build_stream_error_items( + self, + exc: "HTTPException", + responses_so_far: Sequence[Any] | None = None, + ) -> Sequence[Any] | None: + from litellm.proxy.common_request_processing import ( + serialize_http_exception_detail, + ) + from litellm.proxy.guardrails.anthropic_sse import anthropic_sse_error_frames + + message, _ = serialize_http_exception_detail(exc.detail) + return list(anthropic_sse_error_frames(message)) + def _standalone_block_chunks(self, exc: "ModifyResponseException") -> list[bytes]: import uuid diff --git a/litellm/llms/base_llm/guardrail_translation/base_translation.py b/litellm/llms/base_llm/guardrail_translation/base_translation.py index ba96ab3dc99..b07a6b986d7 100644 --- a/litellm/llms/base_llm/guardrail_translation/base_translation.py +++ b/litellm/llms/base_llm/guardrail_translation/base_translation.py @@ -1,8 +1,11 @@ from abc import ABC, abstractmethod +from collections.abc import Sequence from dataclasses import dataclass, field from typing import TYPE_CHECKING, Any, Final, Optional if TYPE_CHECKING: + from fastapi import HTTPException + from litellm.integrations.custom_guardrail import ( CustomGuardrail, ModifyResponseException, @@ -73,6 +76,31 @@ class BaseTranslation(ABC): return transformed + @staticmethod + def merge_user_api_key_metadata_into_request( + request_data: dict[str, Any], # mutable-ok: proxy hooks share and mutate the request payload dict in place + user_api_key_dict: Optional["UserAPIKeyAuth"], + ) -> None: + """ + Add the prefixed ``user_api_key_*`` metadata to the request's resolved + metadata bucket without overwriting existing keys. + + Writes must go through ``get_or_create_metadata_bucket``: creating a + ``litellm_metadata`` key on a route whose bucket is ``metadata`` (chat + completions) flips the bucket for every later metadata write, and spend + logging never sees those writes (e.g. guardrail_information). + """ + from litellm.litellm_core_utils.core_helpers import ( + get_or_create_metadata_bucket, + ) + + user_metadata: Final = BaseTranslation.transform_user_api_key_dict_to_metadata(user_api_key_dict) + if not user_metadata: + return + _, metadata_bucket = get_or_create_metadata_bucket(request_data) + for key, value in user_metadata.items(): + metadata_bucket.setdefault(key, value) + @abstractmethod async def process_input_messages( self, @@ -147,6 +175,26 @@ class BaseTranslation(ABC): """ return None + def build_stream_error_items( + self, + exc: "HTTPException", + responses_so_far: Sequence[Any] | None = None, + ) -> Sequence[Any] | None: + """ + Build the stream items that surface a guardrail HTTPException (a block + with the default exception-on-block config, or a failed scan) after the + response has already started streaming, in this endpoint's wire format. + + Called only once chunks have been sent: the HTTP status is gone, so the + failure must travel as an in-stream error frame. ``responses_so_far`` + holds the chunks the client has already received, for formats whose + error frame continues the stream (e.g. sequence numbers). + + Returns None when the format has no in-stream error frame; the caller + then re-raises ``exc``. Override in endpoint subclasses. + """ + return None + def get_structured_messages(self, data: dict) -> list["AllMessageValues"] | None: """ Convert request data to OpenAI-spec structured messages. diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index e411dc497fc..e61fb719c98 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -14,6 +14,7 @@ Pattern Overview: This pattern can be replicated for other message formats (e.g., Anthropic). """ +from collections.abc import Sequence from typing import TYPE_CHECKING, Any, Final, Union, cast import litellm @@ -46,6 +47,8 @@ from litellm.types.utils import ( ) if TYPE_CHECKING: + from fastapi import HTTPException + from litellm.integrations.custom_guardrail import CustomGuardrail @@ -381,11 +384,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if "response" not in request_data: request_data["response"] = response - # Add user API key metadata with prefixed keys - if "litellm_metadata" not in request_data: - user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) - if user_metadata: - request_data["litellm_metadata"] = user_metadata + self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict) inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) if images_to_check: @@ -554,11 +553,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): if "responses" not in request_data: request_data["responses"] = responses_so_far - # Add user API key metadata with prefixed keys - if "litellm_metadata" not in request_data: - user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) - if user_metadata: - request_data["litellm_metadata"] = user_metadata + self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict) inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) if images_to_check: @@ -590,6 +585,18 @@ class OpenAIChatCompletionsHandler(BaseTranslation): return responses_so_far + def build_stream_error_items( + self, + exc: "HTTPException", + responses_so_far: Sequence[Any] | None = None, + ) -> Sequence[Any] | None: + import json + + from litellm.proxy.common_request_processing import sse_error_payload + + _, error_obj = sse_error_payload(exc) + return [f"data: {json.dumps({'error': error_obj})}\n\n".encode()] + @staticmethod def _accumulate_string_content_by_choice_index( responses_so_far: list["ModelResponseStream"], @@ -652,10 +659,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): request_data = {"responses": responses_so_far} elif "responses" not in request_data: request_data["responses"] = responses_so_far - if "litellm_metadata" not in request_data: - user_metadata: Final = self.transform_user_api_key_dict_to_metadata(user_api_key_dict) - if user_metadata: - request_data["litellm_metadata"] = user_metadata + self.merge_user_api_key_metadata_into_request(request_data, user_api_key_dict) inputs: Final = GenericGuardrailAPIInputs(texts=texts_to_check) if responses_so_far and getattr(responses_so_far[0], "model", None): diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 7c5d8ac99ad..50d7f7452a8 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -48,6 +48,8 @@ from litellm.types.llms.openai import ( AllMessageValues, ChatCompletionToolCallChunk, ChatCompletionToolParam, + ErrorEvent, + ErrorEventError, OpenAIMcpServerTool, ResponsesAPIStreamEvents, ) @@ -59,6 +61,8 @@ from litellm.types.responses.main import ( from litellm.types.utils import GenericGuardrailAPIInputs if TYPE_CHECKING: + from fastapi import HTTPException + from litellm.integrations.custom_guardrail import CustomGuardrail from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.proxy._types import UserAPIKeyAuth @@ -80,6 +84,14 @@ class ResponsesStreamChunk(TypedDict, total=False): text: ReadOnly[str] +def _next_stream_sequence_number(responses_so_far: Sequence[Any] | None) -> int: + sequence_numbers: Final = ( + item.get("sequence_number") if isinstance(item, dict) else getattr(item, "sequence_number", None) + for item in reversed(responses_so_far or []) + ) + return next((n + 1 for n in sequence_numbers if isinstance(n, int)), 0) + + class OpenAIResponsesHandler(BaseTranslation): """ Handler for processing OpenAI Responses API with guardrails. @@ -620,6 +632,29 @@ class OpenAIResponsesHandler(BaseTranslation): } return responses_so_far[-1].get("type") in terminal_types + def build_stream_error_items( + self, + exc: "HTTPException", + responses_so_far: Sequence[Any] | None = None, + ) -> Sequence[Any] | None: + from litellm.proxy.common_request_processing import ( + serialize_http_exception_detail, + ) + + message, _ = serialize_http_exception_detail(exc.detail) + return [ + ErrorEvent( + type=ResponsesAPIStreamEvents.ERROR, + sequence_number=_next_stream_sequence_number(responses_so_far), + error=ErrorEventError( + type="guardrail_error", + code=str(exc.status_code), + message=message, + param=None, + ), + ) + ] + def get_streaming_string_so_far(self, responses_so_far: Sequence[ResponsesStreamChunk]) -> str: """ Get the string so far from the responses so far. diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index ff6c8d1b1f8..6c7c77610b3 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -502,7 +502,7 @@ def _as_success_dispatcher(logging_obj: _DispatchesSuccessHandlers) -> _Dispatch return logging_obj -def _serialize_http_exception_detail( +def serialize_http_exception_detail( detail: object, ) -> tuple[str, dict | None]: """ @@ -803,7 +803,7 @@ async def _buffer_first_chunk_honoring_disconnect( raise _ClientDisconnectedBeforeFirstChunk() -def _sse_error_payload(exc: BaseException) -> tuple[int, Mapping[str, object]]: +def sse_error_payload(exc: BaseException) -> tuple[int, Mapping[str, object]]: """Build the ProxyException-shaped ``{"error": ...}`` body used in SSE error frames. Matches ``ProxyException.to_dict()`` so streaming and non-streaming error frames @@ -812,7 +812,7 @@ def _sse_error_payload(exc: BaseException) -> tuple[int, Mapping[str, object]]: # Preserve status code from HTTPException (e.g. guardrail blocks) error_status: Final = getattr(exc, "status_code", status.HTTP_500_INTERNAL_SERVER_ERROR) raw_detail: Final = _getattr_object(exc, "detail", "Error processing stream start") - message, structured_fields = _serialize_http_exception_detail(raw_detail) + message, structured_fields = serialize_http_exception_detail(raw_detail) existing_fields: Final = getattr(exc, "provider_specific_fields", None) or {} merged_fields: Final = {**existing_fields, **structured_fields} if structured_fields else (existing_fields or None) @@ -927,7 +927,7 @@ async def create_response( # Unexpected error consuming first chunk. verbose_proxy_logger.exception("Error consuming first chunk from generator: %s", e) - error_status, error_obj = _sse_error_payload(e) + error_status, error_obj = sse_error_payload(e) async def error_gen_message() -> AsyncGenerator[str, None]: for frame in _sse_error_frames(error_obj): @@ -1104,7 +1104,7 @@ async def open_sse_before_first_byte( # would never fire and the failure would go unaudited. The hook # also gets to sanitize what reaches the client, by returning or # raising a replacement, so its answer decides the frame. - _, error_obj = _sse_error_payload(await _sanitized_late_failure(exc, on_late_failure)) + _, error_obj = sse_error_payload(await _sanitized_late_failure(exc, on_late_failure)) for frame in _sse_error_frames(error_obj): yield frame.encode() return @@ -2374,7 +2374,7 @@ class ProxyBaseLLMRequestProcessing: logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete elif ( _post_call_guardrails_active - and route_type == "anthropic_messages" + and route_type in ("anthropic_messages", "aresponses") and self._is_streaming_response(response) ): from litellm.litellm_core_utils.logging_worker import ( @@ -3245,7 +3245,7 @@ class ProxyBaseLLMRequestProcessing: if isinstance(e, HTTPException): raw_detail: Final = _getattr_object(e, "detail", str(e)) - message, structured_fields = _serialize_http_exception_detail(raw_detail) + message, structured_fields = serialize_http_exception_detail(raw_detail) existing_fields: Final = getattr(e, "provider_specific_fields", None) or {} if structured_fields: merged_fields: dict | None = {**existing_fields, **structured_fields} diff --git a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py index d7f222db35d..951cbc13290 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py +++ b/litellm/proxy/guardrails/guardrail_hooks/bedrock_guardrails.py @@ -42,7 +42,7 @@ from litellm.llms.custom_httpx.http_handler import ( httpxSpecialProvider, ) from litellm.proxy._types import UserAPIKeyAuth -from litellm.proxy.common_request_processing import _serialize_http_exception_detail +from litellm.proxy.common_request_processing import serialize_http_exception_detail from litellm.proxy.common_utils.sse_keepalive import keepalive_ping_has_fired from litellm.proxy.guardrails.anthropic_sse import ( anthropic_sse_chunks_from_response, @@ -2760,7 +2760,7 @@ class BedrockGuardrail(CustomGuardrail, BaseAWSLLM): ) if not raw_sse or (not is_block and not headers_flushed): raise - block_message, _ = _serialize_http_exception_detail(block_detail) + block_message, _ = serialize_http_exception_detail(block_detail) for error_frame in anthropic_sse_error_frames( block_message if is_block else f"{block_exc.status_code}: {block_message}" ): diff --git a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py index e95e97bfe74..6647ac4c293 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/unified_guardrail/unified_guardrail.py @@ -57,6 +57,9 @@ class _EndpointTranslation(Protocol): @property def build_block_sse_chunks(self) -> "Callable[..., Sequence[bytes] | None]": ... + @property + def build_stream_error_items(self) -> "Callable[..., Sequence[object] | None]": ... + def _as_endpoint_translation(translation: _EndpointTranslation) -> _EndpointTranslation: return translation @@ -408,14 +411,32 @@ class UnifiedLLMGuardrails(CustomLogger): call_type: str | None, responses_so_far: Sequence[object], request_data: dict, + endpoint_translation: _EndpointTranslation | None = None, + stream_started: bool = False, + responses_yielded: Sequence[object] | None = None, ) -> AsyncGenerator[object, None]: - """Surface a mid-stream HTTPException. For A2A call types the response has - already started, so emit an in-stream JSON-RPC error chunk; otherwise - re-raise so the proxy can report it. + """Surface a mid-stream HTTPException (a guardrail block with the default + exception-on-block config, or a failed scan). + + A2A call types emit an in-stream JSON-RPC error chunk. For other call + types, once chunks have already reached the client the HTTP status is + gone, so the failure is delegated to the endpoint translation's + ``build_stream_error_items`` and travels as an in-stream error frame in + that endpoint's wire format. Before the first chunk (or when the format + has no in-stream error frame) the exception is re-raised so the proxy + can report it with a real HTTP status. """ if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES: yield _a2a_jsonrpc_error_chunk(exc, _get_a2a_request_id(responses_so_far, request_data)) return + if stream_started and endpoint_translation is not None: + error_items: Final = endpoint_translation.build_stream_error_items( + exc, responses_so_far=list(responses_yielded) if responses_yielded is not None else None + ) + if error_items is not None: + for error_item in error_items: + yield error_item + return raise exc def _build_transform_chunk( @@ -586,7 +607,15 @@ class UnifiedLLMGuardrails(CustomLogger): yield block_chunk raise _StreamTerminated() except HTTPException as e: - async for error_item in self._emit_streaming_http_error(e, call_type, responses_so_far, request_data): + async for error_item in self._emit_streaming_http_error( + e, + call_type, + responses_so_far, + request_data, + endpoint_translation=endpoint_translation, + stream_started=bool(responses_yielded), + responses_yielded=responses_yielded, + ): yield error_item raise _StreamTerminated() @@ -1070,11 +1099,17 @@ class UnifiedLLMGuardrails(CustomLogger): return except HTTPException as e: # Response already started (we already yielded chunks); cannot send 400. - # For A2A, yield an in-stream JSON-RPC error so the client sees it. - if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES: - yield _a2a_jsonrpc_error_chunk(e, _get_a2a_request_id(responses_so_far, request_data)) - return - raise + async for error_item in self._emit_streaming_http_error( + e, + call_type, + responses_so_far, + request_data, + endpoint_translation=endpoint_translation, + stream_started=chunks_yielded, + responses_yielded=responses_yielded, + ): + yield error_item + return chunks_yielded = True responses_yielded.append(original_item) yield original_item @@ -1133,7 +1168,13 @@ class UnifiedLLMGuardrails(CustomLogger): yield block_chunk return except HTTPException as e: - if call_type is not None and CallTypes(call_type) in A2A_CALL_TYPES: - yield _a2a_jsonrpc_error_chunk(e, _get_a2a_request_id(responses_so_far, request_data)) - else: - raise + async for error_item in self._emit_streaming_http_error( + e, + call_type, + responses_so_far, + request_data, + endpoint_translation=endpoint_translation, + stream_started=bool(responses_yielded), + responses_yielded=responses_yielded, + ): + yield error_item diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 368fd481e63..94c93ae7c6b 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -422,15 +422,20 @@ class BaseResponsesAPIStreamingIterator: end_time: Final = datetime.now() if is_async: - asyncio.create_task( - self.logging_obj.dispatch_success_handlers( - logging_response, - start_time=self.start_time, - end_time=end_time, - cache_hit=self._completed_response_cache_hit, - prefer_async_handlers=True, - ) + logging_coroutine: Final = self.logging_obj.dispatch_success_handlers( + logging_response, + start_time=self.start_time, + end_time=end_time, + cache_hit=self._completed_response_cache_hit, + prefer_async_handlers=True, ) + deferred_dispatch_armed: Final = getattr(self.logging_obj, "_on_deferred_stream_complete", None) is not None + if deferred_dispatch_armed: + # End-of-stream guardrail scans write guardrail_information after + # the terminal event; dispatching now would snapshot metadata early. + self.logging_obj._deferred_stream_complete_args = (logging_coroutine,) + else: + asyncio.create_task(logging_coroutine) else: run_async_function( async_function=self.logging_obj.async_success_handler, diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py index 112bc5e6e49..2b43720a126 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_moderations.py @@ -482,23 +482,25 @@ async def test_openai_moderation_guardrail_streaming_harmful_content(): "metadata": {"guardrails": ["test-openai-moderation"]}, } - # Should raise HTTPException when processing streaming harmful content - from fastapi import HTTPException + # Chunks have already been flushed by end-of-stream moderation, so + # the block surfaces as the in-stream error frame, not a raise. + import json as _json - async def _drain(): - result_chunks = [] - async for chunk in unified_guardrail.async_post_call_streaming_iterator_hook( - user_api_key_dict=user_api_key_dict, - response=mock_stream(), - request_data=request_data, - ): - result_chunks.append(chunk) + result_chunks = [] + async for chunk in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=mock_stream(), + request_data=request_data, + ): + result_chunks.append(chunk) - with pytest.raises(HTTPException) as exc_info: - await _drain() - - assert exc_info.value.status_code == 400 - assert "Violated OpenAI moderation policy" in str(exc_info.value.detail) + frame = result_chunks[-1] + assert isinstance(frame, bytes) + text = frame.decode() + assert text.startswith("data: ") + assert "Violated OpenAI moderation policy" in text + payload = _json.loads(text[len("data: ") :]) + assert payload["error"]["code"] == "400" @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py index 914af0e2368..476d443d8d8 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/openai/test_openai_moderation_streaming.py @@ -161,19 +161,27 @@ async def test_openai_moderation_guardrail_streaming_harmful_content(): "metadata": {"guardrails": ["test-openai-moderation"]}, } - # Should raise HTTPException - with pytest.raises(HTTPException) as exc_info: - async for ( - _ - ) in unified_guardrail.async_post_call_streaming_iterator_hook( - user_api_key_dict=user_api_key_dict, - response=mock_stream(), - request_data=request_data, - ): - pass + # Chunks have already been flushed by end-of-stream moderation, so + # the block surfaces as the in-stream error frame, not a raise. + import json as _json - assert exc_info.value.status_code == 400 - assert "Violated OpenAI moderation policy" in str(exc_info.value.detail) + collected = [] + async for ( + chunk + ) in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=mock_stream(), + request_data=request_data, + ): + collected.append(chunk) + + frame = collected[-1] + assert isinstance(frame, bytes) + text = frame.decode() + assert text.startswith("data: ") + assert "Violated OpenAI moderation policy" in text + payload = _json.loads(text[len("data: ") :]) + assert payload["error"]["code"] == "400" @pytest.mark.asyncio diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py index ca4b0a65e0b..235a5c0c09b 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_bedrock_guardrails.py @@ -5517,3 +5517,76 @@ async def test_masking_keeps_buffered_path_even_when_unbuffered_configured(): assert guardrail._streams_incrementally() is False events = await _run_streaming_hook_recording_order(guardrail) assert events[0] == "scan" + + +@pytest.mark.asyncio +async def test_streaming_end_of_stream_block_emits_error_frame_instead_of_truncating(): + """Regression for PR #38722: a topicPolicy DENY caught by the end-of-stream + scan used to raise after SSE headers were flushed, so the client saw a + silently truncated stream. The unified hook must emit the chat in-stream + error frame instead.""" + from litellm.llms import load_guardrail_translation_mappings + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail import ( + unified_guardrail as unified_module, + ) + from litellm.proxy.guardrails.guardrail_hooks.unified_guardrail.unified_guardrail import ( + UnifiedLLMGuardrails, + ) + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + guardrail = BedrockGuardrail( + guardrailIdentifier="test-guardrail", + guardrailVersion="DRAFT", + streaming_end_of_stream_only=True, + streaming_buffer_until_moderated=False, + guardrail_name="bedrock-eos", + event_hook=GuardrailEventHooks.post_call, + default_on=True, + ) + blocked_response = { + "action": "GUARDRAIL_INTERVENED", + "actionReason": "Guardrail blocked.", + "outputs": [{"text": "Sorry, the model cannot answer this question."}], + "assessments": [ + {"topicPolicy": {"topics": [{"name": "Forbidden topic", "type": "DENY", "action": "BLOCKED"}]}} + ], + } + + def _chunk(content, finish_reason=None): + return ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta={"content": content, "role": "assistant"}, + finish_reason=finish_reason, + ) + ], + ) + + async def _mock_stream(): + yield _chunk("the forbidden ") + yield _chunk("topic answer", finish_reason="stop") + + unified_module.endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings() + try: + with patch.object(guardrail, "make_bedrock_api_request", new_callable=AsyncMock) as mock_api: + mock_api.side_effect = guardrail._get_http_exception_for_blocked_guardrail(blocked_response) + + out = [] + async for item in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="test", request_route="/v1/chat/completions"), + response=_mock_stream(), + request_data={"guardrail_to_apply": guardrail, "model": "gpt-4"}, + ): + out.append(item) + finally: + unified_module.endpoint_guardrail_translation_mappings = None + + assert len(out) == 3 + assert isinstance(out[0], ModelResponseStream) + frame = out[-1] + assert isinstance(frame, bytes) + payload = json.loads(frame.decode()[len("data: ") :]) + assert payload["error"]["message"] == "Violated guardrail policy" + assert payload["error"]["code"] == "400" + assert payload["error"]["provider_specific_fields"]["guardrailIdentifier"] == "test-guardrail" diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py index 8b9ecfbbeee..0b32558a00a 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/unified_guardrails/test_unified_guardrail.py @@ -948,19 +948,24 @@ class TestStreamingTransform: assert streamed == "ABCDEFGHIJ" @pytest.mark.asyncio - async def test_incremental_diff_underflow_raises(self): + async def test_incremental_diff_underflow_emits_error_frame(self): """A transform shorter than what was already streamed cannot retract - bytes: it raises HTTPException(stream_transform_underflow).""" + bytes. Chunks have already been flushed by then, so the underflow + surfaces as the in-stream error frame, not an unraisable HTTPException.""" + import json as _json + # First sample emits "ABCDEF" (6 chars); second sample shrinks to 3. guardrail = _StreamingTextGuardrail(shrink_to="ABC", shrink_after=1) chunks = [_stream_chunk("abcdef"), _stream_chunk("ghij")] - with pytest.raises(unified_module.HTTPException) as exc_info: - await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) - assert exc_info.value.status_code == 400 - assert exc_info.value.detail["error"] == "stream_transform_underflow" + frame = out[-1] + assert isinstance(frame, bytes) + payload = _json.loads(frame.decode()[len("data: ") :]) + assert payload["error"]["message"] == "stream_transform_underflow" + assert payload["error"]["code"] == "400" @pytest.mark.asyncio async def test_incremental_diff_final_chunk_preserves_finish_reason(self): @@ -1747,3 +1752,221 @@ class TestAppliedGuardrailsReflectsExecution: async def test_ordinary_guardrail_is_auto_marked_applied(self): data = await self._run(_AutoLoggingGuardrail()) assert "auto-logging" in _applied_guardrails(data) + + +class _EosHttpBlockingGuardrail(CustomGuardrail): + """Raises the bedrock-shaped block HTTPException at end-of-stream scan time.""" + + def __init__(self): + super().__init__(guardrail_name="eos-http-block") + self.streaming_end_of_stream_only = True + + def should_run_guardrail(self, data, event_type): # type: ignore[override] + return True + + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + raise unified_module.HTTPException( + status_code=400, + detail={ + "error": "Violated guardrail policy", + "bedrock_guardrail_response": "BLOCKED_TOPIC", + }, + ) + + +def _anthropic_sse_event(event_type, data): + import json as _json + + return f"event: {event_type}\ndata: {_json.dumps(data)}\n\n".encode() + + +def _anthropic_message_chunks(texts): + head = [ + _anthropic_sse_event( + "message_start", + { + "type": "message_start", + "message": { + "id": "msg_test", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5", + "content": [], + "stop_reason": None, + "usage": {"input_tokens": 1, "output_tokens": 0}, + }, + }, + ), + _anthropic_sse_event( + "content_block_start", + {"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}, + ), + ] + deltas = [ + _anthropic_sse_event( + "content_block_delta", + {"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}, + ) + for text in texts + ] + tail = [ + _anthropic_sse_event("content_block_stop", {"type": "content_block_stop", "index": 0}), + _anthropic_sse_event( + "message_delta", + { + "type": "message_delta", + "delta": {"stop_reason": "end_turn", "stop_sequence": None}, + "usage": {"output_tokens": 5}, + }, + ), + _anthropic_sse_event("message_stop", {"type": "message_stop"}), + ] + return head + deltas + tail + + +class TestStreamingHttpErrorFrames: + """A post-flush end-of-stream guardrail block (HTTPException) must surface as + the endpoint's in-stream error frame instead of an unhandled raise that + silently truncates the SSE stream (PR #38722 defect 1).""" + + @pytest.fixture(autouse=True) + def _use_real_mappings(self): + unified_module.endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings() + yield + unified_module.endpoint_guardrail_translation_mappings = None + + @pytest.mark.asyncio + async def test_chat_eos_block_emits_data_error_frame(self): + import json as _json + + guardrail = _EosHttpBlockingGuardrail() + chunks = [_stream_chunk("hello "), _stream_chunk("world", finish_reason="stop")] + + out = await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert out[:2] == chunks + frame = out[-1] + assert isinstance(frame, bytes) + text = frame.decode() + assert text.startswith("data: ") + payload = _json.loads(text[len("data: ") :]) + assert payload["error"]["message"] == "Violated guardrail policy" + assert payload["error"]["code"] == "400" + + @pytest.mark.asyncio + async def test_messages_eos_block_emits_anthropic_error_event(self): + guardrail = _EosHttpBlockingGuardrail() + chunks = _anthropic_message_chunks(["hello ", "world"]) + + out = await _drive_stream( + UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/messages" + ) + + raw = b"".join(c for c in out if isinstance(c, bytes)).decode() + assert "hello " in raw + assert "event: error" in raw + assert "Violated guardrail policy" in raw + assert "guardrail_error" in raw + + @pytest.mark.asyncio + async def test_responses_eos_block_emits_error_event_with_next_sequence(self): + guardrail = _EosHttpBlockingGuardrail() + chunks = [ + {"type": "response.created", "sequence_number": 0}, + {"type": "response.output_text.delta", "sequence_number": 1, "delta": "hello"}, + { + "type": "response.completed", + "sequence_number": 2, + "response": { + "model": "gpt-4", + "output": [{"type": "message", "content": [{"type": "output_text", "text": "hello"}]}], + }, + }, + ] + + out = await _drive_stream( + UnifiedLLMGuardrails(), guardrail, chunks, request_route="/v1/responses" + ) + + assert chunks[0] in out and chunks[1] in out + assert chunks[2] not in out + error_event = out[-1] + assert error_event.type == "error" + assert error_event.sequence_number == 2 + assert error_event.error.message == "Violated guardrail policy" + assert error_event.error.code == "400" + assert error_event.error.type == "guardrail_error" + + @pytest.mark.asyncio + async def test_pre_flush_block_still_raises_http_exception(self): + guardrail = _EosHttpBlockingGuardrail() + guardrail.streaming_buffer_until_moderated = True + chunks = [_stream_chunk("hello "), _stream_chunk("world", finish_reason="stop")] + + with pytest.raises(unified_module.HTTPException) as exc_info: + await _drive_stream(UnifiedLLMGuardrails(), guardrail, chunks) + + assert exc_info.value.status_code == 400 + assert exc_info.value.detail["error"] == "Violated guardrail policy" + + +class _AuditRecordingGuardrail(CustomGuardrail): + """Successful scan that records guardrail_information, like a flags-on audit.""" + + def __init__(self): + super().__init__(guardrail_name="audit-recorder") + self.streaming_end_of_stream_only = True + + def should_run_guardrail(self, data, event_type): # type: ignore[override] + return True + + async def apply_guardrail(self, inputs, request_data, input_type, **kwargs): + self.add_standard_logging_guardrail_information_to_request_data( + guardrail_json_response={"action": "NONE"}, + request_data=request_data, + guardrail_status="success", + ) + return inputs + + +class TestStreamingGuardrailInformationBucket: + """guardrail_information written during a chat streaming end-of-stream scan + must land in the request's ``metadata`` bucket that spend logging snapshots. + Regression for PR #38722 defect 2: the chat handler used to plant a + ``litellm_metadata`` key first, flipping the bucket so every later + guardrail_information write was diverted and /spend/logs showed null.""" + + @pytest.fixture(autouse=True) + def _use_real_mappings(self): + unified_module.endpoint_guardrail_translation_mappings = load_guardrail_translation_mappings() + yield + unified_module.endpoint_guardrail_translation_mappings = None + + @pytest.mark.asyncio + async def test_chat_eos_scan_writes_guardrail_information_to_metadata(self): + guardrail = _AuditRecordingGuardrail() + chunks = [_stream_chunk("hello "), _stream_chunk("world", finish_reason="stop")] + + async def _mock_stream(): + for chunk in chunks: + yield chunk + + user_api_key_dict = UserAPIKeyAuth( + api_key="test-key", user_id="user-1", request_route="/v1/chat/completions" + ) + request_data = {"guardrail_to_apply": guardrail, "model": "gpt-4", "metadata": {}} + + out = [] + async for item in UnifiedLLMGuardrails().async_post_call_streaming_iterator_hook( + user_api_key_dict=user_api_key_dict, + response=_mock_stream(), + request_data=request_data, + ): + out.append(item) + + assert "litellm_metadata" not in request_data + recorded = request_data["metadata"]["standard_logging_guardrail_information"] + assert len(recorded) == 1 + assert recorded[0]["guardrail_name"] == "audit-recorder" + assert recorded[0]["guardrail_status"] == "success" + assert request_data["metadata"]["user_api_key_user_id"] == "user-1" diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py index e70fc61de30..b7317d33b84 100644 --- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -1228,3 +1228,72 @@ class TestFireDeferredStreamLogging: assert info is not None, "guardrail_information should be populated" assert len(info) == 1 assert info[0]["guardrail_name"] == "info-writer" + + +class TestResponsesIteratorDeferredLogging: + """Regression for PR #38722 defect 2 on /v1/responses streams: when the + proxy arms _on_deferred_stream_complete, the responses streaming iterator + must store the logging coroutine for ProxyLogging._fire_deferred_stream_logging + (which runs AFTER end-of-stream guardrail scans write guardrail_information) + instead of dispatching immediately with a premature metadata snapshot.""" + + def _iterator(self, logging_obj): + from litellm.responses.streaming_iterator import ( + BaseResponsesAPIStreamingIterator, + ) + + iterator = object.__new__(BaseResponsesAPIStreamingIterator) + iterator.logging_obj = logging_obj + iterator.start_time = None + iterator.completed_response = None + iterator._completed_response_logged = False + iterator._completed_response_cache_hit = None + iterator._persist_completed_response_before_logging = False + return iterator + + def _logging_obj(self): + recorded = {} + + async def dispatch_success_handlers(result=None, **kwargs): + recorded["dispatched"] = True + + logging_obj = MagicMock() + logging_obj.dispatch_success_handlers = dispatch_success_handlers + return logging_obj, recorded + + @pytest.mark.asyncio + async def test_armed_iterator_stores_deferred_coroutine(self): + logging_obj, recorded = self._logging_obj() + logging_obj._on_deferred_stream_complete = MagicMock() + iterator = self._iterator(logging_obj) + + with patch("asyncio.create_task") as mock_create_task: + iterator._log_completed_response(is_async=True) + + mock_create_task.assert_not_called() + args = logging_obj._deferred_stream_complete_args + assert isinstance(args, tuple) and len(args) == 1 + assert "dispatched" not in recorded + await args[0] + assert recorded["dispatched"] is True + + @pytest.mark.asyncio + async def test_unarmed_iterator_dispatches_immediately(self): + logging_obj, recorded = self._logging_obj() + logging_obj._on_deferred_stream_complete = None + iterator = self._iterator(logging_obj) + + created = [] + real_create_task = asyncio.create_task + + def tracking_create_task(coro): + task = real_create_task(coro) + created.append(task) + return task + + with patch("asyncio.create_task", side_effect=tracking_create_task): + iterator._log_completed_response(is_async=True) + + assert len(created) == 1 + await created[0] + assert recorded["dispatched"] is True diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 71d4666416d..99f03d23378 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -1656,32 +1656,32 @@ class TestCommonRequestProcessingHelpers: assert payload["error"]["message"] == "MCP request blocked: no rewritable argument field present" assert payload["error"]["provider_specific_fields"]["error"]["code"] == "panw_prisma_airs_blocked" - async def test_serialize_http_exception_detail_helper(self): + async def testserialize_http_exception_detail_helper(self): """Direct unit coverage for the L1 helper across all branches.""" from litellm.proxy.common_request_processing import ( - _serialize_http_exception_detail, + serialize_http_exception_detail, ) import json as _json - assert _serialize_http_exception_detail("plain") == ("plain", None) + assert serialize_http_exception_detail("plain") == ("plain", None) - msg, fields = _serialize_http_exception_detail({"error": "Violated", "extra": "x"}) + msg, fields = serialize_http_exception_detail({"error": "Violated", "extra": "x"}) assert msg == "Violated" assert fields == {"error": "Violated", "extra": "x"} - msg, fields = _serialize_http_exception_detail({"error": {"message": "blocked", "code": "x"}}) + msg, fields = serialize_http_exception_detail({"error": {"message": "blocked", "code": "x"}}) assert msg == "blocked" assert fields == {"error": {"message": "blocked", "code": "x"}} - msg, fields = _serialize_http_exception_detail({"message": "top-level"}) + msg, fields = serialize_http_exception_detail({"message": "top-level"}) assert msg == "top-level" assert fields == {"message": "top-level"} - msg, fields = _serialize_http_exception_detail({"weird": ["a", "b"]}) + msg, fields = serialize_http_exception_detail({"weird": ["a", "b"]}) assert msg == _json.dumps({"weird": ["a", "b"]}) assert fields == {"weird": ["a", "b"]} - assert _serialize_http_exception_detail(42) == ("42", None) + assert serialize_http_exception_detail(42) == ("42", None) async def test_create_streaming_response_first_chunk_error_string_code(self): """ From f60ccf623460fcf029f3c22179d867eaf5fdd99a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 02:23:12 -0700 Subject: [PATCH 07/49] fix(guardrails): match deferred stream dispatch shape per stream owner and defer passthrough logging until guardrail eos --- litellm/proxy/common_request_processing.py | 132 +++++++++++------ .../streaming_handler.py | 41 ++++-- .../test_deferred_guardrail_logging.py | 123 ++++++++++++++++ .../test_streaming_handler_interrupt.py | 133 ++++++++++++++++-- .../proxy/test_common_request_processing.py | 2 +- .../test_proxy_logging_hook_detection.py | 29 ++-- 6 files changed, 381 insertions(+), 79 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 6c7c77610b3..3fca267e321 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -2341,53 +2341,14 @@ class ProxyBaseLLMRequestProcessing: if requested_model_from_client: self.data["_litellm_client_requested_model"] = requested_model_from_client - # Streaming: attach a closure that fires after all guardrail - # end-of-stream blocks complete. CSW.__anext__ stores the - # assembled response on logging_obj; the outer consumer - # (ProxyLogging._fire_deferred_stream_logging) fires the - # closure after the full streaming pipeline finishes. - # The closure runs non-apply_guardrail hooks on the - # assembled response, then fires success logging. - # Only for CustomStreamWrapper — raw async generators from - # passthrough routes bypass CSW and would orphan the closure. - from litellm.litellm_core_utils.streaming_handler import ( - CustomStreamWrapper, - ) - - if _post_call_guardrails_active and isinstance(response, CustomStreamWrapper): - # Intentionally a live reference (not a copy) — mirrors - # ProxyLogging.post_call_success_hook which also mutates - # data["guardrail_to_apply"] during iteration. - _captured_data: Final = self.data - _captured_user_api_key_dict: Final = user_api_key_dict - _captured_logging_obj: Final = logging_obj - - 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, - captured_logging_obj=_captured_logging_obj, - assembled_response=assembled_response, - cache_hit=cache_hit, - ) - - logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete - elif ( - _post_call_guardrails_active - and route_type in ("anthropic_messages", "aresponses") - and self._is_streaming_response(response) - ): - from litellm.litellm_core_utils.logging_worker import ( - GLOBAL_LOGGING_WORKER, + if _post_call_guardrails_active: + self._arm_deferred_stream_dispatch( + response=response, + route_type=route_type, + user_api_key_dict=user_api_key_dict, + logging_obj=logging_obj, ) - async def _on_deferred_native_stream_complete( - logging_coroutine: Coroutine[object, object, object], - ) -> None: - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) - - logging_obj._on_deferred_stream_complete = _on_deferred_native_stream_complete - if route_type == "allm_passthrough_route": # Check if response is an async generator if self._is_streaming_response(response): @@ -3057,6 +3018,87 @@ class ProxyBaseLLMRequestProcessing: except Exception as e: verbose_proxy_logger.exception("Error firing deferred logging: %s", e) + def _arm_deferred_stream_dispatch( + self, + response: object, + route_type: str, + user_api_key_dict: "UserAPIKeyAuth", + logging_obj: LiteLLMLoggingObj, + ) -> None: + """ + Streaming with post-call guardrails active: attach a closure that + ProxyLogging._fire_deferred_stream_logging fires after all guardrail + end-of-stream blocks complete, so the spend log sees + guardrail_information. + + Three closure shapes, matching who owns logging for the stream: + - CustomStreamWrapper (chat completions) stores + (assembled_response, cache_hit); the closure also runs + non-apply_guardrail post-call hooks via + _run_deferred_stream_guardrails. + - Bridged /v1/responses (LiteLLMCompletionStreamingIterator) shares + its inner CustomStreamWrapper's logging_obj, so it stores the same + (assembled_response, cache_hit) shape; the closure only dispatches + success logging, matching the route's pre-existing hook surface. + - Native anthropic_messages/aresponses iterators store a single + ready-made logging coroutine to enqueue. + + Raw async generators from passthrough routes bypass all three and + would orphan the closure, so they are not armed here. + """ + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + if isinstance(response, CustomStreamWrapper): + # Intentionally a live reference (not a copy) — mirrors + # ProxyLogging.post_call_success_hook which also mutates + # data["guardrail_to_apply"] during iteration. + _captured_data: Final = self.data + _captured_user_api_key_dict: Final = user_api_key_dict + _captured_logging_obj: Final = logging_obj + + 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, + captured_logging_obj=_captured_logging_obj, + assembled_response=assembled_response, + cache_hit=cache_hit, + ) + + logging_obj._on_deferred_stream_complete = _on_deferred_stream_complete + return + + if route_type not in ("anthropic_messages", "aresponses") or not self._is_streaming_response(response): + return + + from litellm.responses.litellm_completion_transformation.streaming_iterator import ( + LiteLLMCompletionStreamingIterator, + ) + + if isinstance(response, LiteLLMCompletionStreamingIterator): + _captured_bridge_logging_obj: Final = logging_obj + + async def _on_deferred_bridged_stream_complete(assembled_response: object, cache_hit: object) -> None: + await _as_success_dispatcher(_captured_bridge_logging_obj).dispatch_success_handlers( + assembled_response, + cache_hit=cache_hit, + start_time=None, + end_time=None, + prefer_async_handlers=True, + ) + + logging_obj._on_deferred_stream_complete = _on_deferred_bridged_stream_complete + return + + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + async def _on_deferred_native_stream_complete( + logging_coroutine: Coroutine[object, object, object], + ) -> None: + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=logging_coroutine) + + logging_obj._on_deferred_stream_complete = _on_deferred_native_stream_complete + @staticmethod async def _run_deferred_stream_guardrails( captured_data: dict, diff --git a/litellm/proxy/pass_through_endpoints/streaming_handler.py b/litellm/proxy/pass_through_endpoints/streaming_handler.py index 5ad41b00890..022a1ecbac4 100644 --- a/litellm/proxy/pass_through_endpoints/streaming_handler.py +++ b/litellm/proxy/pass_through_endpoints/streaming_handler.py @@ -65,6 +65,19 @@ class PassThroughStreamingHandler: route_streaming_logging or PassThroughStreamingHandler._route_streaming_logging_to_handler ) raw_bytes: Final[list[bytes]] = [] + + def _build_logging_coroutine() -> Coroutine[None, None, None]: + return resolved_route_streaming_logging( + litellm_logging_obj=litellm_logging_obj, + passthrough_success_handler_obj=passthrough_success_handler_obj, + url_route=url_route, + request_body=request_body or {}, + endpoint_type=endpoint_type, + start_time=start_time, + raw_bytes=raw_bytes, + end_time=datetime.now(), + ) + logging_scheduled = False model_name: Final = PassThroughStreamingHandler._extract_model_for_cost_injection( request_body=request_body, @@ -114,6 +127,21 @@ class PassThroughStreamingHandler: ) if pending: yield pending + # Stream completed cleanly. When the proxy armed deferred + # dispatch (post-call guardrails active), park the logging + # coroutine on logging_obj instead of enqueueing now, so + # ProxyLogging._fire_deferred_stream_logging fires it after + # guardrail end-of-stream blocks populate guardrail_information. + # Disconnect/exception paths skip this and fall through to the + # immediate enqueue in ``finally`` to keep partial billing + # (LIT-2642). + if ( + getattr(litellm_logging_obj, "_on_deferred_stream_complete", None) is not None + and raw_bytes + and response.status_code < 400 + ): + logging_scheduled = True + litellm_logging_obj._deferred_stream_complete_args = (_build_logging_coroutine(),) except Exception as e: verbose_proxy_logger.error("Error in chunk_processor: %s", e) raise @@ -128,18 +156,7 @@ class PassThroughStreamingHandler: if not logging_scheduled and raw_bytes and response.status_code < 400: logging_scheduled = True try: - GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue( - async_coroutine=resolved_route_streaming_logging( - litellm_logging_obj=litellm_logging_obj, - passthrough_success_handler_obj=passthrough_success_handler_obj, - url_route=url_route, - request_body=request_body or {}, - endpoint_type=endpoint_type, - start_time=start_time, - raw_bytes=raw_bytes, - end_time=datetime.now(), - ) - ) + GLOBAL_LOGGING_WORKER.ensure_initialized_and_enqueue(async_coroutine=_build_logging_coroutine()) except Exception as e: verbose_proxy_logger.error("Error scheduling chunk_processor logging: %s", e) diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py index b7317d33b84..99fae277bdd 100644 --- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -1297,3 +1297,126 @@ class TestResponsesIteratorDeferredLogging: assert len(created) == 1 await created[0] assert recorded["dispatched"] is True + + +class TestArmDeferredStreamDispatch: + """Regression for PR #38722: the closure shape armed on logging_obj must + match the args the stream's logging owner stores. Bridged /v1/responses + (LiteLLMCompletionStreamingIterator) shares its inner CustomStreamWrapper's + logging_obj, which stores (assembled_response, cache_hit); arming the + single-coroutine native closure there made _fire_deferred_stream_logging + raise TypeError inside the streaming hook, leaking an in-stream 500 error + frame on every streamed /v1/responses request.""" + + def _processor(self): + return ProxyBaseLLMRequestProcessing(data={"model": "gpt-test"}) + + def _dispatch_recording_logging_obj(self): + recorded = {} + + async def dispatch_success_handlers( + result=None, start_time=None, end_time=None, cache_hit=None, prefer_async_handlers=False + ): + recorded["result"] = result + recorded["cache_hit"] = cache_hit + recorded["prefer_async_handlers"] = prefer_async_handlers + + logging_obj = MagicMock() + logging_obj.dispatch_success_handlers = dispatch_success_handlers + logging_obj._on_deferred_stream_complete = None + logging_obj._deferred_stream_complete_args = None + return logging_obj, recorded + + @pytest.mark.asyncio + async def test_bridged_responses_iterator_gets_csw_arg_shape(self): + from litellm.responses.litellm_completion_transformation.streaming_iterator import ( + LiteLLMCompletionStreamingIterator, + ) + + logging_obj, recorded = self._dispatch_recording_logging_obj() + bridged = object.__new__(LiteLLMCompletionStreamingIterator) + + self._processor()._arm_deferred_stream_dispatch( + response=bridged, + route_type="aresponses", + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + + assembled = object() + logging_obj._deferred_stream_complete_args = (assembled, False) + ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await asyncio.sleep(0) + + assert recorded["result"] is assembled + assert recorded["cache_hit"] is False + assert recorded["prefer_async_handlers"] is True + + @pytest.mark.asyncio + async def test_native_stream_closure_enqueues_single_coroutine(self): + from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER + + logging_obj, _ = self._dispatch_recording_logging_obj() + + async def _agen(): + yield b"x" + + self._processor()._arm_deferred_stream_dispatch( + response=_agen(), + route_type="anthropic_messages", + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + closure = logging_obj._on_deferred_stream_complete + assert closure is not None + + async def _logging_coroutine(): + return None + + coro = _logging_coroutine() + with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam + GLOBAL_LOGGING_WORKER, "ensure_initialized_and_enqueue" + ) as mock_enqueue: + await closure(coro) + mock_enqueue.assert_called_once_with(async_coroutine=coro) + coro.close() + + @pytest.mark.asyncio + async def test_csw_closure_routes_through_deferred_stream_guardrails(self, monkeypatch): + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + + logging_obj, recorded = self._dispatch_recording_logging_obj() + csw = object.__new__(CustomStreamWrapper) + processor = self._processor() + + monkeypatch.setattr( # test-quality-ok: empty the process-global callback registry so no ambient guardrail runs + litellm, "callbacks", [] + ) + processor._arm_deferred_stream_dispatch( + response=csw, + route_type="acompletion", + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + assembled = object() + await logging_obj._on_deferred_stream_complete(assembled, False) + await asyncio.sleep(0) + + assert recorded["result"] is assembled + assert recorded["cache_hit"] is False + assert recorded["prefer_async_handlers"] is True + + def test_non_native_route_generator_not_armed(self): + logging_obj, _ = self._dispatch_recording_logging_obj() + + async def _agen(): + yield b"x" + + self._processor()._arm_deferred_stream_dispatch( + response=_agen(), + route_type="acompletion", + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + + assert logging_obj._on_deferred_stream_complete is None diff --git a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py index 1d82a5dfc6e..56c89fed79a 100644 --- a/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py +++ b/tests/test_litellm/proxy/pass_through_endpoints/test_streaming_handler_interrupt.py @@ -28,12 +28,21 @@ def _make_streaming_response(chunks): return mock +def _unarmed_logging_obj(): + """Real Logging objects only carry _on_deferred_stream_complete when the + proxy arms deferred dispatch; a bare MagicMock's auto-attribute is truthy + and would spuriously trigger the deferral branch.""" + obj = MagicMock() + obj._on_deferred_stream_complete = None + return obj + + @pytest.mark.asyncio async def test_chunk_processor_logs_on_normal_completion(): chunks = [b"chunk-1", b"chunk-2", b"chunk-3"] response = _make_streaming_response(chunks) - mock_logging_obj = MagicMock() + mock_logging_obj = _unarmed_logging_obj() mock_passthrough_handler = MagicMock() with patch.object( @@ -66,7 +75,7 @@ async def test_chunk_processor_logs_on_client_disconnect(): chunks = [b"event-1", b"event-2", b"event-3"] response = _make_streaming_response(chunks) - mock_logging_obj = MagicMock() + mock_logging_obj = _unarmed_logging_obj() mock_passthrough_handler = MagicMock() with patch.object( @@ -104,7 +113,7 @@ async def test_chunk_processor_does_not_schedule_success_logging_for_upstream_er response = _make_streaming_response(chunks) response.status_code = 403 - mock_logging_obj = MagicMock() + mock_logging_obj = _unarmed_logging_obj() mock_passthrough_handler = MagicMock() with patch.object( @@ -134,7 +143,7 @@ async def test_chunk_processor_does_not_schedule_success_logging_for_upstream_er async def test_chunk_processor_does_not_schedule_logging_when_no_chunks(): response = _make_streaming_response([]) - mock_logging_obj = MagicMock() + mock_logging_obj = _unarmed_logging_obj() mock_passthrough_handler = MagicMock() with patch.object( @@ -189,7 +198,7 @@ async def test_chunk_processor_routes_logging_through_logging_worker(): async for chunk in PassThroughStreamingHandler.chunk_processor( response=response, request_body={"model": "claude-3-haiku"}, - litellm_logging_obj=MagicMock(), + litellm_logging_obj=_unarmed_logging_obj(), endpoint_type=EndpointType.GENERIC, start_time=datetime.now(), passthrough_success_handler_obj=MagicMock(), @@ -230,7 +239,7 @@ async def test_chunk_processor_routes_logging_through_logging_worker_on_disconne gen = PassThroughStreamingHandler.chunk_processor( response=response, request_body={"model": "claude-3-haiku"}, - litellm_logging_obj=MagicMock(), + litellm_logging_obj=_unarmed_logging_obj(), endpoint_type=EndpointType.GENERIC, start_time=datetime.now(), passthrough_success_handler_obj=MagicMock(), @@ -246,7 +255,7 @@ async def test_chunk_processor_routes_logging_through_logging_worker_on_disconne def _logging_obj_with_write_once_cst(): """Build a MagicMock that mirrors the real Logging behavior: _update_completion_start_time latches self.completion_start_time so the write-once guard actually latches.""" - obj = MagicMock() + obj = _unarmed_logging_obj() obj.completion_start_time = None def _update(*, completion_start_time): @@ -301,7 +310,7 @@ async def test_chunk_processor_does_not_reset_completion_start_time_on_later_chu response = _make_streaming_response(chunks) real_first = datetime(2020, 1, 1, 0, 0, 0) - mock_logging_obj = MagicMock() + mock_logging_obj = _unarmed_logging_obj() # Simulate first-chunk stamp having already landed (e.g. under contention or a # prior wrapper that already set it): later chunks must be no-ops. mock_logging_obj.completion_start_time = real_first @@ -387,7 +396,7 @@ async def _collect_openai_passthrough_chunks(chunks, endpoint_type): async for chunk in PassThroughStreamingHandler.chunk_processor( response=response, request_body={"model": "gpt-4o-mini", "stream": True}, - litellm_logging_obj=MagicMock(), + litellm_logging_obj=_unarmed_logging_obj(), endpoint_type=endpoint_type, start_time=datetime.now(), passthrough_success_handler_obj=MagicMock(), @@ -517,3 +526,109 @@ def test_convert_raw_bytes_survives_truncated_multibyte_sequence(): lines = PassThroughStreamingHandler._convert_raw_bytes_to_str_lines(raw_bytes) assert any('"type": "message_delta"' in line for line in lines) + + +@pytest.mark.asyncio +async def test_chunk_processor_defers_logging_until_fire_when_armed(): + """Regression for PR #38722: native /v1/messages streams route through + chunk_processor, which enqueued the spend log the moment the stream ended, + racing the guardrail end-of-stream scan and logging + guardrail_information as null. With deferred dispatch armed, the completed + stream must park the logging coroutine on logging_obj and only enqueue it + when ProxyLogging._fire_deferred_stream_logging fires after the scan.""" + from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing + from litellm.proxy.utils import ProxyLogging + + chunks = [b"event-1", b"event-2"] + response = _make_streaming_response(chunks) + + logging_obj = _unarmed_logging_obj() + logging_obj._deferred_stream_complete_args = None + + enqueued = [] + + def _capture(async_coroutine): + enqueued.append(async_coroutine) + async_coroutine.close() + + with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam + GLOBAL_LOGGING_WORKER, + "ensure_initialized_and_enqueue", + side_effect=_capture, + ) as mock_enqueue: + gen = PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={"model": "claude-3-haiku"}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.ANTHROPIC, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + route_streaming_logging=AsyncMock(), + ) + ProxyBaseLLMRequestProcessing(data={})._arm_deferred_stream_dispatch( + response=gen, + route_type="anthropic_messages", + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + + received = [] + async for chunk in gen: + received.append(chunk) + await asyncio.sleep(0) + + assert received == chunks + mock_enqueue.assert_not_called() + parked = logging_obj._deferred_stream_complete_args + assert isinstance(parked, tuple) and len(parked) == 1 + assert asyncio.iscoroutine(parked[0]) + + ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await asyncio.sleep(0) + + mock_enqueue.assert_called_once() + + +@pytest.mark.asyncio +async def test_chunk_processor_enqueues_immediately_on_disconnect_even_when_armed(): + """Client disconnects never reach _fire_deferred_stream_logging, so parking + the coroutine there would lose the partial-usage spend log (LIT-2642); the + disconnect path must keep enqueueing immediately.""" + chunks = [b"event-1", b"event-2", b"event-3"] + response = _make_streaming_response(chunks) + + logging_obj = _unarmed_logging_obj() + + async def _armed_closure(logging_coroutine): + raise AssertionError("deferred closure must not fire on disconnect") + + logging_obj._on_deferred_stream_complete = _armed_closure + logging_obj._deferred_stream_complete_args = None + + enqueued = [] + + def _capture(async_coroutine): + enqueued.append(async_coroutine) + async_coroutine.close() + + with patch.object( # test-quality-ok: GLOBAL_LOGGING_WORKER is a process-global singleton with no injection seam + GLOBAL_LOGGING_WORKER, + "ensure_initialized_and_enqueue", + side_effect=_capture, + ) as mock_enqueue: + gen = PassThroughStreamingHandler.chunk_processor( + response=response, + request_body={"model": "claude-3-haiku"}, + litellm_logging_obj=logging_obj, + endpoint_type=EndpointType.ANTHROPIC, + start_time=datetime.now(), + passthrough_success_handler_obj=MagicMock(), + url_route="/v1/messages", + route_streaming_logging=AsyncMock(), + ) + await gen.__anext__() + await gen.aclose() + + mock_enqueue.assert_called_once() + assert logging_obj._deferred_stream_complete_args is None diff --git a/tests/test_litellm/proxy/test_common_request_processing.py b/tests/test_litellm/proxy/test_common_request_processing.py index 99f03d23378..0fea239625a 100644 --- a/tests/test_litellm/proxy/test_common_request_processing.py +++ b/tests/test_litellm/proxy/test_common_request_processing.py @@ -1656,7 +1656,7 @@ class TestCommonRequestProcessingHelpers: assert payload["error"]["message"] == "MCP request blocked: no rewritable argument field present" assert payload["error"]["provider_specific_fields"]["error"]["code"] == "panw_prisma_airs_blocked" - async def testserialize_http_exception_detail_helper(self): + async def test_serialize_http_exception_detail_helper(self): """Direct unit coverage for the L1 helper across all branches.""" from litellm.proxy.common_request_processing import ( serialize_http_exception_detail, diff --git a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py index 542572e1e56..9f1321aec2c 100644 --- a/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py +++ b/tests/test_litellm/proxy/test_proxy_logging_hook_detection.py @@ -346,14 +346,14 @@ async def test_post_call_stream_guardrail_keeps_own_iterator_on_chat_completions @pytest.mark.asyncio -async def test_unified_guardrail_iterator_accepts_explicit_guardrail(monkeypatch): +async def test_unified_guardrail_iterator_accepts_explicit_guardrail(): """ The dispatch passes each guardrail explicitly instead of through a shared request_data key, so chaining two unified-routed guardrails cannot drop - all but the last one. + all but the last one. The block fires after the deltas were already + flushed to the client, so it surfaces as a trailing in-stream error frame + rather than a raised HTTPException. """ - from fastapi import HTTPException - from litellm.proxy.utils import unified_guardrail guardrail = _content_filter_guardrail("BLOCK") @@ -367,14 +367,19 @@ async def test_unified_guardrail_iterator_accepts_explicit_guardrail(monkeypatch for chunk in _anthropic_stream_chunks(["the", " zebra runs"]): yield chunk - with pytest.raises(HTTPException): - async for _ in unified_guardrail.async_post_call_streaming_iterator_hook( - user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages"), - response=fake_stream(), - request_data=request_data, - guardrail_to_apply=guardrail, - ): - pass + delivered = [] + async for item in unified_guardrail.async_post_call_streaming_iterator_hook( + user_api_key_dict=UserAPIKeyAuth(api_key="sk-1234", request_route="/v1/messages"), + response=fake_stream(), + request_data=request_data, + guardrail_to_apply=guardrail, + ): + delivered.append(item) + + raw = b"".join(c for c in delivered if isinstance(c, bytes)).decode() + assert "event: error" in raw + assert "guardrail_error" in raw + assert raw.index("guardrail_error") > raw.index(" zebra runs") @pytest.mark.asyncio From 229970c5005fefd540d87c292750fe6c3c7ea6a4 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 02:30:34 -0700 Subject: [PATCH 08/49] fix(guardrails): unwrap HiddenParamsAsyncIteratorWrapper before deferred dispatch class sniffing --- litellm/proxy/common_request_processing.py | 11 ++++-- .../test_deferred_guardrail_logging.py | 34 +++++++++++++++++++ 2 files changed, 43 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/common_request_processing.py b/litellm/proxy/common_request_processing.py index 3fca267e321..69c2cb3f0f0 100644 --- a/litellm/proxy/common_request_processing.py +++ b/litellm/proxy/common_request_processing.py @@ -3045,10 +3045,17 @@ class ProxyBaseLLMRequestProcessing: Raw async generators from passthrough routes bypass all three and would orphan the closure, so they are not armed here. + + The router wraps iterators that cannot carry _hidden_params in + HiddenParamsAsyncIteratorWrapper, so class sniffing runs on the + unwrapped inner iterator. """ from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.router_utils.add_retry_fallback_headers import HiddenParamsAsyncIteratorWrapper - if isinstance(response, CustomStreamWrapper): + unwrapped: Final = response._inner if isinstance(response, HiddenParamsAsyncIteratorWrapper) else response + + if isinstance(unwrapped, CustomStreamWrapper): # Intentionally a live reference (not a copy) — mirrors # ProxyLogging.post_call_success_hook which also mutates # data["guardrail_to_apply"] during iteration. @@ -3075,7 +3082,7 @@ class ProxyBaseLLMRequestProcessing: LiteLLMCompletionStreamingIterator, ) - if isinstance(response, LiteLLMCompletionStreamingIterator): + if isinstance(unwrapped, LiteLLMCompletionStreamingIterator): _captured_bridge_logging_obj: Final = logging_obj async def _on_deferred_bridged_stream_complete(assembled_response: object, cache_hit: object) -> None: diff --git a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py index 99fae277bdd..8fde4cc9d5e 100644 --- a/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py +++ b/tests/test_litellm/proxy/guardrails/test_deferred_guardrail_logging.py @@ -1352,6 +1352,40 @@ class TestArmDeferredStreamDispatch: assert recorded["cache_hit"] is False assert recorded["prefer_async_handlers"] is True + @pytest.mark.asyncio + async def test_router_wrapped_bridged_iterator_gets_csw_arg_shape(self): + """The router wraps iterators without _hidden_params in + HiddenParamsAsyncIteratorWrapper before the proxy arms deferral, so + every production streamed /v1/responses reaches arming wrapped; + sniffing the wrapper instead of the inner iterator armed the 1-arg + native closure against the CSW's 2-arg stored shape and leaked a + TypeError 500 frame into the stream.""" + from litellm.responses.litellm_completion_transformation.streaming_iterator import ( + LiteLLMCompletionStreamingIterator, + ) + from litellm.router_utils.add_retry_fallback_headers import ( + HiddenParamsAsyncIteratorWrapper, + ) + + logging_obj, recorded = self._dispatch_recording_logging_obj() + wrapped = HiddenParamsAsyncIteratorWrapper(object.__new__(LiteLLMCompletionStreamingIterator)) + + self._processor()._arm_deferred_stream_dispatch( + response=wrapped, + route_type="aresponses", + user_api_key_dict=MagicMock(), + logging_obj=logging_obj, + ) + + assembled = object() + logging_obj._deferred_stream_complete_args = (assembled, False) + ProxyLogging._fire_deferred_stream_logging({"litellm_logging_obj": logging_obj}) + await asyncio.sleep(0) + + assert recorded["result"] is assembled + assert recorded["cache_hit"] is False + assert recorded["prefer_async_handlers"] is True + @pytest.mark.asyncio async def test_native_stream_closure_enqueues_single_coroutine(self): from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER From 80e0bc2dc6861222781f3dc87aa5ab72786dae8f Mon Sep 17 00:00:00 2001 From: ryan-crabbe-berri Date: Sat, 29 Aug 2026 10:58:10 -0700 Subject: [PATCH 09/49] docs(claude.md): require tests to check behavior, not code structure Claude-Session: https://claude.ai/code/session_017nZW6omb93ZuAfqqCzKSU5 --- CLAUDE.md | 2 ++ 1 file changed, 2 insertions(+) diff --git a/CLAUDE.md b/CLAUDE.md index 930825aeb89..d9e9e8f1586 100644 --- a/CLAUDE.md +++ b/CLAUDE.md @@ -23,6 +23,8 @@ When adding new features, add meaningful tests. Don't add tests that don't check Same thing for bug fixes. The tests should make it so that this specific bug can never happen again without failing tests (i.e., regression) +Never test structure of code only function of it + `tests/test_litellm/` mirrors `litellm/` in a parallel path (see `tests/test_litellm/readme.md`). Name tests `test_.py`, but always match the existing test file in the directory you touch — many provider dirs use longer descriptive names (e.g. `test_anthropic_chat_transformation.py`) to avoid ambiguity across sibling folders. For bug fixes, extend the existing mapped test file rather than creating a new one. Only create a new test file for a new feature (provider, endpoint, or transformation module) that has no mapped test yet, following that directory's naming convention (or `test_.py` if you're the first test there). One focused regression test beats many shallow ones End-to-end tests belong in `tests/e2e/` and must follow the harness conventions documented in that directory's `CLAUDE.md` From 5024c4d52052e7a04eebb4b33d6615e71ee1cb0c Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sat, 29 Aug 2026 20:56:21 +0000 Subject: [PATCH 10/49] fix: update stale source URLs in model cost map 119 entries pointed at 404ing or permanently-moved pages (Pylon #7777). Replaced with verified working equivalents (200-checked or permanent redirect targets). Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 238 +++++++++--------- model_prices_and_context_window.json | 238 +++++++++--------- 2 files changed, 238 insertions(+), 238 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index ec175025b42..b548de07452 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -3726,7 +3726,7 @@ "output_cost_per_token": 0, "litellm_provider": "azure_ai", "mode": "chat", - "source": "https://azure.microsoft.com/en-us/pricing/details/ai-services/", + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/aoai/", "comment": "Flat cost of $0.14 per M input tokens for Azure AI Foundry Model Router infrastructure. Use pattern: azure_ai/model_router/ where deployment-name is your Azure deployment (e.g., azure-model-router)" }, "azure/eu/gpt-4o-2024-08-06": { @@ -5328,7 +5328,7 @@ "input_cost_per_second": 0.0002833333333333333, "litellm_provider": "azure", "mode": "audio_transcription", - "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/gpt-realtime-whisper", + "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/gpt-realtime-whisper", "supported_endpoints": [ "/v1/realtime", "/v1/realtime/transcription_sessions" @@ -9107,7 +9107,7 @@ "mode": "embedding", "output_cost_per_token": 0.0, "output_vector_size": 1024, - "source": "https://azuremarketplace.microsoft.com/en-us/marketplace/apps/cohere.cohere-embed-v3-english-offer?tab=PlansAndPrice", + "source": "https://marketplace.microsoft.com/en-us/marketplace/apps/cohere.cohere-embed-v3-english-offer?tab=PlansAndPrice", "supports_embedding_image_input": true }, "azure_ai/Cohere-embed-v3-multilingual": { @@ -9118,7 +9118,7 @@ "mode": "embedding", "output_cost_per_token": 0.0, "output_vector_size": 1024, - "source": "https://azuremarketplace.microsoft.com/en-us/marketplace/apps/cohere.cohere-embed-v3-english-offer?tab=PlansAndPrice", + "source": "https://marketplace.microsoft.com/en-us/marketplace/apps/cohere.cohere-embed-v3-english-offer?tab=PlansAndPrice", "supports_embedding_image_input": true }, "azure_ai/FLUX-1.1-pro": { @@ -9134,7 +9134,7 @@ "litellm_provider": "azure_ai", "mode": "image_generation", "output_cost_per_image": 0.04, - "source": "https://azuremarketplace.microsoft.com/pt-br/marketplace/apps/cohere.cohere-embed-4-offer?tab=PlansAndPrice", + "source": "https://marketplace.microsoft.com/pt-br/marketplace/apps/cohere.cohere-embed-4-offer?tab=PlansAndPrice", "supported_endpoints": [ "/v1/images/generations" ] @@ -9467,7 +9467,7 @@ "max_tokens": 2048, "mode": "chat", "output_cost_per_token": 3.7e-07, - "source": "https://azuremarketplace.microsoft.com/en/marketplace/apps/metagenai.meta-llama-3-2-11b-vision-instruct-offer?tab=Overview", + "source": "https://marketplace.microsoft.com/en/marketplace/apps/metagenai.meta-llama-3-2-11b-vision-instruct-offer?tab=Overview", "supports_function_calling": true, "supports_tool_choice": true, "supports_vision": true @@ -9481,7 +9481,7 @@ "max_tokens": 2048, "mode": "chat", "output_cost_per_token": 2.04e-06, - "source": "https://azuremarketplace.microsoft.com/en/marketplace/apps/metagenai.meta-llama-3-2-90b-vision-instruct-offer?tab=Overview", + "source": "https://marketplace.microsoft.com/en/marketplace/apps/metagenai.meta-llama-3-2-90b-vision-instruct-offer?tab=Overview", "supports_function_calling": true, "supports_tool_choice": true, "supports_vision": true @@ -9494,7 +9494,7 @@ "max_tokens": 2048, "mode": "chat", "output_cost_per_token": 7.1e-07, - "source": "https://azuremarketplace.microsoft.com/en/marketplace/apps/metagenai.llama-3-3-70b-instruct-offer?tab=Overview", + "source": "https://marketplace.microsoft.com/en/marketplace/apps/metagenai.llama-3-3-70b-instruct-offer?tab=Overview", "supports_function_calling": true, "supports_tool_choice": true }, @@ -9543,7 +9543,7 @@ "max_tokens": 2048, "mode": "chat", "output_cost_per_token": 1.6e-05, - "source": "https://azuremarketplace.microsoft.com/en-us/marketplace/apps/metagenai.meta-llama-3-1-405b-instruct-offer?tab=PlansAndPrice", + "source": "https://marketplace.microsoft.com/en-us/marketplace/apps/metagenai.meta-llama-3-1-405b-instruct-offer?tab=PlansAndPrice", "supports_tool_choice": true }, "azure_ai/Meta-Llama-3.1-70B-Instruct": { @@ -9554,7 +9554,7 @@ "max_tokens": 2048, "mode": "chat", "output_cost_per_token": 3.54e-06, - "source": "https://azuremarketplace.microsoft.com/en-us/marketplace/apps/metagenai.meta-llama-3-1-70b-instruct-offer?tab=PlansAndPrice", + "source": "https://marketplace.microsoft.com/en-us/marketplace/apps/metagenai.meta-llama-3-1-70b-instruct-offer?tab=PlansAndPrice", "supports_tool_choice": true }, "azure_ai/Meta-Llama-3.1-8B-Instruct": { @@ -9566,7 +9566,7 @@ "max_tokens": 2048, "mode": "chat", "output_cost_per_token": 6.1e-07, - "source": "https://azuremarketplace.microsoft.com/en-us/marketplace/apps/metagenai.meta-llama-3-1-8b-instruct-offer?tab=PlansAndPrice", + "source": "https://marketplace.microsoft.com/en-us/marketplace/apps/metagenai.meta-llama-3-1-8b-instruct-offer?tab=PlansAndPrice", "supports_tool_choice": true }, "azure_ai/Phi-3-medium-128k-instruct": { @@ -9756,7 +9756,7 @@ "supported_endpoints": [ "/v1/ocr" ], - "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/" + "source": "https://ai.azure.com/catalog/models/mistral-document-ai-2512" }, "azure_ai/doc-intelligence/prebuilt-read": { "litellm_provider": "azure_ai", @@ -9967,7 +9967,7 @@ "mode": "embedding", "output_cost_per_token": 0.0, "output_vector_size": 3072, - "source": "https://azuremarketplace.microsoft.com/pt-br/marketplace/apps/cohere.cohere-embed-4-offer?tab=PlansAndPrice", + "source": "https://marketplace.microsoft.com/pt-br/marketplace/apps/cohere.cohere-embed-4-offer?tab=PlansAndPrice", "supported_endpoints": [ "/v1/embeddings" ], @@ -10151,7 +10151,7 @@ "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 0.00971, - "source": "https://azure.microsoft.com/en-us/products/ai-services/ai-foundry/models/jais-30b-chat" + "source": "https://ai.azure.com/catalog/models/jais-30b-chat" }, "azure_ai/jamba-instruct": { "input_cost_per_token": 5e-07, @@ -10208,7 +10208,7 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 4e-08, - "source": "https://azuremarketplace.microsoft.com/en/marketplace/apps/000-000.ministral-3b-2410-offer?tab=Overview", + "source": "https://marketplace.microsoft.com/en/marketplace/apps/000-000.ministral-3b-2410-offer?tab=Overview", "supports_function_calling": true, "supports_tool_choice": true }, @@ -10231,7 +10231,7 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 6e-06, - "source": "https://azuremarketplace.microsoft.com/en/marketplace/apps/000-000.mistral-ai-large-2407-offer?tab=Overview", + "source": "https://marketplace.microsoft.com/en/marketplace/apps/000-000.mistral-ai-large-2407-offer?tab=Overview", "supports_function_calling": true, "supports_tool_choice": true }, @@ -10243,7 +10243,7 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 6e-06, - "source": "https://azuremarketplace.microsoft.com/en/marketplace/apps/000-000.mistral-ai-large-2407-offer?tab=Overview", + "source": "https://marketplace.microsoft.com/en/marketplace/apps/000-000.mistral-ai-large-2407-offer?tab=Overview", "supports_function_calling": true, "supports_tool_choice": true }, @@ -10280,7 +10280,7 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1.5e-07, - "source": "https://azuremarketplace.microsoft.com/en/marketplace/apps/000-000.mistral-nemo-12b-2407?tab=PlansAndPrice", + "source": "https://marketplace.microsoft.com/en/marketplace/apps/000-000.mistral-nemo-12b-2407?tab=PlansAndPrice", "supports_function_calling": true }, "azure_ai/mistral-small": { @@ -22796,7 +22796,7 @@ "output_cost_per_reasoning_token": 3e-06, "output_cost_per_token": 3e-06, "rpm": 2000, - "source": "https://ai.google.dev/pricing/gemini-3", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -22850,7 +22850,7 @@ "output_cost_per_reasoning_token": 9e-06, "output_cost_per_token": 9e-06, "rpm": 2000, - "source": "https://ai.google.dev/pricing/gemini-3", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -22915,7 +22915,7 @@ "output_cost_per_token_batches": 1.875e-06, "output_cost_per_token_flex": 1.875e-06, "rpm": 2000, - "source": "https://ai.google.dev/pricing/gemini-3", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -22974,7 +22974,7 @@ "output_cost_per_token_batches": 1.875e-06, "output_cost_per_token_flex": 1.875e-06, "rpm": 2000, - "source": "https://ai.google.dev/pricing/gemini-3", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -23178,7 +23178,7 @@ "mode": "chat", "output_cost_per_reasoning_token": 3e-06, "output_cost_per_token": 3e-06, - "source": "https://ai.google.dev/pricing/gemini-3", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -23262,7 +23262,7 @@ "mode": "chat", "output_cost_per_reasoning_token": 9e-06, "output_cost_per_token": 9e-06, - "source": "https://ai.google.dev/pricing/gemini-3", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -23325,7 +23325,7 @@ "output_cost_per_token": 3.75e-06, "output_cost_per_token_batches": 1.875e-06, "output_cost_per_token_flex": 1.875e-06, - "source": "https://ai.google.dev/pricing/gemini-3", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -23382,7 +23382,7 @@ "output_cost_per_token": 3.75e-06, "output_cost_per_token_batches": 1.875e-06, "output_cost_per_token_flex": 1.875e-06, - "source": "https://ai.google.dev/pricing/gemini-3", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -26982,7 +26982,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "source": "https://platform.openai.com/docs/models/gpt-5.6-cyber", + "source": "https://developers.openai.com/api/docs/models/gpt-5.6-cyber", "supports_computer_use": true, "supports_parallel_function_calling": true }, @@ -27021,7 +27021,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "source": "https://platform.openai.com/docs/models/daybreak-red-latest", + "source": "https://developers.openai.com/api/docs/models/daybreak-red-latest", "supports_computer_use": true, "supports_parallel_function_calling": true }, @@ -27061,7 +27061,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "source": "https://platform.openai.com/docs/models/daybreak-blue-latest", + "source": "https://developers.openai.com/api/docs/models/daybreak-blue-latest", "supports_parallel_function_calling": true }, "chat-latest": { @@ -27073,7 +27073,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 3e-05, - "source": "https://platform.openai.com/docs/models/chat-latest", + "source": "https://developers.openai.com/api/docs/models/chat-latest", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -30368,7 +30368,7 @@ "search_context_size_low": 0.0025, "search_context_size_medium": 0.0025 }, - "source": "https://dev.meta.ai/docs/getting-started/pricing-rate-limits", + "source": "https://ai.developer.meta.com/docs/pricing-rate-limits", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses", @@ -30409,7 +30409,7 @@ "search_context_size_low": 0.0025, "search_context_size_medium": 0.0025 }, - "source": "https://dev.meta.ai/docs/getting-started/pricing-rate-limits", + "source": "https://ai.developer.meta.com/docs/pricing-rate-limits", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses", @@ -30450,7 +30450,7 @@ "search_context_size_low": 0.0025, "search_context_size_medium": 0.0025 }, - "source": "https://dev.meta.ai/docs/getting-started/pricing-rate-limits", + "source": "https://ai.developer.meta.com/docs/pricing-rate-limits", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses", @@ -30483,7 +30483,7 @@ "max_output_tokens": 4028, "max_tokens": 4028, "mode": "chat", - "source": "https://llama.developer.meta.com/docs/models", + "source": "https://ai.developer.meta.com/docs/models", "supported_modalities": [ "text" ], @@ -30499,7 +30499,7 @@ "max_output_tokens": 4028, "max_tokens": 4028, "mode": "chat", - "source": "https://llama.developer.meta.com/docs/models", + "source": "https://ai.developer.meta.com/docs/models", "supported_modalities": [ "text" ], @@ -30515,7 +30515,7 @@ "max_output_tokens": 4028, "max_tokens": 4028, "mode": "chat", - "source": "https://llama.developer.meta.com/docs/models", + "source": "https://ai.developer.meta.com/docs/models", "supported_modalities": [ "text", "image" @@ -30532,7 +30532,7 @@ "max_output_tokens": 4028, "max_tokens": 4028, "mode": "chat", - "source": "https://llama.developer.meta.com/docs/models", + "source": "https://ai.developer.meta.com/docs/models", "supported_modalities": [ "text", "image" @@ -32477,7 +32477,7 @@ "mode": "chat", "supports_function_calling": true, "supports_reasoning": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/deepseek-ai/DeepSeek-R1-0528": { "max_tokens": 164000, @@ -32489,7 +32489,7 @@ "mode": "chat", "supports_function_calling": true, "supports_reasoning": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/deepseek-ai/DeepSeek-R1-Distill-Llama-70B": { "max_tokens": 128000, @@ -32500,7 +32500,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/deepseek-ai/DeepSeek-V3": { "max_tokens": 128000, @@ -32511,7 +32511,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/deepseek-ai/DeepSeek-V3-0324": { "max_tokens": 128000, @@ -32522,7 +32522,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/google/gemma-3-27b-it": { "max_tokens": 128000, @@ -32534,7 +32534,7 @@ "mode": "chat", "supports_function_calling": true, "supports_vision": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/meta-llama/Llama-3.3-70B-Instruct": { "max_tokens": 128000, @@ -32545,7 +32545,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/meta-llama/Llama-Guard-3-8B": { "max_tokens": 128000, @@ -32555,7 +32555,7 @@ "output_cost_per_token": 6e-08, "litellm_provider": "nebius", "mode": "chat", - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/meta-llama/Meta-Llama-3.1-8B-Instruct": { "max_tokens": 128000, @@ -32566,7 +32566,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/meta-llama/Meta-Llama-3.1-70B-Instruct": { "max_tokens": 128000, @@ -32577,7 +32577,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/meta-llama/Meta-Llama-3.1-405B-Instruct": { "max_tokens": 128000, @@ -32588,7 +32588,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/mistralai/Mistral-Nemo-Instruct-2407": { "max_tokens": 128000, @@ -32599,7 +32599,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/NousResearch/Hermes-3-Llama-3.1-405B": { "max_tokens": 128000, @@ -32610,7 +32610,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/nvidia/Llama-3.1-Nemotron-Ultra-253B-v1": { "max_tokens": 128000, @@ -32621,7 +32621,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/nvidia/Llama-3.3-Nemotron-Super-49B-v1": { "max_tokens": 131072, @@ -32632,7 +32632,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/Qwen/Qwen3-235B-A22B": { "max_tokens": 262144, @@ -32643,7 +32643,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/Qwen/Qwen3-32B": { "max_tokens": 32768, @@ -32654,7 +32654,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/Qwen/Qwen3-30B-A3B": { "max_tokens": 32768, @@ -32665,7 +32665,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/Qwen/Qwen3-14B": { "max_tokens": 32768, @@ -32676,7 +32676,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/Qwen/Qwen3-4B": { "max_tokens": 32768, @@ -32687,7 +32687,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/Qwen/QwQ-32B": { "max_tokens": 32768, @@ -32699,7 +32699,7 @@ "mode": "chat", "supports_function_calling": true, "supports_reasoning": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/Qwen/Qwen2.5-72B-Instruct": { "max_tokens": 128000, @@ -32710,7 +32710,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/Qwen/Qwen2.5-32B-Instruct": { "max_tokens": 128000, @@ -32721,7 +32721,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/Qwen/Qwen2.5-Coder-7B": { "max_tokens": 32768, @@ -32732,7 +32732,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/Qwen/Qwen2.5-VL-72B-Instruct": { "max_tokens": 131072, @@ -32744,7 +32744,7 @@ "mode": "chat", "supports_function_calling": true, "supports_vision": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/Qwen/Qwen2-VL-72B-Instruct": { "max_tokens": 131072, @@ -32756,7 +32756,7 @@ "mode": "chat", "supports_function_calling": true, "supports_vision": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/Qwen/Qwen2-VL-7B-Instruct": { "max_tokens": 131072, @@ -32767,7 +32767,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_vision": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/BAAI/bge-en-icl": { "max_tokens": 32768, @@ -32776,7 +32776,7 @@ "output_cost_per_token": 0.0, "litellm_provider": "nebius", "mode": "embedding", - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/BAAI/bge-multilingual-gemma2": { "max_tokens": 8192, @@ -32785,7 +32785,7 @@ "output_cost_per_token": 0.0, "litellm_provider": "nebius", "mode": "embedding", - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/intfloat/e5-mistral-7b-instruct": { "max_tokens": 32768, @@ -32794,7 +32794,7 @@ "output_cost_per_token": 0.0, "litellm_provider": "nebius", "mode": "embedding", - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nvidia.nemotron-nano-12b-v2": { "input_cost_per_token": 2e-07, @@ -33535,7 +33535,7 @@ "max_tokens": 4000, "mode": "chat", "output_cost_per_token": 1.56e-06, - "source": "https://www.oracle.com/cloud/ai/generative-ai/pricing/", + "source": "https://www.oracle.com/artificial-intelligence/enterprise-ai/cost-estimator/", "supports_function_calling": true, "supports_response_schema": false, "supports_native_streaming": true @@ -33548,7 +33548,7 @@ "max_tokens": 4000, "mode": "chat", "output_cost_per_token": 1.56e-06, - "source": "https://www.oracle.com/cloud/ai/generative-ai/pricing/", + "source": "https://www.oracle.com/artificial-intelligence/enterprise-ai/cost-estimator/", "supports_function_calling": true, "supports_response_schema": false, "supports_native_streaming": true @@ -33561,7 +33561,7 @@ "max_tokens": 4000, "mode": "chat", "output_cost_per_token": 1.56e-06, - "source": "https://www.oracle.com/cloud/ai/generative-ai/pricing/", + "source": "https://www.oracle.com/artificial-intelligence/enterprise-ai/cost-estimator/", "supports_function_calling": true, "supports_response_schema": false, "supports_native_streaming": true @@ -33618,7 +33618,7 @@ "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 1.56e-06, - "source": "https://www.oracle.com/cloud/ai/generative-ai/pricing/", + "source": "https://www.oracle.com/artificial-intelligence/enterprise-ai/cost-estimator/", "supports_function_calling": true, "supports_response_schema": false, "supports_native_streaming": true, @@ -33632,7 +33632,7 @@ "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 1.56e-06, - "source": "https://www.oracle.com/cloud/ai/generative-ai/pricing/", + "source": "https://www.oracle.com/artificial-intelligence/enterprise-ai/cost-estimator/", "supports_function_calling": false, "supports_response_schema": false, "supports_native_streaming": true @@ -33643,7 +33643,7 @@ "max_input_tokens": 512, "mode": "embedding", "output_vector_size": 1024, - "source": "https://www.oracle.com/cloud/ai/generative-ai/pricing/", + "source": "https://www.oracle.com/artificial-intelligence/enterprise-ai/cost-estimator/", "supports_vision": true }, "oci/cohere.command-a-reasoning-08-2025": { @@ -34540,7 +34540,7 @@ "max_tokens": 2048, "mode": "chat", "output_cost_per_token": 2e-07, - "source": "https://openrouter.ai/api/v1/models/bytedance/ui-tars-1.5-7b", + "source": "https://openrouter.ai/bytedance/ui-tars-1.5-7b", "supports_tool_choice": true }, "openrouter/deepseek/deepseek-chat": { @@ -34775,7 +34775,7 @@ "output_cost_per_reasoning_token": 3e-06, "output_cost_per_token": 3e-06, "rpm": 2000, - "source": "https://ai.google.dev/pricing/gemini-3", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -34816,7 +34816,7 @@ "output_cost_per_reasoning_token": 1.5e-06, "output_cost_per_token": 1.5e-06, "rpm": 2000, - "source": "https://ai.google.dev/pricing/gemini-3", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -35901,7 +35901,7 @@ "max_tokens": 131000, "mode": "chat", "output_cost_per_token": 6.7e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/deepseek-r1-distill-llama-70b", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -35915,7 +35915,7 @@ "max_tokens": 131000, "mode": "chat", "output_cost_per_token": 1e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/llama-3-1-8b-instruct", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true @@ -35928,7 +35928,7 @@ "max_tokens": 131000, "mode": "chat", "output_cost_per_token": 6.7e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/meta-llama-3-1-70b-instruct", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": false, "supports_response_schema": false, "supports_tool_choice": false @@ -35941,7 +35941,7 @@ "max_tokens": 131000, "mode": "chat", "output_cost_per_token": 6.7e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/meta-llama-3-3-70b-instruct", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true @@ -35954,7 +35954,7 @@ "max_tokens": 127000, "mode": "chat", "output_cost_per_token": 1e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/mistral-7b-instruct-v0-3", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true @@ -35967,7 +35967,7 @@ "max_tokens": 118000, "mode": "chat", "output_cost_per_token": 1.3e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/mistral-nemo-instruct-2407", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true @@ -35980,7 +35980,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2.8e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/mistral-small-3-2-24b-instruct-2506", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -35994,7 +35994,7 @@ "max_tokens": 32000, "mode": "chat", "output_cost_per_token": 6.3e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/mixtral-8x7b-instruct-v0-1", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": false, "supports_response_schema": true, "supports_tool_choice": false @@ -36007,7 +36007,7 @@ "max_tokens": 32000, "mode": "chat", "output_cost_per_token": 8.7e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/qwen2-5-coder-32b-instruct", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": false, "supports_response_schema": true, "supports_tool_choice": false @@ -36020,7 +36020,7 @@ "max_tokens": 32000, "mode": "chat", "output_cost_per_token": 9.1e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/qwen2-5-vl-72b-instruct", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": false, "supports_response_schema": true, "supports_tool_choice": false, @@ -36034,7 +36034,7 @@ "max_tokens": 32000, "mode": "chat", "output_cost_per_token": 2.3e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/qwen3-32b", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -36048,7 +36048,7 @@ "max_tokens": 131000, "mode": "chat", "output_cost_per_token": 4e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/gpt-oss-120b", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": false, "supports_reasoning": true, "supports_response_schema": true, @@ -36062,7 +36062,7 @@ "max_tokens": 131000, "mode": "chat", "output_cost_per_token": 1.5e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/gpt-oss-20b", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": false, "supports_reasoning": true, "supports_response_schema": true, @@ -36076,7 +36076,7 @@ "max_tokens": 32000, "mode": "chat", "output_cost_per_token": 2.9e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/llava-next-mistral-7b", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": false, "supports_response_schema": true, "supports_tool_choice": false, @@ -36090,7 +36090,7 @@ "max_tokens": 256000, "mode": "chat", "output_cost_per_token": 1.9e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/mamba-codestral-7b-v0-1", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": false, "supports_response_schema": true, "supports_tool_choice": false @@ -38856,7 +38856,7 @@ "max_input_tokens": 262144, "mode": "chat", "output_cost_per_token": 3.6e-06, - "source": "https://www.together.ai/models/Qwen/Qwen3.5-397B-A17B", + "source": "https://www.together.ai/models/qwen3-5-397b-a17b", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -43097,7 +43097,7 @@ "max_tokens": 2000000, "mode": "chat", "output_cost_per_token": 5e-07, - "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)", + "source": "https://docs.x.ai/developers/models", "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -43113,7 +43113,7 @@ "max_tokens": 2000000, "mode": "chat", "output_cost_per_token": 5e-07, - "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)", + "source": "https://docs.x.ai/developers/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -43130,7 +43130,7 @@ "max_tokens": 2000000, "mode": "chat", "output_cost_per_token": 6e-06, - "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)", + "source": "https://docs.x.ai/developers/models", "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -43146,7 +43146,7 @@ "max_tokens": 2000000, "mode": "chat", "output_cost_per_token": 6e-06, - "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)", + "source": "https://docs.x.ai/developers/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -43266,7 +43266,7 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/veo", + "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", "supported_modalities": [ "text" ], @@ -43280,7 +43280,7 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.15, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/veo", + "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", "supported_modalities": [ "text" ], @@ -43295,7 +43295,7 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/veo", + "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", "supported_modalities": [ "text" ], @@ -43310,7 +43310,7 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.15, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/veo", + "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", "supported_modalities": [ "text" ], @@ -44946,7 +44946,7 @@ "litellm_provider": "azure", "mode": "video_generation", "output_cost_per_video_per_second": 0.1, - "source": "https://azure.microsoft.com/en-us/products/ai-services/video-generation", + "source": "https://ai.azure.com/catalog/models/sora-2", "supported_modalities": [ "text" ], @@ -44958,7 +44958,7 @@ "litellm_provider": "azure", "mode": "video_generation", "output_cost_per_video_per_second": 0.3, - "source": "https://azure.microsoft.com/en-us/products/ai-services/video-generation", + "source": "https://ai.azure.com/catalog/models/sora-2-pro", "supported_modalities": [ "text" ], @@ -44970,7 +44970,7 @@ "litellm_provider": "azure", "mode": "video_generation", "output_cost_per_video_per_second": 0.5, - "source": "https://azure.microsoft.com/en-us/products/ai-services/video-generation", + "source": "https://ai.azure.com/catalog/models/sora-2-pro", "supported_modalities": [ "text" ], @@ -49180,7 +49180,7 @@ "input_cost_per_second": 0.0002833333333333333, "litellm_provider": "openai", "mode": "audio_transcription", - "source": "https://platform.openai.com/docs/models/gpt-realtime-whisper", + "source": "https://developers.openai.com/api/docs/models/gpt-realtime-whisper", "supported_endpoints": [ "/v1/realtime", "/v1/realtime/transcription_sessions" @@ -50432,7 +50432,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "source": "https://www.volcengine.com/docs/82379/1330310", + "source": "https://docs.volcengine.com/docs/82379/1330310", "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": false, @@ -50470,7 +50470,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "source": "https://www.volcengine.com/docs/82379/1330310", + "source": "https://docs.volcengine.com/docs/82379/1330310", "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": false, @@ -50508,7 +50508,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "source": "https://www.volcengine.com/docs/82379/1330310", + "source": "https://docs.volcengine.com/docs/82379/1330310", "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": false, @@ -50546,7 +50546,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "source": "https://www.volcengine.com/docs/82379/1330310", + "source": "https://docs.volcengine.com/docs/82379/1330310", "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": false, @@ -51295,7 +51295,7 @@ "supports_function_calling": true, "supports_assistant_prefill": true, "supports_reasoning": true, - "source": "https://pinstripes.io/pricing" + "source": "https://pinstripes.io/" }, "pinstripes/ps/qwen3.6-35b-a3b": { "max_tokens": 131072, @@ -51308,7 +51308,7 @@ "supports_function_calling": true, "supports_assistant_prefill": true, "supports_reasoning": true, - "source": "https://pinstripes.io/pricing" + "source": "https://pinstripes.io/" }, "pinstripes/ps/qwen3-30b-a3b": { "max_tokens": 131072, @@ -51321,7 +51321,7 @@ "supports_function_calling": true, "supports_assistant_prefill": true, "supports_reasoning": true, - "source": "https://pinstripes.io/pricing" + "source": "https://pinstripes.io/" }, "pinstripes/ps/qwen3-coder-30b-a3b": { "max_tokens": 131072, @@ -51334,7 +51334,7 @@ "supports_function_calling": true, "supports_assistant_prefill": true, "supports_reasoning": false, - "source": "https://pinstripes.io/pricing" + "source": "https://pinstripes.io/" }, "pinstripes/ps/deepseek-v4-flash": { "max_tokens": 163840, @@ -51347,7 +51347,7 @@ "supports_function_calling": true, "supports_assistant_prefill": true, "supports_reasoning": true, - "source": "https://pinstripes.io/pricing" + "source": "https://pinstripes.io/" }, "pinstripes/ps/minimax-m2.7": { "max_tokens": 1000192, @@ -51360,7 +51360,7 @@ "supports_function_calling": true, "supports_assistant_prefill": true, "supports_reasoning": false, - "source": "https://pinstripes.io/pricing" + "source": "https://pinstripes.io/" }, "darkbloom/gemma-4-26b": { "input_cost_per_token": 3e-08, @@ -51464,7 +51464,7 @@ "input_cost_per_second": 7.5e-05, "litellm_provider": "openai", "mode": "audio_transcription", - "source": "https://platform.openai.com/docs/models/gpt-transcribe", + "source": "https://developers.openai.com/api/docs/models/gpt-transcribe", "supported_endpoints": [ "/v1/audio/transcriptions", "/v1/realtime/transcription_sessions" @@ -51482,7 +51482,7 @@ "input_cost_per_second": 0.0002833333333333333, "litellm_provider": "openai", "mode": "audio_transcription", - "source": "https://platform.openai.com/docs/models/gpt-live-transcribe", + "source": "https://developers.openai.com/api/docs/models/gpt-live-transcribe", "supported_endpoints": [ "/v1/realtime", "/v1/realtime/transcription_sessions" @@ -51503,7 +51503,7 @@ "max_output_tokens": 2000, "max_tokens": 2000, "mode": "realtime", - "source": "https://platform.openai.com/docs/models/gpt-realtime-translate", + "source": "https://developers.openai.com/api/docs/models/gpt-realtime-translate", "supported_modalities": [ "audio" ], @@ -51531,7 +51531,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "source": "https://docs.claude.com/en/docs/about-claude/models/overview", + "source": "https://platform.claude.com/docs/en/about-claude/models/overview", "supports_adaptive_thinking": true, "thinking_always_on": true, "supports_mid_conversation_system": true, @@ -51570,7 +51570,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "source": "https://docs.claude.com/en/docs/about-claude/models/overview", + "source": "https://platform.claude.com/docs/en/about-claude/models/overview", "supports_adaptive_thinking": true, "thinking_always_on": true, "supports_assistant_prefill": false, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index ec175025b42..b548de07452 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -3726,7 +3726,7 @@ "output_cost_per_token": 0, "litellm_provider": "azure_ai", "mode": "chat", - "source": "https://azure.microsoft.com/en-us/pricing/details/ai-services/", + "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/aoai/", "comment": "Flat cost of $0.14 per M input tokens for Azure AI Foundry Model Router infrastructure. Use pattern: azure_ai/model_router/ where deployment-name is your Azure deployment (e.g., azure-model-router)" }, "azure/eu/gpt-4o-2024-08-06": { @@ -5328,7 +5328,7 @@ "input_cost_per_second": 0.0002833333333333333, "litellm_provider": "azure", "mode": "audio_transcription", - "source": "https://learn.microsoft.com/en-us/azure/ai-foundry/openai/concepts/gpt-realtime-whisper", + "source": "https://learn.microsoft.com/en-us/azure/foundry/openai/concepts/gpt-realtime-whisper", "supported_endpoints": [ "/v1/realtime", "/v1/realtime/transcription_sessions" @@ -9107,7 +9107,7 @@ "mode": "embedding", "output_cost_per_token": 0.0, "output_vector_size": 1024, - "source": "https://azuremarketplace.microsoft.com/en-us/marketplace/apps/cohere.cohere-embed-v3-english-offer?tab=PlansAndPrice", + "source": "https://marketplace.microsoft.com/en-us/marketplace/apps/cohere.cohere-embed-v3-english-offer?tab=PlansAndPrice", "supports_embedding_image_input": true }, "azure_ai/Cohere-embed-v3-multilingual": { @@ -9118,7 +9118,7 @@ "mode": "embedding", "output_cost_per_token": 0.0, "output_vector_size": 1024, - "source": "https://azuremarketplace.microsoft.com/en-us/marketplace/apps/cohere.cohere-embed-v3-english-offer?tab=PlansAndPrice", + "source": "https://marketplace.microsoft.com/en-us/marketplace/apps/cohere.cohere-embed-v3-english-offer?tab=PlansAndPrice", "supports_embedding_image_input": true }, "azure_ai/FLUX-1.1-pro": { @@ -9134,7 +9134,7 @@ "litellm_provider": "azure_ai", "mode": "image_generation", "output_cost_per_image": 0.04, - "source": "https://azuremarketplace.microsoft.com/pt-br/marketplace/apps/cohere.cohere-embed-4-offer?tab=PlansAndPrice", + "source": "https://marketplace.microsoft.com/pt-br/marketplace/apps/cohere.cohere-embed-4-offer?tab=PlansAndPrice", "supported_endpoints": [ "/v1/images/generations" ] @@ -9467,7 +9467,7 @@ "max_tokens": 2048, "mode": "chat", "output_cost_per_token": 3.7e-07, - "source": "https://azuremarketplace.microsoft.com/en/marketplace/apps/metagenai.meta-llama-3-2-11b-vision-instruct-offer?tab=Overview", + "source": "https://marketplace.microsoft.com/en/marketplace/apps/metagenai.meta-llama-3-2-11b-vision-instruct-offer?tab=Overview", "supports_function_calling": true, "supports_tool_choice": true, "supports_vision": true @@ -9481,7 +9481,7 @@ "max_tokens": 2048, "mode": "chat", "output_cost_per_token": 2.04e-06, - "source": "https://azuremarketplace.microsoft.com/en/marketplace/apps/metagenai.meta-llama-3-2-90b-vision-instruct-offer?tab=Overview", + "source": "https://marketplace.microsoft.com/en/marketplace/apps/metagenai.meta-llama-3-2-90b-vision-instruct-offer?tab=Overview", "supports_function_calling": true, "supports_tool_choice": true, "supports_vision": true @@ -9494,7 +9494,7 @@ "max_tokens": 2048, "mode": "chat", "output_cost_per_token": 7.1e-07, - "source": "https://azuremarketplace.microsoft.com/en/marketplace/apps/metagenai.llama-3-3-70b-instruct-offer?tab=Overview", + "source": "https://marketplace.microsoft.com/en/marketplace/apps/metagenai.llama-3-3-70b-instruct-offer?tab=Overview", "supports_function_calling": true, "supports_tool_choice": true }, @@ -9543,7 +9543,7 @@ "max_tokens": 2048, "mode": "chat", "output_cost_per_token": 1.6e-05, - "source": "https://azuremarketplace.microsoft.com/en-us/marketplace/apps/metagenai.meta-llama-3-1-405b-instruct-offer?tab=PlansAndPrice", + "source": "https://marketplace.microsoft.com/en-us/marketplace/apps/metagenai.meta-llama-3-1-405b-instruct-offer?tab=PlansAndPrice", "supports_tool_choice": true }, "azure_ai/Meta-Llama-3.1-70B-Instruct": { @@ -9554,7 +9554,7 @@ "max_tokens": 2048, "mode": "chat", "output_cost_per_token": 3.54e-06, - "source": "https://azuremarketplace.microsoft.com/en-us/marketplace/apps/metagenai.meta-llama-3-1-70b-instruct-offer?tab=PlansAndPrice", + "source": "https://marketplace.microsoft.com/en-us/marketplace/apps/metagenai.meta-llama-3-1-70b-instruct-offer?tab=PlansAndPrice", "supports_tool_choice": true }, "azure_ai/Meta-Llama-3.1-8B-Instruct": { @@ -9566,7 +9566,7 @@ "max_tokens": 2048, "mode": "chat", "output_cost_per_token": 6.1e-07, - "source": "https://azuremarketplace.microsoft.com/en-us/marketplace/apps/metagenai.meta-llama-3-1-8b-instruct-offer?tab=PlansAndPrice", + "source": "https://marketplace.microsoft.com/en-us/marketplace/apps/metagenai.meta-llama-3-1-8b-instruct-offer?tab=PlansAndPrice", "supports_tool_choice": true }, "azure_ai/Phi-3-medium-128k-instruct": { @@ -9756,7 +9756,7 @@ "supported_endpoints": [ "/v1/ocr" ], - "source": "https://azure.microsoft.com/en-us/pricing/details/ai-foundry-models/" + "source": "https://ai.azure.com/catalog/models/mistral-document-ai-2512" }, "azure_ai/doc-intelligence/prebuilt-read": { "litellm_provider": "azure_ai", @@ -9967,7 +9967,7 @@ "mode": "embedding", "output_cost_per_token": 0.0, "output_vector_size": 3072, - "source": "https://azuremarketplace.microsoft.com/pt-br/marketplace/apps/cohere.cohere-embed-4-offer?tab=PlansAndPrice", + "source": "https://marketplace.microsoft.com/pt-br/marketplace/apps/cohere.cohere-embed-4-offer?tab=PlansAndPrice", "supported_endpoints": [ "/v1/embeddings" ], @@ -10151,7 +10151,7 @@ "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 0.00971, - "source": "https://azure.microsoft.com/en-us/products/ai-services/ai-foundry/models/jais-30b-chat" + "source": "https://ai.azure.com/catalog/models/jais-30b-chat" }, "azure_ai/jamba-instruct": { "input_cost_per_token": 5e-07, @@ -10208,7 +10208,7 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 4e-08, - "source": "https://azuremarketplace.microsoft.com/en/marketplace/apps/000-000.ministral-3b-2410-offer?tab=Overview", + "source": "https://marketplace.microsoft.com/en/marketplace/apps/000-000.ministral-3b-2410-offer?tab=Overview", "supports_function_calling": true, "supports_tool_choice": true }, @@ -10231,7 +10231,7 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 6e-06, - "source": "https://azuremarketplace.microsoft.com/en/marketplace/apps/000-000.mistral-ai-large-2407-offer?tab=Overview", + "source": "https://marketplace.microsoft.com/en/marketplace/apps/000-000.mistral-ai-large-2407-offer?tab=Overview", "supports_function_calling": true, "supports_tool_choice": true }, @@ -10243,7 +10243,7 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 6e-06, - "source": "https://azuremarketplace.microsoft.com/en/marketplace/apps/000-000.mistral-ai-large-2407-offer?tab=Overview", + "source": "https://marketplace.microsoft.com/en/marketplace/apps/000-000.mistral-ai-large-2407-offer?tab=Overview", "supports_function_calling": true, "supports_tool_choice": true }, @@ -10280,7 +10280,7 @@ "max_tokens": 4096, "mode": "chat", "output_cost_per_token": 1.5e-07, - "source": "https://azuremarketplace.microsoft.com/en/marketplace/apps/000-000.mistral-nemo-12b-2407?tab=PlansAndPrice", + "source": "https://marketplace.microsoft.com/en/marketplace/apps/000-000.mistral-nemo-12b-2407?tab=PlansAndPrice", "supports_function_calling": true }, "azure_ai/mistral-small": { @@ -22796,7 +22796,7 @@ "output_cost_per_reasoning_token": 3e-06, "output_cost_per_token": 3e-06, "rpm": 2000, - "source": "https://ai.google.dev/pricing/gemini-3", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -22850,7 +22850,7 @@ "output_cost_per_reasoning_token": 9e-06, "output_cost_per_token": 9e-06, "rpm": 2000, - "source": "https://ai.google.dev/pricing/gemini-3", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -22915,7 +22915,7 @@ "output_cost_per_token_batches": 1.875e-06, "output_cost_per_token_flex": 1.875e-06, "rpm": 2000, - "source": "https://ai.google.dev/pricing/gemini-3", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -22974,7 +22974,7 @@ "output_cost_per_token_batches": 1.875e-06, "output_cost_per_token_flex": 1.875e-06, "rpm": 2000, - "source": "https://ai.google.dev/pricing/gemini-3", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -23178,7 +23178,7 @@ "mode": "chat", "output_cost_per_reasoning_token": 3e-06, "output_cost_per_token": 3e-06, - "source": "https://ai.google.dev/pricing/gemini-3", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -23262,7 +23262,7 @@ "mode": "chat", "output_cost_per_reasoning_token": 9e-06, "output_cost_per_token": 9e-06, - "source": "https://ai.google.dev/pricing/gemini-3", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -23325,7 +23325,7 @@ "output_cost_per_token": 3.75e-06, "output_cost_per_token_batches": 1.875e-06, "output_cost_per_token_flex": 1.875e-06, - "source": "https://ai.google.dev/pricing/gemini-3", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -23382,7 +23382,7 @@ "output_cost_per_token": 3.75e-06, "output_cost_per_token_batches": 1.875e-06, "output_cost_per_token_flex": 1.875e-06, - "source": "https://ai.google.dev/pricing/gemini-3", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -26982,7 +26982,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "source": "https://platform.openai.com/docs/models/gpt-5.6-cyber", + "source": "https://developers.openai.com/api/docs/models/gpt-5.6-cyber", "supports_computer_use": true, "supports_parallel_function_calling": true }, @@ -27021,7 +27021,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "source": "https://platform.openai.com/docs/models/daybreak-red-latest", + "source": "https://developers.openai.com/api/docs/models/daybreak-red-latest", "supports_computer_use": true, "supports_parallel_function_calling": true }, @@ -27061,7 +27061,7 @@ "supports_tool_choice": true, "supports_vision": true, "supports_web_search": true, - "source": "https://platform.openai.com/docs/models/daybreak-blue-latest", + "source": "https://developers.openai.com/api/docs/models/daybreak-blue-latest", "supports_parallel_function_calling": true }, "chat-latest": { @@ -27073,7 +27073,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 3e-05, - "source": "https://platform.openai.com/docs/models/chat-latest", + "source": "https://developers.openai.com/api/docs/models/chat-latest", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses" @@ -30368,7 +30368,7 @@ "search_context_size_low": 0.0025, "search_context_size_medium": 0.0025 }, - "source": "https://dev.meta.ai/docs/getting-started/pricing-rate-limits", + "source": "https://ai.developer.meta.com/docs/pricing-rate-limits", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses", @@ -30409,7 +30409,7 @@ "search_context_size_low": 0.0025, "search_context_size_medium": 0.0025 }, - "source": "https://dev.meta.ai/docs/getting-started/pricing-rate-limits", + "source": "https://ai.developer.meta.com/docs/pricing-rate-limits", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses", @@ -30450,7 +30450,7 @@ "search_context_size_low": 0.0025, "search_context_size_medium": 0.0025 }, - "source": "https://dev.meta.ai/docs/getting-started/pricing-rate-limits", + "source": "https://ai.developer.meta.com/docs/pricing-rate-limits", "supported_endpoints": [ "/v1/chat/completions", "/v1/responses", @@ -30483,7 +30483,7 @@ "max_output_tokens": 4028, "max_tokens": 4028, "mode": "chat", - "source": "https://llama.developer.meta.com/docs/models", + "source": "https://ai.developer.meta.com/docs/models", "supported_modalities": [ "text" ], @@ -30499,7 +30499,7 @@ "max_output_tokens": 4028, "max_tokens": 4028, "mode": "chat", - "source": "https://llama.developer.meta.com/docs/models", + "source": "https://ai.developer.meta.com/docs/models", "supported_modalities": [ "text" ], @@ -30515,7 +30515,7 @@ "max_output_tokens": 4028, "max_tokens": 4028, "mode": "chat", - "source": "https://llama.developer.meta.com/docs/models", + "source": "https://ai.developer.meta.com/docs/models", "supported_modalities": [ "text", "image" @@ -30532,7 +30532,7 @@ "max_output_tokens": 4028, "max_tokens": 4028, "mode": "chat", - "source": "https://llama.developer.meta.com/docs/models", + "source": "https://ai.developer.meta.com/docs/models", "supported_modalities": [ "text", "image" @@ -32477,7 +32477,7 @@ "mode": "chat", "supports_function_calling": true, "supports_reasoning": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/deepseek-ai/DeepSeek-R1-0528": { "max_tokens": 164000, @@ -32489,7 +32489,7 @@ "mode": "chat", "supports_function_calling": true, "supports_reasoning": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/deepseek-ai/DeepSeek-R1-Distill-Llama-70B": { "max_tokens": 128000, @@ -32500,7 +32500,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/deepseek-ai/DeepSeek-V3": { "max_tokens": 128000, @@ -32511,7 +32511,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/deepseek-ai/DeepSeek-V3-0324": { "max_tokens": 128000, @@ -32522,7 +32522,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/google/gemma-3-27b-it": { "max_tokens": 128000, @@ -32534,7 +32534,7 @@ "mode": "chat", "supports_function_calling": true, "supports_vision": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/meta-llama/Llama-3.3-70B-Instruct": { "max_tokens": 128000, @@ -32545,7 +32545,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/meta-llama/Llama-Guard-3-8B": { "max_tokens": 128000, @@ -32555,7 +32555,7 @@ "output_cost_per_token": 6e-08, "litellm_provider": "nebius", "mode": "chat", - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/meta-llama/Meta-Llama-3.1-8B-Instruct": { "max_tokens": 128000, @@ -32566,7 +32566,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/meta-llama/Meta-Llama-3.1-70B-Instruct": { "max_tokens": 128000, @@ -32577,7 +32577,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/meta-llama/Meta-Llama-3.1-405B-Instruct": { "max_tokens": 128000, @@ -32588,7 +32588,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/mistralai/Mistral-Nemo-Instruct-2407": { "max_tokens": 128000, @@ -32599,7 +32599,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/NousResearch/Hermes-3-Llama-3.1-405B": { "max_tokens": 128000, @@ -32610,7 +32610,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/nvidia/Llama-3.1-Nemotron-Ultra-253B-v1": { "max_tokens": 128000, @@ -32621,7 +32621,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/nvidia/Llama-3.3-Nemotron-Super-49B-v1": { "max_tokens": 131072, @@ -32632,7 +32632,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/Qwen/Qwen3-235B-A22B": { "max_tokens": 262144, @@ -32643,7 +32643,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/Qwen/Qwen3-32B": { "max_tokens": 32768, @@ -32654,7 +32654,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/Qwen/Qwen3-30B-A3B": { "max_tokens": 32768, @@ -32665,7 +32665,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/Qwen/Qwen3-14B": { "max_tokens": 32768, @@ -32676,7 +32676,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/Qwen/Qwen3-4B": { "max_tokens": 32768, @@ -32687,7 +32687,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/Qwen/QwQ-32B": { "max_tokens": 32768, @@ -32699,7 +32699,7 @@ "mode": "chat", "supports_function_calling": true, "supports_reasoning": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/Qwen/Qwen2.5-72B-Instruct": { "max_tokens": 128000, @@ -32710,7 +32710,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/Qwen/Qwen2.5-32B-Instruct": { "max_tokens": 128000, @@ -32721,7 +32721,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/Qwen/Qwen2.5-Coder-7B": { "max_tokens": 32768, @@ -32732,7 +32732,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_function_calling": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/Qwen/Qwen2.5-VL-72B-Instruct": { "max_tokens": 131072, @@ -32744,7 +32744,7 @@ "mode": "chat", "supports_function_calling": true, "supports_vision": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/Qwen/Qwen2-VL-72B-Instruct": { "max_tokens": 131072, @@ -32756,7 +32756,7 @@ "mode": "chat", "supports_function_calling": true, "supports_vision": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/Qwen/Qwen2-VL-7B-Instruct": { "max_tokens": 131072, @@ -32767,7 +32767,7 @@ "litellm_provider": "nebius", "mode": "chat", "supports_vision": true, - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/BAAI/bge-en-icl": { "max_tokens": 32768, @@ -32776,7 +32776,7 @@ "output_cost_per_token": 0.0, "litellm_provider": "nebius", "mode": "embedding", - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/BAAI/bge-multilingual-gemma2": { "max_tokens": 8192, @@ -32785,7 +32785,7 @@ "output_cost_per_token": 0.0, "litellm_provider": "nebius", "mode": "embedding", - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nebius/intfloat/e5-mistral-7b-instruct": { "max_tokens": 32768, @@ -32794,7 +32794,7 @@ "output_cost_per_token": 0.0, "litellm_provider": "nebius", "mode": "embedding", - "source": "https://nebius.com/prices-ai-studio" + "source": "https://nebius.com/prices" }, "nvidia.nemotron-nano-12b-v2": { "input_cost_per_token": 2e-07, @@ -33535,7 +33535,7 @@ "max_tokens": 4000, "mode": "chat", "output_cost_per_token": 1.56e-06, - "source": "https://www.oracle.com/cloud/ai/generative-ai/pricing/", + "source": "https://www.oracle.com/artificial-intelligence/enterprise-ai/cost-estimator/", "supports_function_calling": true, "supports_response_schema": false, "supports_native_streaming": true @@ -33548,7 +33548,7 @@ "max_tokens": 4000, "mode": "chat", "output_cost_per_token": 1.56e-06, - "source": "https://www.oracle.com/cloud/ai/generative-ai/pricing/", + "source": "https://www.oracle.com/artificial-intelligence/enterprise-ai/cost-estimator/", "supports_function_calling": true, "supports_response_schema": false, "supports_native_streaming": true @@ -33561,7 +33561,7 @@ "max_tokens": 4000, "mode": "chat", "output_cost_per_token": 1.56e-06, - "source": "https://www.oracle.com/cloud/ai/generative-ai/pricing/", + "source": "https://www.oracle.com/artificial-intelligence/enterprise-ai/cost-estimator/", "supports_function_calling": true, "supports_response_schema": false, "supports_native_streaming": true @@ -33618,7 +33618,7 @@ "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 1.56e-06, - "source": "https://www.oracle.com/cloud/ai/generative-ai/pricing/", + "source": "https://www.oracle.com/artificial-intelligence/enterprise-ai/cost-estimator/", "supports_function_calling": true, "supports_response_schema": false, "supports_native_streaming": true, @@ -33632,7 +33632,7 @@ "max_tokens": 8192, "mode": "chat", "output_cost_per_token": 1.56e-06, - "source": "https://www.oracle.com/cloud/ai/generative-ai/pricing/", + "source": "https://www.oracle.com/artificial-intelligence/enterprise-ai/cost-estimator/", "supports_function_calling": false, "supports_response_schema": false, "supports_native_streaming": true @@ -33643,7 +33643,7 @@ "max_input_tokens": 512, "mode": "embedding", "output_vector_size": 1024, - "source": "https://www.oracle.com/cloud/ai/generative-ai/pricing/", + "source": "https://www.oracle.com/artificial-intelligence/enterprise-ai/cost-estimator/", "supports_vision": true }, "oci/cohere.command-a-reasoning-08-2025": { @@ -34540,7 +34540,7 @@ "max_tokens": 2048, "mode": "chat", "output_cost_per_token": 2e-07, - "source": "https://openrouter.ai/api/v1/models/bytedance/ui-tars-1.5-7b", + "source": "https://openrouter.ai/bytedance/ui-tars-1.5-7b", "supports_tool_choice": true }, "openrouter/deepseek/deepseek-chat": { @@ -34775,7 +34775,7 @@ "output_cost_per_reasoning_token": 3e-06, "output_cost_per_token": 3e-06, "rpm": 2000, - "source": "https://ai.google.dev/pricing/gemini-3", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -34816,7 +34816,7 @@ "output_cost_per_reasoning_token": 1.5e-06, "output_cost_per_token": 1.5e-06, "rpm": 2000, - "source": "https://ai.google.dev/pricing/gemini-3", + "source": "https://ai.google.dev/gemini-api/docs/pricing", "supported_endpoints": [ "/v1/chat/completions", "/v1/completions", @@ -35901,7 +35901,7 @@ "max_tokens": 131000, "mode": "chat", "output_cost_per_token": 6.7e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/deepseek-r1-distill-llama-70b", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -35915,7 +35915,7 @@ "max_tokens": 131000, "mode": "chat", "output_cost_per_token": 1e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/llama-3-1-8b-instruct", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true @@ -35928,7 +35928,7 @@ "max_tokens": 131000, "mode": "chat", "output_cost_per_token": 6.7e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/meta-llama-3-1-70b-instruct", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": false, "supports_response_schema": false, "supports_tool_choice": false @@ -35941,7 +35941,7 @@ "max_tokens": 131000, "mode": "chat", "output_cost_per_token": 6.7e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/meta-llama-3-3-70b-instruct", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true @@ -35954,7 +35954,7 @@ "max_tokens": 127000, "mode": "chat", "output_cost_per_token": 1e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/mistral-7b-instruct-v0-3", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true @@ -35967,7 +35967,7 @@ "max_tokens": 118000, "mode": "chat", "output_cost_per_token": 1.3e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/mistral-nemo-instruct-2407", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true @@ -35980,7 +35980,7 @@ "max_tokens": 128000, "mode": "chat", "output_cost_per_token": 2.8e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/mistral-small-3-2-24b-instruct-2506", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -35994,7 +35994,7 @@ "max_tokens": 32000, "mode": "chat", "output_cost_per_token": 6.3e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/mixtral-8x7b-instruct-v0-1", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": false, "supports_response_schema": true, "supports_tool_choice": false @@ -36007,7 +36007,7 @@ "max_tokens": 32000, "mode": "chat", "output_cost_per_token": 8.7e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/qwen2-5-coder-32b-instruct", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": false, "supports_response_schema": true, "supports_tool_choice": false @@ -36020,7 +36020,7 @@ "max_tokens": 32000, "mode": "chat", "output_cost_per_token": 9.1e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/qwen2-5-vl-72b-instruct", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": false, "supports_response_schema": true, "supports_tool_choice": false, @@ -36034,7 +36034,7 @@ "max_tokens": 32000, "mode": "chat", "output_cost_per_token": 2.3e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/qwen3-32b", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -36048,7 +36048,7 @@ "max_tokens": 131000, "mode": "chat", "output_cost_per_token": 4e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/gpt-oss-120b", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": false, "supports_reasoning": true, "supports_response_schema": true, @@ -36062,7 +36062,7 @@ "max_tokens": 131000, "mode": "chat", "output_cost_per_token": 1.5e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/gpt-oss-20b", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": false, "supports_reasoning": true, "supports_response_schema": true, @@ -36076,7 +36076,7 @@ "max_tokens": 32000, "mode": "chat", "output_cost_per_token": 2.9e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/llava-next-mistral-7b", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": false, "supports_response_schema": true, "supports_tool_choice": false, @@ -36090,7 +36090,7 @@ "max_tokens": 256000, "mode": "chat", "output_cost_per_token": 1.9e-07, - "source": "https://endpoints.ai.cloud.ovh.net/models/mamba-codestral-7b-v0-1", + "source": "https://www.ovhcloud.com/en/public-cloud/ai-endpoints/catalog/", "supports_function_calling": false, "supports_response_schema": true, "supports_tool_choice": false @@ -38856,7 +38856,7 @@ "max_input_tokens": 262144, "mode": "chat", "output_cost_per_token": 3.6e-06, - "source": "https://www.together.ai/models/Qwen/Qwen3.5-397B-A17B", + "source": "https://www.together.ai/models/qwen3-5-397b-a17b", "supports_function_calling": true, "supports_parallel_function_calling": true, "supports_prompt_caching": true, @@ -43097,7 +43097,7 @@ "max_tokens": 2000000, "mode": "chat", "output_cost_per_token": 5e-07, - "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)", + "source": "https://docs.x.ai/developers/models", "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -43113,7 +43113,7 @@ "max_tokens": 2000000, "mode": "chat", "output_cost_per_token": 5e-07, - "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)", + "source": "https://docs.x.ai/developers/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -43130,7 +43130,7 @@ "max_tokens": 2000000, "mode": "chat", "output_cost_per_token": 6e-06, - "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)", + "source": "https://docs.x.ai/developers/models", "supports_function_calling": true, "supports_response_schema": true, "supports_tool_choice": true, @@ -43146,7 +43146,7 @@ "max_tokens": 2000000, "mode": "chat", "output_cost_per_token": 6e-06, - "source": "https://docs.x.ai/docs/models (Vertex AI Model Garden)", + "source": "https://docs.x.ai/developers/models", "supports_function_calling": true, "supports_reasoning": true, "supports_response_schema": true, @@ -43266,7 +43266,7 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/veo", + "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", "supported_modalities": [ "text" ], @@ -43280,7 +43280,7 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.15, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/veo", + "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", "supported_modalities": [ "text" ], @@ -43295,7 +43295,7 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.4, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/veo", + "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", "supported_modalities": [ "text" ], @@ -43310,7 +43310,7 @@ "max_tokens": 1024, "mode": "video_generation", "output_cost_per_second": 0.15, - "source": "https://cloud.google.com/vertex-ai/generative-ai/docs/model-reference/veo", + "source": "https://docs.cloud.google.com/gemini-enterprise-agent-platform/models/veo/3-1-generate", "supported_modalities": [ "text" ], @@ -44946,7 +44946,7 @@ "litellm_provider": "azure", "mode": "video_generation", "output_cost_per_video_per_second": 0.1, - "source": "https://azure.microsoft.com/en-us/products/ai-services/video-generation", + "source": "https://ai.azure.com/catalog/models/sora-2", "supported_modalities": [ "text" ], @@ -44958,7 +44958,7 @@ "litellm_provider": "azure", "mode": "video_generation", "output_cost_per_video_per_second": 0.3, - "source": "https://azure.microsoft.com/en-us/products/ai-services/video-generation", + "source": "https://ai.azure.com/catalog/models/sora-2-pro", "supported_modalities": [ "text" ], @@ -44970,7 +44970,7 @@ "litellm_provider": "azure", "mode": "video_generation", "output_cost_per_video_per_second": 0.5, - "source": "https://azure.microsoft.com/en-us/products/ai-services/video-generation", + "source": "https://ai.azure.com/catalog/models/sora-2-pro", "supported_modalities": [ "text" ], @@ -49180,7 +49180,7 @@ "input_cost_per_second": 0.0002833333333333333, "litellm_provider": "openai", "mode": "audio_transcription", - "source": "https://platform.openai.com/docs/models/gpt-realtime-whisper", + "source": "https://developers.openai.com/api/docs/models/gpt-realtime-whisper", "supported_endpoints": [ "/v1/realtime", "/v1/realtime/transcription_sessions" @@ -50432,7 +50432,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "source": "https://www.volcengine.com/docs/82379/1330310", + "source": "https://docs.volcengine.com/docs/82379/1330310", "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": false, @@ -50470,7 +50470,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "source": "https://www.volcengine.com/docs/82379/1330310", + "source": "https://docs.volcengine.com/docs/82379/1330310", "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": false, @@ -50508,7 +50508,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "source": "https://www.volcengine.com/docs/82379/1330310", + "source": "https://docs.volcengine.com/docs/82379/1330310", "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": false, @@ -50546,7 +50546,7 @@ "max_output_tokens": 128000, "max_tokens": 128000, "mode": "chat", - "source": "https://www.volcengine.com/docs/82379/1330310", + "source": "https://docs.volcengine.com/docs/82379/1330310", "supports_function_calling": true, "supports_reasoning": true, "supports_tool_choice": false, @@ -51295,7 +51295,7 @@ "supports_function_calling": true, "supports_assistant_prefill": true, "supports_reasoning": true, - "source": "https://pinstripes.io/pricing" + "source": "https://pinstripes.io/" }, "pinstripes/ps/qwen3.6-35b-a3b": { "max_tokens": 131072, @@ -51308,7 +51308,7 @@ "supports_function_calling": true, "supports_assistant_prefill": true, "supports_reasoning": true, - "source": "https://pinstripes.io/pricing" + "source": "https://pinstripes.io/" }, "pinstripes/ps/qwen3-30b-a3b": { "max_tokens": 131072, @@ -51321,7 +51321,7 @@ "supports_function_calling": true, "supports_assistant_prefill": true, "supports_reasoning": true, - "source": "https://pinstripes.io/pricing" + "source": "https://pinstripes.io/" }, "pinstripes/ps/qwen3-coder-30b-a3b": { "max_tokens": 131072, @@ -51334,7 +51334,7 @@ "supports_function_calling": true, "supports_assistant_prefill": true, "supports_reasoning": false, - "source": "https://pinstripes.io/pricing" + "source": "https://pinstripes.io/" }, "pinstripes/ps/deepseek-v4-flash": { "max_tokens": 163840, @@ -51347,7 +51347,7 @@ "supports_function_calling": true, "supports_assistant_prefill": true, "supports_reasoning": true, - "source": "https://pinstripes.io/pricing" + "source": "https://pinstripes.io/" }, "pinstripes/ps/minimax-m2.7": { "max_tokens": 1000192, @@ -51360,7 +51360,7 @@ "supports_function_calling": true, "supports_assistant_prefill": true, "supports_reasoning": false, - "source": "https://pinstripes.io/pricing" + "source": "https://pinstripes.io/" }, "darkbloom/gemma-4-26b": { "input_cost_per_token": 3e-08, @@ -51464,7 +51464,7 @@ "input_cost_per_second": 7.5e-05, "litellm_provider": "openai", "mode": "audio_transcription", - "source": "https://platform.openai.com/docs/models/gpt-transcribe", + "source": "https://developers.openai.com/api/docs/models/gpt-transcribe", "supported_endpoints": [ "/v1/audio/transcriptions", "/v1/realtime/transcription_sessions" @@ -51482,7 +51482,7 @@ "input_cost_per_second": 0.0002833333333333333, "litellm_provider": "openai", "mode": "audio_transcription", - "source": "https://platform.openai.com/docs/models/gpt-live-transcribe", + "source": "https://developers.openai.com/api/docs/models/gpt-live-transcribe", "supported_endpoints": [ "/v1/realtime", "/v1/realtime/transcription_sessions" @@ -51503,7 +51503,7 @@ "max_output_tokens": 2000, "max_tokens": 2000, "mode": "realtime", - "source": "https://platform.openai.com/docs/models/gpt-realtime-translate", + "source": "https://developers.openai.com/api/docs/models/gpt-realtime-translate", "supported_modalities": [ "audio" ], @@ -51531,7 +51531,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "source": "https://docs.claude.com/en/docs/about-claude/models/overview", + "source": "https://platform.claude.com/docs/en/about-claude/models/overview", "supports_adaptive_thinking": true, "thinking_always_on": true, "supports_mid_conversation_system": true, @@ -51570,7 +51570,7 @@ "search_context_size_low": 0.01, "search_context_size_medium": 0.01 }, - "source": "https://docs.claude.com/en/docs/about-claude/models/overview", + "source": "https://platform.claude.com/docs/en/about-claude/models/overview", "supports_adaptive_thinking": true, "thinking_always_on": true, "supports_assistant_prefill": false, From e2ffb6b01c52764fb31d9e931c64f4c53a14a747 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 26 Aug 2026 14:42:08 -0400 Subject: [PATCH 11/49] feat(ci): close duplicate issues after a 3-day grace period Duplicate detection already labelled and commented on new issues, and then closed them outright at 0.85 title similarity. That gave the reporter no chance to push back, and a title-similarity match is not strong enough evidence to close on its own. Detection now only flags. A new daily sweep closes a flagged issue three days later, and only if nobody engaged with the flag. Replying to it, thumbs-downing it, or applying an opt-out label all keep the issue open. The notice says all of that up front, so the reporter knows what happens and how to stop it. The two workflows hand off through an HTML marker in the comment body rather than its prose, so rewording the notice cannot silently break the sweep. The sweep lists by label instead of walking the whole backlog: 1663 open issues against 23 carrying the label meant a comments request each, which would burn the Actions token's hourly budget for a handful of matches. Candidates are taken as the lowest issue number, not the first one listed. The detector orders by score rather than age, so the first candidate can be newer than the issue being closed, and folding an original report into a later one is backwards. An issue whose only candidates are newer is skipped. Closures use state_reason=duplicate rather than not_planned, which reads as "see the other issue" instead of "we are not doing this". Also drops {{html_url}} from the notice. The detection action only exposes number, title and accuracy, so that placeholder had been rendering empty and every "similar issue" link in the comment pointed nowhere. --- .github/workflows/check_duplicate_issues.yml | 38 ++--- .../close_stale_duplicate_issues.yml | 148 ++++++++++++++++++ 2 files changed, 158 insertions(+), 28 deletions(-) create mode 100644 .github/workflows/close_stale_duplicate_issues.yml diff --git a/.github/workflows/check_duplicate_issues.yml b/.github/workflows/check_duplicate_issues.yml index 78198b2c7bb..a087007bff3 100644 --- a/.github/workflows/check_duplicate_issues.yml +++ b/.github/workflows/check_duplicate_issues.yml @@ -1,5 +1,10 @@ name: Check Duplicate Issues +# Flags newly opened issues that look like existing ones. Flagging only: the actual +# close happens 3 days later in "Close Stale Duplicate Issues", and only if nobody +# replied to the comment posted here. The HTML marker below is the handshake between +# the two workflows, so keep it in the template. + on: issues: types: [opened, edited] @@ -19,35 +24,12 @@ jobs: threshold: 0.6 reaction: eyes comment: | - **⚠️ Potential duplicate detected** + + **Potential duplicate detected** - This issue appears similar to existing issue(s): + This looks similar to: {{#issues}} - - [#{{number}}]({{html_url}}) - {{title}} ({{accuracy}}% similar) + - #{{number}} - {{title}} ({{accuracy}}% similar) {{/issues}} - Please review the linked issue(s) to see if they address your concern. If this is not a duplicate, please provide additional context to help us understand the difference. - - - name: Checkout close script - if: github.event.action == 'opened' - uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 - with: - sparse-checkout: .github/scripts - persist-credentials: false - - - name: Set up Python - if: github.event.action == 'opened' - uses: actions/setup-python@a26af69be951a213d495a4c3e4e4022e16d87065 # v5.6.0 - with: - python-version: "3.12" - - - name: Auto-close if high-confidence duplicate - if: github.event.action == 'opened' - env: - GH_TOKEN: ${{ secrets.GITHUB_TOKEN }} - run: | - python3 .github/scripts/close_duplicate_issues.py \ - --issue-number ${{ github.event.issue.number }} \ - --repo ${{ github.repository }} \ - --threshold 0.85 \ - --close + This issue will close automatically in 3 days unless someone responds. If it is a duplicate, please 👍 the existing issue and follow along there. If it is not, comment here or 👎 this comment and it stays open. diff --git a/.github/workflows/close_stale_duplicate_issues.yml b/.github/workflows/close_stale_duplicate_issues.yml new file mode 100644 index 00000000000..8bc1af90454 --- /dev/null +++ b/.github/workflows/close_stale_duplicate_issues.yml @@ -0,0 +1,148 @@ +name: Close Stale Duplicate Issues + +# Closes issues that "Check Duplicate Issues" flagged and that nobody acknowledged +# within the grace period. Replying to the flag, thumbs-downing it, or applying an +# opt-out label all keep an issue open. +# +# Dry-run preview (touches nothing): +# gh workflow run "Close Stale Duplicate Issues" -f dry_run=true + +on: + schedule: + # Daily at 09:30 UTC, after the midnight stale sweep and off the hour. + - cron: "30 9 * * *" + workflow_dispatch: + inputs: + dry_run: + description: "Report what would close without touching any issue." + required: false + default: "false" + type: choice + options: + - "true" + - "false" + grace_period_days: + description: "Days to wait after the duplicate flag before closing." + required: false + default: "3" + limit: + description: "Maximum number of issues to close in a single run." + required: false + default: "50" + +permissions: + contents: read + issues: write + +jobs: + close-stale-duplicates: + if: github.repository == 'BerriAI/litellm' + runs-on: ubuntu-latest + timeout-minutes: 20 + steps: + - name: Close unacknowledged duplicates + uses: actions/github-script@f28e40c7f34bde8b3046d885e986cb6290c5673b # v7.1.0 + env: + DRY_RUN: ${{ github.event.inputs.dry_run || 'false' }} + GRACE_PERIOD_DAYS: ${{ github.event.inputs.grace_period_days || '3' }} + LIMIT: ${{ github.event.inputs.limit || '50' }} + with: + script: | + const FLAG_MARKER = ''; + const FLAG_LABEL = 'potential-duplicate'; + const OPTOUT_LABELS = ['do not close', 'keep open', 'not a duplicate']; + + const dryRun = process.env.DRY_RUN === 'true'; + const graceDays = Number(process.env.GRACE_PERIOD_DAYS); + const limit = Number(process.env.LIMIT); + const cutoff = Date.now() - graceDays * 86400000; + const { owner, repo } = context.repo; + + // The oldest issue the flag points at, excluding the issue itself. The + // detector orders candidates by score, not age, so the first one listed + // can be newer than the original report. + const canonicalTarget = (body, self) => { + const refs = new Set(); + for (const [, n] of body.matchAll(/#(\d+)/g)) refs.add(Number(n)); + for (const [, n] of body.matchAll(/github\.com\/[^/\s]+\/[^/\s]+\/issues\/(\d+)/g)) refs.add(Number(n)); + refs.delete(self); + return refs.size ? Math.min(...refs) : null; + }; + + // Only issues the detector labelled: scanning the whole open backlog would + // cost one comments request each and exhaust the token's hourly budget. + const issues = await github.paginate(github.rest.issues.listForRepo, { + owner, repo, state: 'open', labels: FLAG_LABEL, per_page: 100, + }); + core.info(`Scanning ${issues.length} open issues labelled '${FLAG_LABEL}' in ${owner}/${repo}.`); + + const closures = []; + for (const issue of issues) { + const skip = (reason) => core.info(` #${issue.number}: skip, ${reason}`); + const labels = issue.labels.map((l) => (l.name || l).toLowerCase()); + const blocking = OPTOUT_LABELS.find((l) => labels.includes(l)); + if (blocking) { skip(`carries opt-out label '${blocking}'`); continue; } + + const comments = await github.paginate(github.rest.issues.listComments, { + owner, repo, issue_number: issue.number, per_page: 100, + }); + + // Author filter matters: "Quote reply" carries the marker into a human + // comment, and treating that as a fresh flag would restart the clock. + const flags = comments.filter((c) => c.user?.type === 'Bot' && c.body?.includes(FLAG_MARKER)); + if (!flags.length) { skip('never flagged as a potential duplicate'); continue; } + + const flag = flags.reduce((a, b) => (new Date(a.created_at) > new Date(b.created_at) ? a : b)); + const flaggedAt = new Date(flag.created_at).getTime(); + if (flaggedAt > cutoff) { skip(`flagged less than ${graceDays}d ago`); continue; } + + if (comments.some((c) => new Date(c.created_at).getTime() > flaggedAt)) { + skip('someone replied after the flag went up'); continue; + } + + const reactions = await github.paginate(github.rest.reactions.listForIssueComment, { + owner, repo, comment_id: flag.id, per_page: 100, + }); + if (reactions.some((r) => r.content === '-1' && r.user?.login === issue.user?.login)) { + skip('author thumbs-downed the flag'); continue; + } + + const target = canonicalTarget(flag.body, issue.number); + if (target === null) { skip('flag comment names no other issue number'); continue; } + if (target > issue.number) { skip(`only candidate #${target} is newer than this issue`); continue; } + + core.info(` #${issue.number}: unacknowledged duplicate of #${target}`); + closures.push({ number: issue.number, title: issue.title, target }); + } + + const actionable = closures.slice(0, limit); + if (closures.length > limit) { + core.info(`Reached limit ${limit}; ${closures.length - limit} further match(es) left for the next run.`); + } + + for (const { number, target } of actionable) { + if (dryRun) { core.info(` WOULD close #${number} as duplicate of #${target}`); continue; } + core.info(` closing #${number} as duplicate of #${target}`); + await github.rest.issues.createComment({ + owner, repo, issue_number: number, + body: `Closing as a duplicate of #${target}.\n\nThe duplicate notice on this issue went ` + + `unanswered for ${graceDays} days, so it is being closed automatically. If that call is ` + + `wrong, reopen the issue and say how it differs from #${target}, and we will pick it back ` + + `up.\n\n`, + }); + await github.rest.issues.addLabels({ owner, repo, issue_number: number, labels: ['duplicate'] }); + await github.rest.issues.update({ + owner, repo, issue_number: number, state: 'closed', state_reason: 'duplicate', + }); + } + + const heading = dryRun ? 'Would close as duplicates (dry run)' : 'Closed as duplicates'; + const rows = actionable.length + ? ['| Issue | Duplicate of |', '| --- | --- |', + ...actionable.map((c) => `| [#${c.number}](https://github.com/${owner}/${repo}/issues/${c.number}) ${c.title} | #${c.target} |`)] + : ['No issue reached the end of its grace period unacknowledged.']; + await core.summary + .addRaw([`## ${heading}`, '', ...rows, '', `Scanned ${issues.length} flagged issues.`].join('\n')) + .write(); + + core.info(`\n${dryRun ? 'Would close' : 'Closed'}: ${actionable.length}`); From af340c02402a3c0f989c388f217d5d974603809f Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 26 Aug 2026 15:04:35 -0400 Subject: [PATCH 12/49] fix(ci): read duplicate candidates from the marker, not the notice prose The notice interpolates each candidate's title, and the sweep scanned the whole comment for issue references and took the lowest. Titles are attacker-controlled, so filing a candidate titled "... see #1" redirected the closure: any later report matching that candidate would be closed as a duplicate of #1 instead. The detector now emits the candidate numbers as a digits-only field inside the marker, built from the API's number field, and the sweep reads only that. Prose is never parsed, so nothing a reporter can type reaches the target selection. --- .github/workflows/check_duplicate_issues.yml | 2 +- .../workflows/close_stale_duplicate_issues.yml | 17 ++++++++++------- 2 files changed, 11 insertions(+), 8 deletions(-) diff --git a/.github/workflows/check_duplicate_issues.yml b/.github/workflows/check_duplicate_issues.yml index a087007bff3..71d2a3b75eb 100644 --- a/.github/workflows/check_duplicate_issues.yml +++ b/.github/workflows/check_duplicate_issues.yml @@ -24,7 +24,7 @@ jobs: threshold: 0.6 reaction: eyes comment: | - + **Potential duplicate detected** This looks similar to: diff --git a/.github/workflows/close_stale_duplicate_issues.yml b/.github/workflows/close_stale_duplicate_issues.yml index 8bc1af90454..c5f904d2b88 100644 --- a/.github/workflows/close_stale_duplicate_issues.yml +++ b/.github/workflows/close_stale_duplicate_issues.yml @@ -48,7 +48,8 @@ jobs: LIMIT: ${{ github.event.inputs.limit || '50' }} with: script: | - const FLAG_MARKER = ''; + const FLAG_MARKER = '/; const FLAG_LABEL = 'potential-duplicate'; const OPTOUT_LABELS = ['do not close', 'keep open', 'not a duplicate']; @@ -58,13 +59,15 @@ jobs: const cutoff = Date.now() - graceDays * 86400000; const { owner, repo } = context.repo; - // The oldest issue the flag points at, excluding the issue itself. The - // detector orders candidates by score, not age, so the first one listed - // can be newer than the original report. + // Read candidates from the marker's digits-only field, never from the prose. + // Titles are user-controlled and get interpolated into this same comment, so + // scanning the body would let an issue titled "... see #1" redirect a closure + // onto an unrelated report. Take the lowest: the detector orders by score, not + // age, so the first candidate listed can be newer than the original report. const canonicalTarget = (body, self) => { - const refs = new Set(); - for (const [, n] of body.matchAll(/#(\d+)/g)) refs.add(Number(n)); - for (const [, n] of body.matchAll(/github\.com\/[^/\s]+\/[^/\s]+\/issues\/(\d+)/g)) refs.add(Number(n)); + const field = body.match(CANDIDATES); + if (!field) return null; + const refs = new Set(field[1].split(',').filter(Boolean).map(Number)); refs.delete(self); return refs.size ? Math.min(...refs) : null; }; From 3ea11b64e65628f07b47a2f0e0853e79b1ff8334 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 26 Aug 2026 15:18:06 -0400 Subject: [PATCH 13/49] refactor(ci): run the duplicate sweep as a Bun TypeScript script Moves the sweep out of inline workflow JavaScript and into scripts/, following the layout anthropics/claude-code uses for the same job: a checked-out repo, a sha-pinned setup-bun step, and `bun run scripts/auto-close-duplicates.ts`. The script mirrors that repo's file shape, keeping the same request helper, interfaces, per-issue debug logging, and top-level catch, so the two read the same way side by side. Two things stay deliberately different. Candidates come from the notice marker's digits-only field rather than a regex over the comment prose, because titles are attacker-controlled and are interpolated into that same comment. The label is also added on its own endpoint instead of alongside the state change, since sending labels with a PATCH replaces every label already on the issue. --- .github/workflows/auto-close-duplicates.yml | 33 ++ .../close_stale_duplicate_issues.yml | 151 --------- scripts/auto-close-duplicates.ts | 308 ++++++++++++++++++ 3 files changed, 341 insertions(+), 151 deletions(-) create mode 100644 .github/workflows/auto-close-duplicates.yml delete mode 100644 .github/workflows/close_stale_duplicate_issues.yml create mode 100644 scripts/auto-close-duplicates.ts diff --git a/.github/workflows/auto-close-duplicates.yml b/.github/workflows/auto-close-duplicates.yml new file mode 100644 index 00000000000..886aeaaa8e6 --- /dev/null +++ b/.github/workflows/auto-close-duplicates.yml @@ -0,0 +1,33 @@ +name: Auto-close duplicate issues +description: Auto-closes issues that are duplicates of existing issues +on: + schedule: + - cron: "0 9 * * *" + workflow_dispatch: + +jobs: + auto-close-duplicates: + if: github.repository == 'BerriAI/litellm' + runs-on: ubuntu-latest + timeout-minutes: 10 + permissions: + contents: read + issues: write + + steps: + - name: Checkout repository + uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Setup Bun + uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2.2.0 (sha-pinned) + with: + bun-version: latest + + - name: Auto-close duplicate issues + run: bun run scripts/auto-close-duplicates.ts + env: + GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} + GITHUB_REPOSITORY_OWNER: ${{ github.repository_owner }} + GITHUB_REPOSITORY_NAME: ${{ github.event.repository.name }} diff --git a/.github/workflows/close_stale_duplicate_issues.yml b/.github/workflows/close_stale_duplicate_issues.yml deleted file mode 100644 index c5f904d2b88..00000000000 --- a/.github/workflows/close_stale_duplicate_issues.yml +++ /dev/null @@ -1,151 +0,0 @@ -name: Close Stale Duplicate Issues - -# Closes issues that "Check Duplicate Issues" flagged and that nobody acknowledged -# within the grace period. Replying to the flag, thumbs-downing it, or applying an -# opt-out label all keep an issue open. -# -# Dry-run preview (touches nothing): -# gh workflow run "Close Stale Duplicate Issues" -f dry_run=true - -on: - schedule: - # Daily at 09:30 UTC, after the midnight stale sweep and off the hour. - - cron: "30 9 * * *" - workflow_dispatch: - inputs: - dry_run: - description: "Report what would close without touching any issue." - required: false - default: "false" - type: choice - options: - - "true" - - "false" - grace_period_days: - description: "Days to wait after the duplicate flag before closing." - required: false - default: "3" - limit: - description: "Maximum number of issues to close in a single run." - required: false - default: "50" - -permissions: - contents: read - issues: write - -jobs: - close-stale-duplicates: - if: github.repository == 'BerriAI/litellm' - runs-on: ubuntu-latest - timeout-minutes: 20 - steps: - - name: Close unacknowledged duplicates - uses: actions/github-script@f28e40c7f34bde8b3046d885e986cb6290c5673b # v7.1.0 - env: - DRY_RUN: ${{ github.event.inputs.dry_run || 'false' }} - GRACE_PERIOD_DAYS: ${{ github.event.inputs.grace_period_days || '3' }} - LIMIT: ${{ github.event.inputs.limit || '50' }} - with: - script: | - const FLAG_MARKER = '/; - const FLAG_LABEL = 'potential-duplicate'; - const OPTOUT_LABELS = ['do not close', 'keep open', 'not a duplicate']; - - const dryRun = process.env.DRY_RUN === 'true'; - const graceDays = Number(process.env.GRACE_PERIOD_DAYS); - const limit = Number(process.env.LIMIT); - const cutoff = Date.now() - graceDays * 86400000; - const { owner, repo } = context.repo; - - // Read candidates from the marker's digits-only field, never from the prose. - // Titles are user-controlled and get interpolated into this same comment, so - // scanning the body would let an issue titled "... see #1" redirect a closure - // onto an unrelated report. Take the lowest: the detector orders by score, not - // age, so the first candidate listed can be newer than the original report. - const canonicalTarget = (body, self) => { - const field = body.match(CANDIDATES); - if (!field) return null; - const refs = new Set(field[1].split(',').filter(Boolean).map(Number)); - refs.delete(self); - return refs.size ? Math.min(...refs) : null; - }; - - // Only issues the detector labelled: scanning the whole open backlog would - // cost one comments request each and exhaust the token's hourly budget. - const issues = await github.paginate(github.rest.issues.listForRepo, { - owner, repo, state: 'open', labels: FLAG_LABEL, per_page: 100, - }); - core.info(`Scanning ${issues.length} open issues labelled '${FLAG_LABEL}' in ${owner}/${repo}.`); - - const closures = []; - for (const issue of issues) { - const skip = (reason) => core.info(` #${issue.number}: skip, ${reason}`); - const labels = issue.labels.map((l) => (l.name || l).toLowerCase()); - const blocking = OPTOUT_LABELS.find((l) => labels.includes(l)); - if (blocking) { skip(`carries opt-out label '${blocking}'`); continue; } - - const comments = await github.paginate(github.rest.issues.listComments, { - owner, repo, issue_number: issue.number, per_page: 100, - }); - - // Author filter matters: "Quote reply" carries the marker into a human - // comment, and treating that as a fresh flag would restart the clock. - const flags = comments.filter((c) => c.user?.type === 'Bot' && c.body?.includes(FLAG_MARKER)); - if (!flags.length) { skip('never flagged as a potential duplicate'); continue; } - - const flag = flags.reduce((a, b) => (new Date(a.created_at) > new Date(b.created_at) ? a : b)); - const flaggedAt = new Date(flag.created_at).getTime(); - if (flaggedAt > cutoff) { skip(`flagged less than ${graceDays}d ago`); continue; } - - if (comments.some((c) => new Date(c.created_at).getTime() > flaggedAt)) { - skip('someone replied after the flag went up'); continue; - } - - const reactions = await github.paginate(github.rest.reactions.listForIssueComment, { - owner, repo, comment_id: flag.id, per_page: 100, - }); - if (reactions.some((r) => r.content === '-1' && r.user?.login === issue.user?.login)) { - skip('author thumbs-downed the flag'); continue; - } - - const target = canonicalTarget(flag.body, issue.number); - if (target === null) { skip('flag comment names no other issue number'); continue; } - if (target > issue.number) { skip(`only candidate #${target} is newer than this issue`); continue; } - - core.info(` #${issue.number}: unacknowledged duplicate of #${target}`); - closures.push({ number: issue.number, title: issue.title, target }); - } - - const actionable = closures.slice(0, limit); - if (closures.length > limit) { - core.info(`Reached limit ${limit}; ${closures.length - limit} further match(es) left for the next run.`); - } - - for (const { number, target } of actionable) { - if (dryRun) { core.info(` WOULD close #${number} as duplicate of #${target}`); continue; } - core.info(` closing #${number} as duplicate of #${target}`); - await github.rest.issues.createComment({ - owner, repo, issue_number: number, - body: `Closing as a duplicate of #${target}.\n\nThe duplicate notice on this issue went ` - + `unanswered for ${graceDays} days, so it is being closed automatically. If that call is ` - + `wrong, reopen the issue and say how it differs from #${target}, and we will pick it back ` - + `up.\n\n`, - }); - await github.rest.issues.addLabels({ owner, repo, issue_number: number, labels: ['duplicate'] }); - await github.rest.issues.update({ - owner, repo, issue_number: number, state: 'closed', state_reason: 'duplicate', - }); - } - - const heading = dryRun ? 'Would close as duplicates (dry run)' : 'Closed as duplicates'; - const rows = actionable.length - ? ['| Issue | Duplicate of |', '| --- | --- |', - ...actionable.map((c) => `| [#${c.number}](https://github.com/${owner}/${repo}/issues/${c.number}) ${c.title} | #${c.target} |`)] - : ['No issue reached the end of its grace period unacknowledged.']; - await core.summary - .addRaw([`## ${heading}`, '', ...rows, '', `Scanned ${issues.length} flagged issues.`].join('\n')) - .write(); - - core.info(`\n${dryRun ? 'Would close' : 'Closed'}: ${actionable.length}`); diff --git a/scripts/auto-close-duplicates.ts b/scripts/auto-close-duplicates.ts new file mode 100644 index 00000000000..1e94b009a1e --- /dev/null +++ b/scripts/auto-close-duplicates.ts @@ -0,0 +1,308 @@ +#!/usr/bin/env bun + +declare global { + var process: { + env: Record; + }; +} + +interface GitHubIssue { + number: number; + title: string; + user: { login: string }; + labels: { name: string }[]; +} + +interface GitHubComment { + id: number; + body: string; + created_at: string; + user: { type: string }; +} + +interface GitHubReaction { + user: { login: string }; + content: string; +} + +const FLAG_LABEL = "potential-duplicate"; +const FLAG_MARKER = "/; +const GRACE_PERIOD_DAYS = 3; + +async function githubRequest( + endpoint: string, + token: string, + method: string = "GET", + body?: any, +): Promise { + const response = await fetch(`https://api.github.com${endpoint}`, { + method, + headers: { + Authorization: `Bearer ${token}`, + Accept: "application/vnd.github.v3+json", + "User-Agent": "auto-close-duplicates-script", + ...(body && { "Content-Type": "application/json" }), + }, + ...(body && { body: JSON.stringify(body) }), + }); + + if (!response.ok) { + throw new Error( + `GitHub API request failed: ${response.status} ${response.statusText}`, + ); + } + + return response.json(); +} + +function extractDuplicateIssueNumber( + commentBody: string, + issueNumber: number, +): number | null { + // Read candidates from the marker's digits-only field, never from the prose. + // Titles are user-controlled and are interpolated into this same comment, so + // scanning the body would let an issue titled "... see #1" redirect a closure + // onto an unrelated report. + const field = commentBody.match(CANDIDATES); + if (!field) { + return null; + } + + const candidates = field[1] + .split(",") + .filter((value) => value !== "") + .map(Number) + .filter((value) => value !== issueNumber); + + // The detector orders candidates by score, not age, so the first one listed can + // be newer than the original report. Duplicates fold into the earliest issue. + return candidates.length > 0 ? Math.min(...candidates) : null; +} + +async function closeIssueAsDuplicate( + owner: string, + repo: string, + issueNumber: number, + duplicateOfNumber: number, + token: string, +): Promise { + await githubRequest( + `/repos/${owner}/${repo}/issues/${issueNumber}/comments`, + token, + "POST", + { + body: `This issue has been automatically closed as a duplicate of #${duplicateOfNumber}. + +The duplicate notice went unanswered for ${GRACE_PERIOD_DAYS} days. If this is incorrect, please re-open this issue and say how it differs from #${duplicateOfNumber}. + +`, + }, + ); + + // Added on its own endpoint rather than in the PATCH below, because sending + // `labels` with the state change replaces every label on the issue. + await githubRequest( + `/repos/${owner}/${repo}/issues/${issueNumber}/labels`, + token, + "POST", + { labels: ["duplicate"] }, + ); + + await githubRequest( + `/repos/${owner}/${repo}/issues/${issueNumber}`, + token, + "PATCH", + { state: "closed", state_reason: "duplicate" }, + ); +} + +async function autoCloseDuplicates(): Promise { + console.log("[DEBUG] Starting auto-close duplicates script"); + + const token = process.env.GITHUB_TOKEN; + if (!token) { + throw new Error("GITHUB_TOKEN environment variable is required"); + } + console.log("[DEBUG] GitHub token found"); + + const owner = process.env.GITHUB_REPOSITORY_OWNER || "BerriAI"; + const repo = process.env.GITHUB_REPOSITORY_NAME || "litellm"; + console.log(`[DEBUG] Repository: ${owner}/${repo}`); + + const threeDaysAgo = new Date(); + threeDaysAgo.setDate(threeDaysAgo.getDate() - GRACE_PERIOD_DAYS); + console.log( + `[DEBUG] Checking for duplicate comments older than: ${threeDaysAgo.toISOString()}`, + ); + + // Only issues the detector labelled. Walking the whole open backlog would cost a + // comments request per issue, which on a four-figure backlog exhausts the Actions + // token's hourly rate limit for a handful of matches. + console.log(`[DEBUG] Fetching open issues labelled '${FLAG_LABEL}'...`); + const allIssues: GitHubIssue[] = []; + let page = 1; + const perPage = 100; + + while (true) { + const pageIssues: GitHubIssue[] = await githubRequest( + `/repos/${owner}/${repo}/issues?state=open&labels=${FLAG_LABEL}&per_page=${perPage}&page=${page}`, + token, + ); + + if (pageIssues.length === 0) break; + + allIssues.push(...pageIssues); + page++; + + // Safety limit to avoid infinite loops + if (page > 20) break; + } + + const issues = allIssues; + console.log(`[DEBUG] Found ${issues.length} flagged issues`); + + let processedCount = 0; + let candidateCount = 0; + + for (const issue of issues) { + processedCount++; + console.log( + `[DEBUG] Processing issue #${issue.number} (${processedCount}/${issues.length}): ${issue.title}`, + ); + + console.log(`[DEBUG] Fetching comments for issue #${issue.number}...`); + const comments: GitHubComment[] = await githubRequest( + `/repos/${owner}/${repo}/issues/${issue.number}/comments?per_page=100`, + token, + ); + console.log( + `[DEBUG] Issue #${issue.number} has ${comments.length} comments`, + ); + + // The author filter matters: GitHub's "Quote reply" carries the HTML marker into + // a human comment, and treating that as a fresh notice restarts the clock. + const dupeComments = comments.filter( + (comment) => + comment.body.includes(FLAG_MARKER) && comment.user.type === "Bot", + ); + console.log( + `[DEBUG] Issue #${issue.number} has ${dupeComments.length} duplicate detection comments`, + ); + + if (dupeComments.length === 0) { + console.log( + `[DEBUG] Issue #${issue.number} - no duplicate comments found, skipping`, + ); + continue; + } + + const lastDupeComment = dupeComments[dupeComments.length - 1]; + const dupeCommentDate = new Date(lastDupeComment.created_at); + console.log( + `[DEBUG] Issue #${issue.number} - most recent duplicate comment from: ${dupeCommentDate.toISOString()}`, + ); + + if (dupeCommentDate > threeDaysAgo) { + console.log( + `[DEBUG] Issue #${issue.number} - duplicate comment is too recent, skipping`, + ); + continue; + } + console.log( + `[DEBUG] Issue #${issue.number} - duplicate comment is old enough (${Math.floor( + (Date.now() - dupeCommentDate.getTime()) / (1000 * 60 * 60 * 24), + )} days)`, + ); + + const commentsAfterDupe = comments.filter( + (comment) => new Date(comment.created_at) > dupeCommentDate, + ); + console.log( + `[DEBUG] Issue #${issue.number} - ${commentsAfterDupe.length} comments after duplicate detection`, + ); + + if (commentsAfterDupe.length > 0) { + console.log( + `[DEBUG] Issue #${issue.number} - has activity after duplicate comment, skipping`, + ); + continue; + } + + console.log( + `[DEBUG] Issue #${issue.number} - checking reactions on duplicate comment...`, + ); + const reactions: GitHubReaction[] = await githubRequest( + `/repos/${owner}/${repo}/issues/comments/${lastDupeComment.id}/reactions?per_page=100`, + token, + ); + console.log( + `[DEBUG] Issue #${issue.number} - duplicate comment has ${reactions.length} reactions`, + ); + + const authorThumbsDown = reactions.some( + (reaction) => + reaction.user.login === issue.user.login && reaction.content === "-1", + ); + console.log( + `[DEBUG] Issue #${issue.number} - author thumbs down reaction: ${authorThumbsDown}`, + ); + + if (authorThumbsDown) { + console.log( + `[DEBUG] Issue #${issue.number} - author disagreed with duplicate detection, skipping`, + ); + continue; + } + + const duplicateIssueNumber = extractDuplicateIssueNumber( + lastDupeComment.body, + issue.number, + ); + if (!duplicateIssueNumber) { + console.log( + `[DEBUG] Issue #${issue.number} - could not extract duplicate issue number from comment, skipping`, + ); + continue; + } + + if (duplicateIssueNumber > issue.number) { + console.log( + `[DEBUG] Issue #${issue.number} - only candidate #${duplicateIssueNumber} is newer, skipping`, + ); + continue; + } + + candidateCount++; + const issueUrl = `https://github.com/${owner}/${repo}/issues/${issue.number}`; + + try { + console.log( + `[INFO] Auto-closing issue #${issue.number} as duplicate of #${duplicateIssueNumber}: ${issueUrl}`, + ); + await closeIssueAsDuplicate( + owner, + repo, + issue.number, + duplicateIssueNumber, + token, + ); + console.log( + `[SUCCESS] Successfully closed issue #${issue.number} as duplicate of #${duplicateIssueNumber}`, + ); + } catch (error) { + console.error( + `[ERROR] Failed to close issue #${issue.number} as duplicate: ${error}`, + ); + } + } + + console.log( + `[DEBUG] Script completed. Processed ${processedCount} issues, found ${candidateCount} candidates for auto-close`, + ); +} + +autoCloseDuplicates().catch(console.error); + +// Make it a module +export {}; From f4542d960511368eaeab50f68580a89aa13903e6 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 26 Aug 2026 15:26:02 -0400 Subject: [PATCH 14/49] fix(ci): pin the Bun runtime instead of tracking latest The setup step ran `bun-version: latest`, carried over from the upstream layout, and the step after it holds an issues: write token. A compromised Bun release would have executed privileged in that job and could rewrite or close issues. Pinned to 1.4.0, the release the passing runs already resolved to. setup-bun takes no checksum input, so pinning the action by sha and the runtime by exact version is as far as this can be hardened without hand-rolling the download. --- .github/workflows/auto-close-duplicates.yml | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/.github/workflows/auto-close-duplicates.yml b/.github/workflows/auto-close-duplicates.yml index 886aeaaa8e6..ff3b5eff7c2 100644 --- a/.github/workflows/auto-close-duplicates.yml +++ b/.github/workflows/auto-close-duplicates.yml @@ -23,7 +23,10 @@ jobs: - name: Setup Bun uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2.2.0 (sha-pinned) with: - bun-version: latest + # Exact version, never latest: the next step holds an issues: write token, + # so a compromised Bun release would run privileged here. setup-bun exposes + # no checksum input, so pinning the action and the version is the ceiling. + bun-version: "1.4.0" - name: Auto-close duplicate issues run: bun run scripts/auto-close-duplicates.ts From f118511f5562eff67d0473346056d3bd6bbf06e1 Mon Sep 17 00:00:00 2001 From: mubashir1osmani Date: Wed, 26 Aug 2026 15:35:36 -0400 Subject: [PATCH 15/49] fix(ci): honour a duplicate-notice thumbs down from anyone The notice tells every reader that a thumbs down keeps the issue open, but the sweep only counted the reaction when it came from the issue author. A maintainer or another affected user could follow the instruction exactly and still watch the issue close, which made the notice a promise the sweep did not keep. Any thumbs down now spares the issue. That buys back nothing an abuser did not already have: a plain comment stops the clock for anyone, so restricting the reaction only ever penalised people who did what they were told. Drops the issue and reaction author fields, since nothing reads them now. --- scripts/auto-close-duplicates.ts | 16 +++++++--------- 1 file changed, 7 insertions(+), 9 deletions(-) diff --git a/scripts/auto-close-duplicates.ts b/scripts/auto-close-duplicates.ts index 1e94b009a1e..6e361af3e24 100644 --- a/scripts/auto-close-duplicates.ts +++ b/scripts/auto-close-duplicates.ts @@ -9,7 +9,6 @@ declare global { interface GitHubIssue { number: number; title: string; - user: { login: string }; labels: { name: string }[]; } @@ -21,7 +20,6 @@ interface GitHubComment { } interface GitHubReaction { - user: { login: string }; content: string; } @@ -240,17 +238,17 @@ async function autoCloseDuplicates(): Promise { `[DEBUG] Issue #${issue.number} - duplicate comment has ${reactions.length} reactions`, ); - const authorThumbsDown = reactions.some( - (reaction) => - reaction.user.login === issue.user.login && reaction.content === "-1", - ); + // Any thumbs down, not just the author's. The notice tells every reader that a + // 👎 keeps the issue open, and anyone can already stop the clock by commenting, + // so honouring only the author would make the notice a lie without buying safety. + const thumbsDown = reactions.some((reaction) => reaction.content === "-1"); console.log( - `[DEBUG] Issue #${issue.number} - author thumbs down reaction: ${authorThumbsDown}`, + `[DEBUG] Issue #${issue.number} - thumbs down reaction: ${thumbsDown}`, ); - if (authorThumbsDown) { + if (thumbsDown) { console.log( - `[DEBUG] Issue #${issue.number} - author disagreed with duplicate detection, skipping`, + `[DEBUG] Issue #${issue.number} - someone disagreed with duplicate detection, skipping`, ); continue; } From 539bc8ef929e93603e2f99d8b5529cb5eb13a14c Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 14:02:07 -0700 Subject: [PATCH 16/49] fix(ci): close only identical-title duplicates, dry-run the sweep, reopen on reply The merged detector's 0.6 flag threshold had become the close bar, and 6 of the 7 real flagged pairs at 85% or more were not duplicates. The sweep now closes only when an older open issue has the identical normalized title, measures the grace period from the latest bot notice, and leaves the issue open when anyone replies or gives the notice a thumbs down. A reporter cannot reopen an issue the bot closed, so a reporter comment after the automatic close reopens it, drops the duplicate label, and asks for a human look. Manual dispatch defaults to a dry run and takes a grace_period_days input, the runner supplies the repository, the dead python closer is gone, and the decision core has bun tests on a PR-triggered job. --- .github/scripts/close_duplicate_issues.py | 230 -------- .github/workflows/auto-close-duplicates.yml | 59 +- .github/workflows/check_duplicate_issues.yml | 14 +- scripts/auto-close-duplicates.test.ts | 327 +++++++++++ scripts/auto-close-duplicates.ts | 543 +++++++++---------- 5 files changed, 647 insertions(+), 526 deletions(-) delete mode 100755 .github/scripts/close_duplicate_issues.py create mode 100644 scripts/auto-close-duplicates.test.ts diff --git a/.github/scripts/close_duplicate_issues.py b/.github/scripts/close_duplicate_issues.py deleted file mode 100755 index ec522af4f88..00000000000 --- a/.github/scripts/close_duplicate_issues.py +++ /dev/null @@ -1,230 +0,0 @@ -#!/usr/bin/env python3 -""" -Detect and close duplicate GitHub issues using title similarity. - -Modes: - --scan Compare all open issues against each other (batch) - --issue-number N Check a single issue against older open issues - -Requires the `gh` CLI to be authenticated. -""" - -import argparse -import difflib -import json -import re -import subprocess -import sys - - -def normalize_title(title: str) -> str: - """Strip common prefixes, lowercase, and collapse whitespace.""" - title = re.sub( - r"^\[?(bug|feature request|enhancement|question|docs)[:\]]?\s*", - "", - title, - flags=re.IGNORECASE, - ) - return " ".join(title.lower().split()) - - -def gh(*args: str) -> str: - """Run a gh CLI command and return stdout.""" - result = subprocess.run( - ["gh", *args], - capture_output=True, - text=True, - check=True, - ) - return result.stdout - - -def fetch_open_issues(repo: str | None) -> list[dict]: - """Fetch all open issues (excluding PRs) via gh api --paginate.""" - if repo: - endpoint = ( - f"repos/{repo}/issues?state=open&per_page=100&sort=created&direction=asc" - ) - else: - endpoint = "repos/{owner}/{repo}/issues?state=open&per_page=100&sort=created&direction=asc" - cmd = ["api", "--paginate", endpoint] - - raw = gh(*cmd) - # gh --paginate concatenates JSON arrays, so we may get multiple arrays - issues = [] - for line in raw.strip().splitlines(): - line = line.strip() - if not line: - continue - parsed = json.loads(line) - if isinstance(parsed, list): - issues.extend(parsed) - else: - issues.append(parsed) - - # Filter out pull requests (they also appear in the issues endpoint) - return [i for i in issues if "pull_request" not in i] - - -def close_as_duplicate( - issue_number: int, duplicate_of: int, repo: str | None, dry_run: bool -) -> None: - """Close an issue as duplicate of another, adding a comment and label.""" - repo_args = ["--repo", repo] if repo else [] - - if dry_run: - print( - f" [DRY RUN] Would close #{issue_number} as duplicate of #{duplicate_of}" - ) - return - - # Add comment - comment_body = ( - f"Closing as duplicate of #{duplicate_of}.\n\n" - "If you believe this is not a duplicate, please reopen and add context " - "explaining how this differs." - ) - gh("issue", "comment", str(issue_number), "--body", comment_body, *repo_args) - - # Add label - gh("issue", "edit", str(issue_number), "--add-label", "duplicate", *repo_args) - - # Close with not_planned reason - gh( - "api", - f"repos/{repo or '{owner}/{repo}'}/issues/{issue_number}", - "-X", - "PATCH", - "-f", - "state=closed", - "-f", - "state_reason=not_planned", - ) - - print(f" Closed #{issue_number} as duplicate of #{duplicate_of}") - - -def find_duplicate( - issue: dict, candidates: list[dict], threshold: float -) -> dict | None: - """Return the first candidate whose normalized title is above threshold.""" - norm = normalize_title(issue["title"]) - for candidate in candidates: - if candidate["number"] == issue["number"]: - continue - cand_norm = normalize_title(candidate["title"]) - ratio = difflib.SequenceMatcher(None, norm, cand_norm).ratio() - if ratio >= threshold: - return candidate - return None - - -def scan_all( - issues: list[dict], threshold: float, repo: str | None, dry_run: bool -) -> int: - """Compare every issue against all older issues. Returns count of duplicates found.""" - # Sort oldest first - issues.sort(key=lambda i: i["number"]) - closed_count = 0 - - for idx, issue in enumerate(issues): - older = issues[:idx] - if not older: - continue - dup = find_duplicate(issue, older, threshold) - if dup: - ratio = difflib.SequenceMatcher( - None, - normalize_title(issue["title"]), - normalize_title(dup["title"]), - ).ratio() - print( - f"#{issue['number']}: \"{issue['title']}\"\n" - f" -> duplicate of #{dup['number']}: \"{dup['title']}\" " - f"({ratio:.0%} similar)" - ) - close_as_duplicate(issue["number"], dup["number"], repo, dry_run) - closed_count += 1 - - return closed_count - - -def check_single( - issue_number: int, - issues: list[dict], - threshold: float, - repo: str | None, - dry_run: bool, -) -> bool: - """Check a single issue against all older open issues. Returns True if duplicate found.""" - target = None - for i in issues: - if i["number"] == issue_number: - target = i - break - - if target is None: - print(f"Issue #{issue_number} not found among open issues.") - return False - - older = [i for i in issues if i["number"] < issue_number] - dup = find_duplicate(target, older, threshold) - if dup: - ratio = difflib.SequenceMatcher( - None, - normalize_title(target["title"]), - normalize_title(dup["title"]), - ).ratio() - print( - f"#{target['number']}: \"{target['title']}\"\n" - f" -> duplicate of #{dup['number']}: \"{dup['title']}\" " - f"({ratio:.0%} similar)" - ) - close_as_duplicate(issue_number, dup["number"], repo, dry_run) - return True - - print(f"#{issue_number}: no duplicate found above threshold {threshold}") - return False - - -def main() -> None: - parser = argparse.ArgumentParser( - description="Detect and close duplicate GitHub issues" - ) - mode = parser.add_mutually_exclusive_group(required=True) - mode.add_argument("--scan", action="store_true", help="Scan all open issues") - mode.add_argument("--issue-number", type=int, help="Check a single issue number") - parser.add_argument( - "--threshold", type=float, default=0.85, help="Similarity threshold (0-1)" - ) - parser.add_argument( - "--close", - action="store_true", - help="Actually close duplicates (default is dry-run)", - ) - parser.add_argument( - "--repo", type=str, help="Repository (owner/repo). Auto-detected if omitted." - ) - args = parser.parse_args() - - dry_run = not args.close - - if dry_run: - print("=== DRY RUN MODE (pass --close to actually close issues) ===\n") - - print("Fetching open issues...") - issues = fetch_open_issues(args.repo) - print(f"Found {len(issues)} open issues.\n") - - if args.scan: - count = scan_all(issues, args.threshold, args.repo, dry_run) - print(f"\nTotal duplicates {'found' if dry_run else 'closed'}: {count}") - else: - found = check_single( - args.issue_number, issues, args.threshold, args.repo, dry_run - ) - sys.exit(0 if found else 0) # Always exit 0; finding no dup is not an error - - -if __name__ == "__main__": - main() diff --git a/.github/workflows/auto-close-duplicates.yml b/.github/workflows/auto-close-duplicates.yml index ff3b5eff7c2..d8256917805 100644 --- a/.github/workflows/auto-close-duplicates.yml +++ b/.github/workflows/auto-close-duplicates.yml @@ -1,19 +1,33 @@ name: Auto-close duplicate issues -description: Auto-closes issues that are duplicates of existing issues + on: schedule: - cron: "0 9 * * *" workflow_dispatch: + inputs: + dry_run: + description: Log which issues would close without closing anything + type: boolean + default: true + grace_period_days: + description: Days a duplicate notice must go unanswered before the close + type: number + default: 3 + pull_request: + paths: + - .github/workflows/auto-close-duplicates.yml + - scripts/auto-close-duplicates.ts + - scripts/auto-close-duplicates.test.ts + +permissions: {} jobs: - auto-close-duplicates: - if: github.repository == 'BerriAI/litellm' + test: + if: github.event_name == 'pull_request' runs-on: ubuntu-latest - timeout-minutes: 10 + timeout-minutes: 5 permissions: contents: read - issues: write - steps: - name: Checkout repository uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 @@ -21,16 +35,35 @@ jobs: persist-credentials: false - name: Setup Bun - uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2.2.0 (sha-pinned) + uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2.2.0 with: - # Exact version, never latest: the next step holds an issues: write token, - # so a compromised Bun release would run privileged here. setup-bun exposes - # no checksum input, so pinning the action and the version is the ceiling. bun-version: "1.4.0" - - name: Auto-close duplicate issues + - name: Test the sweep + run: bun test scripts/auto-close-duplicates.test.ts + + sweep: + if: github.event_name != 'pull_request' && github.repository == 'BerriAI/litellm' + runs-on: ubuntu-latest + timeout-minutes: 10 + permissions: + contents: read + issues: write + steps: + - name: Checkout repository + uses: actions/checkout@08eba0b27e820071cde6df949e0beb9ba4906955 # v4.3.0 + with: + persist-credentials: false + + - name: Setup Bun + uses: oven-sh/setup-bun@0c5077e51419868618aeaa5fe8019c62421857d6 # v2.2.0 + with: + # Exact version, never latest: the next step holds an issues: write token + bun-version: "1.4.0" + + - name: Close unanswered duplicates, reopen ones the reporter answered run: bun run scripts/auto-close-duplicates.ts env: GITHUB_TOKEN: ${{ secrets.GITHUB_TOKEN }} - GITHUB_REPOSITORY_OWNER: ${{ github.repository_owner }} - GITHUB_REPOSITORY_NAME: ${{ github.event.repository.name }} + DRY_RUN: ${{ inputs.dry_run == true }} + GRACE_PERIOD_DAYS: ${{ inputs.grace_period_days }} diff --git a/.github/workflows/check_duplicate_issues.yml b/.github/workflows/check_duplicate_issues.yml index 71d2a3b75eb..41ec43a1d9b 100644 --- a/.github/workflows/check_duplicate_issues.yml +++ b/.github/workflows/check_duplicate_issues.yml @@ -1,17 +1,19 @@ name: Check Duplicate Issues -# Flags newly opened issues that look like existing ones. Flagging only: the actual -# close happens 3 days later in "Close Stale Duplicate Issues", and only if nobody -# replied to the comment posted here. The HTML marker below is the handshake between -# the two workflows, so keep it in the template. +# Flagging only. "Auto-close duplicate issues" closes a flagged issue 3 days later, +# and only when its title is identical to an older open issue and nobody replied. +# The HTML marker below is the handshake between the two, so keep it in the template. on: issues: types: [opened, edited] +permissions: {} + jobs: check-duplicate: runs-on: ubuntu-latest + timeout-minutes: 5 permissions: issues: write contents: read @@ -29,7 +31,7 @@ jobs: This looks similar to: {{#issues}} - - #{{number}} - {{title}} ({{accuracy}}% similar) + - #{{number}} - {{title}} {{/issues}} - This issue will close automatically in 3 days unless someone responds. If it is a duplicate, please 👍 the existing issue and follow along there. If it is not, comment here or 👎 this comment and it stays open. + If this is a duplicate, add a thumbs-up reaction to the existing issue and follow along there. When the title is identical to an older open issue, this issue closes automatically in 3 days unless someone responds. If it is not a duplicate, comment here or add a thumbs-down reaction to this comment and it stays open. diff --git a/scripts/auto-close-duplicates.test.ts b/scripts/auto-close-duplicates.test.ts new file mode 100644 index 00000000000..6a7a3a507bd --- /dev/null +++ b/scripts/auto-close-duplicates.test.ts @@ -0,0 +1,327 @@ +import { describe, expect, test } from "bun:test"; + +import { + CLOSED_MARKER, + REOPEN_COMMENT, + candidateNumbers, + duplicateTarget, + normalizeTitle, + pendingNotice, + readConfig, + reopenTarget, + sweepClosedIssue, + sweepIssue, + type Comment, + type GitHubApi, + type Issue, + type SweepConfig, +} from "./auto-close-duplicates"; + +const NOW = new Date("2026-09-04T09:00:00Z"); +const DAY_MS = 24 * 60 * 60 * 1000; +const daysAgo = (days: number): string => new Date(NOW.getTime() - days * DAY_MS).toISOString(); + +const issue = (number: number, title: string, overrides: Partial = {}): Issue => ({ + number, + title, + state: "open", + user: { login: "reporter" }, + ...overrides, +}); + +const notice = (candidates: readonly number[], createdAt: string, overrides: Partial = {}): Comment => ({ + id: 900, + body: `\n**Potential duplicate detected**`, + created_at: createdAt, + user: { type: "Bot", login: "github-actions[bot]" }, + ...overrides, +}); + +const humanComment = (createdAt: string, body = "It is not the same thing", login = "reporter"): Comment => ({ + id: 901, + body, + created_at: createdAt, + user: { type: "User", login }, +}); + +const config: SweepConfig = { repo: "BerriAI/litellm", graceDays: 3, dryRun: false, now: NOW }; + +describe("normalizeTitle", () => { + test("drops the template prefix, case, and punctuation", () => { + expect(normalizeTitle("[Bug]: Gemma 4-e4b fails on Vertex!")).toBe("gemma 4 e4b fails on vertex"); + expect(normalizeTitle("[Feature]: ")).toBe(""); + }); +}); + +describe("candidateNumbers", () => { + test("reads only the marker field, keeps older issues, sorted ascending and deduplicated", () => { + const body = "\n- #1 - see #1 (100% similar)"; + expect(candidateNumbers(body, 35)).toEqual([10, 30]); + }); + + test("returns nothing without the marker", () => { + expect(candidateNumbers("- #1 - looks like #1", 35)).toEqual([]); + }); +}); + +describe("pendingNotice", () => { + test("waits out the grace period from the latest notice", () => { + const fresh = pendingNotice(issue(35, "t"), [notice([10], daysAgo(2.9))], config); + expect(fresh.kind).toBe("skip"); + const aged = pendingNotice(issue(35, "t"), [notice([10], daysAgo(3.1))], config); + expect(aged.kind).toBe("pending"); + const reposted = pendingNotice( + issue(35, "t"), + [notice([10], daysAgo(6)), notice([10], daysAgo(1), { id: 902 })], + config, + ); + expect(reposted.kind).toBe("skip"); + }); + + test("a zero-day grace period acts on the notice at once", () => { + const verdict = pendingNotice(issue(35, "t"), [notice([10], daysAgo(0.01))], { ...config, graceDays: 0 }); + expect(verdict.kind).toBe("pending"); + }); + + test("a human reply after the notice keeps the issue open, a bot reply does not", () => { + const human = pendingNotice(issue(35, "t"), [notice([10], daysAgo(5)), humanComment(daysAgo(4))], config); + expect(human).toEqual({ kind: "skip", reason: "someone replied after the notice" }); + const bot = pendingNotice( + issue(35, "t"), + [notice([10], daysAgo(5)), { id: 903, body: "triage", created_at: daysAgo(4), user: { type: "Bot", login: "triage[bot]" } }], + config, + ); + expect(bot.kind).toBe("pending"); + }); + + test("a human quoting the marker is not a notice", () => { + const quoted = pendingNotice(issue(35, "t"), [notice([10], daysAgo(5), { user: { type: "User", login: "reporter" } })], config); + expect(quoted).toEqual({ kind: "skip", reason: "carries no duplicate notice" }); + }); + + test("never closes an issue twice: a reopened issue is left alone", () => { + const reopened = pendingNotice( + issue(35, "t"), + [notice([10], daysAgo(9)), { id: 904, body: `Closed automatically\n\n${CLOSED_MARKER}`, created_at: daysAgo(5), user: { type: "Bot", login: "github-actions[bot]" } }], + config, + ); + expect(reopened).toEqual({ kind: "skip", reason: "was reopened after an automatic close" }); + }); + + test("skips pull requests and issues whose only candidates are newer", () => { + expect(pendingNotice(issue(35, "t", { pull_request: {} }), [notice([10], daysAgo(5))], config).kind).toBe("skip"); + expect(pendingNotice(issue(35, "t"), [notice([40], daysAgo(5))], config)).toEqual({ + kind: "skip", + reason: "no candidate is older than this issue", + }); + }); +}); + +describe("duplicateTarget", () => { + const reporter = issue(35, "[Bug]: Gemma 4-e4b fails on Vertex"); + + test("closes only against the earliest open issue with the identical normalized title", () => { + const verdict = duplicateTarget( + reporter, + [issue(10, "[Bug]: Gemma 4-e4n fails on Vertex"), issue(20, "[bug]: gemma 4-e4b fails on vertex"), issue(30, "[Bug]: Gemma 4-e4b fails on Vertex")], + [], + ); + expect(verdict).toEqual({ kind: "close", duplicateOf: 20 }); + }); + + test("a near miss in the title is not a duplicate", () => { + const verdict = duplicateTarget(reporter, [issue(10, "[Bug]: Gemma 4-e4n fails on Vertex")], []); + expect(verdict).toEqual({ kind: "skip", reason: "no older open issue has the identical title" }); + }); + + test("bare template titles never match each other", () => { + const verdict = duplicateTarget(issue(35, "[Bug]: "), [issue(10, "[Bug]: ")], []); + expect(verdict.kind).toBe("skip"); + expect(verdict.kind === "skip" && verdict.reason).toContain("too short"); + }); + + test("a closed candidate or a pull request is never the target", () => { + expect(duplicateTarget(reporter, [issue(10, reporter.title, { state: "closed" })], []).kind).toBe("skip"); + expect(duplicateTarget(reporter, [issue(10, reporter.title, { pull_request: {} })], []).kind).toBe("skip"); + }); + + test("a thumbs down on the notice keeps the issue open", () => { + const verdict = duplicateTarget(reporter, [issue(10, reporter.title)], [{ content: "+1" }, { content: "-1" }]); + expect(verdict).toEqual({ kind: "skip", reason: "someone gave the notice a thumbs down" }); + }); +}); + +describe("sweepIssue", () => { + const reporter = issue(35, "[Bug]: Gemma 4-e4b fails on Vertex"); + const original = issue(10, "[Bug]: Gemma 4-e4b fails on Vertex"); + + function fakeApi(): { readonly api: GitHubApi; readonly writes: readonly string[] } { + const writes: string[] = []; + const api: GitHubApi = { + request: async (method: string, path: string, body?: object): Promise => { + if (method !== "GET") { + writes.push(`${method} ${path} ${JSON.stringify(body)}`); + return {} as T; + } + if (path.startsWith("/repos/BerriAI/litellm/issues/35/comments")) { + return [notice([10], daysAgo(5))] as T; + } + if (path.startsWith("/repos/BerriAI/litellm/issues/comments/900/reactions")) { + return [] as T; + } + if (path === "/repos/BerriAI/litellm/issues/10") { + return original as T; + } + throw new Error(`unexpected GET ${path}`); + }, + }; + return { api, writes }; + } + + test("a dry run reports the close and writes nothing", async () => { + const { api, writes } = fakeApi(); + const verdict = await sweepIssue(api, { ...config, dryRun: true }, reporter); + expect(verdict).toEqual({ kind: "close", duplicateOf: 10 }); + expect(writes).toEqual([]); + }); + + test("a real run comments, labels, then closes with the duplicate reason", async () => { + const { api, writes } = fakeApi(); + const verdict = await sweepIssue(api, config, reporter); + expect(verdict).toEqual({ kind: "close", duplicateOf: 10 }); + expect(writes.map((write) => write.split(" ").slice(0, 2).join(" "))).toEqual([ + "POST /repos/BerriAI/litellm/issues/35/comments", + "POST /repos/BerriAI/litellm/issues/35/labels", + "PATCH /repos/BerriAI/litellm/issues/35", + ]); + expect(writes[0]).toContain("duplicate of #10"); + expect(writes[0]).toContain("unanswered for 3 days"); + expect(writes[0]).toContain(CLOSED_MARKER); + expect(writes[1]).toContain('{"labels":["duplicate"]}'); + expect(writes[2]).toContain('{"state":"closed","state_reason":"duplicate"}'); + }); +}); + +describe("reopenTarget", () => { + const closedByBot = (overrides: Partial = {}): Issue => + issue(35, "t", { state: "closed", closed_by: { type: "Bot" }, ...overrides }); + const closeMarker = (createdAt: string): Comment => ({ + id: 905, + body: `Closed automatically as a duplicate of #10.\n\n${CLOSED_MARKER}`, + created_at: createdAt, + user: { type: "Bot", login: "github-actions[bot]" }, + }); + + test("a reporter reply after the automatic close reopens", () => { + const verdict = reopenTarget(closedByBot(), [closeMarker(daysAgo(2)), humanComment(daysAgo(1))]); + expect(verdict).toEqual({ kind: "reopen" }); + }); + + test("an issue closed by a person stays closed", () => { + const verdict = reopenTarget(closedByBot({ closed_by: { type: "User" } }), [ + closeMarker(daysAgo(2)), + humanComment(daysAgo(1)), + ]); + expect(verdict).toEqual({ kind: "skip", reason: "was closed by a person" }); + }); + + test("without the automatic-close marker nothing reopens", () => { + const verdict = reopenTarget(closedByBot(), [humanComment(daysAgo(1))]); + expect(verdict).toEqual({ kind: "skip", reason: "carries no automatic-close marker" }); + }); + + test("a maintainer reply alone does not reopen", () => { + const verdict = reopenTarget(closedByBot(), [ + closeMarker(daysAgo(2)), + humanComment(daysAgo(1), "Confirmed duplicate", "maintainer"), + ]); + expect(verdict).toEqual({ kind: "skip", reason: "the reporter has not replied since the close" }); + }); + + test("a reporter comment from before the close does not reopen", () => { + const verdict = reopenTarget(closedByBot(), [humanComment(daysAgo(3)), closeMarker(daysAgo(2))]); + expect(verdict).toEqual({ kind: "skip", reason: "the reporter has not replied since the close" }); + }); + + test("a pull request never reopens", () => { + const verdict = reopenTarget(closedByBot({ pull_request: {} }), [closeMarker(daysAgo(2)), humanComment(daysAgo(1))]); + expect(verdict).toEqual({ kind: "skip", reason: "is a pull request" }); + }); +}); + +describe("sweepClosedIssue", () => { + function fakeApi(issueBody: Issue, comments: readonly Comment[]): { readonly api: GitHubApi; readonly writes: readonly string[] } { + const writes: string[] = []; + const api: GitHubApi = { + request: async (method: string, path: string, body?: object): Promise => { + if (method !== "GET") { + writes.push(`${method} ${path} ${JSON.stringify(body)}`); + return {} as T; + } + if (path.startsWith("/repos/BerriAI/litellm/issues/35/comments")) { + return comments as T; + } + if (path === "/repos/BerriAI/litellm/issues/35") { + return issueBody as T; + } + throw new Error(`unexpected GET ${path}`); + }, + }; + return { api, writes }; + } + + const closedByBot = issue(35, "t", { state: "closed", closed_by: { type: "Bot" } }); + const closeMarker: Comment = { + id: 905, + body: `Closed automatically as a duplicate of #10.\n\n${CLOSED_MARKER}`, + created_at: daysAgo(2), + user: { type: "Bot", login: "github-actions[bot]" }, + }; + + test("a real run unlabels, reopens, then explains", async () => { + const { api, writes } = fakeApi(closedByBot, [closeMarker, humanComment(daysAgo(1))]); + const verdict = await sweepClosedIssue(api, config, 35); + expect(verdict).toEqual({ kind: "reopen" }); + expect(writes).toEqual([ + "DELETE /repos/BerriAI/litellm/issues/35/labels/duplicate undefined", + 'PATCH /repos/BerriAI/litellm/issues/35 {"state":"open"}', + `POST /repos/BerriAI/litellm/issues/35/comments {"body":"${REOPEN_COMMENT}"}`, + ]); + }); + + test("a dry run reports the reopen and writes nothing", async () => { + const { api, writes } = fakeApi(closedByBot, [closeMarker, humanComment(daysAgo(1))]); + const verdict = await sweepClosedIssue(api, { ...config, dryRun: true }, 35); + expect(verdict).toEqual({ kind: "reopen" }); + expect(writes).toEqual([]); + }); +}); + +describe("readConfig", () => { + test("defaults to a real run with a 3-day grace period", () => { + const parsed = readConfig({ GITHUB_TOKEN: "t", GITHUB_REPOSITORY: "BerriAI/litellm" }, NOW); + expect(parsed).toEqual({ token: "t", repo: "BerriAI/litellm", graceDays: 3, dryRun: false, now: NOW }); + }); + + test("honors DRY_RUN and GRACE_PERIOD_DAYS overrides", () => { + const parsed = readConfig( + { GITHUB_TOKEN: "t", GITHUB_REPOSITORY: "o/r", DRY_RUN: "true", GRACE_PERIOD_DAYS: "0" }, + NOW, + ); + expect(parsed.dryRun).toBe(true); + expect(parsed.graceDays).toBe(0); + }); + + test("an empty GRACE_PERIOD_DAYS, as a schedule run renders it, means the default", () => { + const parsed = readConfig({ GITHUB_TOKEN: "t", GITHUB_REPOSITORY: "o/r", GRACE_PERIOD_DAYS: "" }, NOW); + expect(parsed.graceDays).toBe(3); + }); + + test("refuses a missing token, a malformed repository, or a bad grace period", () => { + expect(() => readConfig({ GITHUB_REPOSITORY: "o/r" }, NOW)).toThrow("GITHUB_TOKEN"); + expect(() => readConfig({ GITHUB_TOKEN: "t", GITHUB_REPOSITORY: "litellm" }, NOW)).toThrow("owner/repo"); + expect(() => readConfig({ GITHUB_TOKEN: "t", GITHUB_REPOSITORY: "o/r", GRACE_PERIOD_DAYS: "-1" }, NOW)).toThrow( + "GRACE_PERIOD_DAYS", + ); + }); +}); diff --git a/scripts/auto-close-duplicates.ts b/scripts/auto-close-duplicates.ts index 6e361af3e24..941f281efe6 100644 --- a/scripts/auto-close-duplicates.ts +++ b/scripts/auto-close-duplicates.ts @@ -1,306 +1,295 @@ #!/usr/bin/env bun -declare global { - var process: { - env: Record; - }; +declare const process: { readonly env: Readonly> }; + +export interface Issue { + readonly number: number; + readonly title: string; + readonly state: string; + readonly user: { readonly login: string }; + readonly closed_by?: { readonly type: string } | null; + readonly pull_request?: unknown; } -interface GitHubIssue { - number: number; - title: string; - labels: { name: string }[]; +export interface Comment { + readonly id: number; + readonly body: string; + readonly created_at: string; + readonly user: { readonly type: string; readonly login: string }; } -interface GitHubComment { - id: number; - body: string; - created_at: string; - user: { type: string }; +export interface Reaction { + readonly content: string; } -interface GitHubReaction { - content: string; +export interface GitHubApi { + readonly request: (method: "GET" | "POST" | "PATCH" | "DELETE", path: string, body?: object) => Promise; } -const FLAG_LABEL = "potential-duplicate"; -const FLAG_MARKER = "/; -const GRACE_PERIOD_DAYS = 3; - -async function githubRequest( - endpoint: string, - token: string, - method: string = "GET", - body?: any, -): Promise { - const response = await fetch(`https://api.github.com${endpoint}`, { - method, - headers: { - Authorization: `Bearer ${token}`, - Accept: "application/vnd.github.v3+json", - "User-Agent": "auto-close-duplicates-script", - ...(body && { "Content-Type": "application/json" }), - }, - ...(body && { body: JSON.stringify(body) }), - }); - - if (!response.ok) { - throw new Error( - `GitHub API request failed: ${response.status} ${response.statusText}`, - ); - } - - return response.json(); +export interface SweepConfig { + readonly repo: string; + readonly graceDays: number; + readonly dryRun: boolean; + readonly now: Date; } -function extractDuplicateIssueNumber( - commentBody: string, - issueNumber: number, -): number | null { - // Read candidates from the marker's digits-only field, never from the prose. - // Titles are user-controlled and are interpolated into this same comment, so - // scanning the body would let an issue titled "... see #1" redirect a closure - // onto an unrelated report. - const field = commentBody.match(CANDIDATES); +export type NoticeVerdict = + | { readonly kind: "pending"; readonly notice: Comment; readonly candidates: readonly number[] } + | { readonly kind: "skip"; readonly reason: string }; + +export type CloseVerdict = + | { readonly kind: "close"; readonly duplicateOf: number } + | { readonly kind: "skip"; readonly reason: string }; + +export type ReopenVerdict = + | { readonly kind: "reopen" } + | { readonly kind: "skip"; readonly reason: string }; + +export const FLAG_LABEL = "potential-duplicate"; +export const CLOSED_MARKER = ""; +export const DEFAULT_GRACE_DAYS = 3; +export const REOPEN_COMMENT = + "Reopened automatically: the reporter replied after the duplicate close, so this needs a human look."; +const NOTICE_MARKER = //; +const MIN_TITLE_WORDS = 3; +const PAGE_SIZE = 100; +const DAY_MS = 24 * 60 * 60 * 1000; +const REOPEN_LOOKBACK_DAYS = 30; + +const skip = (reason: string): { readonly kind: "skip"; readonly reason: string } => ({ kind: "skip", reason }); + +export function normalizeTitle(title: string): string { + return title + .toLowerCase() + .replace(/^\s*\[[^\]]*\]\s*:?/, "") + .replace(/[^a-z0-9]+/g, " ") + .trim(); +} + +export function candidateNumbers(noticeBody: string, issueNumber: number): readonly number[] { + const field = noticeBody.match(NOTICE_MARKER); if (!field) { - return null; + return []; } - - const candidates = field[1] + const older = field[1] .split(",") .filter((value) => value !== "") .map(Number) - .filter((value) => value !== issueNumber); - - // The detector orders candidates by score, not age, so the first one listed can - // be newer than the original report. Duplicates fold into the earliest issue. - return candidates.length > 0 ? Math.min(...candidates) : null; + .filter((candidate) => candidate < issueNumber); + return [...new Set(older)].sort((a, b) => a - b); } -async function closeIssueAsDuplicate( - owner: string, - repo: string, +export function pendingNotice( + issue: Issue, + comments: readonly Comment[], + config: Pick, +): NoticeVerdict { + if (issue.pull_request !== undefined) { + return skip("is a pull request"); + } + if (comments.some((comment) => comment.body.includes(CLOSED_MARKER))) { + return skip("was reopened after an automatic close"); + } + const notices = comments.filter((comment) => comment.user.type === "Bot" && NOTICE_MARKER.test(comment.body)); + const notice = notices[notices.length - 1]; + if (notice === undefined) { + return skip("carries no duplicate notice"); + } + const noticeAt = new Date(notice.created_at); + const ageDays = (config.now.getTime() - noticeAt.getTime()) / DAY_MS; + if (ageDays < config.graceDays) { + return skip(`notice is ${ageDays.toFixed(1)} days old, grace period is ${config.graceDays}`); + } + if (comments.some((comment) => comment.user.type !== "Bot" && new Date(comment.created_at) > noticeAt)) { + return skip("someone replied after the notice"); + } + const candidates = candidateNumbers(notice.body, issue.number); + if (candidates.length === 0) { + return skip("no candidate is older than this issue"); + } + return { kind: "pending", notice, candidates }; +} + +export function duplicateTarget( + issue: Issue, + candidates: readonly Issue[], + reactions: readonly Reaction[], +): CloseVerdict { + if (reactions.some((reaction) => reaction.content === "-1")) { + return skip("someone gave the notice a thumbs down"); + } + const title = normalizeTitle(issue.title); + if (title.split(" ").length < MIN_TITLE_WORDS) { + return skip(`title "${issue.title}" is too short to match on`); + } + const original = candidates.find( + (candidate) => + candidate.state === "open" && candidate.pull_request === undefined && normalizeTitle(candidate.title) === title, + ); + if (original === undefined) { + return skip("no older open issue has the identical title"); + } + return { kind: "close", duplicateOf: original.number }; +} + +export function reopenTarget(issue: Issue, comments: readonly Comment[]): ReopenVerdict { + if (issue.pull_request !== undefined) { + return skip("is a pull request"); + } + if (issue.closed_by?.type !== "Bot") { + return skip("was closed by a person"); + } + const marker = comments.find((comment) => comment.body.includes(CLOSED_MARKER)); + if (marker === undefined) { + return skip("carries no automatic-close marker"); + } + const markerAt = new Date(marker.created_at); + if (!comments.some((comment) => comment.user.login === issue.user.login && new Date(comment.created_at) > markerAt)) { + return skip("the reporter has not replied since the close"); + } + return { kind: "reopen" }; +} + +export function closingComment(duplicateOf: number, graceDays: number): string { + return `Closed automatically as a duplicate of #${duplicateOf}. Its title is identical to that older open issue and the duplicate notice above went unanswered for ${graceDays} days. If this is wrong, comment here with how it differs from #${duplicateOf} and this issue will be reopened automatically within a day. + +${CLOSED_MARKER}`; +} + +async function listAll(api: GitHubApi, path: string, page = 1): Promise { + const separator = path.includes("?") ? "&" : "?"; + const batch = await api.request("GET", `${path}${separator}per_page=${PAGE_SIZE}&page=${page}`); + return batch.length < PAGE_SIZE ? batch : [...batch, ...(await listAll(api, path, page + 1))]; +} + +async function closeAsDuplicate( + api: GitHubApi, + config: SweepConfig, issueNumber: number, - duplicateOfNumber: number, - token: string, + duplicateOf: number, ): Promise { - await githubRequest( - `/repos/${owner}/${repo}/issues/${issueNumber}/comments`, - token, - "POST", - { - body: `This issue has been automatically closed as a duplicate of #${duplicateOfNumber}. + const issuePath = `/repos/${config.repo}/issues/${issueNumber}`; + await api.request("POST", `${issuePath}/comments`, { body: closingComment(duplicateOf, config.graceDays) }); + await api.request("POST", `${issuePath}/labels`, { labels: ["duplicate"] }); + await api.request("PATCH", issuePath, { state: "closed", state_reason: "duplicate" }); +} -The duplicate notice went unanswered for ${GRACE_PERIOD_DAYS} days. If this is incorrect, please re-open this issue and say how it differs from #${duplicateOfNumber}. +async function reopenForReporter(api: GitHubApi, config: SweepConfig, issueNumber: number): Promise { + const issuePath = `/repos/${config.repo}/issues/${issueNumber}`; + await api.request("DELETE", `${issuePath}/labels/duplicate`); + await api.request("PATCH", issuePath, { state: "open" }); + await api.request("POST", `${issuePath}/comments`, { body: REOPEN_COMMENT }); +} -`, +export async function sweepClosedIssue(api: GitHubApi, config: SweepConfig, issueNumber: number): Promise { + const issue = await api.request("GET", `/repos/${config.repo}/issues/${issueNumber}`); + const comments = await listAll(api, `/repos/${config.repo}/issues/${issueNumber}/comments`); + const verdict = reopenTarget(issue, comments); + if (verdict.kind === "reopen" && !config.dryRun) { + await reopenForReporter(api, config, issueNumber); + } + return verdict; +} + +export async function sweepIssue(api: GitHubApi, config: SweepConfig, issue: Issue): Promise { + const comments = await listAll(api, `/repos/${config.repo}/issues/${issue.number}/comments`); + const pending = pendingNotice(issue, comments, config); + if (pending.kind === "skip") { + return pending; + } + const reactions = await listAll(api, `/repos/${config.repo}/issues/comments/${pending.notice.id}/reactions`); + const candidates = await Promise.all( + pending.candidates.map((candidate) => api.request("GET", `/repos/${config.repo}/issues/${candidate}`)), + ); + const verdict = duplicateTarget(issue, candidates, reactions); + if (verdict.kind === "close" && !config.dryRun) { + await closeAsDuplicate(api, config, issue.number, verdict.duplicateOf); + } + return verdict; +} + +function describe(issue: Issue, verdict: CloseVerdict, dryRun: boolean): string { + if (verdict.kind === "skip") { + return `#${issue.number}: skipped, ${verdict.reason}`; + } + return `#${issue.number}: ${dryRun ? "would close" : "closed"} as a duplicate of #${verdict.duplicateOf}`; +} + +export async function sweep(api: GitHubApi, config: SweepConfig): Promise { + const issues = await listAll(api, `/repos/${config.repo}/issues?state=open&labels=${FLAG_LABEL}`); + console.log(`${issues.length} open issues carry the ${FLAG_LABEL} label in ${config.repo}${config.dryRun ? " (dry run)" : ""}`); + return issues.reduce>(async (previous, issue) => { + const verdicts = await previous; + const verdict = await sweepIssue(api, config, issue); + console.log(describe(issue, verdict, config.dryRun)); + return [...verdicts, verdict]; + }, Promise.resolve([])); +} + +function describeReopen(issueNumber: number, verdict: ReopenVerdict, dryRun: boolean): string { + if (verdict.kind === "skip") { + return `#${issueNumber}: skipped, ${verdict.reason}`; + } + return `#${issueNumber}: ${dryRun ? "would reopen" : "reopened"} for the reporter's reply`; +} + +export async function reopenSweep(api: GitHubApi, config: SweepConfig): Promise { + const since = new Date(config.now.getTime() - REOPEN_LOOKBACK_DAYS * DAY_MS).toISOString(); + const closedPath = `/repos/${config.repo}/issues?state=closed&labels=duplicate,${FLAG_LABEL}&since=${encodeURIComponent(since)}`; + const issues = await listAll(api, closedPath); + console.log(`${issues.length} recently closed issues carry the duplicate and ${FLAG_LABEL} labels in ${config.repo}${config.dryRun ? " (dry run)" : ""}`); + return issues.reduce>(async (previous, issue) => { + const verdicts = await previous; + const verdict = await sweepClosedIssue(api, config, issue.number); + console.log(describeReopen(issue.number, verdict, config.dryRun)); + return [...verdicts, verdict]; + }, Promise.resolve([])); +} + +export function readConfig(env: Readonly>, now: Date): SweepConfig & { readonly token: string } { + const token = env.GITHUB_TOKEN; + const repo = env.GITHUB_REPOSITORY; + if (!token || !repo || !/^[\w.-]+\/[\w.-]+$/.test(repo)) { + throw new Error("GITHUB_TOKEN and GITHUB_REPOSITORY (owner/repo) are required"); + } + const rawGraceDays = env.GRACE_PERIOD_DAYS?.trim(); + const graceDays = rawGraceDays === undefined || rawGraceDays === "" ? DEFAULT_GRACE_DAYS : Number(rawGraceDays); + if (!Number.isFinite(graceDays) || graceDays < 0) { + throw new Error(`GRACE_PERIOD_DAYS must be a non-negative number, got "${env.GRACE_PERIOD_DAYS}"`); + } + return { token, repo, graceDays, dryRun: env.DRY_RUN === "true", now }; +} + +export function githubApi(token: string): GitHubApi { + return { + request: async (method: "GET" | "POST" | "PATCH" | "DELETE", path: string, body?: object): Promise => { + const response = await fetch(`https://api.github.com${path}`, { + method, + headers: { + Authorization: `Bearer ${token}`, + Accept: "application/vnd.github+json", + "X-GitHub-Api-Version": "2022-11-28", + "User-Agent": "litellm-auto-close-duplicates", + ...(body === undefined ? {} : { "Content-Type": "application/json" }), + }, + body: body === undefined ? undefined : JSON.stringify(body), + }); + if (!response.ok) { + throw new Error(`${method} ${path} failed: ${response.status} ${response.statusText}`); + } + return (await response.json()) as T; }, - ); - - // Added on its own endpoint rather than in the PATCH below, because sending - // `labels` with the state change replaces every label on the issue. - await githubRequest( - `/repos/${owner}/${repo}/issues/${issueNumber}/labels`, - token, - "POST", - { labels: ["duplicate"] }, - ); - - await githubRequest( - `/repos/${owner}/${repo}/issues/${issueNumber}`, - token, - "PATCH", - { state: "closed", state_reason: "duplicate" }, - ); + }; } -async function autoCloseDuplicates(): Promise { - console.log("[DEBUG] Starting auto-close duplicates script"); - - const token = process.env.GITHUB_TOKEN; - if (!token) { - throw new Error("GITHUB_TOKEN environment variable is required"); - } - console.log("[DEBUG] GitHub token found"); - - const owner = process.env.GITHUB_REPOSITORY_OWNER || "BerriAI"; - const repo = process.env.GITHUB_REPOSITORY_NAME || "litellm"; - console.log(`[DEBUG] Repository: ${owner}/${repo}`); - - const threeDaysAgo = new Date(); - threeDaysAgo.setDate(threeDaysAgo.getDate() - GRACE_PERIOD_DAYS); +if (import.meta.main) { + const { token, ...config } = readConfig(process.env, new Date()); + const api = githubApi(token); + const closeVerdicts = await sweep(api, config); + const reopenVerdicts = await reopenSweep(api, config); + const closed = closeVerdicts.filter((verdict) => verdict.kind === "close").length; + const reopened = reopenVerdicts.filter((verdict) => verdict.kind === "reopen").length; console.log( - `[DEBUG] Checking for duplicate comments older than: ${threeDaysAgo.toISOString()}`, - ); - - // Only issues the detector labelled. Walking the whole open backlog would cost a - // comments request per issue, which on a four-figure backlog exhausts the Actions - // token's hourly rate limit for a handful of matches. - console.log(`[DEBUG] Fetching open issues labelled '${FLAG_LABEL}'...`); - const allIssues: GitHubIssue[] = []; - let page = 1; - const perPage = 100; - - while (true) { - const pageIssues: GitHubIssue[] = await githubRequest( - `/repos/${owner}/${repo}/issues?state=open&labels=${FLAG_LABEL}&per_page=${perPage}&page=${page}`, - token, - ); - - if (pageIssues.length === 0) break; - - allIssues.push(...pageIssues); - page++; - - // Safety limit to avoid infinite loops - if (page > 20) break; - } - - const issues = allIssues; - console.log(`[DEBUG] Found ${issues.length} flagged issues`); - - let processedCount = 0; - let candidateCount = 0; - - for (const issue of issues) { - processedCount++; - console.log( - `[DEBUG] Processing issue #${issue.number} (${processedCount}/${issues.length}): ${issue.title}`, - ); - - console.log(`[DEBUG] Fetching comments for issue #${issue.number}...`); - const comments: GitHubComment[] = await githubRequest( - `/repos/${owner}/${repo}/issues/${issue.number}/comments?per_page=100`, - token, - ); - console.log( - `[DEBUG] Issue #${issue.number} has ${comments.length} comments`, - ); - - // The author filter matters: GitHub's "Quote reply" carries the HTML marker into - // a human comment, and treating that as a fresh notice restarts the clock. - const dupeComments = comments.filter( - (comment) => - comment.body.includes(FLAG_MARKER) && comment.user.type === "Bot", - ); - console.log( - `[DEBUG] Issue #${issue.number} has ${dupeComments.length} duplicate detection comments`, - ); - - if (dupeComments.length === 0) { - console.log( - `[DEBUG] Issue #${issue.number} - no duplicate comments found, skipping`, - ); - continue; - } - - const lastDupeComment = dupeComments[dupeComments.length - 1]; - const dupeCommentDate = new Date(lastDupeComment.created_at); - console.log( - `[DEBUG] Issue #${issue.number} - most recent duplicate comment from: ${dupeCommentDate.toISOString()}`, - ); - - if (dupeCommentDate > threeDaysAgo) { - console.log( - `[DEBUG] Issue #${issue.number} - duplicate comment is too recent, skipping`, - ); - continue; - } - console.log( - `[DEBUG] Issue #${issue.number} - duplicate comment is old enough (${Math.floor( - (Date.now() - dupeCommentDate.getTime()) / (1000 * 60 * 60 * 24), - )} days)`, - ); - - const commentsAfterDupe = comments.filter( - (comment) => new Date(comment.created_at) > dupeCommentDate, - ); - console.log( - `[DEBUG] Issue #${issue.number} - ${commentsAfterDupe.length} comments after duplicate detection`, - ); - - if (commentsAfterDupe.length > 0) { - console.log( - `[DEBUG] Issue #${issue.number} - has activity after duplicate comment, skipping`, - ); - continue; - } - - console.log( - `[DEBUG] Issue #${issue.number} - checking reactions on duplicate comment...`, - ); - const reactions: GitHubReaction[] = await githubRequest( - `/repos/${owner}/${repo}/issues/comments/${lastDupeComment.id}/reactions?per_page=100`, - token, - ); - console.log( - `[DEBUG] Issue #${issue.number} - duplicate comment has ${reactions.length} reactions`, - ); - - // Any thumbs down, not just the author's. The notice tells every reader that a - // 👎 keeps the issue open, and anyone can already stop the clock by commenting, - // so honouring only the author would make the notice a lie without buying safety. - const thumbsDown = reactions.some((reaction) => reaction.content === "-1"); - console.log( - `[DEBUG] Issue #${issue.number} - thumbs down reaction: ${thumbsDown}`, - ); - - if (thumbsDown) { - console.log( - `[DEBUG] Issue #${issue.number} - someone disagreed with duplicate detection, skipping`, - ); - continue; - } - - const duplicateIssueNumber = extractDuplicateIssueNumber( - lastDupeComment.body, - issue.number, - ); - if (!duplicateIssueNumber) { - console.log( - `[DEBUG] Issue #${issue.number} - could not extract duplicate issue number from comment, skipping`, - ); - continue; - } - - if (duplicateIssueNumber > issue.number) { - console.log( - `[DEBUG] Issue #${issue.number} - only candidate #${duplicateIssueNumber} is newer, skipping`, - ); - continue; - } - - candidateCount++; - const issueUrl = `https://github.com/${owner}/${repo}/issues/${issue.number}`; - - try { - console.log( - `[INFO] Auto-closing issue #${issue.number} as duplicate of #${duplicateIssueNumber}: ${issueUrl}`, - ); - await closeIssueAsDuplicate( - owner, - repo, - issue.number, - duplicateIssueNumber, - token, - ); - console.log( - `[SUCCESS] Successfully closed issue #${issue.number} as duplicate of #${duplicateIssueNumber}`, - ); - } catch (error) { - console.error( - `[ERROR] Failed to close issue #${issue.number} as duplicate: ${error}`, - ); - } - } - - console.log( - `[DEBUG] Script completed. Processed ${processedCount} issues, found ${candidateCount} candidates for auto-close`, + `${config.dryRun ? "Would close" : "Closed"} ${closed} of ${closeVerdicts.length} flagged issues, ${config.dryRun ? "would reopen" : "reopened"} ${reopened}`, ); } - -autoCloseDuplicates().catch(console.error); - -// Make it a module -export {}; From ed5761daef4ae17152446d182c860630c38b7268 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sat, 29 Aug 2026 14:28:47 -0700 Subject: [PATCH 17/49] fix(ci): keep earlier objections when the duplicate notice is re-posted The detector fires on issue edits and posts a fresh notice each time, so the sweep now counts replies from the first notice on and a thumbs down on any notice --- scripts/auto-close-duplicates.test.ts | 29 +++++++++++++++++++++++---- scripts/auto-close-duplicates.ts | 23 ++++++++++++--------- 2 files changed, 39 insertions(+), 13 deletions(-) diff --git a/scripts/auto-close-duplicates.test.ts b/scripts/auto-close-duplicates.test.ts index 6a7a3a507bd..b49bf05cbc2 100644 --- a/scripts/auto-close-duplicates.test.ts +++ b/scripts/auto-close-duplicates.test.ts @@ -14,6 +14,7 @@ import { type Comment, type GitHubApi, type Issue, + type Reaction, type SweepConfig, } from "./auto-close-duplicates"; @@ -78,6 +79,15 @@ describe("pendingNotice", () => { expect(reposted.kind).toBe("skip"); }); + test("an objection posted before a re-posted notice still keeps the issue open", () => { + const verdict = pendingNotice( + issue(35, "t"), + [notice([10], daysAgo(10)), humanComment(daysAgo(7)), notice([10], daysAgo(4), { id: 902 })], + config, + ); + expect(verdict).toEqual({ kind: "skip", reason: "someone replied after the notice" }); + }); + test("a zero-day grace period acts on the notice at once", () => { const verdict = pendingNotice(issue(35, "t"), [notice([10], daysAgo(0.01))], { ...config, graceDays: 0 }); expect(verdict.kind).toBe("pending"); @@ -155,7 +165,10 @@ describe("sweepIssue", () => { const reporter = issue(35, "[Bug]: Gemma 4-e4b fails on Vertex"); const original = issue(10, "[Bug]: Gemma 4-e4b fails on Vertex"); - function fakeApi(): { readonly api: GitHubApi; readonly writes: readonly string[] } { + function fakeApi( + comments: readonly Comment[] = [notice([10], daysAgo(5))], + reactionsByNotice: Readonly> = {}, + ): { readonly api: GitHubApi; readonly writes: readonly string[] } { const writes: string[] = []; const api: GitHubApi = { request: async (method: string, path: string, body?: object): Promise => { @@ -164,10 +177,11 @@ describe("sweepIssue", () => { return {} as T; } if (path.startsWith("/repos/BerriAI/litellm/issues/35/comments")) { - return [notice([10], daysAgo(5))] as T; + return comments as T; } - if (path.startsWith("/repos/BerriAI/litellm/issues/comments/900/reactions")) { - return [] as T; + const reactionsPath = path.match(/^\/repos\/BerriAI\/litellm\/issues\/comments\/(\d+)\/reactions/); + if (reactionsPath) { + return (reactionsByNotice[Number(reactionsPath[1])] ?? []) as T; } if (path === "/repos/BerriAI/litellm/issues/10") { return original as T; @@ -185,6 +199,13 @@ describe("sweepIssue", () => { expect(writes).toEqual([]); }); + test("a thumbs down on an earlier notice still keeps the issue open", async () => { + const { api, writes } = fakeApi([notice([10], daysAgo(9)), notice([10], daysAgo(5), { id: 902 })], { 900: [{ content: "-1" }] }); + const verdict = await sweepIssue(api, config, reporter); + expect(verdict).toEqual({ kind: "skip", reason: "someone gave the notice a thumbs down" }); + expect(writes).toEqual([]); + }); + test("a real run comments, labels, then closes with the duplicate reason", async () => { const { api, writes } = fakeApi(); const verdict = await sweepIssue(api, config, reporter); diff --git a/scripts/auto-close-duplicates.ts b/scripts/auto-close-duplicates.ts index 941f281efe6..c595104d886 100644 --- a/scripts/auto-close-duplicates.ts +++ b/scripts/auto-close-duplicates.ts @@ -34,7 +34,7 @@ export interface SweepConfig { } export type NoticeVerdict = - | { readonly kind: "pending"; readonly notice: Comment; readonly candidates: readonly number[] } + | { readonly kind: "pending"; readonly notices: readonly Comment[]; readonly candidates: readonly number[] } | { readonly kind: "skip"; readonly reason: string }; export type CloseVerdict = @@ -91,23 +91,24 @@ export function pendingNotice( return skip("was reopened after an automatic close"); } const notices = comments.filter((comment) => comment.user.type === "Bot" && NOTICE_MARKER.test(comment.body)); - const notice = notices[notices.length - 1]; - if (notice === undefined) { + const first = notices[0]; + const latest = notices[notices.length - 1]; + if (first === undefined || latest === undefined) { return skip("carries no duplicate notice"); } - const noticeAt = new Date(notice.created_at); - const ageDays = (config.now.getTime() - noticeAt.getTime()) / DAY_MS; + const ageDays = (config.now.getTime() - new Date(latest.created_at).getTime()) / DAY_MS; if (ageDays < config.graceDays) { return skip(`notice is ${ageDays.toFixed(1)} days old, grace period is ${config.graceDays}`); } - if (comments.some((comment) => comment.user.type !== "Bot" && new Date(comment.created_at) > noticeAt)) { + const firstNoticeAt = new Date(first.created_at); + if (comments.some((comment) => comment.user.type !== "Bot" && new Date(comment.created_at) > firstNoticeAt)) { return skip("someone replied after the notice"); } - const candidates = candidateNumbers(notice.body, issue.number); + const candidates = candidateNumbers(latest.body, issue.number); if (candidates.length === 0) { return skip("no candidate is older than this issue"); } - return { kind: "pending", notice, candidates }; + return { kind: "pending", notices, candidates }; } export function duplicateTarget( @@ -197,7 +198,11 @@ export async function sweepIssue(api: GitHubApi, config: SweepConfig, issue: Iss if (pending.kind === "skip") { return pending; } - const reactions = await listAll(api, `/repos/${config.repo}/issues/comments/${pending.notice.id}/reactions`); + const reactions = ( + await Promise.all( + pending.notices.map((notice) => listAll(api, `/repos/${config.repo}/issues/comments/${notice.id}/reactions`)), + ) + ).flat(); const candidates = await Promise.all( pending.candidates.map((candidate) => api.request("GET", `/repos/${config.repo}/issues/${candidate}`)), ); From 3ea4b715ba8d6c1666205443a12475d74a1af040 Mon Sep 17 00:00:00 2001 From: siyoon Date: Sun, 30 Aug 2026 14:22:41 +0900 Subject: [PATCH 18/49] feat(friendli): add zai-org/GLM-5.3-Flash model pricing Per https://api.friendli.ai/serverless/v1/models: - $0.15 input / $0.50 output / $0.03 cached input per MTok - 1M context, 1M max output, reasoning with effort low/high/max (per HF chat_template.jinja: low/high honored, anything else -> max) - tool calling, parallel tool calls, structured output, prompt caching - image + video input (native multimodal) --- model_prices_and_context_window.json | 25 +++++++++++++ ...t_friendli_glm_5_3_flash_model_metadata.py | 37 +++++++++++++++++++ 2 files changed, 62 insertions(+) create mode 100644 tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 05c1cfd3179..fd6b9e11adf 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -19564,6 +19564,31 @@ "supports_system_messages": true, "supports_tool_choice": true }, + "friendliai/zai-org/GLM-5.3-Flash": { + "litellm_provider": "friendliai", + "supports_reasoning": true, + "supports_function_calling": true, + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "max_output_tokens": 1048576, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 5e-07, + "cache_read_input_token_cost": 3e-08, + "supports_prompt_caching": true, + "supports_max_reasoning_effort": true, + "supports_low_reasoning_effort": true, + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "mode": "chat", + "comment": "Native multimodal GLM model for efficient coding and long-horizon agent tasks", + "source": "https://api.friendli.ai/serverless/v1/models", + "supports_vision": true, + "supports_image_input": true, + "supports_video_input": true + }, "ft:babbage-002": { "deprecation_date": "2026-10-23", "input_cost_per_token": 1.6e-06, diff --git a/tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py b/tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py new file mode 100644 index 00000000000..5110307e5bb --- /dev/null +++ b/tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py @@ -0,0 +1,37 @@ +import json +from pathlib import Path + +from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider + + +def test_friendli_glm_5_3_flash_model_info(): + model = "friendliai/zai-org/GLM-5.3-Flash" + json_path = Path(__file__).parents[2] / "model_prices_and_context_window.json" + with open(json_path) as f: + model_cost = json.load(f) + + info = model_cost.get(model) + assert ( + info is not None + ), f"{model} not found in model_prices_and_context_window.json" + assert info["litellm_provider"] == "friendliai" + assert info["mode"] == "chat" + assert info["input_cost_per_token"] == 1.5e-07 + assert info["output_cost_per_token"] == 5e-07 + assert info["cache_read_input_token_cost"] == 3e-08 + assert info["max_input_tokens"] == 1048576 + assert info["max_output_tokens"] == 1048576 + assert info["supports_function_calling"] is True + assert info["supports_reasoning"] is True + assert info["supports_low_reasoning_effort"] is True + assert info["supports_max_reasoning_effort"] is True + assert info["supports_tool_choice"] is True + assert info["supports_prompt_caching"] is True + # Friendli serves image and video input + assert info["supports_vision"] is True + assert info["supports_image_input"] is True + assert info["supports_video_input"] is True + + routed_model, provider, _, _ = get_llm_provider(model=model) + assert routed_model == "zai-org/GLM-5.3-Flash" + assert provider == "friendliai" From e9f1af84730ffd1c11b9f77a10c38d0bf20351be Mon Sep 17 00:00:00 2001 From: siyoon Date: Sun, 30 Aug 2026 14:29:28 +0900 Subject: [PATCH 19/49] chore(tests): drop redundant capability comment Per greptile review + CLAUDE.md comment policy: the comment restated the immediately following assertions without adding value. --- tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py | 1 - 1 file changed, 1 deletion(-) diff --git a/tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py b/tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py index 5110307e5bb..0acd750e49a 100644 --- a/tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py +++ b/tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py @@ -27,7 +27,6 @@ def test_friendli_glm_5_3_flash_model_info(): assert info["supports_max_reasoning_effort"] is True assert info["supports_tool_choice"] is True assert info["supports_prompt_caching"] is True - # Friendli serves image and video input assert info["supports_vision"] is True assert info["supports_image_input"] is True assert info["supports_video_input"] is True From 4fb1440747d704140beb558b54a21520aa3c57d3 Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 30 Aug 2026 07:40:56 +0000 Subject: [PATCH 20/49] docs(proxy): clarify spend semantics on /v2/user/info and /user/daily/activity Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- .../internal_user_endpoints.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 9a98bdbb6b1..73e993b37a1 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -997,6 +997,13 @@ async def user_info_v2( This is the v2 replacement for /user/info, designed to avoid the "god endpoint" problem where the old endpoint loaded all keys and teams into memory. + Note on `spend`: this is the user's running budget counter, which is zeroed by the + budget reset job whenever `budget_reset_at` elapses (see `budget_duration`). It is NOT + lifetime or per-period historical spend. For historical spend over a date range, use + `/user/daily/activity` or `/user/daily/activity/aggregated`, which read immutable daily + spend records that are never reset. The two values are expected to diverge once a + budget reset has occurred within the queried period. + Access control: - Proxy admins can query any user - Team admins can query users within their teams @@ -2687,6 +2694,10 @@ async def get_user_daily_activity( Meant to optimize querying spend data for analytics for a user. + Reads immutable daily spend records, which are never affected by budget resets. + This can legitimately exceed the `spend` field returned by `/v2/user/info`, which + is a running budget counter zeroed on every budget reset. + Returns: (by date) - spend @@ -2800,6 +2811,10 @@ async def get_user_daily_activity_aggregated( """ Aggregated analytics for a user's daily activity without pagination. Returns the same response shape as the paginated endpoint with page metadata set to single-page. + + Reads immutable daily spend records, which are never affected by budget resets. + This can legitimately exceed the `spend` field returned by `/v2/user/info`, which + is a running budget counter zeroed on every budget reset. """ from litellm.proxy.proxy_server import prisma_client From 36c53e1288051aaef506f8bbfde0722aabc3fa3b Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Sun, 30 Aug 2026 07:46:45 +0000 Subject: [PATCH 21/49] chore(ui): regenerate schema.d.ts for updated endpoint descriptions Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ui/litellm-dashboard/src/lib/http/schema.d.ts | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 137c67e837c..b4526834b98 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16210,6 +16210,10 @@ export interface paths { * * Meant to optimize querying spend data for analytics for a user. * + * Reads immutable daily spend records, which are never affected by budget resets. + * This can legitimately exceed the `spend` field returned by `/v2/user/info`, which + * is a running budget counter zeroed on every budget reset. + * * Returns: * (by date) * - spend @@ -16241,6 +16245,10 @@ export interface paths { * Get User Daily Activity Aggregated * @description Aggregated analytics for a user's daily activity without pagination. * Returns the same response shape as the paginated endpoint with page metadata set to single-page. + * + * Reads immutable daily spend records, which are never affected by budget resets. + * This can legitimately exceed the `spend` field returned by `/v2/user/info`, which + * is a running budget counter zeroed on every budget reset. */ get: operations["get_user_daily_activity_aggregated_user_daily_activity_aggregated_get"]; put?: never; @@ -21004,6 +21012,13 @@ export interface paths { * This is the v2 replacement for /user/info, designed to avoid the "god endpoint" problem * where the old endpoint loaded all keys and teams into memory. * + * Note on `spend`: this is the user's running budget counter, which is zeroed by the + * budget reset job whenever `budget_reset_at` elapses (see `budget_duration`). It is NOT + * lifetime or per-period historical spend. For historical spend over a date range, use + * `/user/daily/activity` or `/user/daily/activity/aggregated`, which read immutable daily + * spend records that are never reset. The two values are expected to diverge once a + * budget reset has occurred within the queried period. + * * Access control: * - Proxy admins can query any user * - Team admins can query users within their teams From 14f392bb9b10bdcee7e134e8132f1b070e913e0a Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Sun, 30 Aug 2026 13:11:32 -0700 Subject: [PATCH 22/49] fix(docker): bump wolfi-base for glibc 2.44 and pin apk python to 3.13 --- Dockerfile | 15 ++++++++------- backend/Dockerfile | 12 ++++++------ docker/Dockerfile.database | 15 ++++++++------- docker/Dockerfile.non_root | 17 +++++++++-------- gateway/Dockerfile | 12 ++++++------ 5 files changed, 37 insertions(+), 34 deletions(-) diff --git a/Dockerfile b/Dockerfile index 700b0d6525e..b3ee85e9ed1 100644 --- a/Dockerfile +++ b/Dockerfile @@ -1,10 +1,10 @@ # syntax=docker/dockerfile:1.7 # Base image for building -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d # Runtime image -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a # Pinned by digest like the other base images; bump explicitly on Node upgrades. ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43 @@ -40,8 +40,8 @@ COPY --from=uvbin /uvx /usr/local/bin/uvx RUN apk add --no-cache \ bash \ gcc \ - python3 \ - python3-dev \ + python-3.13 \ + python-3.13-dev \ rust \ openssl \ openssl-dev \ @@ -51,6 +51,7 @@ RUN apk add --no-cache \ ENV UV_PROJECT_ENVIRONMENT=/app/.venv \ UV_LINK_MODE=copy \ + UV_PYTHON_DOWNLOADS=0 \ PATH="/app/.venv/bin:${PATH}" # Copy dependency metadata first for layer caching @@ -65,7 +66,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr --extra extra_proxy \ --extra semantic-router \ --extra saml \ - --python python3 + --python python3.13 # Copy full source tree COPY . . @@ -86,7 +87,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \ --extra extra_proxy \ --extra semantic-router \ --extra saml \ - --python python3 + --python python3.13 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ npm_config_cache=/root/.npm \ @@ -101,7 +102,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime USER root # node (without npm) is required by the prisma CLI at runtime -RUN apk add --no-cache bash openssl tzdata nodejs python3 libsndfile +RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile WORKDIR /app ENV PATH="/app/.venv/bin:${PATH}" \ diff --git a/backend/Dockerfile b/backend/Dockerfile index 4ca40944606..aa01b9fba8b 100644 --- a/backend/Dockerfile +++ b/backend/Dockerfile @@ -1,5 +1,5 @@ -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72 -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a FROM $UV_IMAGE AS uvbin @@ -16,7 +16,7 @@ COPY --from=uvbin /uv /uvx /usr/local/bin/ # instead of nodeenv downloading one whose dynamic deps may not be in Wolfi # (e.g. Node 26.2.0 needs libatomic). Retry for transient apk.cgr.dev flakes. RUN for i in 1 2 3; do \ - apk add --no-cache bash gcc python3 python3-dev openssl openssl-dev libsndfile nodejs npm && break; \ + apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl openssl-dev libsndfile nodejs npm && break; \ [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ sleep 5; \ done @@ -46,7 +46,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --extra proxy-runtime \ --extra extra_proxy \ --extra semantic-router \ - --python python3 + --python python3.13 # Stage 2 — copy source and install the project + workspace members. COPY . . @@ -57,7 +57,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --extra proxy-runtime \ --extra extra_proxy \ --extra semantic-router \ - --python python3 + --python python3.13 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ npm_config_cache=/root/.npm \ @@ -71,7 +71,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime USER root RUN for i in 1 2 3; do \ - apk add --no-cache bash openssl tzdata python3 libsndfile libatomic && break; \ + apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic && break; \ [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ sleep 5; \ done diff --git a/docker/Dockerfile.database b/docker/Dockerfile.database index f0d6d02fccf..c1348f68231 100644 --- a/docker/Dockerfile.database +++ b/docker/Dockerfile.database @@ -1,10 +1,10 @@ # syntax=docker/dockerfile:1.7 # Base image for building -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d # Runtime image -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72 +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a # Pinned by digest like the other base images; bump explicitly on Node upgrades. ARG UI_BUILD_IMAGE=node:24.19-alpine3.24@sha256:d32cdf619f63fe0471182d08996dd516c6275bb5fd31ae06e55a570bd9e1ad43 @@ -39,8 +39,8 @@ COPY --from=uvbin /uvx /usr/local/bin/uvx RUN apk add --no-cache \ bash \ gcc \ - python3 \ - python3-dev \ + python-3.13 \ + python-3.13-dev \ openssl \ openssl-dev \ nodejs \ @@ -49,6 +49,7 @@ RUN apk add --no-cache \ ENV UV_PROJECT_ENVIRONMENT=/app/.venv \ UV_LINK_MODE=copy \ + UV_PYTHON_DOWNLOADS=0 \ PATH="/app/.venv/bin:${PATH}" # Copy dependency metadata first for layer caching @@ -63,7 +64,7 @@ RUN uv sync --frozen --no-install-project --no-install-workspace --no-default-gr --extra extra_proxy \ --extra semantic-router \ --extra saml \ - --python python3 + --python python3.13 # Copy full source tree COPY . . @@ -84,7 +85,7 @@ RUN uv sync --frozen --no-default-groups --no-editable \ --extra extra_proxy \ --extra semantic-router \ --extra saml \ - --python python3 + --python python3.13 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ npm_config_cache=/root/.npm \ @@ -98,7 +99,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime USER root # node (without npm) is required by the prisma CLI at runtime -RUN apk add --no-cache bash openssl tzdata nodejs python3 libsndfile +RUN apk add --no-cache bash openssl tzdata nodejs python-3.13 libsndfile WORKDIR /app ENV PATH="/app/.venv/bin:${PATH}" \ diff --git a/docker/Dockerfile.non_root b/docker/Dockerfile.non_root index 4a5df6ecd69..2221435a83a 100644 --- a/docker/Dockerfile.non_root +++ b/docker/Dockerfile.non_root @@ -1,8 +1,8 @@ # syntax=docker/dockerfile:1.7 # Base images -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72 -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d ARG PROXY_EXTRAS_SOURCE=published ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a # Pinned by digest like the other base images; bump explicitly on Node upgrades. @@ -37,8 +37,8 @@ COPY --from=uvbin /uvx /usr/local/bin/uvx RUN for i in 1 2 3; do \ apk add --no-cache \ - python3 \ - python3-dev \ + python-3.13 \ + python-3.13-dev \ gcc \ rust \ bash \ @@ -52,6 +52,7 @@ RUN for i in 1 2 3; do \ ENV UV_PROJECT_ENVIRONMENT=/app/.venv \ UV_LINK_MODE=copy \ + UV_PYTHON_DOWNLOADS=0 \ PATH="/app/.venv/bin:${PATH}" \ LITELLM_NON_ROOT=true \ XDG_CACHE_HOME=/app/.cache @@ -69,7 +70,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ --extra extra_proxy \ --extra semantic-router \ --extra saml \ - --python python3 + --python python3.13 # Copy full source tree COPY . . @@ -96,7 +97,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ --extra extra_proxy \ --extra semantic-router \ --extra saml \ - --python python3 \ + --python python3.13 \ --no-sources-package litellm-proxy-extras; \ else \ uv sync --frozen --no-default-groups --no-editable \ @@ -105,7 +106,7 @@ RUN --mount=type=cache,target=/app/.cache/uv,id=litellm-uv-cache \ --extra extra_proxy \ --extra semantic-router \ --extra saml \ - --python python3; \ + --python python3.13; \ fi RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ @@ -124,7 +125,7 @@ RUN for i in 1 2 3; do \ apk upgrade --no-cache && break || sleep 5; \ done && \ for i in 1 2 3; do \ - apk add --no-cache python3 bash openssl tzdata libsndfile nodejs && break || sleep 5; \ + apk add --no-cache python-3.13 bash openssl tzdata libsndfile nodejs && break || sleep 5; \ done # Copy only what runtime needs. The application is installed inside the venv; diff --git a/gateway/Dockerfile b/gateway/Dockerfile index 4a2e32e186e..308d70a6b26 100644 --- a/gateway/Dockerfile +++ b/gateway/Dockerfile @@ -1,5 +1,5 @@ -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72 -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a FROM $UV_IMAGE AS uvbin @@ -16,7 +16,7 @@ COPY --from=uvbin /uv /uvx /usr/local/bin/ # instead of nodeenv downloading one whose dynamic deps may not be in Wolfi # (e.g. Node 26.2.0 needs libatomic). Retry for transient apk.cgr.dev flakes. RUN for i in 1 2 3; do \ - apk add --no-cache bash gcc python3 python3-dev openssl openssl-dev libsndfile nodejs npm && break; \ + apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl openssl-dev libsndfile nodejs npm && break; \ [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ sleep 5; \ done @@ -47,7 +47,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --extra extra_proxy \ --extra semantic-router \ --extra bedrock-realtime \ - --python python3 + --python python3.13 # Stage 2 — copy source and install the project + workspace members. COPY . . @@ -59,7 +59,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ --extra extra_proxy \ --extra semantic-router \ --extra bedrock-realtime \ - --python python3 + --python python3.13 RUN HOME=/opt/prisma XDG_CACHE_HOME=/opt/prisma/.cache PRISMA_BINARY_CACHE_DIR=/opt/prisma/binaries \ npm_config_cache=/root/.npm \ @@ -73,7 +73,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime USER root RUN for i in 1 2 3; do \ - apk add --no-cache bash openssl tzdata python3 libsndfile libatomic && break; \ + apk add --no-cache bash openssl tzdata python-3.13 libsndfile libatomic && break; \ [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ sleep 5; \ done From cafdfda8ba586429dcaf0c6c1dcf7e7c5caf70a7 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 31 Aug 2026 05:52:17 +0000 Subject: [PATCH 23/49] feat(cli): set ENABLE_TOOL_SEARCH=true for lite claude Claude Code turns tool search off when ANTHROPIC_BASE_URL is a proxy. lite claude, lite up, login --config-claude, and autoroute now force ENABLE_TOOL_SEARCH=true so MCP tools stay deferred through the proxy Co-authored-by: Mateo Wang --- litellm/proxy/client/cli/README.md | 6 +++--- litellm/proxy/client/cli/commands/agents.py | 7 ++++++- .../client/cli/commands/autoroute/settings.py | 3 +++ .../proxy/client/cli/commands/claude_settings.py | 15 +++++++++++---- .../proxy/client/cli/autoroute/test_commands.py | 1 + .../proxy/client/cli/autoroute/test_settings.py | 7 +++++++ .../test_litellm/proxy/client/cli/test_agents.py | 14 ++++++++++++++ .../proxy/client/cli/test_auth_commands.py | 1 + .../proxy/client/cli/test_claude_settings.py | 1 + .../proxy/client/cli/test_up_commands.py | 12 +++++++++++- 10 files changed, 58 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index fe417396317..0da219df700 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -489,7 +489,7 @@ lite codex exec "summarize the repo" Each command resolves your LiteLLM key (logging in via SSO when none is stored and you are at a terminal; otherwise it expects `LITELLM_PROXY_API_KEY` or `--api-key`), checks the key against the proxy so bad credentials fail immediately instead of deep inside the agent, exports the environment variables the agent reads, then replaces itself with the agent process. -The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). +The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). Options (these belong to the wrapper, so put them before the agent's own flags): @@ -505,7 +505,7 @@ The credential is short-lived by design (default 24h, configurable via `LITELLM_ ### Route Every Claude Code Session Through the Proxy -`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it. +`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true`, and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it. Two things need to already be true: you've run `lite login` (or `lite login --pkce`, whose key the helper renews on its own), since the apiKeyHelper depends on that stored token, and the proxy is already reachable, since `lite up` does not start one for you. @@ -529,7 +529,7 @@ Cursor is not supported: it has no equivalent file-based config to hot-patch thi lite --base-url https://your-proxy.example.com login --config-claude ``` -It writes the same two settings `lite up` does, `env.ANTHROPIC_BASE_URL` and `apiKeyHelper`, but persistently: there is no backup, nothing to restore, and no foreground process to keep alive. Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag. +It writes the same settings `lite up` does, `env.ANTHROPIC_BASE_URL`, `env.ENABLE_TOOL_SEARCH`, and `apiKeyHelper`, but persistently: there is no backup, nothing to restore, and no foreground process to keep alive. Every other key in `~/.claude/settings.json` is preserved, the file is created if it does not exist, and it is written atomically with owner-only permissions. Plain `lite login` is unchanged; nothing happens to your Claude Code config unless you pass the flag. Because the credential is reached through `apiKeyHelper` rather than copied into the file, a later `lite login` refreshes it with no further action: Claude Code re-runs the helper on every request and picks up whatever token the most recent login stored. Nothing secret is written to `settings.json`. diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index e05e85ae483..4d3a441272e 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -13,6 +13,8 @@ from .auth import context_secret_vault, get_stored_api_key, login ANTHROPIC_BASE_URL_ENV: Final = "ANTHROPIC_BASE_URL" ANTHROPIC_AUTH_TOKEN_ENV: Final = "ANTHROPIC_AUTH_TOKEN" ANTHROPIC_API_KEY_ENV: Final = "ANTHROPIC_API_KEY" +ENABLE_TOOL_SEARCH_ENV: Final = "ENABLE_TOOL_SEARCH" +ENABLE_TOOL_SEARCH_VALUE: Final = "true" OPENAI_BASE_URL_ENV: Final = "OPENAI_BASE_URL" OPENAI_API_KEY_ENV: Final = "OPENAI_API_KEY" @@ -61,7 +63,9 @@ def build_agent_env( Anthropic clients (Claude Code) append /v1/messages to ANTHROPIC_BASE_URL, so it stays the bare proxy root; OpenAI clients (Codex, OpenCode) expect the /v1 suffix on OPENAI_BASE_URL. ANTHROPIC_API_KEY is dropped so a stray - Anthropic key cannot win over the bearer token we set. + Anthropic key cannot win over the bearer token we set. ENABLE_TOOL_SEARCH is + forced on because Claude Code turns tool search off when ANTHROPIC_BASE_URL + is not a first-party Anthropic host. """ env: Final = dict(base_env) root: Final = base_url.rstrip("/") @@ -69,6 +73,7 @@ def build_agent_env( env[ANTHROPIC_BASE_URL_ENV] = root env[ANTHROPIC_AUTH_TOKEN_ENV] = api_key env.pop(ANTHROPIC_API_KEY_ENV, None) + env[ENABLE_TOOL_SEARCH_ENV] = ENABLE_TOOL_SEARCH_VALUE if PROFILE_OPENAI in profiles: env[OPENAI_BASE_URL_ENV] = root + "/v1" env[OPENAI_API_KEY_ENV] = api_key diff --git a/litellm/proxy/client/cli/commands/autoroute/settings.py b/litellm/proxy/client/cli/commands/autoroute/settings.py index 9fcb11a585b..8331d41ae78 100644 --- a/litellm/proxy/client/cli/commands/autoroute/settings.py +++ b/litellm/proxy/client/cli/commands/autoroute/settings.py @@ -9,6 +9,8 @@ API_KEY_HELPER_KEY: Final = "apiKeyHelper" ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY" ANTHROPIC_AUTH_TOKEN_KEY: Final = "ANTHROPIC_AUTH_TOKEN" ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL" +ENABLE_TOOL_SEARCH_KEY: Final = "ENABLE_TOOL_SEARCH" +ENABLE_TOOL_SEARCH_VALUE: Final = "true" # Force every one of Claude Code's own model tiers to request the auto-router by name. # Router's auto-router registry is keyed by the literal requested model string # (litellm/router.py:10711-10717) with no wildcard/pattern resolution, so a bare "*" @@ -37,6 +39,7 @@ def merge_claude_settings_static_token( **base_env, ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"), ANTHROPIC_AUTH_TOKEN_KEY: auth_token, + ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE, **{key: AUTOROUTER_MODEL_NAME for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS}, } env.pop(ANTHROPIC_API_KEY_KEY, None) diff --git a/litellm/proxy/client/cli/commands/claude_settings.py b/litellm/proxy/client/cli/commands/claude_settings.py index e18e5b1b7ee..6d88c69bbd4 100644 --- a/litellm/proxy/client/cli/commands/claude_settings.py +++ b/litellm/proxy/client/cli/commands/claude_settings.py @@ -21,6 +21,8 @@ ENV_KEY: Final = "env" API_KEY_HELPER_KEY: Final = "apiKeyHelper" ANTHROPIC_BASE_URL_KEY: Final = "ANTHROPIC_BASE_URL" ANTHROPIC_API_KEY_KEY: Final = "ANTHROPIC_API_KEY" +ENABLE_TOOL_SEARCH_KEY: Final = "ENABLE_TOOL_SEARCH" +ENABLE_TOOL_SEARCH_VALUE: Final = "true" CLAUDE_SETTINGS_PATH: Final = Path.home() / ".claude" / "settings.json" BACKUP_PATH: Final = Path.home() / ".litellm" / "claude_settings_backup.json" @@ -68,16 +70,19 @@ def merge_claude_settings( ) -> dict[str, JsonValue]: """Return a new settings dict wired to route Claude Code through the proxy. - Only env.ANTHROPIC_BASE_URL and the top-level apiKeyHelper are overridden; a - stray env.ANTHROPIC_API_KEY is dropped so it cannot outrank the helper-issued - token (same reasoning as build_agent_env in agents.py). Every other key is - preserved untouched. + Only env.ANTHROPIC_BASE_URL, env.ENABLE_TOOL_SEARCH, and the top-level + apiKeyHelper are overridden; a stray env.ANTHROPIC_API_KEY is dropped so it + cannot outrank the helper-issued token (same reasoning as build_agent_env + in agents.py). ENABLE_TOOL_SEARCH is forced on because Claude Code turns + tool search off when ANTHROPIC_BASE_URL is not a first-party Anthropic host. + Every other key is preserved untouched. """ raw_env: Final = settings.get(ENV_KEY, {}) base_env: Final = raw_env if isinstance(raw_env, dict) else {} env: Final = { **{key: value for key, value in base_env.items() if key != ANTHROPIC_API_KEY_KEY}, ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"), + ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE, } return {**settings, ENV_KEY: env, API_KEY_HELPER_KEY: api_key_helper} @@ -141,6 +146,8 @@ __all__ = ( "ANTHROPIC_API_KEY_KEY", "ANTHROPIC_BASE_URL_KEY", "API_KEY_HELPER_KEY", + "ENABLE_TOOL_SEARCH_KEY", + "ENABLE_TOOL_SEARCH_VALUE", "AUTOROUTE_BACKUP_PATH", "BACKUP_PATH", "CLAUDE_SETTINGS_PATH", diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py b/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py index 74bf1c95777..4a3b3ef22c4 100644 --- a/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_commands.py @@ -157,6 +157,7 @@ class TestUpCommand: assert captured["settings"]["theme"] == "dark" assert captured["settings"]["env"]["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:5483" assert captured["settings"]["env"]["ANTHROPIC_AUTH_TOKEN"] == "fixed-master-key" + assert captured["settings"]["env"]["ENABLE_TOOL_SEARCH"] == "true" assert "apiKeyHelper" not in captured["settings"] assert captured["settings_mode"] == 0o600 diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py b/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py index 40d3e7f2aee..ded8f0808a0 100644 --- a/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py @@ -19,6 +19,13 @@ def test_sets_base_url_and_auth_token(): merged = merge_claude_settings_static_token({}, "http://127.0.0.1:4000/", "token-abc") assert merged["env"]["ANTHROPIC_BASE_URL"] == "http://127.0.0.1:4000" assert merged["env"]["ANTHROPIC_AUTH_TOKEN"] == "token-abc" + assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true" + + +def test_forces_tool_search_on(): + settings = {"env": {"ENABLE_TOOL_SEARCH": "false"}} + merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") + assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true" def test_drops_stray_api_key(): diff --git a/tests/test_litellm/proxy/client/cli/test_agents.py b/tests/test_litellm/proxy/client/cli/test_agents.py index 32dfb8d521d..d03c5d9db2d 100644 --- a/tests/test_litellm/proxy/client/cli/test_agents.py +++ b/tests/test_litellm/proxy/client/cli/test_agents.py @@ -77,9 +77,19 @@ class TestBuildAgentEnv: ) assert env["ANTHROPIC_BASE_URL"] == "http://localhost:4000" assert env["ANTHROPIC_AUTH_TOKEN"] == "sk-key" + assert env["ENABLE_TOOL_SEARCH"] == "true" assert "OPENAI_BASE_URL" not in env assert "OPENAI_API_KEY" not in env + def test_anthropic_profile_forces_tool_search_on(self): + env = build_agent_env( + {"ENABLE_TOOL_SEARCH": "false"}, + "http://localhost:4000", + "sk-key", + frozenset({"anthropic"}), + ) + assert env["ENABLE_TOOL_SEARCH"] == "true" + def test_anthropic_profile_drops_existing_api_key(self): env = build_agent_env( {"ANTHROPIC_API_KEY": "real-key"}, @@ -96,6 +106,7 @@ class TestBuildAgentEnv: assert env["OPENAI_BASE_URL"] == "http://localhost:4000/v1" assert env["OPENAI_API_KEY"] == "sk-key" assert "ANTHROPIC_BASE_URL" not in env + assert "ENABLE_TOOL_SEARCH" not in env def test_both_profiles_set_everything(self): env = build_agent_env( @@ -105,6 +116,7 @@ class TestBuildAgentEnv: assert env["OPENAI_BASE_URL"] == "http://localhost:4000/v1" assert env["ANTHROPIC_AUTH_TOKEN"] == "sk-key" assert env["OPENAI_API_KEY"] == "sk-key" + assert env["ENABLE_TOOL_SEARCH"] == "true" def test_preserves_unrelated_env_and_does_not_mutate_input(self): base = {"PATH": "/usr/bin", "ANTHROPIC_API_KEY": "real-key"} @@ -201,6 +213,7 @@ class TestRunAgent: env = calls["env"] assert env["ANTHROPIC_BASE_URL"] == "http://localhost:4000" assert env["ANTHROPIC_AUTH_TOKEN"] == "sk-key" + assert env["ENABLE_TOOL_SEARCH"] == "true" assert "ANTHROPIC_API_KEY" not in env assert "OPENAI_BASE_URL" not in env @@ -218,6 +231,7 @@ class TestRunAgent: assert calls["env"]["OPENAI_BASE_URL"] == "http://localhost:4000/v1" assert calls["env"]["OPENAI_API_KEY"] == "sk-key" assert "ANTHROPIC_BASE_URL" not in calls["env"] + assert "ENABLE_TOOL_SEARCH" not in calls["env"] def test_codex_injects_proxy_provider_args_before_user_args(self): calls = {} diff --git a/tests/test_litellm/proxy/client/cli/test_auth_commands.py b/tests/test_litellm/proxy/client/cli/test_auth_commands.py index 85a4d90abf9..1d0a99b8e0a 100644 --- a/tests/test_litellm/proxy/client/cli/test_auth_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_auth_commands.py @@ -1373,6 +1373,7 @@ class TestLoginConfigClaude: assert result.exit_code == 0 written = json.loads(settings_path.read_text()) assert written["env"]["ANTHROPIC_BASE_URL"] == "https://test.example.com" + assert written["env"]["ENABLE_TOOL_SEARCH"] == "true" assert written["apiKeyHelper"] == "/usr/local/bin/lite --base-url https://test.example.com auth print-token" assert "Configured Claude Code" in result.output diff --git a/tests/test_litellm/proxy/client/cli/test_claude_settings.py b/tests/test_litellm/proxy/client/cli/test_claude_settings.py index 9010fb4c022..898f9ab1ed7 100644 --- a/tests/test_litellm/proxy/client/cli/test_claude_settings.py +++ b/tests/test_litellm/proxy/client/cli/test_claude_settings.py @@ -48,6 +48,7 @@ class TestWriteClaudeSettings: written = json.loads(settings_path.read_text()) assert written["env"]["ANTHROPIC_BASE_URL"] == "https://proxy.example.com" + assert written["env"]["ENABLE_TOOL_SEARCH"] == "true" assert written["apiKeyHelper"] == "/usr/local/bin/lite --base-url https://proxy.example.com auth print-token" def test_updates_an_existing_file_preserving_unrelated_settings(self, paths, lite_on_path): diff --git a/tests/test_litellm/proxy/client/cli/test_up_commands.py b/tests/test_litellm/proxy/client/cli/test_up_commands.py index 9958286884b..a1aa9eeeb8a 100644 --- a/tests/test_litellm/proxy/client/cli/test_up_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_up_commands.py @@ -55,8 +55,14 @@ class TestMergeClaudeSettings: } merged = merge_claude_settings(settings, "http://localhost:4000/", "new-helper") assert merged["env"]["ANTHROPIC_BASE_URL"] == "http://localhost:4000" + assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true" assert merged["apiKeyHelper"] == "new-helper" + def test_forces_tool_search_on(self): + settings = {"env": {"ENABLE_TOOL_SEARCH": "false"}} + merged = merge_claude_settings(settings, "http://localhost:4000", "helper") + assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true" + def test_drops_stray_api_key(self): settings = {"env": {"ANTHROPIC_API_KEY": "leaked-key"}} merged = merge_claude_settings(settings, "http://localhost:4000", "helper") @@ -64,7 +70,10 @@ class TestMergeClaudeSettings: def test_works_from_empty_settings(self): merged = merge_claude_settings({}, "http://localhost:4000", "helper") - assert merged["env"] == {"ANTHROPIC_BASE_URL": "http://localhost:4000"} + assert merged["env"] == { + "ANTHROPIC_BASE_URL": "http://localhost:4000", + "ENABLE_TOOL_SEARCH": "true", + } assert merged["apiKeyHelper"] == "helper" def test_does_not_mutate_input(self): @@ -486,6 +495,7 @@ class TestUpCommand: assert captured["backup_existed"] is True assert captured["settings"]["theme"] == "dark" assert captured["settings"]["env"]["ANTHROPIC_BASE_URL"] == "http://localhost:4000" + assert captured["settings"]["env"]["ENABLE_TOOL_SEARCH"] == "true" assert captured["settings"]["apiKeyHelper"] == "/usr/local/bin/lite auth print-token" assert json.loads(settings_path.read_text()) == original assert not backup_path.exists() From 096e016baf04fc833ce617025eace2f8937cefe2 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 31 Aug 2026 05:58:00 +0000 Subject: [PATCH 24/49] fix(cli): sort claude_settings __all__ for ruff RUF022 Co-authored-by: Mateo Wang --- litellm/proxy/client/cli/commands/claude_settings.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/client/cli/commands/claude_settings.py b/litellm/proxy/client/cli/commands/claude_settings.py index 6d88c69bbd4..41ba25447e1 100644 --- a/litellm/proxy/client/cli/commands/claude_settings.py +++ b/litellm/proxy/client/cli/commands/claude_settings.py @@ -146,11 +146,11 @@ __all__ = ( "ANTHROPIC_API_KEY_KEY", "ANTHROPIC_BASE_URL_KEY", "API_KEY_HELPER_KEY", - "ENABLE_TOOL_SEARCH_KEY", - "ENABLE_TOOL_SEARCH_VALUE", "AUTOROUTE_BACKUP_PATH", "BACKUP_PATH", "CLAUDE_SETTINGS_PATH", + "ENABLE_TOOL_SEARCH_KEY", + "ENABLE_TOOL_SEARCH_VALUE", "ENV_KEY", "SETTINGS_FILE_OWNERS", "ClaudeSettingsError", From 43c64e24715cdd1aa7595544f1d4aba55f7c2512 Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 31 Aug 2026 06:00:45 +0000 Subject: [PATCH 25/49] fix(cli): keep an existing ENABLE_TOOL_SEARCH value Default remains true so lite claude turns tool search back on through a proxy. An explicit false or auto in the env or settings is left alone Co-authored-by: Mateo Wang --- litellm/proxy/client/cli/README.md | 4 ++-- litellm/proxy/client/cli/commands/agents.py | 10 ++++++---- .../client/cli/commands/autoroute/settings.py | 2 +- .../proxy/client/cli/commands/claude_settings.py | 14 +++++++------- .../proxy/client/cli/autoroute/test_settings.py | 4 ++-- tests/test_litellm/proxy/client/cli/test_agents.py | 4 ++-- .../proxy/client/cli/test_up_commands.py | 4 ++-- 7 files changed, 22 insertions(+), 20 deletions(-) diff --git a/litellm/proxy/client/cli/README.md b/litellm/proxy/client/cli/README.md index 0da219df700..3ddce35b53d 100644 --- a/litellm/proxy/client/cli/README.md +++ b/litellm/proxy/client/cli/README.md @@ -489,7 +489,7 @@ lite codex exec "summarize the repo" Each command resolves your LiteLLM key (logging in via SSO when none is stored and you are at a terminal; otherwise it expects `LITELLM_PROXY_API_KEY` or `--api-key`), checks the key against the proxy so bad credentials fail immediately instead of deep inside the agent, exports the environment variables the agent reads, then replaces itself with the agent process. -The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). +The right variables are picked per agent. Claude Code gets `ANTHROPIC_BASE_URL` (the proxy root, so it appends `/v1/messages`) and `ANTHROPIC_AUTH_TOKEN`, with any stray `ANTHROPIC_API_KEY` cleared so the proxy token wins, and `ENABLE_TOOL_SEARCH=true` (unless you already set it) so Claude Code keeps tool search on even though the base URL is a proxy rather than a first-party Anthropic host. Codex and OpenCode get `OPENAI_BASE_URL` (the proxy plus `/v1`) and `OPENAI_API_KEY`. Codex ignores `OPENAI_BASE_URL`, so it is additionally pointed at the proxy through a custom provider passed as `-c` config overrides (HTTP/SSE Responses transport, since the proxy does not speak the Responses WebSocket protocol). Options (these belong to the wrapper, so put them before the agent's own flags): @@ -505,7 +505,7 @@ The credential is short-lived by design (default 24h, configurable via `LITELLM_ ### Route Every Claude Code Session Through the Proxy -`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true`, and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it. +`lite claude` wraps a single invocation, but `lite up` goes further: it patches `~/.claude/settings.json`, Claude Code's own config file, so that every Claude Code session started afterward -- from any terminal, launched normally with just `claude`, no wrapper needed -- routes through your LiteLLM proxy. It sets `env.ANTHROPIC_BASE_URL` to the proxy URL, `env.ENABLE_TOOL_SEARCH` to `true` when that key is missing, and `apiKeyHelper` to a `lite auth print-token` invocation, drops any stray static `ANTHROPIC_API_KEY` so the helper-issued token wins, and leaves every other setting in the file untouched. It backs up the original file before patching it. Two things need to already be true: you've run `lite login` (or `lite login --pkce`, whose key the helper renews on its own), since the apiKeyHelper depends on that stored token, and the proxy is already reachable, since `lite up` does not start one for you. diff --git a/litellm/proxy/client/cli/commands/agents.py b/litellm/proxy/client/cli/commands/agents.py index 4d3a441272e..45e05d353fb 100644 --- a/litellm/proxy/client/cli/commands/agents.py +++ b/litellm/proxy/client/cli/commands/agents.py @@ -63,9 +63,10 @@ def build_agent_env( Anthropic clients (Claude Code) append /v1/messages to ANTHROPIC_BASE_URL, so it stays the bare proxy root; OpenAI clients (Codex, OpenCode) expect the /v1 suffix on OPENAI_BASE_URL. ANTHROPIC_API_KEY is dropped so a stray - Anthropic key cannot win over the bearer token we set. ENABLE_TOOL_SEARCH is - forced on because Claude Code turns tool search off when ANTHROPIC_BASE_URL - is not a first-party Anthropic host. + Anthropic key cannot win over the bearer token we set. ENABLE_TOOL_SEARCH + defaults to true because Claude Code turns tool search off when + ANTHROPIC_BASE_URL is not a first-party Anthropic host; a value already in + the environment is left alone. """ env: Final = dict(base_env) root: Final = base_url.rstrip("/") @@ -73,7 +74,8 @@ def build_agent_env( env[ANTHROPIC_BASE_URL_ENV] = root env[ANTHROPIC_AUTH_TOKEN_ENV] = api_key env.pop(ANTHROPIC_API_KEY_ENV, None) - env[ENABLE_TOOL_SEARCH_ENV] = ENABLE_TOOL_SEARCH_VALUE + if ENABLE_TOOL_SEARCH_ENV not in env: + env[ENABLE_TOOL_SEARCH_ENV] = ENABLE_TOOL_SEARCH_VALUE if PROFILE_OPENAI in profiles: env[OPENAI_BASE_URL_ENV] = root + "/v1" env[OPENAI_API_KEY_ENV] = api_key diff --git a/litellm/proxy/client/cli/commands/autoroute/settings.py b/litellm/proxy/client/cli/commands/autoroute/settings.py index 8331d41ae78..60729b5410d 100644 --- a/litellm/proxy/client/cli/commands/autoroute/settings.py +++ b/litellm/proxy/client/cli/commands/autoroute/settings.py @@ -36,10 +36,10 @@ def merge_claude_settings_static_token( raw_env: Final = settings.get(ENV_KEY, {}) base_env: Final = raw_env if isinstance(raw_env, dict) else {} env: Final[dict[str, JsonValue]] = { + ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE, **base_env, ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"), ANTHROPIC_AUTH_TOKEN_KEY: auth_token, - ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE, **{key: AUTOROUTER_MODEL_NAME for key in ANTHROPIC_DEFAULT_MODEL_ENV_KEYS}, } env.pop(ANTHROPIC_API_KEY_KEY, None) diff --git a/litellm/proxy/client/cli/commands/claude_settings.py b/litellm/proxy/client/cli/commands/claude_settings.py index 41ba25447e1..ea1fa019c83 100644 --- a/litellm/proxy/client/cli/commands/claude_settings.py +++ b/litellm/proxy/client/cli/commands/claude_settings.py @@ -70,19 +70,19 @@ def merge_claude_settings( ) -> dict[str, JsonValue]: """Return a new settings dict wired to route Claude Code through the proxy. - Only env.ANTHROPIC_BASE_URL, env.ENABLE_TOOL_SEARCH, and the top-level - apiKeyHelper are overridden; a stray env.ANTHROPIC_API_KEY is dropped so it - cannot outrank the helper-issued token (same reasoning as build_agent_env - in agents.py). ENABLE_TOOL_SEARCH is forced on because Claude Code turns - tool search off when ANTHROPIC_BASE_URL is not a first-party Anthropic host. - Every other key is preserved untouched. + Only env.ANTHROPIC_BASE_URL and the top-level apiKeyHelper are overridden; a + stray env.ANTHROPIC_API_KEY is dropped so it cannot outrank the helper-issued + token (same reasoning as build_agent_env in agents.py). ENABLE_TOOL_SEARCH + defaults to true because Claude Code turns tool search off when + ANTHROPIC_BASE_URL is not a first-party Anthropic host; an existing value is + left alone. Every other key is preserved untouched. """ raw_env: Final = settings.get(ENV_KEY, {}) base_env: Final = raw_env if isinstance(raw_env, dict) else {} env: Final = { + ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE, **{key: value for key, value in base_env.items() if key != ANTHROPIC_API_KEY_KEY}, ANTHROPIC_BASE_URL_KEY: base_url.rstrip("/"), - ENABLE_TOOL_SEARCH_KEY: ENABLE_TOOL_SEARCH_VALUE, } return {**settings, ENV_KEY: env, API_KEY_HELPER_KEY: api_key_helper} diff --git a/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py b/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py index ded8f0808a0..87a33c79a79 100644 --- a/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py +++ b/tests/test_litellm/proxy/client/cli/autoroute/test_settings.py @@ -22,10 +22,10 @@ def test_sets_base_url_and_auth_token(): assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true" -def test_forces_tool_search_on(): +def test_preserves_existing_tool_search(): settings = {"env": {"ENABLE_TOOL_SEARCH": "false"}} merged = merge_claude_settings_static_token(settings, "http://127.0.0.1:4000", "token-abc") - assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true" + assert merged["env"]["ENABLE_TOOL_SEARCH"] == "false" def test_drops_stray_api_key(): diff --git a/tests/test_litellm/proxy/client/cli/test_agents.py b/tests/test_litellm/proxy/client/cli/test_agents.py index d03c5d9db2d..0191dad3d94 100644 --- a/tests/test_litellm/proxy/client/cli/test_agents.py +++ b/tests/test_litellm/proxy/client/cli/test_agents.py @@ -81,14 +81,14 @@ class TestBuildAgentEnv: assert "OPENAI_BASE_URL" not in env assert "OPENAI_API_KEY" not in env - def test_anthropic_profile_forces_tool_search_on(self): + def test_anthropic_profile_preserves_existing_tool_search(self): env = build_agent_env( {"ENABLE_TOOL_SEARCH": "false"}, "http://localhost:4000", "sk-key", frozenset({"anthropic"}), ) - assert env["ENABLE_TOOL_SEARCH"] == "true" + assert env["ENABLE_TOOL_SEARCH"] == "false" def test_anthropic_profile_drops_existing_api_key(self): env = build_agent_env( diff --git a/tests/test_litellm/proxy/client/cli/test_up_commands.py b/tests/test_litellm/proxy/client/cli/test_up_commands.py index a1aa9eeeb8a..053c90c36b4 100644 --- a/tests/test_litellm/proxy/client/cli/test_up_commands.py +++ b/tests/test_litellm/proxy/client/cli/test_up_commands.py @@ -58,10 +58,10 @@ class TestMergeClaudeSettings: assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true" assert merged["apiKeyHelper"] == "new-helper" - def test_forces_tool_search_on(self): + def test_preserves_existing_tool_search(self): settings = {"env": {"ENABLE_TOOL_SEARCH": "false"}} merged = merge_claude_settings(settings, "http://localhost:4000", "helper") - assert merged["env"]["ENABLE_TOOL_SEARCH"] == "true" + assert merged["env"]["ENABLE_TOOL_SEARCH"] == "false" def test_drops_stray_api_key(self): settings = {"env": {"ANTHROPIC_API_KEY": "leaked-key"}} From 4291afbfa507e08a0ea270b77966b08e1a735b6b Mon Sep 17 00:00:00 2001 From: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 31 Aug 2026 13:08:06 +0000 Subject: [PATCH 26/49] fix(registry): correct OpenAI preview shutdown dates, add whisper/transcribe and Bedrock/Vertex deprecation dates Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- ...odel_prices_and_context_window_backup.json | 29 +++++++++++-------- model_prices_and_context_window.json | 29 +++++++++++-------- 2 files changed, 34 insertions(+), 24 deletions(-) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 9def855bc51..7071eaa0807 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -12315,7 +12315,8 @@ "output_cost_per_token": 1e-05, "supported_endpoints": [ "/v1/audio/transcriptions" - ] + ], + "deprecation_date": "2027-02-26" }, "claude-haiku-4-5-20251001": { "deprecation_date": "2026-10-15", @@ -17767,7 +17768,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "deprecation_date": "2027-01-08" }, "eu.anthropic.claude-opus-4-20250514-v1:0": { "cache_creation_input_token_cost": 1.875e-05, @@ -25327,7 +25329,7 @@ "supports_vision": true }, "gpt-4o-audio-preview": { - "deprecation_date": "2027-01-20", + "deprecation_date": "2026-05-07", "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 2.5e-06, "litellm_provider": "openai", @@ -25650,7 +25652,7 @@ "supports_vision": true }, "gpt-4o-mini-audio-preview": { - "deprecation_date": "2027-01-20", + "deprecation_date": "2026-05-07", "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 1.5e-07, "litellm_provider": "openai", @@ -25688,7 +25690,7 @@ "gpt-4o-mini-realtime-preview": { "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 3e-07, - "deprecation_date": "2027-01-20", + "deprecation_date": "2026-05-07", "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 6e-07, "litellm_provider": "openai", @@ -25787,7 +25789,8 @@ "output_cost_per_token": 5e-06, "supported_endpoints": [ "/v1/audio/transcriptions" - ] + ], + "deprecation_date": "2027-02-26" }, "gpt-4o-mini-tts": { "input_cost_per_token": 2.5e-06, @@ -25809,7 +25812,7 @@ }, "gpt-4o-realtime-preview": { "cache_read_input_token_cost": 2.5e-06, - "deprecation_date": "2027-01-20", + "deprecation_date": "2026-05-07", "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 5e-06, "litellm_provider": "openai", @@ -25828,7 +25831,7 @@ }, "gpt-4o-realtime-preview-2024-12-17": { "cache_read_input_token_cost": 2.5e-06, - "deprecation_date": "2027-01-20", + "deprecation_date": "2026-05-07", "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 5e-06, "litellm_provider": "openai", @@ -25847,7 +25850,7 @@ }, "gpt-4o-realtime-preview-2025-06-03": { "cache_read_input_token_cost": 2.5e-06, - "deprecation_date": "2027-01-20", + "deprecation_date": "2026-05-07", "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 5e-06, "litellm_provider": "openai", @@ -25926,7 +25929,8 @@ "output_cost_per_token": 1e-05, "supported_endpoints": [ "/v1/audio/transcriptions" - ] + ], + "deprecation_date": "2027-02-26" }, "gpt-image-1.5": { "cache_read_input_token_cost": 1.25e-06, @@ -38321,7 +38325,7 @@ "source": "https://docs.mistral.ai/capabilities/code_generation/" }, "text-embedding-004": { - "deprecation_date": "2026-01-14", + "deprecation_date": "2027-04-01", "input_cost_per_character": 2.5e-08, "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-embedding-models", @@ -44100,7 +44104,8 @@ "output_cost_per_second": 0.0001, "supported_endpoints": [ "/v1/audio/transcriptions" - ] + ], + "deprecation_date": "2027-02-26" }, "xai/grok-3": { "cache_read_input_token_cost": 2e-07, diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index 9def855bc51..7071eaa0807 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -12315,7 +12315,8 @@ "output_cost_per_token": 1e-05, "supported_endpoints": [ "/v1/audio/transcriptions" - ] + ], + "deprecation_date": "2027-02-26" }, "claude-haiku-4-5-20251001": { "deprecation_date": "2026-10-15", @@ -17767,7 +17768,8 @@ "supports_response_schema": true, "supports_tool_choice": true, "supports_vision": true, - "prompt_cache_min_tokens": 1024 + "prompt_cache_min_tokens": 1024, + "deprecation_date": "2027-01-08" }, "eu.anthropic.claude-opus-4-20250514-v1:0": { "cache_creation_input_token_cost": 1.875e-05, @@ -25327,7 +25329,7 @@ "supports_vision": true }, "gpt-4o-audio-preview": { - "deprecation_date": "2027-01-20", + "deprecation_date": "2026-05-07", "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 2.5e-06, "litellm_provider": "openai", @@ -25650,7 +25652,7 @@ "supports_vision": true }, "gpt-4o-mini-audio-preview": { - "deprecation_date": "2027-01-20", + "deprecation_date": "2026-05-07", "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 1.5e-07, "litellm_provider": "openai", @@ -25688,7 +25690,7 @@ "gpt-4o-mini-realtime-preview": { "cache_creation_input_audio_token_cost": 3e-07, "cache_read_input_token_cost": 3e-07, - "deprecation_date": "2027-01-20", + "deprecation_date": "2026-05-07", "input_cost_per_audio_token": 1e-05, "input_cost_per_token": 6e-07, "litellm_provider": "openai", @@ -25787,7 +25789,8 @@ "output_cost_per_token": 5e-06, "supported_endpoints": [ "/v1/audio/transcriptions" - ] + ], + "deprecation_date": "2027-02-26" }, "gpt-4o-mini-tts": { "input_cost_per_token": 2.5e-06, @@ -25809,7 +25812,7 @@ }, "gpt-4o-realtime-preview": { "cache_read_input_token_cost": 2.5e-06, - "deprecation_date": "2027-01-20", + "deprecation_date": "2026-05-07", "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 5e-06, "litellm_provider": "openai", @@ -25828,7 +25831,7 @@ }, "gpt-4o-realtime-preview-2024-12-17": { "cache_read_input_token_cost": 2.5e-06, - "deprecation_date": "2027-01-20", + "deprecation_date": "2026-05-07", "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 5e-06, "litellm_provider": "openai", @@ -25847,7 +25850,7 @@ }, "gpt-4o-realtime-preview-2025-06-03": { "cache_read_input_token_cost": 2.5e-06, - "deprecation_date": "2027-01-20", + "deprecation_date": "2026-05-07", "input_cost_per_audio_token": 4e-05, "input_cost_per_token": 5e-06, "litellm_provider": "openai", @@ -25926,7 +25929,8 @@ "output_cost_per_token": 1e-05, "supported_endpoints": [ "/v1/audio/transcriptions" - ] + ], + "deprecation_date": "2027-02-26" }, "gpt-image-1.5": { "cache_read_input_token_cost": 1.25e-06, @@ -38321,7 +38325,7 @@ "source": "https://docs.mistral.ai/capabilities/code_generation/" }, "text-embedding-004": { - "deprecation_date": "2026-01-14", + "deprecation_date": "2027-04-01", "input_cost_per_character": 2.5e-08, "input_cost_per_token": 1e-07, "litellm_provider": "vertex_ai-embedding-models", @@ -44100,7 +44104,8 @@ "output_cost_per_second": 0.0001, "supported_endpoints": [ "/v1/audio/transcriptions" - ] + ], + "deprecation_date": "2027-02-26" }, "xai/grok-3": { "cache_read_input_token_cost": 2e-07, From f079e4061bf986d6bb368da864f6c9dec22a0cac Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 31 Aug 2026 09:22:36 -0700 Subject: [PATCH 27/49] fix(proxy): deliver budget alerts on webhook-only alerting and accept ALERTING_WEBHOOK_URL (#38441) * fix(proxy): deliver budget alerts on webhook-only alerting and accept ALERTING_WEBHOOK_URL ProxyLogging.budget_alerts forwarded to the alerting pipeline only when 'slack' was in general_settings.alerting, so alerting: ['webhook'] plus WEBHOOK_URL silently never delivered a budget alert (the config /health/services?service=webhook exists to test). Forward when 'webhook' is present too; SlackAlerting.send_alert already fans out per channel. Also accept a provider-neutral ALERTING_WEBHOOK_URL env fallback for the Slack-format channel (any Slack-compatible receiver works), mark it as a sensitive var, and de-brand the admin UI alerting copy. Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(ui): format settings.tsx with prettier Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * chore(ui): regenerate schema.d.ts for updated alerting description Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * ci: retrigger checks after ALERTING_WEBHOOK_URL docs merged Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/constants.py | 1 + .../SlackAlerting/slack_alerting.py | 8 +-- litellm/proxy/_types.py | 2 +- litellm/proxy/proxy_server.py | 3 +- litellm/proxy/utils.py | 6 +- .../SlackAlerting/test_slack_alerting.py | 55 ++++++++++++++++++- .../test_slack_alerting_digest.py | 17 ++++++ .../utils/proxy_logging/test_alerting.py | 29 ++++++++++ .../src/components/settings.tsx | 5 +- ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 +- 10 files changed, 116 insertions(+), 12 deletions(-) diff --git a/litellm/constants.py b/litellm/constants.py index cc6db6c10cc..9872783bfab 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1728,6 +1728,7 @@ SENTRY_DENYLIST: Final = [ "jwt_token", "private_key", "SLACK_WEBHOOK_URL", + "ALERTING_WEBHOOK_URL", "webhook_url", "LANGFUSE_SECRET_KEY", # Email Configuration diff --git a/litellm/integrations/SlackAlerting/slack_alerting.py b/litellm/integrations/SlackAlerting/slack_alerting.py index d7d06387d85..94d734546be 100644 --- a/litellm/integrations/SlackAlerting/slack_alerting.py +++ b/litellm/integrations/SlackAlerting/slack_alerting.py @@ -1485,9 +1485,9 @@ Model Info: elif self.default_webhook_url is not None: _digest_webhook = self.default_webhook_url else: - _digest_webhook = os.getenv("SLACK_WEBHOOK_URL", None) + _digest_webhook = os.getenv("SLACK_WEBHOOK_URL") or os.getenv("ALERTING_WEBHOOK_URL") if _digest_webhook is None: - raise ValueError("Missing SLACK_WEBHOOK_URL from environment") + raise ValueError("Missing SLACK_WEBHOOK_URL / ALERTING_WEBHOOK_URL from environment") digest_key: Final = f"{alert_type_name_str}:{request_model or ''}:{api_base or ''}" @@ -1516,10 +1516,10 @@ Model Info: elif self.default_webhook_url is not None: slack_webhook_url = self.default_webhook_url else: - slack_webhook_url = os.getenv("SLACK_WEBHOOK_URL", None) + slack_webhook_url = os.getenv("SLACK_WEBHOOK_URL") or os.getenv("ALERTING_WEBHOOK_URL") if slack_webhook_url is None: - raise ValueError("Missing SLACK_WEBHOOK_URL from environment") + raise ValueError("Missing SLACK_WEBHOOK_URL / ALERTING_WEBHOOK_URL from environment") payload: Final = {"text": formatted_message} headers: Final = {"Content-type": "application/json"} diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 5ba5e8fa1aa..0f2d97b8b1c 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -2541,7 +2541,7 @@ class ConfigGeneralSettings(LiteLLMPydanticObjectBase): ) alerting: list | None = Field( None, - description="List of alerting integrations. Today, just slack - `alerting: ['slack']`", + description="List of alerting integrations - e.g. `alerting: ['slack', 'webhook', 'email']`. 'slack' posts Slack-format messages to any Slack-compatible webhook (Slack, Rocket.Chat, Mattermost); 'webhook' posts structured JSON budget alerts to WEBHOOK_URL", ) alert_types: list[AlertType] | None = Field( None, diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 887716a383a..da2e09bd7f6 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -1363,7 +1363,7 @@ _OPENAPI_HTTP_METHODS: Final = { # the UI. Kept here at module scope to match the analogous descriptor # `is_secret` flags in litellm.proxy.config_resolvers and the # `_CACHE_SENSITIVE_FIELDS` constant in the cache endpoint file. -_ALERTING_SENSITIVE_VARS: Final[set[str]] = {"SLACK_WEBHOOK_URL", "SMTP_PASSWORD"} +_ALERTING_SENSITIVE_VARS: Final[set[str]] = {"ALERTING_WEBHOOK_URL", "SLACK_WEBHOOK_URL", "SMTP_PASSWORD"} def _strip_operation_id_method_suffix(operation_id: str) -> str: @@ -16566,6 +16566,7 @@ async def create_config_audit_log( _EXTRA_SECRET_CALLBACK_ENV_VARS: Final = frozenset( { + "ALERTING_WEBHOOK_URL", "GALILEO_USERNAME", "GENERIC_LOGGER_HEADERS", "OTEL_HEADERS", diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index cf56fc0b1dd..051d36c4d0f 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -645,7 +645,7 @@ class ProxyLogging: self.max_parallel_request_limiter = _PROXY_MaxParallelRequestsHandler(self.internal_usage_cache) self.max_budget_limiter = _PROXY_MaxBudgetLimiter() self.cache_control_check = _PROXY_CacheControlCheck() - self.alerting: list | None = None + self.alerting: list[str] | None = None self.alerting_threshold: float = 300 # default to 5 min. threshold self.alert_types: list[AlertType] = DEFAULT_ALERT_TYPES self.alert_to_webhook_url: dict | None = None @@ -2364,7 +2364,9 @@ class ProxyLogging: # do nothing if alerting is not switched on (unless it's a soft_budget alert with team-specific emails) return - if self.alerting is not None and ("slack" in self.alerting or "ms_teams" in self.alerting): + if self.alerting is not None and ( + "slack" in self.alerting or "ms_teams" in self.alerting or "webhook" in self.alerting + ): if self.slack_alerting_instance is not None: await self.slack_alerting_instance.budget_alerts( type=type, diff --git a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py index cfbd3e76a88..55e2dcdc270 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py +++ b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting.py @@ -12,7 +12,7 @@ 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 +from litellm.types.integrations.slack_alerting import AlertType, SlackAlertingCacheKeys class TestSlackAlerting(unittest.TestCase): @@ -366,3 +366,56 @@ async def test_scheduled_daily_report_threads_the_pod_lock_manager_through(): _, kwargs = slack_alerting._run_scheduler_helper.await_args assert kwargs["pod_lock_manager"] is pod_lock_manager + + +def _slack_alerting_with_env_resolution() -> SlackAlerting: + slack_alerting: Final = SlackAlerting(alerting=["slack"], internal_usage_cache=DualCache()) + slack_alerting.periodic_started = True + return slack_alerting + + +@pytest.mark.asyncio +async def test_send_alert_falls_back_to_alerting_webhook_url_env(monkeypatch): + monkeypatch.delenv("SLACK_WEBHOOK_URL", raising=False) + monkeypatch.setenv("ALERTING_WEBHOOK_URL", "https://chat.example.com/hooks/abc") + slack_alerting: Final = _slack_alerting_with_env_resolution() + + await slack_alerting.send_alert( + message="budget crossed", + level="High", + alert_type=AlertType.budget_alerts, + alerting_metadata={}, + ) + + assert slack_alerting.log_queue[0]["url"] == "https://chat.example.com/hooks/abc" + + +@pytest.mark.asyncio +async def test_send_alert_prefers_slack_webhook_url_over_fallback(monkeypatch): + monkeypatch.setenv("SLACK_WEBHOOK_URL", "https://hooks.slack.com/services/T0/B0/X0") + monkeypatch.setenv("ALERTING_WEBHOOK_URL", "https://chat.example.com/hooks/abc") + slack_alerting: Final = _slack_alerting_with_env_resolution() + + await slack_alerting.send_alert( + message="budget crossed", + level="High", + alert_type=AlertType.budget_alerts, + alerting_metadata={}, + ) + + assert slack_alerting.log_queue[0]["url"] == "https://hooks.slack.com/services/T0/B0/X0" + + +@pytest.mark.asyncio +async def test_send_alert_raises_when_no_webhook_url_configured(monkeypatch): + monkeypatch.delenv("SLACK_WEBHOOK_URL", raising=False) + monkeypatch.delenv("ALERTING_WEBHOOK_URL", raising=False) + slack_alerting: Final = _slack_alerting_with_env_resolution() + + with pytest.raises(ValueError, match="SLACK_WEBHOOK_URL / ALERTING_WEBHOOK_URL"): + await slack_alerting.send_alert( + message="budget crossed", + level="High", + alert_type=AlertType.budget_alerts, + alerting_metadata={}, + ) diff --git a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting_digest.py b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting_digest.py index edce5c5f3a2..d614823c0ef 100644 --- a/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting_digest.py +++ b/tests/test_litellm/integrations/SlackAlerting/test_slack_alerting_digest.py @@ -79,6 +79,23 @@ class TestDigestMode(unittest.IsolatedAsyncioTestCase): self.assertEqual(len(self.slack_alerting.digest_buckets), 2) + async def test_digest_falls_back_to_alerting_webhook_url_env(self): + """With SLACK_WEBHOOK_URL unset, the digest entry resolves ALERTING_WEBHOOK_URL instead.""" + env = {k: v for k, v in os.environ.items() if k != "SLACK_WEBHOOK_URL"} + env["ALERTING_WEBHOOK_URL"] = "https://chat.example.com/hooks/abc" + with unittest.mock.patch.dict(os.environ, env, clear=True): + await self.slack_alerting.send_alert( + message="`Requests are hanging`", + level="Medium", + alert_type=AlertType.llm_requests_hanging, + alerting_metadata={}, + request_model="gemini-2.5-flash", + api_base="None", + ) + + bucket = list(self.slack_alerting.digest_buckets.values())[0] + self.assertEqual(bucket["webhook_url"], "https://chat.example.com/hooks/abc") + async def test_non_digest_alert_goes_to_queue(self): """Alert types without digest enabled should go straight to the log queue.""" message = "Budget exceeded" diff --git a/tests/test_litellm/proxy/utils/proxy_logging/test_alerting.py b/tests/test_litellm/proxy/utils/proxy_logging/test_alerting.py index cede859cb38..77c0f71dbf9 100644 --- a/tests/test_litellm/proxy/utils/proxy_logging/test_alerting.py +++ b/tests/test_litellm/proxy/utils/proxy_logging/test_alerting.py @@ -115,6 +115,35 @@ async def test_budget_alerts_slack_when_slack_alerting(proxy_logging): assert snapshot == {"type": "user_budget", "user_info_is_callinfo": True, "user_id": "u1"} +@pytest.mark.asyncio +async def test_budget_alerts_webhook_only_forwards_to_slack_alerting_instance(proxy_logging): + proxy_logging.alerting = ["webhook"] + captured: Dict[str, Any] = {} + + async def fake_alert(**kwargs): + captured.update(kwargs) + + proxy_logging.slack_alerting_instance = MagicMock(budget_alerts=fake_alert) + proxy_logging.email_logging_instance = None + await proxy_logging.budget_alerts(type="user_budget", user_info=_user_info()) + snapshot = { + "type": captured["type"], + "user_info_is_callinfo": isinstance(captured["user_info"], CallInfo), + "user_id": captured["user_info"].user_id, + } + assert snapshot == {"type": "user_budget", "user_info_is_callinfo": True, "user_id": "u1"} + + +@pytest.mark.asyncio +async def test_budget_alerts_email_only_skips_slack_alerting_instance(proxy_logging): + proxy_logging.alerting = ["email"] + proxy_logging.slack_alerting_instance = MagicMock(budget_alerts=AsyncMock()) + proxy_logging.email_logging_instance = MagicMock(budget_alerts=AsyncMock()) + await proxy_logging.budget_alerts(type="user_budget", user_info=_user_info()) + proxy_logging.slack_alerting_instance.budget_alerts.assert_not_called() + proxy_logging.email_logging_instance.budget_alerts.assert_called_once() + + @pytest.mark.asyncio async def test_budget_alerts_soft_budget_with_alert_emails_bypasses_global(proxy_logging): proxy_logging.alerting = None diff --git a/ui/litellm-dashboard/src/components/settings.tsx b/ui/litellm-dashboard/src/components/settings.tsx index 05e66985e6f..72da01d0918 100644 --- a/ui/litellm-dashboard/src/components/settings.tsx +++ b/ui/litellm-dashboard/src/components/settings.tsx @@ -522,7 +522,8 @@ const Settings: React.FC = ({ accessToken, userRole, userID,

- Alerts are only supported for Slack Webhook URLs. Get your webhook urls from{" "} + Alerts are sent to any Slack-compatible incoming webhook URL (Slack, Rocket.Chat, Mattermost, etc.). Get + Slack webhook urls from{" "} here @@ -532,7 +533,7 @@ const Settings: React.FC = ({ accessToken, userRole, userID, - Slack Webhook URL + Webhook URL (Slack-compatible) diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 137c67e837c..8a564a07489 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -25182,7 +25182,7 @@ export interface components { alert_types?: components["schemas"]["AlertType"][] | null; /** * Alerting - * @description List of alerting integrations. Today, just slack - `alerting: ['slack']` + * @description List of alerting integrations - e.g. `alerting: ['slack', 'webhook', 'email']`. 'slack' posts Slack-format messages to any Slack-compatible webhook (Slack, Rocket.Chat, Mattermost); 'webhook' posts structured JSON budget alerts to WEBHOOK_URL */ alerting?: unknown[] | null; /** From 39473745ddb6759b14397df3e9f9499e7b49ce0b Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 09:53:01 -0700 Subject: [PATCH 28/49] fix(docker): bump wolfi-base for glibc 2.44 and pin apk python to 3.13 in migrations image --- migrations/Dockerfile | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/migrations/Dockerfile b/migrations/Dockerfile index 6335e6f6bd8..c6d1b0cc46e 100644 --- a/migrations/Dockerfile +++ b/migrations/Dockerfile @@ -1,5 +1,5 @@ -ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72 -ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:a31344ab2cb8618db84f535eec56f76f6178b142cb92cb2e48676cc2dcebea72 +ARG LITELLM_BUILD_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d +ARG LITELLM_RUNTIME_IMAGE=cgr.dev/chainguard/wolfi-base@sha256:e624c5d5e42382ce7165ddafcbbf8e6769a24cbd02ea6114b880b05ae5ba2a8d ARG UV_IMAGE=ghcr.io/astral-sh/uv:0.11.7@sha256:240fb85ab0f263ef12f492d8476aa3a2e4e1e333f7d67fbdd923d00a506a516a FROM $UV_IMAGE AS uvbin @@ -35,7 +35,7 @@ COPY --from=uvbin /uv /uvx /usr/local/bin/ # instead of nodeenv downloading one whose dynamic deps may not be in Wolfi # (e.g. Node 26.2.0 needs libatomic). Retry for transient apk.cgr.dev flakes. RUN for i in 1 2 3; do \ - apk add --no-cache bash gcc python3 python3-dev openssl openssl-dev libsndfile nodejs npm && break; \ + apk add --no-cache bash gcc python-3.13 python-3.13-dev openssl openssl-dev libsndfile nodejs npm && break; \ [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ sleep 5; \ done @@ -56,7 +56,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ uv sync --frozen --no-install-project --no-install-workspace --no-default-groups --no-editable \ --extra proxy \ --extra extra_proxy \ - --python python3 + --python python3.13 # Stage 2 — copy source and install the project + workspace members. COPY . . @@ -65,7 +65,7 @@ RUN --mount=type=cache,target=/root/.cache/uv \ uv sync --frozen --no-default-groups --no-editable \ --extra proxy \ --extra extra_proxy \ - --python python3 + --python python3.13 COPY migrations/run.py /app/run.py @@ -87,7 +87,7 @@ FROM $LITELLM_RUNTIME_IMAGE AS runtime USER root RUN for i in 1 2 3; do \ - apk add --no-cache bash openssl tzdata python3 nodejs libsndfile libatomic && break; \ + apk add --no-cache bash openssl tzdata python-3.13 nodejs libsndfile libatomic && break; \ [ $i = 3 ] && { echo "apk add failed after 3 retries" >&2; exit 1; }; \ sleep 5; \ done From 30f32285103cea768c55f2f380a6422792363c91 Mon Sep 17 00:00:00 2001 From: yucheng-berri Date: Mon, 31 Aug 2026 09:58:35 -0700 Subject: [PATCH 29/49] test(newrelic): cover static default_team_settings per-team routing (#38857) * test(newrelic): cover static default_team_settings per-team routing The dynamic POST /team/{team_id}/callback path for New Relic is tested, but the static default_team_settings twin had no regression coverage. Add a test that drives default_team_settings -> add_team_based_callbacks_from_config and asserts the resolved trusted vars dispatch to BOTH the per-team metrics logger (cost/usage) and the trace logger (LLM/agent spans), so a config-file customer gets the same per-team routing as the API customer. Also correct the /team/callback docstring: callback_name is a str validated against the credential-capable callbacks, not a fixed langfuse/langsmith/gcs Literal, and document the newrelic_api_key / newrelic_region vars. * chore(ui): sync schema.d.ts with the /team/callback docstring Regenerate the dashboard OpenAPI types for the add_team_callbacks description change: callback_name is a validated str (not a langfuse/langsmith/gcs Literal) and the newrelic_api_key / newrelic_region vars are documented. * docs(newrelic): note LITELLM_OTEL_V2 prerequisite, trim test comments Address review: team-scoped New Relic config is rejected with a 400 unless the proxy runs with LITELLM_OTEL_V2=true, so document that in the /team/callback endpoint and sync schema.d.ts. Drop the narrative setup comments in the new test per the repo comment convention; the test name and docstring already say why. --- .../team_callback_endpoints.py | 4 +- tests/proxy_unit_tests/test_proxy_utils.py | 56 +++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 4 +- 3 files changed, 62 insertions(+), 2 deletions(-) diff --git a/litellm/proxy/management_endpoints/team_callback_endpoints.py b/litellm/proxy/management_endpoints/team_callback_endpoints.py index 08346983f32..c2f5dbb4032 100644 --- a/litellm/proxy/management_endpoints/team_callback_endpoints.py +++ b/litellm/proxy/management_endpoints/team_callback_endpoints.py @@ -252,7 +252,7 @@ async def add_team_callbacks( Use this if if you want different teams to have different success/failure callbacks Parameters: - - callback_name (Literal["langfuse", "langsmith", "gcs"], required): The name of the callback to add + - callback_name (str, required): The name of the callback to add, e.g. "langfuse", "langsmith", "gcs", "newrelic". The value is validated against the callbacks that support team-scoped credentials - callback_type (Literal["success", "failure", "success_and_failure"], required): The type of callback to add. One of: - "success": Callback for successful LLM calls - "failure": Callback for failed LLM calls @@ -268,6 +268,8 @@ async def add_team_callbacks( - langsmith_api_key: The API key for the Langsmith callback - langsmith_project: The project for the Langsmith callback - langsmith_base_url: The base URL for the Langsmith callback + - newrelic_api_key: The ingest license key for the team's New Relic account; routes both LLM/agent traces and cost metrics to that account. Requires the proxy to run with LITELLM_OTEL_V2=true, otherwise this callback is rejected with a 400 + - newrelic_region: The New Relic region for the team's account ("us" or "eu"), riding the team's own key Example curl: ``` diff --git a/tests/proxy_unit_tests/test_proxy_utils.py b/tests/proxy_unit_tests/test_proxy_utils.py index 3bde72ccd49..35de9961054 100644 --- a/tests/proxy_unit_tests/test_proxy_utils.py +++ b/tests/proxy_unit_tests/test_proxy_utils.py @@ -1317,6 +1317,62 @@ def test_proxy_config_state_post_init_callback_call(monkeypatch): assert config["litellm_settings"]["default_team_settings"][0]["team_id"] == "test" +@pytest.mark.asyncio +async def test_default_team_settings_newrelic_resolves_traces_and_metrics(): + """Static `default_team_settings` is the config-file twin of POST /team/callback. + + A team pinned to New Relic through `default_team_settings` must reach the + same two loggers the dynamic path does: the per-team metrics logger (cost + and usage) and the trace logger (LLM/agent spans). This proves the static + path resolves both, not just one, so the config-file customer gets the + same per-team routing as the API customer. + """ + from litellm.litellm_core_utils.litellm_logging import Logging + from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup + from litellm.proxy.proxy_server import ProxyConfig + + pc = ProxyConfig() + pc.config = { + "litellm_settings": { + "default_team_settings": [ + { + "team_id": "team-a", + "success_callback": ["newrelic"], + "newrelic_api_key": "team-a-ingest-key", + "newrelic_region": "eu", + } + ] + } + } + + callback_metadata = LiteLLMProxyRequestSetup.add_team_based_callbacks_from_config( + team_id="team-a", + proxy_config=pc, + ) + + assert callback_metadata is not None + assert callback_metadata.success_callback == ["newrelic"] + assert callback_metadata.callback_vars == { + "newrelic_api_key": "team-a-ingest-key", + "newrelic_region": "eu", + } + + logging_obj = Logging( + model="gpt-3.5-turbo", + messages=[{"role": "user", "content": "hi"}], + stream=False, + call_type="completion", + start_time=None, + litellm_call_id="static-nr-1", + function_id="static-nr-1", + ) + logging_obj._trusted_callback_vars = tuple(callback_metadata.callback_vars.items()) + + resolved = logging_obj._resolve_dynamic_callback_string("newrelic") + resolved_names = {type(logger).__name__ for logger in resolved} + assert resolved_names == {"NewRelicMetricsLogger", "NewRelicLogger"} + + def test_proxy_config_state_get_config_state_error(): """ Ensures that get_config_state does not raise an error when the config is not a valid dictionary diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 8a564a07489..81a2fefe32e 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -15505,7 +15505,7 @@ export interface paths { * Use this if if you want different teams to have different success/failure callbacks * * Parameters: - * - callback_name (Literal["langfuse", "langsmith", "gcs"], required): The name of the callback to add + * - callback_name (str, required): The name of the callback to add, e.g. "langfuse", "langsmith", "gcs", "newrelic". The value is validated against the callbacks that support team-scoped credentials * - callback_type (Literal["success", "failure", "success_and_failure"], required): The type of callback to add. One of: * - "success": Callback for successful LLM calls * - "failure": Callback for failed LLM calls @@ -15521,6 +15521,8 @@ export interface paths { * - langsmith_api_key: The API key for the Langsmith callback * - langsmith_project: The project for the Langsmith callback * - langsmith_base_url: The base URL for the Langsmith callback + * - newrelic_api_key: The ingest license key for the team's New Relic account; routes both LLM/agent traces and cost metrics to that account. Requires the proxy to run with LITELLM_OTEL_V2=true, otherwise this callback is rejected with a 400 + * - newrelic_region: The New Relic region for the team's account ("us" or "eu"), riding the team's own key * * Example curl: * ``` From e938e89d138eec2ef20a998e09d200dfb44b71a1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 11:19:11 -0700 Subject: [PATCH 30/49] docs(proxy): account for budget rollover and daily upserts in spend wording --- .../internal_user_endpoints.py | 25 +++++++++++-------- ui/litellm-dashboard/src/lib/http/schema.d.ts | 25 +++++++++++-------- 2 files changed, 28 insertions(+), 22 deletions(-) diff --git a/litellm/proxy/management_endpoints/internal_user_endpoints.py b/litellm/proxy/management_endpoints/internal_user_endpoints.py index 73e993b37a1..6edb75eacb4 100644 --- a/litellm/proxy/management_endpoints/internal_user_endpoints.py +++ b/litellm/proxy/management_endpoints/internal_user_endpoints.py @@ -997,12 +997,13 @@ async def user_info_v2( This is the v2 replacement for /user/info, designed to avoid the "god endpoint" problem where the old endpoint loaded all keys and teams into memory. - Note on `spend`: this is the user's running budget counter, which is zeroed by the - budget reset job whenever `budget_reset_at` elapses (see `budget_duration`). It is NOT + Note on `spend`: this is the user's running budget counter, which the budget reset job + resets whenever `budget_reset_at` elapses (see `budget_duration`): to zero by default, + or to the overage above `max_budget` when `budget_rollover` is enabled. It is NOT lifetime or per-period historical spend. For historical spend over a date range, use - `/user/daily/activity` or `/user/daily/activity/aggregated`, which read immutable daily - spend records that are never reset. The two values are expected to diverge once a - budget reset has occurred within the queried period. + `/user/daily/activity` or `/user/daily/activity/aggregated`, which read daily spend + records that only ever accumulate and are never reset. The two values are expected to + diverge once a budget reset has occurred within the queried period. Access control: - Proxy admins can query any user @@ -2694,9 +2695,10 @@ async def get_user_daily_activity( Meant to optimize querying spend data for analytics for a user. - Reads immutable daily spend records, which are never affected by budget resets. - This can legitimately exceed the `spend` field returned by `/v2/user/info`, which - is a running budget counter zeroed on every budget reset. + Reads daily spend records that only ever accumulate and are never affected by budget + resets. Their total can legitimately exceed the `spend` field returned by + `/v2/user/info`, which is a running budget counter that every budget reset sets back + to zero (or to the overage above `max_budget` when `budget_rollover` is enabled). Returns: (by date) @@ -2812,9 +2814,10 @@ async def get_user_daily_activity_aggregated( Aggregated analytics for a user's daily activity without pagination. Returns the same response shape as the paginated endpoint with page metadata set to single-page. - Reads immutable daily spend records, which are never affected by budget resets. - This can legitimately exceed the `spend` field returned by `/v2/user/info`, which - is a running budget counter zeroed on every budget reset. + Reads daily spend records that only ever accumulate and are never affected by budget + resets. Their total can legitimately exceed the `spend` field returned by + `/v2/user/info`, which is a running budget counter that every budget reset sets back + to zero (or to the overage above `max_budget` when `budget_rollover` is enabled). """ from litellm.proxy.proxy_server import prisma_client diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index b4526834b98..1599db10c7c 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -16210,9 +16210,10 @@ export interface paths { * * Meant to optimize querying spend data for analytics for a user. * - * Reads immutable daily spend records, which are never affected by budget resets. - * This can legitimately exceed the `spend` field returned by `/v2/user/info`, which - * is a running budget counter zeroed on every budget reset. + * Reads daily spend records that only ever accumulate and are never affected by budget + * resets. Their total can legitimately exceed the `spend` field returned by + * `/v2/user/info`, which is a running budget counter that every budget reset sets back + * to zero (or to the overage above `max_budget` when `budget_rollover` is enabled). * * Returns: * (by date) @@ -16246,9 +16247,10 @@ export interface paths { * @description Aggregated analytics for a user's daily activity without pagination. * Returns the same response shape as the paginated endpoint with page metadata set to single-page. * - * Reads immutable daily spend records, which are never affected by budget resets. - * This can legitimately exceed the `spend` field returned by `/v2/user/info`, which - * is a running budget counter zeroed on every budget reset. + * Reads daily spend records that only ever accumulate and are never affected by budget + * resets. Their total can legitimately exceed the `spend` field returned by + * `/v2/user/info`, which is a running budget counter that every budget reset sets back + * to zero (or to the overage above `max_budget` when `budget_rollover` is enabled). */ get: operations["get_user_daily_activity_aggregated_user_daily_activity_aggregated_get"]; put?: never; @@ -21012,12 +21014,13 @@ export interface paths { * This is the v2 replacement for /user/info, designed to avoid the "god endpoint" problem * where the old endpoint loaded all keys and teams into memory. * - * Note on `spend`: this is the user's running budget counter, which is zeroed by the - * budget reset job whenever `budget_reset_at` elapses (see `budget_duration`). It is NOT + * Note on `spend`: this is the user's running budget counter, which the budget reset job + * resets whenever `budget_reset_at` elapses (see `budget_duration`): to zero by default, + * or to the overage above `max_budget` when `budget_rollover` is enabled. It is NOT * lifetime or per-period historical spend. For historical spend over a date range, use - * `/user/daily/activity` or `/user/daily/activity/aggregated`, which read immutable daily - * spend records that are never reset. The two values are expected to diverge once a - * budget reset has occurred within the queried period. + * `/user/daily/activity` or `/user/daily/activity/aggregated`, which read daily spend + * records that only ever accumulate and are never reset. The two values are expected to + * diverge once a budget reset has occurred within the queried period. * * Access control: * - Proxy admins can query any user From cf1b431d58264399c0cf6f7a53e7bfd73b8560b3 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 11:50:45 -0700 Subject: [PATCH 31/49] fix(bedrock): stop duplicating Converse config blocks inside inferenceConfig --- .../bedrock/chat/converse_transformation.py | 9 ++++-- .../chat/test_converse_transformation.py | 30 ++++++++++++++----- 2 files changed, 29 insertions(+), 10 deletions(-) diff --git a/litellm/llms/bedrock/chat/converse_transformation.py b/litellm/llms/bedrock/chat/converse_transformation.py index db9c8a5cedd..395d99a4caa 100644 --- a/litellm/llms/bedrock/chat/converse_transformation.py +++ b/litellm/llms/bedrock/chat/converse_transformation.py @@ -1631,6 +1631,11 @@ class AmazonConverseConfig(BaseConfig): bedrock_tool_config["toolChoice"] = tool_choice_values self._drop_tool_choice_type_conflicting_with_tool_config(additional_request_params) + config_block_entries: Final = tuple( + (config_name, config_class, inference_params.pop(config_name, None)) + for config_name, config_class in self.get_config_blocks().items() + ) + data: Final[CommonRequestObject] = { "inferenceConfig": self._transform_inference_params(inference_params=inference_params), } @@ -1641,9 +1646,7 @@ class AmazonConverseConfig(BaseConfig): if system_content_blocks: data["system"] = system_content_blocks - # Handle all config blocks - for config_name, config_class in self.get_config_blocks().items(): - config_value = inference_params.pop(config_name, None) + for config_name, config_class, config_value in config_block_entries: if config_value is not None: data[config_name] = config_class(**config_value) 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 226bba6826a..63f895e1819 100644 --- a/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py +++ b/tests/test_litellm/llms/bedrock/chat/test_converse_transformation.py @@ -957,6 +957,28 @@ def test_transform_request_helper_includes_anthropic_beta_and_tools(): assert fields["tools"][0]["type"] == "computer_20250124" +def test_config_blocks_do_not_leak_into_inference_config(): + """Regression: inferenceConfig was built before the config blocks were popped, so a dead + nested copy of each block (guardrailConfig, performanceConfig, serviceTier) rode inside + inferenceConfig alongside the real top-level one.""" + data = AmazonConverseConfig()._transform_request_helper( + model="anthropic.claude-haiku-4-5-20251001-v1:0", + system_content_blocks=[], + optional_params={ + "maxTokens": 100, + "guardrailConfig": {"guardrailIdentifier": "gr-id", "guardrailVersion": "DRAFT"}, + "performanceConfig": {"latency": "optimized"}, + "serviceTier": {"type": "priority"}, + }, + messages=[{"role": "user", "content": "hi"}], + ) + + assert data["inferenceConfig"] == {"maxTokens": 100} + assert data["guardrailConfig"] == {"guardrailIdentifier": "gr-id", "guardrailVersion": "DRAFT"} + assert data["performanceConfig"] == {"latency": "optimized"} + assert data["serviceTier"] == {"type": "priority"} + + def test_parallel_tool_calls_config_kept_for_sonnet_5(monkeypatch): old_env = os.environ.get("LITELLM_LOCAL_MODEL_COST_MAP") old_cost = litellm.model_cost @@ -2853,17 +2875,11 @@ def test_guarded_text_guardrail_config_preserved(): headers={}, ) - # GuardrailConfig should be present at top level assert "guardrailConfig" in result assert result["guardrailConfig"]["guardrailIdentifier"] == "gr-abc123" - # GuardrailConfig should also be in inferenceConfig assert "inferenceConfig" in result - assert "guardrailConfig" in result["inferenceConfig"] - assert ( - result["inferenceConfig"]["guardrailConfig"]["guardrailIdentifier"] - == "gr-abc123" - ) + assert "guardrailConfig" not in result["inferenceConfig"] def test_auto_convert_last_user_message_to_guarded_text(): From ae945f4fa31d415c5ef7e56be90911ca0120f6d1 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 11:54:04 -0700 Subject: [PATCH 32/49] feat(openai): support workload identity federation (OIDC token exchange) --- basedpyright-code-budget.json | 2 +- litellm/llms/openai/common_utils.py | 1 + litellm/llms/openai/openai.py | 61 ++++-- .../llms/openai/responses/transformation.py | 11 + litellm/llms/openai/workload_identity.py | 92 +++++++++ .../openai/test_openai_workload_identity.py | 188 ++++++++++++++++++ 6 files changed, 335 insertions(+), 20 deletions(-) create mode 100644 litellm/llms/openai/workload_identity.py create mode 100644 tests/test_litellm/llms/openai/test_openai_workload_identity.py diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index d60c3e9c0af..229d1eca3e8 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -3,7 +3,7 @@ "limit": 16171 }, "reportArgumentType": { - "limit": 2226 + "limit": 2224 }, "reportAssignmentType": { "limit": 319 diff --git a/litellm/llms/openai/common_utils.py b/litellm/llms/openai/common_utils.py index 1b1ab80e85d..4d774f6f165 100644 --- a/litellm/llms/openai/common_utils.py +++ b/litellm/llms/openai/common_utils.py @@ -268,6 +268,7 @@ class BaseOpenAILLM: "max_retries", "organization", "api_base", + "workload_identity_config", ) openai_client_fields: Final = ( BaseOpenAILLM.get_openai_client_initialization_param_fields(client_type=client_type) diff --git a/litellm/llms/openai/openai.py b/litellm/llms/openai/openai.py index 6e66c998acf..16fa0017b23 100644 --- a/litellm/llms/openai/openai.py +++ b/litellm/llms/openai/openai.py @@ -51,6 +51,7 @@ from .common_utils import ( drop_params_from_unprocessable_entity_error, is_output_token_limit_error, ) +from .workload_identity import resolve_openai_workload_identity_config openaiOSeriesConfig: Final = OpenAIOSeriesConfig() openAIGPT5Config: Final = OpenAIGPT5Config() @@ -349,6 +350,7 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): client: OpenAI | AsyncOpenAI | None = None, shared_session: Optional["ClientSession"] = None, ) -> OpenAI | AsyncOpenAI | None: + workload_identity_config: Final = resolve_openai_workload_identity_config(api_key=api_key, api_base=api_base) client_initialization_params: Final[dict] = locals() if client is None: if not isinstance(max_retries, int): @@ -364,28 +366,49 @@ class OpenAIChatCompletion(BaseLLM, BaseOpenAILLM): if cached_client: if isinstance(cached_client, OpenAI) or isinstance(cached_client, AsyncOpenAI): return cached_client - http_client: Final[httpx.Client | httpx.AsyncClient | None] = ( - OpenAIChatCompletion._get_async_http_client(shared_session=shared_session) - if is_async - else OpenAIChatCompletion._get_sync_http_client() - ) if is_async: - _new_client: OpenAI | AsyncOpenAI = AsyncOpenAI( - api_key=api_key, - base_url=api_base, - http_client=http_client, - timeout=timeout, - max_retries=max_retries, - organization=organization, + async_http_client: Final = OpenAIChatCompletion._get_async_http_client(shared_session=shared_session) + http_client: httpx.Client | httpx.AsyncClient | None = async_http_client + _new_client: OpenAI | AsyncOpenAI = ( + AsyncOpenAI( + workload_identity=workload_identity_config.to_sdk_workload_identity(), + base_url=api_base, + http_client=async_http_client, + timeout=timeout, + max_retries=max_retries, + organization=organization, + ) + if workload_identity_config is not None + else AsyncOpenAI( + api_key=api_key, + base_url=api_base, + http_client=async_http_client, + timeout=timeout, + max_retries=max_retries, + organization=organization, + ) ) else: - _new_client = OpenAI( - api_key=api_key, - base_url=api_base, - http_client=http_client, - timeout=timeout, - max_retries=max_retries, - organization=organization, + sync_http_client: Final = OpenAIChatCompletion._get_sync_http_client() + http_client = sync_http_client + _new_client = ( + OpenAI( + workload_identity=workload_identity_config.to_sdk_workload_identity(), + base_url=api_base, + http_client=sync_http_client, + timeout=timeout, + max_retries=max_retries, + organization=organization, + ) + if workload_identity_config is not None + else OpenAI( + api_key=api_key, + base_url=api_base, + http_client=sync_http_client, + timeout=timeout, + max_retries=max_retries, + organization=organization, + ) ) ## SAVE CACHE KEY diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index eac844a790d..1479c378014 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -21,6 +21,7 @@ from litellm.types.router import GenericLiteLLMParams from litellm.types.utils import LlmProviders from ..common_utils import OpenAIError +from ..workload_identity import get_workload_identity_bearer_token, resolve_openai_workload_identity_config OPENAI_RESPONSES_API_MIN_MAX_OUTPUT_TOKENS: Final = 16 @@ -392,6 +393,16 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): litellm_params = litellm_params or GenericLiteLLMParams() api_key = litellm_params.api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") headers.setdefault("Content-Type", "application/json") + workload_identity_config: Final = resolve_openai_workload_identity_config( + api_key=api_key, + api_base=litellm_params.api_base + or litellm.api_base + or get_secret_str("OPENAI_BASE_URL") + or get_secret_str("OPENAI_API_BASE"), + ) + if workload_identity_config is not None: + headers["Authorization"] = f"Bearer {get_workload_identity_bearer_token(workload_identity_config)}" + return headers headers["Authorization"] = f"Bearer {api_key}" return headers diff --git a/litellm/llms/openai/workload_identity.py b/litellm/llms/openai/workload_identity.py new file mode 100644 index 00000000000..15105d67957 --- /dev/null +++ b/litellm/llms/openai/workload_identity.py @@ -0,0 +1,92 @@ +from __future__ import annotations + +from dataclasses import dataclass +from functools import lru_cache +from typing import TYPE_CHECKING, Final +from urllib.parse import urlparse + +from litellm.secret_managers.main import get_secret_str + +from .common_utils import OpenAIError + +if TYPE_CHECKING: + from collections.abc import Callable + + from openai.auth import SubjectTokenProvider, WorkloadIdentity, WorkloadIdentityAuth + +OPENAI_WIF_CLIENT_ID: Final = "litellm" +_OPENAI_API_HOST: Final = "api.openai.com" +_SDK_UPGRADE_MESSAGE: Final = ( + "OpenAI workload identity federation requires openai>=2.32.0. " + "Upgrade the installed openai package to use OPENAI_IDENTITY_PROVIDER_ID / " + "OPENAI_SERVICE_ACCOUNT_ID / OPENAI_IDENTITY_TOKEN_FILE." +) + + +@dataclass(frozen=True, slots=True) +class OpenAIWorkloadIdentityConfig: + identity_provider_id: str + service_account_id: str + token_file: str + + def to_sdk_workload_identity(self) -> WorkloadIdentity: + k8s_token_provider: Final = _load_sdk_k8s_token_provider() + workload_identity: Final[WorkloadIdentity] = { + "client_id": OPENAI_WIF_CLIENT_ID, + "identity_provider_id": self.identity_provider_id, + "service_account_id": self.service_account_id, + "provider": k8s_token_provider(self.token_file), + } + return workload_identity + + +def resolve_openai_workload_identity_config( + api_key: str | None, + api_base: str | None, +) -> OpenAIWorkloadIdentityConfig | None: + if api_key is not None: + return None + if not _targets_openai_api(api_base): + return None + identity_provider_id: Final = get_secret_str("OPENAI_IDENTITY_PROVIDER_ID") + service_account_id: Final = get_secret_str("OPENAI_SERVICE_ACCOUNT_ID") + token_file: Final = get_secret_str("OPENAI_IDENTITY_TOKEN_FILE") + if not identity_provider_id or not service_account_id or not token_file: + return None + return OpenAIWorkloadIdentityConfig( + identity_provider_id=identity_provider_id, + service_account_id=service_account_id, + token_file=token_file, + ) + + +def get_workload_identity_bearer_token(config: OpenAIWorkloadIdentityConfig) -> str: + return _workload_identity_auth(config).get_token() + + +def _targets_openai_api(api_base: str | None) -> bool: + if api_base is None: + return True + return urlparse(api_base).hostname == _OPENAI_API_HOST + + +@lru_cache(maxsize=16) +def _workload_identity_auth(config: OpenAIWorkloadIdentityConfig) -> WorkloadIdentityAuth: + sdk_workload_identity_auth: Final = _load_sdk_workload_identity_auth() + return sdk_workload_identity_auth(workload_identity=config.to_sdk_workload_identity()) + + +def _load_sdk_workload_identity_auth() -> type[WorkloadIdentityAuth]: + try: + from openai.auth import WorkloadIdentityAuth as sdk_workload_identity_auth + except ImportError as e: + raise OpenAIError(status_code=500, message=_SDK_UPGRADE_MESSAGE) from e + return sdk_workload_identity_auth + + +def _load_sdk_k8s_token_provider() -> Callable[[str], SubjectTokenProvider]: + try: + from openai.auth import k8s_service_account_token_provider + except ImportError as e: + raise OpenAIError(status_code=500, message=_SDK_UPGRADE_MESSAGE) from e + return k8s_service_account_token_provider diff --git a/tests/test_litellm/llms/openai/test_openai_workload_identity.py b/tests/test_litellm/llms/openai/test_openai_workload_identity.py new file mode 100644 index 00000000000..b8cc8c9c80f --- /dev/null +++ b/tests/test_litellm/llms/openai/test_openai_workload_identity.py @@ -0,0 +1,188 @@ +import json +import sys +from pathlib import Path +from typing import Final + +import httpx +import pytest +import respx +from openai import AsyncOpenAI, OpenAI + +import litellm +from litellm.llms.openai.common_utils import BaseOpenAILLM, OpenAIError +from litellm.llms.openai.openai import OpenAIChatCompletion +from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig +from litellm.llms.openai.workload_identity import ( + OpenAIWorkloadIdentityConfig, + _workload_identity_auth, + get_workload_identity_bearer_token, + resolve_openai_workload_identity_config, +) +from litellm.types.router import GenericLiteLLMParams + +TOKEN_EXCHANGE_URL: Final = "https://auth.openai.com/oauth/token" + + +@pytest.fixture +def wif_env(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> OpenAIWorkloadIdentityConfig: + token_file: Final = tmp_path / "subject_token.jwt" + token_file.write_text("subject-token-from-file") + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + monkeypatch.setenv("OPENAI_IDENTITY_PROVIDER_ID", "idp_test123") + monkeypatch.setenv("OPENAI_SERVICE_ACCOUNT_ID", "user-test456") + monkeypatch.setenv("OPENAI_IDENTITY_TOKEN_FILE", str(token_file)) + _workload_identity_auth.cache_clear() + litellm.in_memory_llm_clients_cache.flush_cache() + return OpenAIWorkloadIdentityConfig( + identity_provider_id="idp_test123", + service_account_id="user-test456", + token_file=str(token_file), + ) + + +def mock_token_exchange(access_token: str = "exchanged-bearer-token") -> respx.Route: + return respx.post(TOKEN_EXCHANGE_URL).mock( + return_value=httpx.Response(200, json={"access_token": access_token, "expires_in": 3600}) + ) + + +class TestResolveConfig: + def test_resolves_from_env(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: + assert resolve_openai_workload_identity_config(api_key=None, api_base=None) == wif_env + + def test_static_api_key_wins(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: + assert resolve_openai_workload_identity_config(api_key="sk-static", api_base=None) is None + + def test_foreign_api_base_disables(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: + assert resolve_openai_workload_identity_config(api_key=None, api_base="https://my-vllm.internal/v1") is None + + def test_openai_api_base_allows(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: + assert resolve_openai_workload_identity_config(api_key=None, api_base="https://api.openai.com/v1") == wif_env + + @pytest.mark.parametrize( + "missing_var", + ["OPENAI_IDENTITY_PROVIDER_ID", "OPENAI_SERVICE_ACCOUNT_ID", "OPENAI_IDENTITY_TOKEN_FILE"], + ) + def test_partial_env_disables( + self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch, missing_var: str + ) -> None: + monkeypatch.delenv(missing_var) + assert resolve_openai_workload_identity_config(api_key=None, api_base=None) is None + + +class TestTokenExchange: + @respx.mock + def test_exchanges_subject_token_for_bearer(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: + route: Final = mock_token_exchange() + assert get_workload_identity_bearer_token(wif_env) == "exchanged-bearer-token" + request_body: Final = json.loads(route.calls.last.request.content) + assert request_body["grant_type"] == "urn:ietf:params:oauth:grant-type:token-exchange" + assert request_body["subject_token"] == "subject-token-from-file" + assert request_body["subject_token_type"] == "urn:ietf:params:oauth:token-type:jwt" + assert request_body["identity_provider_id"] == "idp_test123" + assert request_body["service_account_id"] == "user-test456" + + @respx.mock + def test_token_cached_across_mints(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: + route: Final = mock_token_exchange() + first: Final = get_workload_identity_bearer_token(wif_env) + second: Final = get_workload_identity_bearer_token(wif_env) + assert first == second == "exchanged-bearer-token" + assert route.call_count == 1 + + def test_old_sdk_raises_upgrade_error( + self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch + ) -> None: + import openai as openai_module + + monkeypatch.delattr(openai_module, "auth", raising=False) + monkeypatch.setitem(sys.modules, "openai.auth", None) + with pytest.raises(OpenAIError, match=r"openai>=2\.32\.0"): + wif_env.to_sdk_workload_identity() + + +class TestClientConstruction: + def test_sync_client_uses_workload_identity(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: + client: Final = OpenAIChatCompletion()._get_openai_client(is_async=False, api_key=None, api_base=None) + assert isinstance(client, OpenAI) + assert client.api_key == "workload-identity-auth" + assert client._workload_identity_auth is not None + + def test_async_client_uses_workload_identity(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: + client: Final = OpenAIChatCompletion()._get_openai_client(is_async=True, api_key=None, api_base=None) + assert isinstance(client, AsyncOpenAI) + assert client.api_key == "workload-identity-auth" + assert client._workload_identity_auth is not None + + def test_static_key_client_unaffected(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: + client: Final = OpenAIChatCompletion()._get_openai_client(is_async=False, api_key="sk-static", api_base=None) + assert isinstance(client, OpenAI) + assert client.api_key == "sk-static" + assert client._workload_identity_auth is None + + def test_cache_key_separates_wif_identities(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: + other_config: Final = OpenAIWorkloadIdentityConfig( + identity_provider_id="idp_other", + service_account_id="user-other", + token_file=wif_env.token_file, + ) + keys: Final = tuple( + BaseOpenAILLM.get_openai_client_cache_key( + client_initialization_params={"api_key": None, "is_async": False, "workload_identity_config": config}, + client_type="openai", + ) + for config in (wif_env, other_config, None) + ) + assert len(set(keys)) == 3 + + @respx.mock + def test_request_carries_exchanged_bearer(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: + mock_token_exchange() + completion_route: Final = respx.post("https://api.openai.com/v1/chat/completions").mock( + return_value=httpx.Response( + 200, + json={ + "id": "chatcmpl-wif", + "object": "chat.completion", + "created": 1, + "model": "gpt-4o-mini", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": "ok"}, + "finish_reason": "stop", + } + ], + "usage": {"prompt_tokens": 1, "completion_tokens": 1, "total_tokens": 2}, + }, + ) + ) + client = OpenAIChatCompletion()._get_openai_client(is_async=False, api_key=None, api_base=None) + assert isinstance(client, OpenAI) + client.chat.completions.create(model="gpt-4o-mini", messages=[{"role": "user", "content": "hi"}]) + auth_header: Final = completion_route.calls.last.request.headers["Authorization"] + assert auth_header == "Bearer exchanged-bearer-token" + + +class TestResponsesValidateEnvironment: + @respx.mock + def test_mints_bearer_when_wif_configured(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: + mock_token_exchange() + headers: Final = OpenAIResponsesAPIConfig().validate_environment( + headers={}, model="gpt-4o-mini", litellm_params=GenericLiteLLMParams() + ) + assert headers["Authorization"] == "Bearer exchanged-bearer-token" + + def test_static_key_wins(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: + headers: Final = OpenAIResponsesAPIConfig().validate_environment( + headers={}, model="gpt-4o-mini", litellm_params=GenericLiteLLMParams(api_key="sk-responses") + ) + assert headers["Authorization"] == "Bearer sk-responses" + + def test_foreign_api_base_skips_wif(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: + headers: Final = OpenAIResponsesAPIConfig().validate_environment( + headers={}, + model="gpt-4o-mini", + litellm_params=GenericLiteLLMParams(api_base="https://my-vllm.internal/v1"), + ) + assert headers["Authorization"] == "Bearer None" From c7c382402a6ec0e26d32b45b30a2ee1c10ad64d5 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 12:04:44 -0700 Subject: [PATCH 33/49] feat(proxy): add /v1/responses/input_tokens token counting endpoint --- litellm/proxy/_types.py | 3 + .../proxy/response_api_endpoints/endpoints.py | 157 ++++++++++++- .../response_api_endpoints/test_endpoints.py | 215 +++++++++++++----- ui/litellm-dashboard/src/lib/http/schema.d.ts | 153 +++++++++++++ 4 files changed, 467 insertions(+), 61 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 0f2d97b8b1c..9574eb36ee5 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -422,6 +422,9 @@ class LiteLLMRoutes(enum.Enum): "/responses/{response_id}/cancel", "/v1/responses/{response_id}/cancel", "/openai/v1/responses/{response_id}/cancel", + "/responses/input_tokens", + "/v1/responses/input_tokens", + "/openai/v1/responses/input_tokens", # vector stores "/vector_stores", "/v1/vector_stores", diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 5e56e822484..100b42a9e2a 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1,14 +1,18 @@ import asyncio import json import time -from collections.abc import AsyncIterator, Mapping +from collections.abc import AsyncIterator, Awaitable, Mapping +from enum import Enum from types import MappingProxyType -from typing import TYPE_CHECKING, Any, Final, NamedTuple, cast, get_args +from typing import TYPE_CHECKING, Any, Final, NamedTuple, Protocol, cast, get_args from uuid import uuid4 import fastapi from fastapi import APIRouter, Depends, HTTPException, Request, Response +from fastapi.responses import JSONResponse +from openai.types.responses.response_create_params import ResponseInputParam from starlette.websockets import WebSocket, WebSocketDisconnect +from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger from litellm.integrations.custom_guardrail import ModifyResponseException @@ -26,8 +30,13 @@ from litellm.proxy.common_utils.http_parsing_utils import ( _read_request_body, _safe_set_request_parsed_body, ) -from litellm.types.llms.openai import REASONING_EFFORT, ResponsesAPIResponse +from litellm.types.llms.openai import ( + REASONING_EFFORT, + ResponsesAPIOptionalRequestParams, + ResponsesAPIResponse, +) from litellm.types.responses.main import DeleteResponseResult +from litellm.types.utils import TokenCountResponse if TYPE_CHECKING: from litellm.router import Router @@ -35,7 +44,7 @@ if TYPE_CHECKING: router: Final = APIRouter() _user_api_key_auth_dep: Final = Depends(user_api_key_auth) -_RESPONSES_TAGS: Final = ["responses"] # mutable-ok: fastapi's route signature requires List[str] tags +_RESPONSES_TAGS: Final[list[str | Enum]] = ["responses"] # mutable-ok: fastapi's route signature requires list tags _TOOL_PAYLOAD_KEYS: Final[Mapping[str, tuple[str, ...]]] = MappingProxyType( { @@ -1017,6 +1026,146 @@ async def compact_response( ) +class _ResponsesApiErrorDetail(TypedDict): + message: ReadOnly[str] + type: ReadOnly[str] + param: ReadOnly[str | None] + code: ReadOnly[str | None] + + +class _ResponsesApiErrorBody(TypedDict): + error: ReadOnly[_ResponsesApiErrorDetail] + + +class _ResponsesInputTokensResult(TypedDict): + object: ReadOnly[str] + input_tokens: ReadOnly[int] + + +class _TokenCountPayload(TypedDict): + model: ReadOnly[str] + messages: ReadOnly[tuple[Mapping[str, object], ...]] + tools: ReadOnly[object] + + +class _TokenCounter(Protocol): + def __call__(self, request: TokenCountRequest, call_endpoint: bool) -> Awaitable[TokenCountResponse]: ... + + +def _proxy_token_counter() -> _TokenCounter: + from litellm.proxy.proxy_server import token_counter + + return token_counter + + +_token_counter_dep: Final = Depends(_proxy_token_counter) + + +def _responses_invalid_request_response(message: str, param: str | None, code: str | None) -> JSONResponse: + body: Final[_ResponsesApiErrorBody] = { + "error": { + "message": message, + "type": "invalid_request_error", + "param": param, + "code": code, + } + } + return JSONResponse(status_code=400, content=body) + + +def _missing_responses_param_response(param: str) -> JSONResponse: + return _responses_invalid_request_response( + message=f"Missing required parameter: '{param}'.", + param=param, + code="missing_required_parameter", + ) + + +def _responses_input_as_token_count_messages( + input_value: str | ResponseInputParam, + instructions: str | None, +) -> tuple[Mapping[str, object], ...]: + from litellm.responses.litellm_completion_transformation.transformation import ( + LiteLLMCompletionResponsesConfig, + ) + + request_params: Final[ResponsesAPIOptionalRequestParams] = {"instructions": instructions} + transformed: Final = LiteLLMCompletionResponsesConfig.transform_responses_api_input_to_messages( + input=input_value, + responses_api_request=request_params, + ) + return tuple( + message if isinstance(message, dict) else message.model_dump(exclude_none=True) for message in transformed + ) + + +@router.post( + "/v1/responses/input_tokens", + dependencies=(_user_api_key_auth_dep,), + tags=_RESPONSES_TAGS, +) +@router.post( + "/responses/input_tokens", + dependencies=(_user_api_key_auth_dep,), + tags=_RESPONSES_TAGS, +) +@router.post( + "/openai/v1/responses/input_tokens", + dependencies=(_user_api_key_auth_dep,), + tags=_RESPONSES_TAGS, +) +async def responses_input_tokens( + request: Request, + token_counter: _TokenCounter = _token_counter_dep, +): + """ + Count the input tokens of a Responses API request without calling the model. + + Follows the OpenAI Responses API spec: https://platform.openai.com/docs/api-reference/responses/input-tokens + + ```bash + curl -X POST http://localhost:4000/v1/responses/input_tokens \ + -H "Content-Type: application/json" \ + -H "Authorization: Bearer sk-1234" \ + -d '{ + "model": "gpt-4o", + "input": "Hello, how are you?" + }' + ``` + + Returns: `{"object": "response.input_tokens", "input_tokens": }` + """ + data: Final = await _read_request_body(request=request) + model_name: Final = data.get("model") + input_value: Final = data.get("input") + if not isinstance(model_name, str) or not model_name: + return _missing_responses_param_response("model") + if input_value is None: + return _missing_responses_param_response("input") + + try: + payload: Final[_TokenCountPayload] = { + "model": model_name, + "messages": _responses_input_as_token_count_messages( + input_value=input_value, + instructions=data.get("instructions"), + ), + "tools": data.get("tools"), + } + token_request: Final = TokenCountRequest.model_validate(payload) + except Exception as e: + return _responses_invalid_request_response( + message=f"Invalid request for token counting: {e}", param=None, code=None + ) + + token_response: Final = await token_counter(request=token_request, call_endpoint=True) + result: Final[_ResponsesInputTokensResult] = { + "object": "response.input_tokens", + "input_tokens": token_response.total_tokens, + } + return result + + @router.post( "/v1/responses/{response_id}/cancel", dependencies=[Depends(user_api_key_auth)], diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index 791d64c6428..dc43e7f5c06 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -82,11 +82,7 @@ class TestResponsesAPIEndpoints(unittest.TestCase): ResponseOutputMessage( type="message", role="assistant", - content=[ - ResponseOutputText( - type="output_text", text="Hello from Cursor!" - ) - ], + content=[ResponseOutputText(type="output_text", text="Hello from Cursor!")], ) ], ) @@ -121,9 +117,7 @@ class TestResponsesAPIEndpoints(unittest.TestCase): @pytest.mark.asyncio @patch("litellm.proxy.proxy_server.llm_router") @patch("litellm.proxy.proxy_server.user_api_key_auth") - async def test_responses_api_key_spend_header_includes_response_cost( - self, mock_auth, mock_router - ): + async def test_responses_api_key_spend_header_includes_response_cost(self, mock_auth, mock_router): """ Test that x-litellm-key-spend header includes the current request's response_cost for /v1/responses endpoint. @@ -159,9 +153,7 @@ class TestResponsesAPIEndpoints(unittest.TestCase): ResponseOutputMessage( type="message", role="assistant", - content=[ - ResponseOutputText(type="output_text", text="Test response") - ], + content=[ResponseOutputText(type="output_text", text="Test response")], ) ], ) @@ -356,6 +348,7 @@ class TestWSModelExtraction: from litellm.proxy.response_api_endpoints.endpoints import ( _extract_model_from_first_ws_event, ) + event = {"type": "response.create", "model": "gpt-4o", "input": "hello"} assert _extract_model_from_first_ws_event(event) == "gpt-4o" @@ -363,6 +356,7 @@ class TestWSModelExtraction: from litellm.proxy.response_api_endpoints.endpoints import ( _extract_model_from_first_ws_event, ) + event = {"type": "response.create", "response": {"model": "gpt-4o", "input": "hello"}} assert _extract_model_from_first_ws_event(event) == "gpt-4o" @@ -370,6 +364,7 @@ class TestWSModelExtraction: from litellm.proxy.response_api_endpoints.endpoints import ( _extract_model_from_first_ws_event, ) + event = { "type": "response.create", "model": "flat-model", @@ -381,6 +376,7 @@ class TestWSModelExtraction: from litellm.proxy.response_api_endpoints.endpoints import ( _extract_model_from_first_ws_event, ) + event = {"type": "response.create", "input": "hello"} assert _extract_model_from_first_ws_event(event) is None @@ -400,9 +396,7 @@ class TestResponsesWSFirstFrameValidation: ) ws = MagicMock() - ws.receive_text = AsyncMock( - return_value=json.dumps({"type": "session.update", "model": "gpt-4o"}) - ) + ws.receive_text = AsyncMock(return_value=json.dumps({"type": "session.update", "model": "gpt-4o"})) ws.send_text = AsyncMock() ws.close = AsyncMock() @@ -412,10 +406,7 @@ class TestResponsesWSFirstFrameValidation: ws.send_text.assert_awaited_once() ws.close.assert_awaited_once_with(code=1008, reason="Invalid first message") error_payload = json.loads(ws.send_text.await_args.args[0]) - assert ( - error_payload["error"]["message"] - == "First message must be a response.create JSON object." - ) + assert error_payload["error"]["message"] == "First message must be a response.create JSON object." @pytest.mark.asyncio async def test_rejects_non_object_json_first_frame(self): @@ -484,16 +475,12 @@ class TestResponsesWSFirstFrameModelAuth: ws.url = "ws://testserver/v1/responses" ws.accept = AsyncMock() ws.receive_text = AsyncMock( - return_value=json.dumps( - {"type": "response.create", "model": "gpt-4o-mini", "input": []} - ) + return_value=json.dumps({"type": "response.create", "model": "gpt-4o-mini", "input": []}) ) ws.close = AsyncMock() processor = MagicMock() - processor.common_processing_pre_call_logic = AsyncMock( - return_value=({"model": "gpt-4o-mini"}, MagicMock()) - ) + processor.common_processing_pre_call_logic = AsyncMock(return_value=({"model": "gpt-4o-mini"}, MagicMock())) async def fake_llm_call(): return None @@ -529,9 +516,7 @@ class TestResponsesWSFirstFrameModelAuth: _enforce_responses_ws_first_frame_model_auth, ) - request = Request( - {"type": "http", "method": "POST", "path": "/v1/responses", "headers": []} - ) + request = Request({"type": "http", "method": "POST", "path": "/v1/responses", "headers": []}) user_api_key_dict = MagicMock() llm_router = MagicMock() @@ -593,9 +578,7 @@ class TestReadWSModelFromFirstFrameErrors: assert result is None ws.send_text.assert_not_awaited() - ws.close.assert_awaited_once_with( - code=1008, reason="Timed out waiting for first message" - ) + ws.close.assert_awaited_once_with(code=1008, reason="Timed out waiting for first message") @pytest.mark.asyncio async def test_invalid_json_sends_error_and_closes(self): @@ -613,9 +596,7 @@ class TestReadWSModelFromFirstFrameErrors: assert result is None payload = json.loads(ws.send_text.await_args.args[0]) assert payload["error"]["message"] == "First message is not valid JSON." - ws.close.assert_awaited_once_with( - code=1008, reason="Invalid JSON in first message" - ) + ws.close.assert_awaited_once_with(code=1008, reason="Invalid JSON in first message") @pytest.mark.asyncio async def test_missing_model_sends_error_and_closes(self): @@ -624,9 +605,7 @@ class TestReadWSModelFromFirstFrameErrors: ) ws = MagicMock() - ws.receive_text = AsyncMock( - return_value=json.dumps({"type": "response.create", "input": []}) - ) + ws.receive_text = AsyncMock(return_value=json.dumps({"type": "response.create", "input": []})) ws.send_text = AsyncMock() ws.close = AsyncMock() @@ -679,10 +658,7 @@ class TestManagedResponsesSameProvider: assert self._handler("gpt-4o")._same_provider("gpt-4o-mini") is True def test_different_provider_is_not_same(self): - assert ( - self._handler("gpt-4o")._same_provider("vertex_ai/gemini-2.0-flash") - is False - ) + assert self._handler("gpt-4o")._same_provider("vertex_ai/gemini-2.0-flash") is False def test_inject_credentials_keeps_provider_for_same_provider_model(self): handler = self._handler("gpt-4o", custom_llm_provider="openai") @@ -697,18 +673,14 @@ class TestManagedResponsesSameProvider: assert "custom_llm_provider" not in call_kwargs def test_unresolvable_connection_model_falls_back_to_custom_provider(self): - handler = self._handler( - "my-custom-deployment", custom_llm_provider="openai" - ) + handler = self._handler("my-custom-deployment", custom_llm_provider="openai") assert handler._same_provider("gpt-4o-mini") is True call_kwargs: dict = {} handler._inject_credentials(call_kwargs, model="gpt-4o-mini") assert call_kwargs["custom_llm_provider"] == "openai" def test_unresolvable_connection_model_still_drops_cross_provider(self): - handler = self._handler( - "my-custom-deployment", custom_llm_provider="openai" - ) + handler = self._handler("my-custom-deployment", custom_llm_provider="openai") call_kwargs: dict = {} handler._inject_credentials(call_kwargs, model="vertex_ai/gemini-2.0-flash") assert "custom_llm_provider" not in call_kwargs @@ -840,9 +812,7 @@ def test_cursor_chat_completions_input_body_uses_responses_pipeline_and_strips_s type="message", role="assistant", status="completed", - content=[ - ResponseOutputText(type="output_text", text="agent reply", annotations=[]) - ], + content=[ResponseOutputText(type="output_text", text="agent reply", annotations=[])], ) ], ) @@ -851,9 +821,12 @@ def test_cursor_chat_completions_input_body_uses_responses_pipeline_and_strips_s app.dependency_overrides[user_api_key_auth] = _auth_override try: - with patch.object(ps, "llm_router", mock_router), patch( - "litellm.proxy.response_api_endpoints.endpoints._read_request_body", - side_effect=capturing_read_request_body, + with ( + patch.object(ps, "llm_router", mock_router), + patch( + "litellm.proxy.response_api_endpoints.endpoints._read_request_body", + side_effect=capturing_read_request_body, + ), ): client = TestClient(app) response = client.post( @@ -1488,8 +1461,8 @@ def _router_serving_only(base_model: str) -> MagicMock: mock_router.router_general_settings.pass_through_all_models = False mock_router.default_deployment = None mock_router.pattern_router.patterns = {base_model: ["anthropic/*"]} - mock_router.pattern_router.get_pattern.side_effect = ( - lambda model: [{"model_name": "anthropic/*"}] if model == base_model else None + mock_router.pattern_router.get_pattern.side_effect = lambda model: ( + [{"model_name": "anthropic/*"}] if model == base_model else None ) return mock_router @@ -1739,9 +1712,7 @@ class TestCursorGateRecognizesRoutingGroups: from litellm.proxy.response_api_endpoints.endpoints import _resolve_cursor_model_variant router = Router( - model_list=[ - {"model_name": "member-fast", "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"}} - ], + model_list=[{"model_name": "member-fast", "litellm_params": {"model": "openai/gpt-4o", "api_key": "fake"}}], routing_groups=[ {"group_name": "grouped-thinking-high", "models": ["member-fast"], "routing_strategy": "simple-shuffle"} ], @@ -1836,3 +1807,133 @@ class TestGuardrailBlockedResponsesUsage: assert usage["input_tokens"] == 0 assert usage["output_tokens"] == 0 assert usage["total_tokens"] == 0 + + +class TestResponsesInputTokens: + """Regression tests for POST /v1/responses/input_tokens. + + The docs promise OpenAI-format token counting on the proxy, but the route was + never registered, so the POST fell through to the GET/DELETE-only + /v1/responses/{response_id} route and returned 405.""" + + def _post_input_tokens(self, body, path="/v1/responses/input_tokens", counter=None): + from litellm.proxy._types import UserAPIKeyAuth + from litellm.proxy.auth.user_api_key_auth import user_api_key_auth + from litellm.proxy.response_api_endpoints.endpoints import _proxy_token_counter + from litellm.types.utils import TokenCountResponse + + token_counter_mock = ( + counter + if counter is not None + else AsyncMock( + return_value=TokenCountResponse( + total_tokens=13, + request_model=body.get("model", ""), + model_used=body.get("model", ""), + tokenizer_type="openai_api", + ) + ) + ) + app.dependency_overrides[user_api_key_auth] = lambda: UserAPIKeyAuth(api_key="sk-test", request_route=path) + app.dependency_overrides[_proxy_token_counter] = lambda: token_counter_mock + try: + client = TestClient(app) + response = client.post(path, json=body, headers={"Authorization": "Bearer sk-1234"}) + return response, token_counter_mock + finally: + app.dependency_overrides.pop(user_api_key_auth, None) + app.dependency_overrides.pop(_proxy_token_counter, None) + + def test_string_input_returns_openai_input_tokens_shape(self): + response, counter = self._post_input_tokens({"model": "gpt-4o", "input": "Hello, how are you?"}) + + assert response.status_code == 200, response.text + assert response.json() == {"object": "response.input_tokens", "input_tokens": 13} + counter.assert_awaited_once() + assert counter.call_args.kwargs["call_endpoint"] is True + token_request = counter.call_args.kwargs["request"] + assert token_request.model == "gpt-4o" + assert token_request.messages == [{"role": "user", "content": "Hello, how are you?"}] + + def test_every_route_alias_is_registered(self): + for path in ("/v1/responses/input_tokens", "/responses/input_tokens", "/openai/v1/responses/input_tokens"): + response, _ = self._post_input_tokens({"model": "gpt-4o", "input": "hi"}, path=path) + assert response.status_code == 200, f"{path}: {response.status_code} {response.text}" + + def test_input_items_instructions_and_tools_are_forwarded(self): + tools = [ + { + "type": "function", + "name": "get_weather", + "description": "Get weather for a city", + "parameters": {"type": "object", "properties": {"city": {"type": "string"}}}, + } + ] + response, counter = self._post_input_tokens( + { + "model": "gpt-4o", + "input": [{"role": "user", "content": "What is the weather in Paris?"}], + "instructions": "You are terse.", + "tools": tools, + } + ) + + assert response.status_code == 200, response.text + token_request = counter.call_args.kwargs["request"] + assert token_request.messages == [ + {"role": "system", "content": "You are terse."}, + {"role": "user", "content": "What is the weather in Paris?"}, + ] + assert token_request.tools == tools + + def test_missing_model_returns_openai_400(self): + response, counter = self._post_input_tokens({"input": "Hello"}) + + assert response.status_code == 400, response.text + assert response.json() == { + "error": { + "message": "Missing required parameter: 'model'.", + "type": "invalid_request_error", + "param": "model", + "code": "missing_required_parameter", + } + } + counter.assert_not_awaited() + + def test_missing_input_returns_openai_400(self): + response, counter = self._post_input_tokens({"model": "gpt-4o"}) + + assert response.status_code == 400, response.text + assert response.json() == { + "error": { + "message": "Missing required parameter: 'input'.", + "type": "invalid_request_error", + "param": "input", + "code": "missing_required_parameter", + } + } + counter.assert_not_awaited() + + def test_invalid_tools_returns_openai_400(self): + response, counter = self._post_input_tokens({"model": "gpt-4o", "input": "hi", "tools": "not-a-list"}) + + assert response.status_code == 400, response.text + error = response.json()["error"] + assert error["type"] == "invalid_request_error" + counter.assert_not_awaited() + + def test_provider_error_maps_status_code(self): + from litellm.proxy._types import ProxyException + + failing_counter = AsyncMock( + side_effect=ProxyException( + message="rate limited", + type="token_counting_error", + param="model", + code="429", + ) + ) + response, _ = self._post_input_tokens({"model": "gpt-4o", "input": "hi"}, counter=failing_counter) + + assert response.status_code == 429, response.text + assert response.json()["error"]["message"] == "rate limited" diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 81a2fefe32e..01a1b516d11 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -9700,6 +9700,37 @@ export interface paths { patch?: never; trace?: never; }; + "/openai/v1/responses/input_tokens": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Responses Input Tokens + * @description Count the input tokens of a Responses API request without calling the model. + * + * Follows the OpenAI Responses API spec: https://platform.openai.com/docs/api-reference/responses/input-tokens + * + * ```bash + * curl -X POST http://localhost:4000/v1/responses/input_tokens -H "Content-Type: application/json" -H "Authorization: Bearer sk-1234" -d '{ + * "model": "gpt-4o", + * "input": "Hello, how are you?" + * }' + * ``` + * + * Returns: `{"object": "response.input_tokens", "input_tokens": }` + */ + post: operations["responses_input_tokens_openai_v1_responses_input_tokens_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/openai/v1/responses/{response_id}": { parameters: { query?: never; @@ -12619,6 +12650,37 @@ export interface paths { patch?: never; trace?: never; }; + "/responses/input_tokens": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Responses Input Tokens + * @description Count the input tokens of a Responses API request without calling the model. + * + * Follows the OpenAI Responses API spec: https://platform.openai.com/docs/api-reference/responses/input-tokens + * + * ```bash + * curl -X POST http://localhost:4000/v1/responses/input_tokens -H "Content-Type: application/json" -H "Authorization: Bearer sk-1234" -d '{ + * "model": "gpt-4o", + * "input": "Hello, how are you?" + * }' + * ``` + * + * Returns: `{"object": "response.input_tokens", "input_tokens": }` + */ + post: operations["responses_input_tokens_responses_input_tokens_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/responses/{response_id}": { parameters: { query?: never; @@ -19184,6 +19246,37 @@ export interface paths { patch?: never; trace?: never; }; + "/v1/responses/input_tokens": { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + get?: never; + put?: never; + /** + * Responses Input Tokens + * @description Count the input tokens of a Responses API request without calling the model. + * + * Follows the OpenAI Responses API spec: https://platform.openai.com/docs/api-reference/responses/input-tokens + * + * ```bash + * curl -X POST http://localhost:4000/v1/responses/input_tokens -H "Content-Type: application/json" -H "Authorization: Bearer sk-1234" -d '{ + * "model": "gpt-4o", + * "input": "Hello, how are you?" + * }' + * ``` + * + * Returns: `{"object": "response.input_tokens", "input_tokens": }` + */ + post: operations["responses_input_tokens_v1_responses_input_tokens_post"]; + delete?: never; + options?: never; + head?: never; + patch?: never; + trace?: never; + }; "/v1/responses/{response_id}": { parameters: { query?: never; @@ -51476,6 +51569,26 @@ export interface operations { }; }; }; + responses_input_tokens_openai_v1_responses_input_tokens_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; get_response_openai_v1_responses__response_id__get: { parameters: { query?: never; @@ -54440,6 +54553,26 @@ export interface operations { }; }; }; + responses_input_tokens_responses_input_tokens_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; get_response_responses__response_id__get: { parameters: { query?: never; @@ -62862,6 +62995,26 @@ export interface operations { }; }; }; + responses_input_tokens_v1_responses_input_tokens_post: { + parameters: { + query?: never; + header?: never; + path?: never; + cookie?: never; + }; + requestBody?: never; + responses: { + /** @description Successful Response */ + 200: { + headers: { + [name: string]: unknown; + }; + content: { + "application/json": unknown; + }; + }; + }; + }; get_response_v1_responses__response_id__get: { parameters: { query?: never; From 9f67a58198ec2b2d992a88812e2fae3c304bfd2e Mon Sep 17 00:00:00 2001 From: davida-ps Date: Mon, 31 Aug 2026 22:05:57 +0300 Subject: [PATCH 34/49] fix(guardrails): configure Prompt Security file timeout policy (#38083) * fix(guardrails): fail open on Prompt Security file timeouts * fix(guardrails): configure Prompt Security timeout policy --- .../prompt_security/__init__.py | 1 + .../prompt_security/prompt_security.py | 49 +++++++++ .../guardrail_hooks/prompt_security.py | 4 + .../test_prompt_security_guardrails.py | 103 ++++++++++++++++-- 4 files changed, 148 insertions(+), 9 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/__init__.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/__init__.py index fa1f9f3d36d..0aaba4016cd 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/__init__.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/__init__.py @@ -20,6 +20,7 @@ def initialize_guardrail(litellm_params: "LitellmParams", guardrail: "Guardrail" guardrail_name=guardrail.get("guardrail_name", ""), event_hook=litellm_params.mode, default_on=litellm_params.default_on, + file_sanitization_fail_open=getattr(litellm_params, "file_sanitization_fail_open", None), ) litellm.logging_callback_manager.add_litellm_callback(_prompt_security_callback) diff --git a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py index 809d5e0fb31..84c4f118b00 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py +++ b/litellm/proxy/guardrails/guardrail_hooks/prompt_security/prompt_security.py @@ -4,10 +4,12 @@ import os from collections.abc import Mapping, Sequence from typing import TYPE_CHECKING, Final, Literal, Optional +import httpx from fastapi import HTTPException from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger +from litellm.exceptions import Timeout as LiteLLMTimeout from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, @@ -24,6 +26,9 @@ if TYPE_CHECKING: from litellm.types.proxy.guardrails.guardrail_hooks.base import GuardrailConfigModel +_SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS: Final = 30.0 + + class PromptSecurityGuardrailMissingSecrets(Exception): pass @@ -63,6 +68,13 @@ class _SanitizeStatusResponse(TypedDict, total=False): metadata: ReadOnly[_SanitizeMetadata] +class _SanitizeResult(TypedDict): + action: ReadOnly[str] + content: ReadOnly[str | None] + metadata: ReadOnly[_SanitizeMetadata] + violations: ReadOnly[Sequence[str]] + + class PromptSecurityGuardrail(CustomGuardrail): @classmethod def get_supported_event_hooks(cls) -> list[GuardrailEventHooks]: @@ -79,6 +91,8 @@ class PromptSecurityGuardrail(CustomGuardrail): user: str | None = None, system_prompt: str | None = None, check_tool_results: bool | None = None, + file_sanitization_timeout: float = _SANITIZE_FILE_FAIL_OPEN_TIMEOUT_SECONDS, + file_sanitization_fail_open: bool | None = None, **kwargs, ): kwargs.setdefault("supported_event_hooks", list(self.get_supported_event_hooks())) @@ -108,6 +122,8 @@ class PromptSecurityGuardrail(CustomGuardrail): # Configuration for file sanitization self.max_poll_attempts = 30 # Maximum number of polling attempts self.poll_interval = 2 # Seconds between polling attempts + self.file_sanitization_timeout = file_sanitization_timeout + self.file_sanitization_fail_open = file_sanitization_fail_open is not False super().__init__(**kwargs) @@ -397,6 +413,39 @@ class PromptSecurityGuardrail(CustomGuardrail): Sanitize file content using Prompt Security API. Returns: dict with keys 'action', 'content', 'metadata' """ + try: + return await asyncio.wait_for( + self._sanitize_file_content(file_data, filename, user_api_key_alias), + timeout=self.file_sanitization_timeout, + ) + except (asyncio.TimeoutError, httpx.TimeoutException, LiteLLMTimeout) as exc: + if not self.file_sanitization_fail_open: + verbose_proxy_logger.error( + "Prompt Security Guardrail: file sanitization for %s timed out with %s; failing closed", + filename, + type(exc).__name__, + ) + raise HTTPException(status_code=408, detail="File sanitization timeout") from exc + + verbose_proxy_logger.error( + "Prompt Security Guardrail: file sanitization for %s timed out with %s; failing open", + filename, + type(exc).__name__, + ) + fail_open_result: Final[_SanitizeResult] = { + "action": "allow", + "content": None, + "metadata": {}, + "violations": (), + } + return fail_open_result + + async def _sanitize_file_content( + self, + file_data: bytes, + filename: str, + user_api_key_alias: str | None, + ) -> _SanitizeResult: headers: Final = {"APP-ID": self.api_key} if user_api_key_alias: headers["X-LiteLLM-Key-Alias"] = user_api_key_alias diff --git a/litellm/types/proxy/guardrails/guardrail_hooks/prompt_security.py b/litellm/types/proxy/guardrails/guardrail_hooks/prompt_security.py index 6e64f0f47a5..94f8161f44e 100644 --- a/litellm/types/proxy/guardrails/guardrail_hooks/prompt_security.py +++ b/litellm/types/proxy/guardrails/guardrail_hooks/prompt_security.py @@ -12,6 +12,10 @@ class PromptSecurityGuardrailConfigModel(GuardrailConfigModel): default=None, description="The API base for the Prompt Security guardrail. If not provided, the `PROMPT_SECURITY_API_BASE` environment variable is used.", ) + file_sanitization_fail_open: bool = Field( + default=True, + description="Whether file sanitization timeouts allow the original file through instead of blocking the request.", + ) @staticmethod def ui_friendly_name() -> str: diff --git a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py index 26beaa78a46..ab4e15ff423 100644 --- a/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py +++ b/tests/test_litellm/proxy/guardrails/test_prompt_security_guardrails.py @@ -1,16 +1,16 @@ -from fastapi.exceptions import HTTPException -from unittest.mock import patch, AsyncMock -from httpx import Response, Request +import asyncio import base64 +from unittest.mock import AsyncMock, patch import pytest - -from litellm.proxy.guardrails.guardrail_hooks.prompt_security.prompt_security import ( - PromptSecurityGuardrailMissingSecrets, - PromptSecurityGuardrail, -) +from fastapi.exceptions import HTTPException +from httpx import ReadTimeout, Request, Response import litellm +from litellm.proxy.guardrails.guardrail_hooks.prompt_security.prompt_security import ( + PromptSecurityGuardrail, + PromptSecurityGuardrailMissingSecrets, +) from litellm.proxy.guardrails.init_guardrails import init_guardrails_v2 @@ -30,6 +30,7 @@ def test_prompt_security_guard_config(monkeypatch: pytest.MonkeyPatch): "guardrail": "prompt_security", "mode": "during_call", "default_on": True, + "file_sanitization_fail_open": False, }, } ], @@ -41,6 +42,10 @@ def test_prompt_security_guard_config(monkeypatch: pytest.MonkeyPatch): assert registered[0].guardrail_name == "prompt_security" assert registered[0].default_on is True assert registered[0].event_hook == "during_call" + assert registered[0].file_sanitization_fail_open is False + config_model = registered[0].get_config_model() + assert config_model is not None + assert config_model().file_sanitization_fail_open is True def test_prompt_security_guard_config_no_api_key(monkeypatch: pytest.MonkeyPatch): @@ -374,6 +379,86 @@ async def test_file_sanitization(monkeypatch: pytest.MonkeyPatch): assert result is not None +@pytest.mark.asyncio +@pytest.mark.parametrize( + "timeout", + ( + litellm.Timeout( + message="Prompt Security upload timed out", + model="default-model-name", + llm_provider="litellm-httpx-handler", + ), + ReadTimeout( + "Prompt Security poll timed out", + request=Request(method="GET", url="https://test.prompt.security/api/sanitizeFile"), + ), + ), + ids=("litellm", "httpx"), +) +@pytest.mark.parametrize("fail_open", (True, False), ids=("fail-open", "fail-closed")) +async def test_file_sanitization_request_timeout_policy( + monkeypatch: pytest.MonkeyPatch, timeout: Exception, fail_open: bool +): + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") + + guardrail = PromptSecurityGuardrail( + guardrail_name="test-guard", + event_hook="pre_call", + default_on=True, + file_sanitization_fail_open=fail_open, + ) + + with patch.object(guardrail.async_handler, "post", AsyncMock(side_effect=timeout)): + if not fail_open: + with pytest.raises(HTTPException) as exc_info: + await guardrail.sanitize_file_content(b"file-content", "document.pdf") + assert exc_info.value.status_code == 408 + assert exc_info.value.detail == "File sanitization timeout" + return + + result = await guardrail.sanitize_file_content(b"file-content", "document.pdf") + + assert result == { + "action": "allow", + "content": None, + "metadata": {}, + "violations": (), + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("fail_open", (True, False), ids=("fail-open", "fail-closed")) +async def test_file_sanitization_overall_timeout_policy(monkeypatch: pytest.MonkeyPatch, fail_open: bool): + monkeypatch.setenv("PROMPT_SECURITY_API_KEY", "test-key") + monkeypatch.setenv("PROMPT_SECURITY_API_BASE", "https://test.prompt.security") + + guardrail = PromptSecurityGuardrail( + guardrail_name="test-guard", + event_hook="pre_call", + default_on=True, + file_sanitization_timeout=0.01, + file_sanitization_fail_open=fail_open, + ) + + async def hanging_post(*_args: object, **_kwargs: object) -> None: + await asyncio.sleep(60) + raise AssertionError("sanitization request should have been cancelled") + + with patch.object(guardrail.async_handler, "post", side_effect=hanging_post): + if not fail_open: + with pytest.raises(HTTPException) as exc_info: + await guardrail.sanitize_file_content(b"file-content", "document.pdf") + assert exc_info.value.status_code == 408 + assert exc_info.value.detail == "File sanitization timeout" + return + + result = await guardrail.sanitize_file_content(b"file-content", "document.pdf") + + assert result["action"] == "allow" + assert result["content"] is None + + @pytest.mark.asyncio async def test_file_sanitization_block(monkeypatch: pytest.MonkeyPatch): """Test that file sanitization blocks malicious files""" @@ -544,7 +629,7 @@ async def test_role_filtering(monkeypatch: pytest.MonkeyPatch): return mock_response with patch.object(guardrail.async_handler, "post", side_effect=mock_post): - result = await guardrail.apply_guardrail( + await guardrail.apply_guardrail( inputs=inputs, request_data=request_data, input_type="request", From 72adeda9ce3bb5cf61ca6816f01fc5693b85a1e2 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 12:15:52 -0700 Subject: [PATCH 35/49] fix(openai): scope workload identity to the openai provider and env-resolved base/key --- .../llms/openai/responses/transformation.py | 10 +++--- litellm/llms/openai/workload_identity.py | 8 +++-- .../openai/test_openai_workload_identity.py | 34 +++++++++++++++++++ 3 files changed, 44 insertions(+), 8 deletions(-) diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 1479c378014..eadc087383a 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -393,12 +393,10 @@ class OpenAIResponsesAPIConfig(BaseResponsesAPIConfig): litellm_params = litellm_params or GenericLiteLLMParams() api_key = litellm_params.api_key or litellm.api_key or litellm.openai_key or get_secret_str("OPENAI_API_KEY") headers.setdefault("Content-Type", "application/json") - workload_identity_config: Final = resolve_openai_workload_identity_config( - api_key=api_key, - api_base=litellm_params.api_base - or litellm.api_base - or get_secret_str("OPENAI_BASE_URL") - or get_secret_str("OPENAI_API_BASE"), + workload_identity_config: Final = ( + resolve_openai_workload_identity_config(api_key=api_key, api_base=litellm_params.api_base) + if self.custom_llm_provider is LlmProviders.OPENAI + else None ) if workload_identity_config is not None: headers["Authorization"] = f"Bearer {get_workload_identity_bearer_token(workload_identity_config)}" diff --git a/litellm/llms/openai/workload_identity.py b/litellm/llms/openai/workload_identity.py index 15105d67957..be4d015af8a 100644 --- a/litellm/llms/openai/workload_identity.py +++ b/litellm/llms/openai/workload_identity.py @@ -5,6 +5,7 @@ from functools import lru_cache from typing import TYPE_CHECKING, Final from urllib.parse import urlparse +import litellm from litellm.secret_managers.main import get_secret_str from .common_utils import OpenAIError @@ -44,9 +45,12 @@ def resolve_openai_workload_identity_config( api_key: str | None, api_base: str | None, ) -> OpenAIWorkloadIdentityConfig | None: - if api_key is not None: + if api_key is not None or get_secret_str("OPENAI_API_KEY") is not None: return None - if not _targets_openai_api(api_base): + effective_api_base: Final = ( + api_base or litellm.api_base or get_secret_str("OPENAI_BASE_URL") or get_secret_str("OPENAI_API_BASE") + ) + if not _targets_openai_api(effective_api_base): return None identity_provider_id: Final = get_secret_str("OPENAI_IDENTITY_PROVIDER_ID") service_account_id: Final = get_secret_str("OPENAI_SERVICE_ACCOUNT_ID") diff --git a/tests/test_litellm/llms/openai/test_openai_workload_identity.py b/tests/test_litellm/llms/openai/test_openai_workload_identity.py index b8cc8c9c80f..d7257d6af89 100644 --- a/tests/test_litellm/llms/openai/test_openai_workload_identity.py +++ b/tests/test_litellm/llms/openai/test_openai_workload_identity.py @@ -9,6 +9,7 @@ import respx from openai import AsyncOpenAI, OpenAI import litellm +from litellm.llms.litellm_proxy.responses.transformation import LiteLLMProxyResponsesAPIConfig from litellm.llms.openai.common_utils import BaseOpenAILLM, OpenAIError from litellm.llms.openai.openai import OpenAIChatCompletion from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig @@ -28,6 +29,9 @@ def wif_env(monkeypatch: pytest.MonkeyPatch, tmp_path: Path) -> OpenAIWorkloadId token_file: Final = tmp_path / "subject_token.jwt" token_file.write_text("subject-token-from-file") monkeypatch.delenv("OPENAI_API_KEY", raising=False) + monkeypatch.delenv("OPENAI_BASE_URL", raising=False) + monkeypatch.delenv("OPENAI_API_BASE", raising=False) + monkeypatch.setattr(litellm, "api_base", None) monkeypatch.setenv("OPENAI_IDENTITY_PROVIDER_ID", "idp_test123") monkeypatch.setenv("OPENAI_SERVICE_ACCOUNT_ID", "user-test456") monkeypatch.setenv("OPENAI_IDENTITY_TOKEN_FILE", str(token_file)) @@ -53,12 +57,36 @@ class TestResolveConfig: def test_static_api_key_wins(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: assert resolve_openai_workload_identity_config(api_key="sk-static", api_base=None) is None + def test_env_openai_api_key_wins( + self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "sk-from-env") + assert resolve_openai_workload_identity_config(api_key=None, api_base=None) is None + def test_foreign_api_base_disables(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: assert resolve_openai_workload_identity_config(api_key=None, api_base="https://my-vllm.internal/v1") is None def test_openai_api_base_allows(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: assert resolve_openai_workload_identity_config(api_key=None, api_base="https://api.openai.com/v1") == wif_env + def test_foreign_env_base_url_disables( + self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("OPENAI_BASE_URL", "https://my-vllm.internal/v1") + assert resolve_openai_workload_identity_config(api_key=None, api_base=None) is None + + def test_openai_env_base_url_allows( + self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("OPENAI_BASE_URL", "https://api.openai.com/v1") + assert resolve_openai_workload_identity_config(api_key=None, api_base=None) == wif_env + + def test_foreign_litellm_api_base_disables( + self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "api_base", "https://my-vllm.internal/v1") + assert resolve_openai_workload_identity_config(api_key=None, api_base=None) is None + @pytest.mark.parametrize( "missing_var", ["OPENAI_IDENTITY_PROVIDER_ID", "OPENAI_SERVICE_ACCOUNT_ID", "OPENAI_IDENTITY_TOKEN_FILE"], @@ -186,3 +214,9 @@ class TestResponsesValidateEnvironment: litellm_params=GenericLiteLLMParams(api_base="https://my-vllm.internal/v1"), ) assert headers["Authorization"] == "Bearer None" + + def test_litellm_proxy_subclass_never_mints_wif(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: + headers: Final = LiteLLMProxyResponsesAPIConfig().validate_environment( + headers={}, model="gpt-4o-mini", litellm_params=GenericLiteLLMParams() + ) + assert headers["Authorization"] == "Bearer None" From 6b7159323bebc5bb5f2ad6a7b8680942cbfb8ab8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 12:25:45 -0700 Subject: [PATCH 36/49] fix(proxy): match OpenAI on empty input and skip budget reservation for token counting /v1/responses/input_tokens returned 200 with a count for an empty "input" ("" or []), while OpenAI returns a 400 missing_required_parameter. The route also went through optimistic budget reservation, which is only released by LLM success/failure callbacks that a token count never reaches, so every call leaked a reservation until TTL expiry and could 429 real traffic. Both routes plus the /openai alias now join /utils/token_counter in the reservation exemption set. --- .../proxy/response_api_endpoints/endpoints.py | 6 +++ .../spend_tracking/budget_reservation.py | 9 +++- .../response_api_endpoints/test_endpoints.py | 15 ++++++ .../spend_tracking/test_budget_reservation.py | 48 +++++++++++++++++++ 4 files changed, 77 insertions(+), 1 deletion(-) create mode 100644 tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index 100b42a9e2a..aa7595ed13d 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1142,6 +1142,12 @@ async def responses_input_tokens( return _missing_responses_param_response("model") if input_value is None: return _missing_responses_param_response("input") + if isinstance(input_value, (str, list)) and not input_value: + return _responses_invalid_request_response( + message="""One of "input" or "previous_response_id" or 'prompt' or 'conversation' must be provided.""", + param=None, + code="missing_required_parameter", + ) try: payload: Final[_TokenCountPayload] = { diff --git a/litellm/proxy/spend_tracking/budget_reservation.py b/litellm/proxy/spend_tracking/budget_reservation.py index 8b2a5dd9312..2d113cfe355 100644 --- a/litellm/proxy/spend_tracking/budget_reservation.py +++ b/litellm/proxy/spend_tracking/budget_reservation.py @@ -172,7 +172,14 @@ async def reserve_budget_for_request( ) -> dict | None: if valid_token is None or not RouteChecks.is_llm_api_route(route=route): return None - if route in {"/models", "/v1/models", "/utils/token_counter"}: + if route in { + "/models", + "/v1/models", + "/utils/token_counter", + "/responses/input_tokens", + "/v1/responses/input_tokens", + "/openai/v1/responses/input_tokens", + }: return None if get_model_from_request(request_body, route, llm_router=llm_router) is None: return None diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index dc43e7f5c06..b31c53c14a8 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -1914,6 +1914,21 @@ class TestResponsesInputTokens: } counter.assert_not_awaited() + @pytest.mark.parametrize("empty_input", ["", []]) + def test_empty_input_returns_openai_400(self, empty_input): + response, counter = self._post_input_tokens({"model": "gpt-4o", "input": empty_input}) + + assert response.status_code == 400, response.text + assert response.json() == { + "error": { + "message": """One of "input" or "previous_response_id" or 'prompt' or 'conversation' must be provided.""", + "type": "invalid_request_error", + "param": None, + "code": "missing_required_parameter", + } + } + counter.assert_not_awaited() + def test_invalid_tools_returns_openai_400(self): response, counter = self._post_input_tokens({"model": "gpt-4o", "input": "hi", "tools": "not-a-list"}) diff --git a/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py new file mode 100644 index 00000000000..f65f68812a2 --- /dev/null +++ b/tests/test_litellm/proxy/spend_tracking/test_budget_reservation.py @@ -0,0 +1,48 @@ +from typing import Final + +import pytest + +from litellm.caching import DualCache +from litellm.proxy._types import UserAPIKeyAuth +from litellm.proxy.common_utils.user_api_key_cache import UserApiKeyCache +from litellm.proxy.spend_tracking.budget_reservation import reserve_budget_for_request +from litellm.proxy.utils import ProxyLogging + +TOKEN_COUNTING_ROUTES: Final = ( + "/responses/input_tokens", + "/v1/responses/input_tokens", + "/openai/v1/responses/input_tokens", + "/utils/token_counter", +) + + +def _budgeted_token() -> UserAPIKeyAuth: + return UserAPIKeyAuth(api_key="sk-test", token="hashed-token", max_budget=100.0, spend=0.0) + + +async def _reserve(route: str) -> dict | None: + return await reserve_budget_for_request( + request_body={"model": "gpt-4o", "input": "hello"}, + route=route, + llm_router=None, + valid_token=_budgeted_token(), + team_object=None, + user_object=None, + prisma_client=None, + user_api_key_cache=UserApiKeyCache(), + proxy_logging_obj=ProxyLogging(user_api_key_cache=DualCache()), + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("route", TOKEN_COUNTING_ROUTES) +async def test_token_counting_routes_are_exempt_from_budget_reservation(route): + assert await _reserve(route) is None + + +@pytest.mark.asyncio +async def test_non_exempt_llm_route_still_reserves_budget(): + reservation: Final = await _reserve("/v1/responses") + + assert reservation is not None + assert reservation["reserved_cost"] > 0 From ef72e7b37dc66d0d755af8dd67934d6f2ae1824d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 12:32:00 -0700 Subject: [PATCH 37/49] fix(openai): require https for workload identity api_base targets --- litellm/llms/openai/workload_identity.py | 3 ++- .../test_litellm/llms/openai/test_openai_workload_identity.py | 3 +++ 2 files changed, 5 insertions(+), 1 deletion(-) diff --git a/litellm/llms/openai/workload_identity.py b/litellm/llms/openai/workload_identity.py index be4d015af8a..8a914039bd0 100644 --- a/litellm/llms/openai/workload_identity.py +++ b/litellm/llms/openai/workload_identity.py @@ -71,7 +71,8 @@ def get_workload_identity_bearer_token(config: OpenAIWorkloadIdentityConfig) -> def _targets_openai_api(api_base: str | None) -> bool: if api_base is None: return True - return urlparse(api_base).hostname == _OPENAI_API_HOST + parsed: Final = urlparse(api_base) + return parsed.scheme == "https" and parsed.hostname == _OPENAI_API_HOST @lru_cache(maxsize=16) diff --git a/tests/test_litellm/llms/openai/test_openai_workload_identity.py b/tests/test_litellm/llms/openai/test_openai_workload_identity.py index d7257d6af89..852e499dc54 100644 --- a/tests/test_litellm/llms/openai/test_openai_workload_identity.py +++ b/tests/test_litellm/llms/openai/test_openai_workload_identity.py @@ -69,6 +69,9 @@ class TestResolveConfig: def test_openai_api_base_allows(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: assert resolve_openai_workload_identity_config(api_key=None, api_base="https://api.openai.com/v1") == wif_env + def test_plaintext_http_api_base_disables(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: + assert resolve_openai_workload_identity_config(api_key=None, api_base="http://api.openai.com/v1") is None + def test_foreign_env_base_url_disables( self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch ) -> None: From 73ab647b1c26c8b1fcff137c87733e46e8d90326 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 12:43:17 -0700 Subject: [PATCH 38/49] fix(count_tokens): preserve image inputs when counting Responses API tokens The chat-to-Responses reverse transform kept only text blocks, so an image input was dropped before the count went to OpenAI. A 256x256 image request counted 13 tokens instead of 268. --- .../responses/count_tokens/transformation.py | 73 ++++++++++++-- ...test_openai_count_tokens_transformation.py | 97 +++++++++++++++++++ .../response_api_endpoints/test_endpoints.py | 9 +- 3 files changed, 169 insertions(+), 10 deletions(-) diff --git a/litellm/llms/openai/responses/count_tokens/transformation.py b/litellm/llms/openai/responses/count_tokens/transformation.py index 6b2f4535df1..72038ae6f6c 100644 --- a/litellm/llms/openai/responses/count_tokens/transformation.py +++ b/litellm/llms/openai/responses/count_tokens/transformation.py @@ -4,7 +4,69 @@ OpenAI Responses API token counting transformation logic. This module handles the transformation of requests to OpenAI's /v1/responses/input_tokens endpoint. """ -from typing import Any, Final +from collections.abc import Mapping, Sequence +from typing import Any, Final, Literal + +from typing_extensions import ReadOnly, TypedDict + + +class ResponsesInputTextPart(TypedDict): + type: ReadOnly[Literal["input_text"]] + text: ReadOnly[str] + + +class ResponsesInputImagePart(TypedDict): + type: ReadOnly[Literal["input_image"]] + image_url: ReadOnly[str] + detail: ReadOnly[str] + + +ResponsesInputPart = ResponsesInputTextPart | ResponsesInputImagePart + + +def _chat_image_block_to_responses_part(image_url: object) -> ResponsesInputImagePart | None: + url: Final = image_url.get("url") if isinstance(image_url, Mapping) else image_url + if not isinstance(url, str) or not url: + return None + detail: Final = image_url.get("detail") if isinstance(image_url, Mapping) else None + part: Final[ResponsesInputImagePart] = { + "type": "input_image", + "image_url": url, + "detail": detail if isinstance(detail, str) and detail else "auto", + } + return part + + +def _chat_block_to_responses_part(block: object) -> ResponsesInputPart | None: + if isinstance(block, str): + bare: Final[ResponsesInputTextPart] = {"type": "input_text", "text": block} + return bare + if not isinstance(block, Mapping): + return None + match block.get("type"): + case "text": + text_value: Final = block.get("text") + text: Final[ResponsesInputTextPart] = { + "type": "input_text", + "text": text_value if isinstance(text_value, str) else "", + } + return text + case "image_url": + return _chat_image_block_to_responses_part(block.get("image_url")) + case _: + return None + + +def chat_content_blocks_to_responses_content( + content: Sequence[object], +) -> str | tuple[ResponsesInputPart, ...]: + """Text-only content collapses to a joined string, so text-only counts stay unchanged.""" + parts: Final = tuple( + part for part in (_chat_block_to_responses_part(block) for block in content) if part is not None + ) + if any(part["type"] != "input_text" for part in parts): + return parts + return "\n".join(part["text"] for part in parts if part["type"] == "input_text") class OpenAICountTokensConfig: @@ -120,14 +182,7 @@ class OpenAICountTokensConfig: instructions_parts.append("\n".join(text_parts)) elif role == "user": if isinstance(content, list): - # Extract text from content blocks for Responses API - text_parts = [] - for block in content: - if isinstance(block, dict) and block.get("type") == "text": - text_parts.append(block.get("text", "")) - elif isinstance(block, str): - text_parts.append(block) - content = "\n".join(text_parts) + content = chat_content_blocks_to_responses_content(content) input_items.append({"role": "user", "content": content}) elif role == "assistant": # Map tool_calls to Responses API function_call items diff --git a/tests/test_litellm/llms/openai/responses/test_openai_count_tokens_transformation.py b/tests/test_litellm/llms/openai/responses/test_openai_count_tokens_transformation.py index e1cc6a92927..ca9e7ab52c3 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_count_tokens_transformation.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_count_tokens_transformation.py @@ -163,6 +163,103 @@ def test_messages_to_responses_input_with_tool(): } +def test_messages_to_responses_input_preserves_images(): + """An image block must survive the round trip, or OpenAI counts only the text. + + A 256x256 image is worth 255 tokens to OpenAI's counting API; dropping it + turned a 268-token request into a 13-token one. + """ + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is in this image?"}, + { + "type": "image_url", + "image_url": {"url": "data:image/png;base64,iVBORw0KGgo=", "detail": "high"}, + }, + ], + } + ] + + input_items, instructions = OpenAICountTokensConfig.messages_to_responses_input(messages) + + assert instructions is None + assert input_items == [ + { + "role": "user", + "content": ( + {"type": "input_text", "text": "What is in this image?"}, + { + "type": "input_image", + "image_url": "data:image/png;base64,iVBORw0KGgo=", + "detail": "high", + }, + ), + } + ] + + +def test_messages_to_responses_input_image_without_detail_defaults_to_auto(): + messages = [ + { + "role": "user", + "content": [{"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}}], + } + ] + + input_items, _ = OpenAICountTokensConfig.messages_to_responses_input(messages) + + assert input_items[0]["content"] == ( + {"type": "input_image", "image_url": "https://example.com/cat.png", "detail": "auto"}, + ) + + +def test_messages_to_responses_input_bare_string_image_url_is_preserved(): + messages = [{"role": "user", "content": [{"type": "image_url", "image_url": "https://example.com/cat.png"}]}] + + input_items, _ = OpenAICountTokensConfig.messages_to_responses_input(messages) + + assert input_items[0]["content"] == ( + {"type": "input_image", "image_url": "https://example.com/cat.png", "detail": "auto"}, + ) + + +def test_messages_to_responses_input_text_only_blocks_stay_a_joined_string(): + """Text-only content must keep collapsing to a string so existing counts do not shift.""" + messages = [ + { + "role": "user", + "content": [{"type": "text", "text": "first"}, {"type": "text", "text": "second"}], + } + ] + + input_items, _ = OpenAICountTokensConfig.messages_to_responses_input(messages) + + assert input_items == [{"role": "user", "content": "first\nsecond"}] + + +def test_messages_to_responses_input_drops_unmappable_blocks(): + """A block with no Responses API equivalent is skipped, never forwarded verbatim.""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "hi"}, + {"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}}, + {"type": "input_audio", "input_audio": {"data": "AAAA", "format": "wav"}}, + ], + } + ] + + input_items, _ = OpenAICountTokensConfig.messages_to_responses_input(messages) + + assert input_items[0]["content"] == ( + {"type": "input_text", "text": "hi"}, + {"type": "input_image", "image_url": "https://example.com/cat.png", "detail": "auto"}, + ) + + def test_validate_request_valid(): """Test that valid requests pass validation.""" config = OpenAICountTokensConfig() diff --git a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py index b31c53c14a8..d7010de6405 100644 --- a/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/test_litellm/proxy/response_api_endpoints/test_endpoints.py @@ -3,10 +3,12 @@ Test for response_api_endpoints/endpoints.py """ import unittest +from typing import Any from unittest.mock import AsyncMock, MagicMock, patch import pytest from fastapi.testclient import TestClient +from httpx import Response import litellm from litellm.proxy.proxy_server import app @@ -1816,7 +1818,12 @@ class TestResponsesInputTokens: never registered, so the POST fell through to the GET/DELETE-only /v1/responses/{response_id} route and returned 405.""" - def _post_input_tokens(self, body, path="/v1/responses/input_tokens", counter=None): + def _post_input_tokens( + self, + body: dict[str, Any], + path: str = "/v1/responses/input_tokens", + counter: AsyncMock | None = None, + ) -> tuple[Response, AsyncMock]: from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.response_api_endpoints.endpoints import _proxy_token_counter From 9f9236e8d5f2485c6e4f1638f1cb1e3b471262e2 Mon Sep 17 00:00:00 2001 From: Ashton Sidhu Date: Mon, 31 Aug 2026 15:50:42 -0400 Subject: [PATCH 39/49] fix(guardrails): exclude images from HiddenLayer v1 scans (#29210) * Don't scan images * Fix failing tests * Fix lint: typed image-part filter, restore monkeypatch-based tests --------- Co-authored-by: Yucheng Zhu --- .../hiddenlayer/hiddenlayer.py | 27 ++++++++++++++++++- .../guardrail_hooks/test_hiddenlayer.py | 7 ++--- 2 files changed, 30 insertions(+), 4 deletions(-) diff --git a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py index a6ea2e09583..68914a1989e 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py +++ b/litellm/proxy/guardrails/guardrail_hooks/hiddenlayer/hiddenlayer.py @@ -156,6 +156,31 @@ def _header_value(headers: Mapping[str, str], key: str, default: str) -> str: return headers.get(key, default) +def _is_image_part(item: object) -> bool: + """Whether a structured-message content part carries an image rather than text.""" + + if not isinstance(item, Mapping): + return False + + part: Final[Mapping[object, object]] = item + return part.get("type") == "image_url" + + +def _scannable_text(content: object) -> str: + """Flatten a structured message's content into the single string the v1 detection endpoint takes. + + Image parts are dropped: the endpoint accepts one string, so an image would only reach it as + its stringified source (a base64 blob or a URL), which is not text the scanner can evaluate. + """ + + if not isinstance(content, list): + return str(content or "") + + parts: Final[Sequence[object]] = content + text_parts: Final = [item for item in parts if not _is_image_part(item)] # mutable-ok: sent as a list repr + return str(text_parts or "") + + def is_saas(host: str) -> bool: """Checks whether the connection is to the SaaS platform""" @@ -270,7 +295,7 @@ class HiddenlayerGuardrail(CustomGuardrail): "messages": [ { "role": last_msg.get("role", "user"), - "content": str(last_msg.get("content", "")), + "content": _scannable_text(last_msg.get("content")), } ] }, diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py index b140082a3bf..f5d51a601d7 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_hiddenlayer.py @@ -432,7 +432,7 @@ class TestHiddenlayerGuardrail: @pytest.mark.asyncio async def test_apply_guardrail_request_with_image(self, monkeypatch: pytest.MonkeyPatch): - """Test apply_guardrail sends multimodal content (image) to HiddenLayer v1.""" + """Test apply_guardrail strips images from multimodal content before sending to HiddenLayer v1.""" monkeypatch.setenv("HIDDENLAYER_API_BASE", "https://my.hiddenlayer") guardrail = HiddenlayerGuardrail( @@ -485,12 +485,13 @@ class TestHiddenlayerGuardrail: logging_obj=logging_obj, ) - # v1 API requires string content — multimodal list is stringified + # v1 API requires string content — image_url items are stripped and the + # remaining (text-only) content is stringified before being sent. mock_post.assert_called_once() call_kwargs = mock_post.call_args.kwargs sent_content = call_kwargs["json"]["input"]["messages"][0]["content"] assert isinstance(sent_content, str) - assert sent_content == str(multimodal_content) + assert sent_content == str([{"type": "text", "text": "how much is on this receipt?"}]) # Result should be returned without error assert result is not None From 0c21b30cb72aab7f56ab88bda47c00243aab1e0c Mon Sep 17 00:00:00 2001 From: Mateo Wang <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 12:52:34 -0700 Subject: [PATCH 40/49] feat(spend_tracking): persist router metadata in spend logs for internal router models (#39001) * feat(spend_tracking): persist router metadata in spend logs for internal router models * test(spend_tracking): expect router_metadata key in exact-payload tests, type the routed-kwargs helper --- litellm/proxy/_types.py | 14 ++++ .../spend_tracking/spend_tracking_utils.py | 60 ++++++++++++---- litellm/types/router.py | 5 ++ .../test_spend_management_endpoints.py | 6 +- .../test_spend_tracking_utils.py | 69 +++++++++++++++++++ ui/litellm-dashboard/src/lib/http/schema.d.ts | 2 + 6 files changed, 139 insertions(+), 17 deletions(-) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index 0f2d97b8b1c..a84fae7fd23 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3549,6 +3549,19 @@ class AllCallbacks(LiteLLMPydanticObjectBase): ) +class SpendLogsRouterMetadata(TypedDict): + """ + Router provenance stamped on spend logs for deployments flagged with + model_info.internal_router_model, correlating the requested model group + with the provider deployment that served the call + """ + + requested_model: ReadOnly[str | None] + selected_model: ReadOnly[str | None] + selected_provider: ReadOnly[str | None] + router_correlation_id: ReadOnly[str | None] + + class SpendLogsMetadata(TypedDict): """ Specific metadata k,v pairs logged to spendlogs for easier cost tracking @@ -3591,6 +3604,7 @@ class SpendLogsMetadata(TypedDict): compression_savings: CompressionSavingsMetadata | None autorouter_savings: ReadOnly[float | None] # stamped by the logging payload; None = not auto-routed litellm_gateway_injected_cache: ReadOnly[str | None] + router_metadata: ReadOnly[SpendLogsRouterMetadata | None] # None = deployment not flagged internal_router_model class SpendLogsPayload(TypedDict): diff --git a/litellm/proxy/spend_tracking/spend_tracking_utils.py b/litellm/proxy/spend_tracking/spend_tracking_utils.py index 43709e4e6ff..9f718b7d20d 100644 --- a/litellm/proxy/spend_tracking/spend_tracking_utils.py +++ b/litellm/proxy/spend_tracking/spend_tracking_utils.py @@ -30,7 +30,7 @@ from litellm.litellm_core_utils.litellm_logging import ( request_model_access_groups_from_litellm_params, ) from litellm.litellm_core_utils.safe_json_dumps import safe_dumps, strip_null_bytes -from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload +from litellm.proxy._types import SpendLogsMetadata, SpendLogsPayload, SpendLogsRouterMetadata from litellm.proxy.spend_tracking.spend_log_error_logger import spend_log_error from litellm.proxy.utils import PrismaClient, hash_token from litellm.types.utils import ( @@ -93,6 +93,24 @@ def _redact_logged_api_key(value: str | None, *, already_redacted: bool = False) return hash_token(stripped) +def _get_router_metadata_for_spend_log( + metadata: Mapping[str, object] | None, + requested_model: str | None, + selected_model: str | None, + selected_provider: str | None, + router_correlation_id: str | None, +) -> SpendLogsRouterMetadata | None: + model_info: Final = metadata.get("model_info") if metadata is not None else None + if not isinstance(model_info, Mapping) or model_info.get("internal_router_model") is not True: + return None + return SpendLogsRouterMetadata( + requested_model=requested_model or None, + selected_model=selected_model or None, + selected_provider=selected_provider or None, + router_correlation_id=router_correlation_id, + ) + + def _get_spend_logs_metadata( metadata: dict | None, applied_guardrails: list[str] | None = None, @@ -109,6 +127,7 @@ def _get_spend_logs_metadata( cost_breakdown: CostBreakdown | None = None, litellm_call_id: str | None = None, autorouter_savings: float | None = None, + router_metadata: SpendLogsRouterMetadata | None = None, ) -> SpendLogsMetadata: if metadata is None: return SpendLogsMetadata( @@ -148,13 +167,17 @@ def _get_spend_logs_metadata( autorouter_savings=autorouter_savings, litellm_gateway_injected_cache=None, litellm_call_id=litellm_call_id, + router_metadata=router_metadata, ) verbose_proxy_logger.debug( "getting payload for SpendLogs, available keys in metadata: " + str(list(metadata.keys())) ) # Filter the metadata dictionary to include only the specified keys - clean_metadata: Final = SpendLogsMetadata(**{key: metadata.get(key) for key in SpendLogsMetadata.__annotations__}) + clean_metadata: Final = SpendLogsMetadata( + **{key: metadata.get(key) for key in SpendLogsMetadata.__annotations__ if key != "router_metadata"}, + router_metadata=router_metadata, + ) _raw_key: Final = clean_metadata.get("user_api_key") _trusted_hash: Final = metadata.get("user_api_key_hash") _already_redacted: Final = ( @@ -375,6 +398,20 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs hidden_params: Final = standard_logging_payload.get("hidden_params", {}) litellm_overhead_time_ms = hidden_params.get("litellm_overhead_time_ms") + custom_llm_provider: Final = ( + kwargs.get("custom_llm_provider") + or _sl_attribution_fallback(standard_logging_payload, "custom_llm_provider") + or None + ) + raw_model: Final = cast(str, kwargs.get("model") or "") + model_name: Final = ( + standard_logging_payload.get("model") if standard_logging_payload is not None else None + ) or reconstruct_model_name(raw_model, custom_llm_provider, metadata or {}) + litellm_call_id: Final = cast( + str | None, + kwargs.get("litellm_call_id") or litellm_params.get("litellm_call_id"), + ) + # clean up litellm metadata clean_metadata = _get_spend_logs_metadata( metadata, @@ -433,9 +470,13 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs autorouter_savings=( standard_logging_payload.get("autorouter_savings", None) if standard_logging_payload is not None else None ), - litellm_call_id=cast( - str | None, - kwargs.get("litellm_call_id") or litellm_params.get("litellm_call_id"), + litellm_call_id=litellm_call_id, + router_metadata=_get_router_metadata_for_spend_log( + metadata=metadata, + requested_model=_model_group, + selected_model=model_name, + selected_provider=custom_llm_provider, + router_correlation_id=litellm_call_id, ), ) @@ -480,15 +521,6 @@ def get_logging_payload(kwargs, response_obj, start_time, end_time) -> SpendLogs # Extract agent_id for A2A requests (set directly on model_call_details) agent_id: Final[str | None] = kwargs.get("agent_id") or metadata.get("agent_id") - custom_llm_provider: Final = ( - kwargs.get("custom_llm_provider") - or _sl_attribution_fallback(standard_logging_payload, "custom_llm_provider") - or None - ) - raw_model: Final = cast(str, kwargs.get("model") or "") - model_name: Final = ( - standard_logging_payload.get("model") if standard_logging_payload is not None else None - ) or reconstruct_model_name(raw_model, custom_llm_provider, metadata or {}) try: payload: Final[SpendLogsPayload] = SpendLogsPayload( diff --git a/litellm/types/router.py b/litellm/types/router.py index 97bd93f3f47..ab6c807ba20 100644 --- a/litellm/types/router.py +++ b/litellm/types/router.py @@ -189,6 +189,11 @@ class ModelInfo(MirroredPricingParams): # router-wide default. enable_tag_filtering: bool | None = None + # when True, calls routed to this deployment persist a router_metadata block + # (requested model group, selected model + provider, router correlation id) + # in the spend log row's metadata. Set it on every deployment of the group. + internal_router_model: 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/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py index 10c3e5fecf8..a0dcbf802ef 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_management_endpoints.py @@ -2865,7 +2865,7 @@ class TestSpendLogsPayload: "model": "gpt-4o", "user": "", "team_id": "", - "metadata": '{"applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}', + "metadata": '{"applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "router_metadata": null, "usage_object": {"completion_tokens": 20, "prompt_tokens": 10, "total_tokens": 30, "completion_tokens_details": null, "prompt_tokens_details": null}, "model_map_information": {"model_map_key": "gpt-4o", "model_map_value": {"key": "gpt-4o", "max_tokens": 16384, "max_input_tokens": 128000, "max_output_tokens": 16384, "input_cost_per_token": 2.5e-06, "cache_creation_input_token_cost": null, "cache_read_input_token_cost": 1.25e-06, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": 1.25e-06, "output_cost_per_token_batches": 5e-06, "output_cost_per_token": 1e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_reasoning_token": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "openai", "mode": "chat", "supports_system_messages": true, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": false, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": false, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": true, "supports_reasoning": false, "search_context_cost_per_query": {"search_context_size_low": 0.03, "search_context_size_medium": 0.035, "search_context_size_high": 0.05}, "tpm": null, "rpm": null, "supported_openai_params": ["frequency_penalty", "logit_bias", "logprobs", "top_logprobs", "max_tokens", "max_completion_tokens", "modalities", "prediction", "n", "presence_penalty", "seed", "stop", "stream", "stream_options", "temperature", "top_p", "tools", "tool_choice", "function_call", "functions", "max_retries", "extra_headers", "parallel_tool_calls", "audio", "response_format", "user"]}}, "additional_usage_values": {"completion_tokens_details": null, "prompt_tokens_details": null}}', "cache_key": "Cache OFF", "spend": 0.00022500000000000002, "total_tokens": 30, @@ -2961,7 +2961,7 @@ class TestSpendLogsPayload: "model": "claude-4-sonnet-20250514", "user": "", "team_id": "", - "metadata": '{"applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', + "metadata": '{"applied_guardrails": [], "attempted_fallbacks": null, "original_model_group": null, "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "router_metadata": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', "cache_key": "Cache OFF", "spend": 0.01383, "total_tokens": 2598, @@ -3055,7 +3055,7 @@ class TestSpendLogsPayload: "model": "claude-4-sonnet-20250514", "user": "", "team_id": "", - "metadata": '{"applied_guardrails": [], "attempted_fallbacks": 0, "original_model_group": "my-anthropic-model-group", "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', + "metadata": '{"applied_guardrails": [], "attempted_fallbacks": 0, "original_model_group": "my-anthropic-model-group", "batch_models": null, "batch_successful_requests": null, "batch_failed_requests": null, "mcp_tool_call_metadata": null, "vector_store_request_metadata": null, "routing_decision": null, "internal_call_origin": null, "guardrail_information": null, "compression_savings": null, "litellm_gateway_injected_cache": null, "router_metadata": null, "usage_object": {"completion_tokens": 503, "prompt_tokens": 2095, "total_tokens": 2598, "completion_tokens_details": null, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}, "model_map_information": {"model_map_key": "claude-4-sonnet-20250514", "model_map_value": {"key": "claude-4-sonnet-20250514", "max_tokens": 128000, "max_input_tokens": 200000, "max_output_tokens": 128000, "input_cost_per_token": 3e-06, "cache_creation_input_token_cost": 3.75e-06, "cache_read_input_token_cost": 3e-07, "input_cost_per_character": null, "input_cost_per_token_above_128k_tokens": null, "input_cost_per_token_above_200k_tokens": null, "input_cost_per_query": null, "input_cost_per_second": null, "input_cost_per_audio_token": null, "input_cost_per_token_batches": null, "output_cost_per_token_batches": null, "output_cost_per_token": 1.5e-05, "output_cost_per_audio_token": null, "output_cost_per_character": null, "output_cost_per_token_above_128k_tokens": null, "output_cost_per_character_above_128k_tokens": null, "output_cost_per_token_above_200k_tokens": null, "output_cost_per_second": null, "output_cost_per_image": null, "output_vector_size": null, "litellm_provider": "anthropic", "mode": "chat", "supports_system_messages": null, "supports_response_schema": true, "supports_vision": true, "supports_function_calling": true, "supports_tool_choice": true, "supports_assistant_prefill": true, "supports_prompt_caching": true, "supports_audio_input": false, "supports_audio_output": false, "supports_pdf_input": true, "supports_embedding_image_input": false, "supports_native_streaming": null, "supports_web_search": false, "supports_reasoning": true, "search_context_cost_per_query": null, "tpm": null, "rpm": null, "supported_openai_params": ["stream", "stop", "temperature", "top_p", "max_tokens", "max_completion_tokens", "tools", "tool_choice", "extra_headers", "parallel_tool_calls", "response_format", "user", "reasoning_effort", "thinking"]}}, "additional_usage_values": {"completion_tokens_details": {"accepted_prediction_tokens": null, "audio_tokens": null, "reasoning_tokens": null, "rejected_prediction_tokens": null, "text_tokens": 503, "image_tokens": null}, "prompt_tokens_details": {"audio_tokens": null, "cached_tokens": 0, "text_tokens": null, "image_tokens": null}, "cache_creation_input_tokens": 0, "cache_read_input_tokens": 0}}', "cache_key": "Cache OFF", "spend": 0.01383, "total_tokens": 2598, diff --git a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py index 5022dab32be..9e5917637a8 100644 --- a/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py +++ b/tests/test_litellm/proxy/spend_tracking/test_spend_tracking_utils.py @@ -2,6 +2,7 @@ import asyncio import datetime import json from datetime import timezone +from collections.abc import Mapping from typing import Any, Final, cast from unittest.mock import AsyncMock, MagicMock, patch @@ -3956,3 +3957,71 @@ def test_passthrough_caching_carries_no_injection_marker(): ) metadata = json.loads(payload["metadata"]) assert metadata["litellm_gateway_injected_cache"] is None + + +def _routed_call_kwargs(model_info: Mapping[str, object]) -> dict[str, object]: + return { + "model": "claude-haiku-4-5", + "custom_llm_provider": "azure_ai", + "litellm_call_id": "router-corr-123", + "litellm_params": { + "metadata": { + "user_api_key": "test-key", + "model_group": "internal-router/gpt-5.4", + "deployment": "azure_ai/claude-haiku-4-5", + "model_info": model_info, + } + }, + } + + +def test_router_metadata_stamped_for_internal_router_model_deployment(): + """A deployment flagged model_info.internal_router_model gets a router_metadata + block correlating the requested model group with the selected deployment.""" + payload = get_logging_payload( + kwargs=_routed_call_kwargs({"id": "mi-1", "internal_router_model": True}), + response_obj=litellm.ModelResponse(id="chatcmpl-router-meta", choices=[], usage=litellm.Usage()), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + metadata = json.loads(payload["metadata"]) + assert metadata["router_metadata"] == { + "requested_model": "internal-router/gpt-5.4", + "selected_model": "azure_ai/claude-haiku-4-5", + "selected_provider": "azure_ai", + "router_correlation_id": "router-corr-123", + } + + +def test_router_metadata_absent_without_internal_router_model_flag(): + payload = get_logging_payload( + kwargs=_routed_call_kwargs({"id": "mi-1"}), + response_obj=litellm.ModelResponse(id="chatcmpl-unflagged", choices=[], usage=litellm.Usage()), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + metadata = json.loads(payload["metadata"]) + assert metadata["router_metadata"] is None + + +@pytest.mark.parametrize("bucket", ["metadata", "litellm_metadata"]) +def test_caller_forged_router_metadata_is_discarded(bucket): + """The raw request bucket is client-writable and _get_spend_logs_metadata projects + every SpendLogsMetadata key from it, so the server-derived value must overwrite + unconditionally or a caller could plant router provenance the router never produced.""" + payload = get_logging_payload( + kwargs={ + "model": "gpt-4o-mini", + "litellm_params": { + bucket: { + "user_api_key": "test-key", + "router_metadata": {"requested_model": "forged", "router_correlation_id": "forged-id"}, + } + }, + }, + response_obj=litellm.ModelResponse(id="chatcmpl-forged-router-meta", choices=[], usage=litellm.Usage()), + start_time=datetime.datetime.now(timezone.utc), + end_time=datetime.datetime.now(timezone.utc), + ) + metadata = json.loads(payload["metadata"]) + assert metadata["router_metadata"] is None diff --git a/ui/litellm-dashboard/src/lib/http/schema.d.ts b/ui/litellm-dashboard/src/lib/http/schema.d.ts index 0ddb26ba035..1ed5366dac1 100644 --- a/ui/litellm-dashboard/src/lib/http/schema.d.ts +++ b/ui/litellm-dashboard/src/lib/http/schema.d.ts @@ -38438,6 +38438,8 @@ export interface components { input_cost_per_character?: number | null; /** Input Cost Per Token */ input_cost_per_token?: number | null; + /** Internal Router Model */ + internal_router_model?: boolean | null; /** Output Cost Per Character */ output_cost_per_character?: number | null; /** Output Cost Per Token */ From e7dc0213bdd1d88dccf2596aa349e8d771d74311 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 12:52:55 -0700 Subject: [PATCH 41/49] fix(openai): treat empty api key values as unset for workload identity --- litellm/llms/openai/workload_identity.py | 2 +- .../llms/openai/test_openai_workload_identity.py | 9 +++++++++ 2 files changed, 10 insertions(+), 1 deletion(-) diff --git a/litellm/llms/openai/workload_identity.py b/litellm/llms/openai/workload_identity.py index 8a914039bd0..e9ac26e26a2 100644 --- a/litellm/llms/openai/workload_identity.py +++ b/litellm/llms/openai/workload_identity.py @@ -45,7 +45,7 @@ def resolve_openai_workload_identity_config( api_key: str | None, api_base: str | None, ) -> OpenAIWorkloadIdentityConfig | None: - if api_key is not None or get_secret_str("OPENAI_API_KEY") is not None: + if api_key or get_secret_str("OPENAI_API_KEY"): return None effective_api_base: Final = ( api_base or litellm.api_base or get_secret_str("OPENAI_BASE_URL") or get_secret_str("OPENAI_API_BASE") diff --git a/tests/test_litellm/llms/openai/test_openai_workload_identity.py b/tests/test_litellm/llms/openai/test_openai_workload_identity.py index 852e499dc54..228d591b47a 100644 --- a/tests/test_litellm/llms/openai/test_openai_workload_identity.py +++ b/tests/test_litellm/llms/openai/test_openai_workload_identity.py @@ -63,6 +63,15 @@ class TestResolveConfig: monkeypatch.setenv("OPENAI_API_KEY", "sk-from-env") assert resolve_openai_workload_identity_config(api_key=None, api_base=None) is None + def test_empty_env_openai_api_key_counts_as_unset( + self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "") + assert resolve_openai_workload_identity_config(api_key=None, api_base=None) == wif_env + + def test_empty_api_key_param_counts_as_unset(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: + assert resolve_openai_workload_identity_config(api_key="", api_base=None) == wif_env + def test_foreign_api_base_disables(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: assert resolve_openai_workload_identity_config(api_key=None, api_base="https://my-vllm.internal/v1") is None From ae83444a3e0eca7536f49101557bc8ec044b501a Mon Sep 17 00:00:00 2001 From: Cursor Agent Date: Mon, 31 Aug 2026 19:54:45 +0000 Subject: [PATCH 42/49] fix(openai): treat empty api_key as unset for WIF resolution --- litellm/llms/openai/workload_identity.py | 7 +++++-- .../llms/openai/test_openai_workload_identity.py | 16 ++++++++++------ 2 files changed, 15 insertions(+), 8 deletions(-) diff --git a/litellm/llms/openai/workload_identity.py b/litellm/llms/openai/workload_identity.py index e9ac26e26a2..ecec161ed46 100644 --- a/litellm/llms/openai/workload_identity.py +++ b/litellm/llms/openai/workload_identity.py @@ -6,7 +6,7 @@ from typing import TYPE_CHECKING, Final from urllib.parse import urlparse import litellm -from litellm.secret_managers.main import get_secret_str +from litellm.secret_managers.main import get_secret_str, normalize_nonempty_secret_str from .common_utils import OpenAIError @@ -45,7 +45,10 @@ def resolve_openai_workload_identity_config( api_key: str | None, api_base: str | None, ) -> OpenAIWorkloadIdentityConfig | None: - if api_key or get_secret_str("OPENAI_API_KEY"): + static_api_key: Final = normalize_nonempty_secret_str(api_key) or normalize_nonempty_secret_str( + get_secret_str("OPENAI_API_KEY") + ) + if static_api_key is not None: return None effective_api_base: Final = ( api_base or litellm.api_base or get_secret_str("OPENAI_BASE_URL") or get_secret_str("OPENAI_API_BASE") diff --git a/tests/test_litellm/llms/openai/test_openai_workload_identity.py b/tests/test_litellm/llms/openai/test_openai_workload_identity.py index 228d591b47a..d8d9936e9a1 100644 --- a/tests/test_litellm/llms/openai/test_openai_workload_identity.py +++ b/tests/test_litellm/llms/openai/test_openai_workload_identity.py @@ -63,14 +63,18 @@ class TestResolveConfig: monkeypatch.setenv("OPENAI_API_KEY", "sk-from-env") assert resolve_openai_workload_identity_config(api_key=None, api_base=None) is None - def test_empty_env_openai_api_key_counts_as_unset( - self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch + @pytest.mark.parametrize("empty_key", ["", " "]) + def test_empty_api_key_arg_does_not_disable_wif( + self, wif_env: OpenAIWorkloadIdentityConfig, empty_key: str ) -> None: - monkeypatch.setenv("OPENAI_API_KEY", "") - assert resolve_openai_workload_identity_config(api_key=None, api_base=None) == wif_env + assert resolve_openai_workload_identity_config(api_key=empty_key, api_base=None) == wif_env - def test_empty_api_key_param_counts_as_unset(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: - assert resolve_openai_workload_identity_config(api_key="", api_base=None) == wif_env + @pytest.mark.parametrize("empty_key", ["", " "]) + def test_empty_env_openai_api_key_does_not_disable_wif( + self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch, empty_key: str + ) -> None: + monkeypatch.setenv("OPENAI_API_KEY", empty_key) + assert resolve_openai_workload_identity_config(api_key=None, api_base=None) == wif_env def test_foreign_api_base_disables(self, wif_env: OpenAIWorkloadIdentityConfig) -> None: assert resolve_openai_workload_identity_config(api_key=None, api_base="https://my-vllm.internal/v1") is None From 1249f84b10e39f1ad7ddfffb0fe11069abe0d2f1 Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 31 Aug 2026 12:56:34 -0700 Subject: [PATCH 43/49] fix(vertex_ai): graft default vertex path when api_base has a version-only path (#38986) * fix(vertex_ai): graft default vertex path when api_base has a version-only path Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(vertex_ai): keep query and fragment placement when grafting vertex path Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(vertex_ai): merge alt=sse into existing query when streaming Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: yassin Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/llms/vertex_ai/vertex_llm_base.py | 20 +++- .../llms/vertex_ai/test_vertex_llm_base.py | 110 ++++++++++++++++++ 2 files changed, 127 insertions(+), 3 deletions(-) diff --git a/litellm/llms/vertex_ai/vertex_llm_base.py b/litellm/llms/vertex_ai/vertex_llm_base.py index 75098515deb..aca257dc095 100644 --- a/litellm/llms/vertex_ai/vertex_llm_base.py +++ b/litellm/llms/vertex_ai/vertex_llm_base.py @@ -27,6 +27,15 @@ from .common_utils import ( get_vertex_base_url, ) + +def _graft_default_vertex_path(api_base: str, default_url: str) -> str: + parsed_api_base: Final = urlparse(api_base) + default_segments: Final = urlparse(default_url).path.lstrip("/").split("/") + graft_segments: Final = default_segments[1:] if default_segments[0] in ("v1", "v1beta1") else default_segments + grafted_path: Final = parsed_api_base.path.rstrip("/") + "/" + "/".join(graft_segments) + return parsed_api_base._replace(path=grafted_path).geturl() + + GOOGLE_IMPORT_ERROR_MESSAGE: Final = ( "Google Cloud SDK not found. Install it with: pip install 'litellm[google]' or pip install google-cloud-aiplatform" ) @@ -621,8 +630,9 @@ class VertexBase: Handles custom api_base for: 1. Gemini (Google AI Studio) - constructs /models/{model}:{endpoint} - 2. Vertex AI with standard proxies - constructs {api_base}:{endpoint}; - if api_base has no path (bare host), grafts the default vertex URL path onto it + 2. Vertex AI with standard proxies - grafts the default vertex URL path onto the + api_base when its path is empty or only an API version (/v1, /v1beta1); + otherwise constructs {api_base}:{endpoint} 3. Vertex AI with PSC endpoints - constructs full path structure {api_base}/v1/projects/{project}/locations/{location}/endpoints/{model}:{endpoint} (only when use_psc_endpoint_format=True) @@ -669,10 +679,14 @@ class VertexBase: ) elif urlparse(api_base).path in ("", "/"): url = api_base.rstrip("/") + urlparse(url).path + elif urlparse(api_base).path.rstrip("/") in ("/v1", "/v1beta1") and "/projects/" in urlparse(url).path: + url = _graft_default_vertex_path(api_base=api_base, default_url=url) else: url = f"{api_base}:{endpoint}" if stream is True: - url = url + "?alt=sse" + parsed_stream_url: Final = urlparse(url) + stream_query: Final = f"{parsed_stream_url.query}&alt=sse" if parsed_stream_url.query else "alt=sse" + url = parsed_stream_url._replace(query=stream_query).geturl() return auth_header, url def _get_token_and_url( diff --git a/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py b/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py index 29d22e844a5..a4d67606698 100644 --- a/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py +++ b/tests/test_litellm/llms/vertex_ai/test_vertex_llm_base.py @@ -982,6 +982,116 @@ class TestVertexBase: assert result_url == f"{gateway_api_base}:embedContent" + def test_check_custom_proxy_vertex_api_base_with_version_path_grafts_default_path(self): + vertex_base = VertexBase() + + _, result_url = vertex_base._check_custom_proxy( + api_base="https://aiplatform.googleapis.com/v1beta1", + custom_llm_provider="vertex_ai", + gemini_api_key=None, + endpoint="generateContent", + stream=None, + auth_header="Bearer token123", + url="https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite:generateContent", + model="gemini-3.5-flash-lite", + ) + + assert ( + result_url + == "https://aiplatform.googleapis.com/v1beta1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite:generateContent" + ) + + def test_check_custom_proxy_vertex_api_base_with_version_path_trailing_slash_grafts_default_path(self): + vertex_base = VertexBase() + + _, result_url = vertex_base._check_custom_proxy( + api_base="https://internal-gateway.example.com/v1/", + custom_llm_provider="vertex_ai", + gemini_api_key=None, + endpoint="generateContent", + stream=None, + auth_header="Bearer token123", + url="https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite:generateContent", + model="gemini-3.5-flash-lite", + ) + + assert ( + result_url + == "https://internal-gateway.example.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite:generateContent" + ) + + def test_check_custom_proxy_vertex_api_base_with_version_path_and_query_grafts_before_query(self): + vertex_base = VertexBase() + + _, result_url = vertex_base._check_custom_proxy( + api_base="https://internal-gateway.example.com/v1beta1?key=abc", + custom_llm_provider="vertex_ai", + gemini_api_key=None, + endpoint="generateContent", + stream=None, + auth_header="Bearer token123", + url="https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite:generateContent", + model="gemini-3.5-flash-lite", + ) + + assert ( + result_url + == "https://internal-gateway.example.com/v1beta1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite:generateContent?key=abc" + ) + + def test_check_custom_proxy_vertex_api_base_with_version_path_and_query_streaming_appends_alt_sse(self): + vertex_base = VertexBase() + + _, result_url = vertex_base._check_custom_proxy( + api_base="https://internal-gateway.example.com/v1beta1?key=abc", + custom_llm_provider="vertex_ai", + gemini_api_key=None, + endpoint="streamGenerateContent", + stream=True, + auth_header="Bearer token123", + url="https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite:streamGenerateContent", + model="gemini-3.5-flash-lite", + ) + + assert ( + result_url + == "https://internal-gateway.example.com/v1beta1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite:streamGenerateContent?key=abc&alt=sse" + ) + + def test_check_custom_proxy_vertex_api_base_with_non_version_path_keeps_endpoint_append(self): + vertex_base = VertexBase() + gateway_api_base = "https://gateway.example.com/vertex-proxy" + + _, result_url = vertex_base._check_custom_proxy( + api_base=gateway_api_base, + custom_llm_provider="vertex_ai", + gemini_api_key=None, + endpoint="generateContent", + stream=None, + auth_header="Bearer token123", + url="https://us-central1-aiplatform.googleapis.com/v1/projects/test-project/locations/us-central1/publishers/google/models/gemini-3.5-flash-lite:generateContent", + model="gemini-3.5-flash-lite", + ) + + assert result_url == f"{gateway_api_base}:generateContent" + + def test_check_custom_proxy_vertex_api_base_without_projects_in_default_url_keeps_endpoint_append(self): + vertex_base = VertexBase() + gemma_api_base = "https://example.com/custom/gemma-deployment" + + _, result_url = vertex_base._check_custom_proxy( + api_base=gemma_api_base, + custom_llm_provider="vertex_ai", + gemini_api_key=None, + endpoint="predict", + stream=False, + auth_header=None, + url=gemma_api_base, + model="gemma-3-27b-it", + ) + + assert result_url == f"{gemma_api_base}:predict" + def test_check_custom_proxy_vertex_bare_host_streaming_keeps_single_alt_sse(self): vertex_base = VertexBase() From fe90c6f6fca6440f302e21bb31f5816225b96770 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 12:59:40 -0700 Subject: [PATCH 44/49] fix(count_tokens): keep assistant turns on the provider counting API Assistant list content was forwarded to /v1/responses/input_tokens as chat `text` blocks, which the Responses API rejects (it accepts only output_text and refusal inside an assistant turn). The 400 sent the whole request to the local tokenizer, so any conversation with an assistant turn silently lost provider-exact counting, including the image counting added in 73ab647b1c. Assistant content now collapses to the plain string the Responses API counts identically, and image parts are kept to user turns where they are legal. --- .../responses/count_tokens/transformation.py | 19 ++++-- ...test_openai_count_tokens_transformation.py | 63 +++++++++++++++++++ 2 files changed, 77 insertions(+), 5 deletions(-) diff --git a/litellm/llms/openai/responses/count_tokens/transformation.py b/litellm/llms/openai/responses/count_tokens/transformation.py index 72038ae6f6c..62dbebb4fe6 100644 --- a/litellm/llms/openai/responses/count_tokens/transformation.py +++ b/litellm/llms/openai/responses/count_tokens/transformation.py @@ -23,6 +23,8 @@ class ResponsesInputImagePart(TypedDict): ResponsesInputPart = ResponsesInputTextPart | ResponsesInputImagePart +ResponsesContentRole = Literal["user", "assistant"] + def _chat_image_block_to_responses_part(image_url: object) -> ResponsesInputImagePart | None: url: Final = image_url.get("url") if isinstance(image_url, Mapping) else image_url @@ -37,7 +39,7 @@ def _chat_image_block_to_responses_part(image_url: object) -> ResponsesInputImag return part -def _chat_block_to_responses_part(block: object) -> ResponsesInputPart | None: +def _chat_block_to_responses_part(block: object, role: ResponsesContentRole) -> ResponsesInputPart | None: if isinstance(block, str): bare: Final[ResponsesInputTextPart] = {"type": "input_text", "text": block} return bare @@ -51,7 +53,7 @@ def _chat_block_to_responses_part(block: object) -> ResponsesInputPart | None: "text": text_value if isinstance(text_value, str) else "", } return text - case "image_url": + case "image_url" if role == "user": return _chat_image_block_to_responses_part(block.get("image_url")) case _: return None @@ -59,10 +61,15 @@ def _chat_block_to_responses_part(block: object) -> ResponsesInputPart | None: def chat_content_blocks_to_responses_content( content: Sequence[object], + role: ResponsesContentRole, ) -> str | tuple[ResponsesInputPart, ...]: - """Text-only content collapses to a joined string, so text-only counts stay unchanged.""" + """Text-only content collapses to a joined string, which every role accepts and counts identically. + + Only a user turn may carry an image part: the Responses API rejects any part but + output_text and refusal inside an assistant turn. + """ parts: Final = tuple( - part for part in (_chat_block_to_responses_part(block) for block in content) if part is not None + part for part in (_chat_block_to_responses_part(block, role) for block in content) if part is not None ) if any(part["type"] != "input_text" for part in parts): return parts @@ -182,11 +189,13 @@ class OpenAICountTokensConfig: instructions_parts.append("\n".join(text_parts)) elif role == "user": if isinstance(content, list): - content = chat_content_blocks_to_responses_content(content) + content = chat_content_blocks_to_responses_content(content, "user") input_items.append({"role": "user", "content": content}) elif role == "assistant": # Map tool_calls to Responses API function_call items tool_calls = msg.get("tool_calls") + if isinstance(content, list): + content = chat_content_blocks_to_responses_content(content, "assistant") if content: input_items.append({"role": "assistant", "content": content}) if tool_calls: diff --git a/tests/test_litellm/llms/openai/responses/test_openai_count_tokens_transformation.py b/tests/test_litellm/llms/openai/responses/test_openai_count_tokens_transformation.py index ca9e7ab52c3..22761272321 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_count_tokens_transformation.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_count_tokens_transformation.py @@ -260,6 +260,69 @@ def test_messages_to_responses_input_drops_unmappable_blocks(): ) +def test_messages_to_responses_input_assistant_blocks_collapse_to_a_string(): + """An assistant turn must never forward chat `text` blocks. + + The Responses API only accepts output_text and refusal inside an assistant turn, so + forwarding them 400s the whole request and silently drops the count back to the local + tokenizer, which is exactly what defeats the image fix above. + """ + messages = [ + {"role": "user", "content": [{"type": "text", "text": "What is the capital of France?"}]}, + {"role": "assistant", "content": [{"type": "text", "text": "Paris."}]}, + ] + + input_items, _ = OpenAICountTokensConfig.messages_to_responses_input(messages) + + assert input_items == [ + {"role": "user", "content": "What is the capital of France?"}, + {"role": "assistant", "content": "Paris."}, + ] + + +def test_messages_to_responses_input_assistant_image_block_is_dropped(): + """An image part is illegal inside an assistant turn, so it must not reach the provider.""" + messages = [ + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Here it is"}, + {"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}}, + ], + } + ] + + input_items, _ = OpenAICountTokensConfig.messages_to_responses_input(messages) + + assert input_items == [{"role": "assistant", "content": "Here it is"}] + + +def test_messages_to_responses_input_keeps_user_image_alongside_an_assistant_turn(): + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "What is in this image?"}, + {"type": "image_url", "image_url": {"url": "https://example.com/cat.png"}}, + ], + }, + {"role": "assistant", "content": [{"type": "text", "text": "A cat."}]}, + ] + + input_items, _ = OpenAICountTokensConfig.messages_to_responses_input(messages) + + assert input_items == [ + { + "role": "user", + "content": ( + {"type": "input_text", "text": "What is in this image?"}, + {"type": "input_image", "image_url": "https://example.com/cat.png", "detail": "auto"}, + ), + }, + {"role": "assistant", "content": "A cat."}, + ] + + def test_validate_request_valid(): """Test that valid requests pass validation.""" config = OpenAICountTokensConfig() From c9908ffabb732f902f3409a7d763c5ea6a1541a8 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 13:17:43 -0700 Subject: [PATCH 45/49] fix(responses): count input_file tokens instead of silently dropping the file The Responses-to-chat transform dropped the filename OpenAI requires next to file_data, so a request carrying an inline PDF counted 13 tokens instead of 36 and a real completion through the chat bridge got a 400. --- .../responses/count_tokens/transformation.py | 28 ++++++- .../transformation.py | 2 + ...test_openai_count_tokens_transformation.py | 74 +++++++++++++++++++ .../test_litellm_completion_responses.py | 19 +++++ 4 files changed, 121 insertions(+), 2 deletions(-) diff --git a/litellm/llms/openai/responses/count_tokens/transformation.py b/litellm/llms/openai/responses/count_tokens/transformation.py index 62dbebb4fe6..88f04c59e01 100644 --- a/litellm/llms/openai/responses/count_tokens/transformation.py +++ b/litellm/llms/openai/responses/count_tokens/transformation.py @@ -21,7 +21,13 @@ class ResponsesInputImagePart(TypedDict): detail: ReadOnly[str] -ResponsesInputPart = ResponsesInputTextPart | ResponsesInputImagePart +class ResponsesInputFilePart(TypedDict): + type: ReadOnly[Literal["input_file"]] + filename: ReadOnly[str] + file_data: ReadOnly[str] + + +ResponsesInputPart = ResponsesInputTextPart | ResponsesInputImagePart | ResponsesInputFilePart ResponsesContentRole = Literal["user", "assistant"] @@ -39,6 +45,22 @@ def _chat_image_block_to_responses_part(image_url: object) -> ResponsesInputImag return part +def _chat_file_block_to_responses_part(file_value: object) -> ResponsesInputFilePart | None: + """Only an inline file round trips: OpenAI rejects `file_data` without the `filename` beside it.""" + if not isinstance(file_value, Mapping): + return None + filename: Final = file_value.get("filename") + file_data: Final = file_value.get("file_data") + if not isinstance(filename, str) or not filename or not isinstance(file_data, str) or not file_data: + return None + part: Final[ResponsesInputFilePart] = { + "type": "input_file", + "filename": filename, + "file_data": file_data, + } + return part + + def _chat_block_to_responses_part(block: object, role: ResponsesContentRole) -> ResponsesInputPart | None: if isinstance(block, str): bare: Final[ResponsesInputTextPart] = {"type": "input_text", "text": block} @@ -55,6 +77,8 @@ def _chat_block_to_responses_part(block: object, role: ResponsesContentRole) -> return text case "image_url" if role == "user": return _chat_image_block_to_responses_part(block.get("image_url")) + case "file" if role == "user": + return _chat_file_block_to_responses_part(block.get("file")) case _: return None @@ -65,7 +89,7 @@ def chat_content_blocks_to_responses_content( ) -> str | tuple[ResponsesInputPart, ...]: """Text-only content collapses to a joined string, which every role accepts and counts identically. - Only a user turn may carry an image part: the Responses API rejects any part but + Only a user turn may carry an image or file part: the Responses API rejects any part but output_text and refusal inside an assistant turn. """ parts: Final = tuple( diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index f39df38d069..4a416f61e1b 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -1629,6 +1629,8 @@ class LiteLLMCompletionResponsesConfig: file_dict["file_id"] = file_id if item.get("file_data"): file_dict["file_data"] = item["file_data"] + if item.get("filename"): + file_dict["filename"] = item["filename"] new_item: Final[dict[str, object]] = {"type": "file", "file": file_dict} if "cache_control" in item: diff --git a/tests/test_litellm/llms/openai/responses/test_openai_count_tokens_transformation.py b/tests/test_litellm/llms/openai/responses/test_openai_count_tokens_transformation.py index 22761272321..c2efc1acdb9 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_count_tokens_transformation.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_count_tokens_transformation.py @@ -323,6 +323,80 @@ def test_messages_to_responses_input_keeps_user_image_alongside_an_assistant_tur ] +def test_messages_to_responses_input_preserves_inline_files(): + """An inline file must survive the round trip, or the count silently drops the file. + + A small PDF is worth 36 tokens to OpenAI's counting API; dropping it left the same + request counting 13, the text-only total. + """ + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize this file."}, + { + "type": "file", + "file": {"filename": "report.pdf", "file_data": "data:application/pdf;base64,JVBERi0="}, + }, + ], + } + ] + + input_items, _ = OpenAICountTokensConfig.messages_to_responses_input(messages) + + assert input_items == [ + { + "role": "user", + "content": ( + {"type": "input_text", "text": "Summarize this file."}, + { + "type": "input_file", + "filename": "report.pdf", + "file_data": "data:application/pdf;base64,JVBERi0=", + }, + ), + } + ] + + +def test_messages_to_responses_input_drops_a_file_with_no_inline_data(): + """OpenAI rejects `file_data` without a `filename`, and a rejected request loses the whole count.""" + messages = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize this file."}, + {"type": "file", "file": {"file_data": "data:application/pdf;base64,JVBERi0="}}, + {"type": "file", "file": {"file_id": "file-abc123"}}, + ], + } + ] + + input_items, _ = OpenAICountTokensConfig.messages_to_responses_input(messages) + + assert input_items == [{"role": "user", "content": "Summarize this file."}] + + +def test_messages_to_responses_input_assistant_file_block_is_dropped(): + """A file part is illegal inside an assistant turn, so it must not reach the provider.""" + messages = [ + { + "role": "assistant", + "content": [ + {"type": "text", "text": "Here it is"}, + { + "type": "file", + "file": {"filename": "report.pdf", "file_data": "data:application/pdf;base64,JVBERi0="}, + }, + ], + } + ] + + input_items, _ = OpenAICountTokensConfig.messages_to_responses_input(messages) + + assert input_items == [{"role": "assistant", "content": "Here it is"}] + + def test_validate_request_valid(): """Test that valid requests pass validation.""" config = OpenAICountTokensConfig() diff --git a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py index b96d2eb5322..bb59e576568 100644 --- a/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py +++ b/tests/test_litellm/responses/litellm_completion_transformation/test_litellm_completion_responses.py @@ -124,6 +124,25 @@ class TestLiteLLMCompletionResponsesConfig: assert "extra_field" not in result["file"] assert "another_field" not in result["file"] + def test_transform_input_file_item_to_file_item_keeps_filename(self): + """OpenAI rejects file_data with no filename beside it, so dropping it 400s the request""" + result = ( + LiteLLMCompletionResponsesConfig._transform_input_file_item_to_file_item( + { + "type": "input_file", + "filename": "report.pdf", + "file_data": "data:application/pdf;base64,JVBERi0=", + } + ) + ) + assert result == { + "type": "file", + "file": { + "file_data": "data:application/pdf;base64,JVBERi0=", + "filename": "report.pdf", + }, + } + def test_transform_input_file_item_to_file_item_with_file_url(self): """file_url should be mapped to file_id for downstream URL handling""" result = ( From a90fb538bfd0a71e39d82b344c5c666200d01a63 Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 13:25:51 -0700 Subject: [PATCH 46/49] fix(friendli): declare GLM-5.3-Flash reasoning efforts as explicit levels --- model_prices_and_context_window.json | 7 +++++-- .../test_friendli_glm_5_3_flash_model_metadata.py | 3 +-- 2 files changed, 6 insertions(+), 4 deletions(-) diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index fd6b9e11adf..f0e95d77018 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -19575,8 +19575,11 @@ "output_cost_per_token": 5e-07, "cache_read_input_token_cost": 3e-08, "supports_prompt_caching": true, - "supports_max_reasoning_effort": true, - "supports_low_reasoning_effort": true, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ], "supports_parallel_function_calling": true, "supports_response_schema": true, "supports_native_structured_output": true, diff --git a/tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py b/tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py index 0acd750e49a..7e94205fb09 100644 --- a/tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py +++ b/tests/test_litellm/test_friendli_glm_5_3_flash_model_metadata.py @@ -23,8 +23,7 @@ def test_friendli_glm_5_3_flash_model_info(): assert info["max_output_tokens"] == 1048576 assert info["supports_function_calling"] is True assert info["supports_reasoning"] is True - assert info["supports_low_reasoning_effort"] is True - assert info["supports_max_reasoning_effort"] is True + assert info["reasoning_effort_levels"] == ["low", "high", "max"] assert info["supports_tool_choice"] is True assert info["supports_prompt_caching"] is True assert info["supports_vision"] is True From b7da4717843e61cf6bd1aeb10b1cbaf70e946e1d Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 13:53:22 -0700 Subject: [PATCH 47/49] fix(count_tokens): price an inline file block instead of raising on it `ChatCompletionFileObject` is in the union `_count_content_list` accepts, but `file` was missing from its match, so every local count of a Responses `input_file` raised `Invalid content item type: file`. On /v1/responses/input_tokens that surfaced as an opaque 500 whenever the model's provider counting API refused the block and the local tokenizer took over. Count it the way the module already counts the same thing in Anthropic's dialect: the filename like a document title, the inline bytes through the image pricer. --- litellm/litellm_core_utils/token_counter.py | 28 ++++++++++++++- .../litellm_core_utils/test_token_counter.py | 35 +++++++++++++++++++ 2 files changed, 62 insertions(+), 1 deletion(-) diff --git a/litellm/litellm_core_utils/token_counter.py b/litellm/litellm_core_utils/token_counter.py index 256bee7b348..c350bc5569e 100644 --- a/litellm/litellm_core_utils/token_counter.py +++ b/litellm/litellm_core_utils/token_counter.py @@ -693,6 +693,26 @@ def _count_document_tokens( ) +def _count_file_tokens( + file_value: object, + count_function: TokenCounterFunction, + use_default_image_token_count: bool, +) -> int: + """An OpenAI `file` block is the chat-completions spelling of a document, so it prices like one.""" + if not isinstance(file_value, Mapping): + return 0 + filename: Final = file_value.get("filename") + file_data: Final = file_value.get("file_data") + name_tokens: Final = count_function(filename) if isinstance(filename, str) and filename else 0 + if not isinstance(file_data, str) or not file_data: + return name_tokens + return name_tokens + calculate_img_tokens( + data=file_data, + mode="auto", + use_default_image_token_count=use_default_image_token_count, + ) + + def _count_anthropic_content( content: Mapping[str, Any], count_function: TokenCounterFunction, @@ -778,6 +798,12 @@ def _count_content_list( use_default_image_token_count, default_token_count, ) + elif c["type"] == "file": + num_tokens += _count_file_tokens( + c.get("file"), + count_function, + use_default_image_token_count, + ) elif c["type"] in ("tool_use", "tool_result"): num_tokens += _count_anthropic_content( c, @@ -807,7 +833,7 @@ def _count_content_list( raise ValueError( f"Invalid content item type: {content_type}. " f"Expected str or dict with 'type' field " - f"(text, image_url, image, document, tool_use, tool_result, thinking, tool_reference)." + f"(text, image_url, image, document, file, tool_use, tool_result, thinking, tool_reference)." ) return num_tokens except Exception as e: diff --git a/tests/test_litellm/litellm_core_utils/test_token_counter.py b/tests/test_litellm/litellm_core_utils/test_token_counter.py index 572b505e94c..4694fa8fbed 100644 --- a/tests/test_litellm/litellm_core_utils/test_token_counter.py +++ b/tests/test_litellm/litellm_core_utils/test_token_counter.py @@ -1377,3 +1377,38 @@ def test_anthropic_document_title_and_context_add_their_tokens(): {"type": "document", "source": source}, ] ) + + +def test_openai_file_block_prices_like_the_equivalent_anthropic_document(): + """An inline `file` is a `document` in the chat-completions dialect, so it must price identically, not raise. + + Before the fix `file` was missing from the content-block match even though `ChatCompletionFileObject` + is in the union this counter accepts, so every local count of a Responses `input_file` raised + `Invalid content item type: file` and surfaced as a 500 on /v1/responses/input_tokens. + """ + prompt = {"type": "text", "text": "Summarize this file."} + inline_file = { + "type": "file", + "file": {"filename": "report.pdf", "file_data": "data:application/pdf;base64,JVBERi0xLjQK"}, + } + document = { + "type": "document", + "title": "report.pdf", + "source": {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"}, + } + + assert _count_user_content([prompt, inline_file]) == _count_user_content([prompt, document]) + assert _count_user_content([prompt, inline_file]) > _count_user_content([prompt]) + + +def test_openai_file_block_without_inline_bytes_counts_what_it_carries(): + """A `file` block naming an uploaded file has no bytes to price, so it adds only the filename's tokens.""" + prompt = {"type": "text", "text": "Summarize this file."} + + by_id = {"type": "file", "file": {"file_id": "file-abc123"}} + assert _count_user_content([prompt, by_id]) == _count_user_content([prompt]) + + named = {"type": "file", "file": {"file_id": "file-abc123", "filename": "report.pdf"}} + assert _count_user_content([prompt, named]) == _count_user_content( + [prompt, {"type": "text", "text": "report.pdf"}] + ) From 2bbf5135a5782a07e62d51d6df1fc2a774ec501e Mon Sep 17 00:00:00 2001 From: mateo-berri <277851410+mateo-berri@users.noreply.github.com> Date: Mon, 31 Aug 2026 13:55:39 -0700 Subject: [PATCH 48/49] fix(friendli): ship GLM-5.3-Flash in the bundled backup cost map --- ...odel_prices_and_context_window_backup.json | 28 +++++++++++++++++++ 1 file changed, 28 insertions(+) diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index 05c1cfd3179..f0e95d77018 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -19564,6 +19564,34 @@ "supports_system_messages": true, "supports_tool_choice": true }, + "friendliai/zai-org/GLM-5.3-Flash": { + "litellm_provider": "friendliai", + "supports_reasoning": true, + "supports_function_calling": true, + "max_input_tokens": 1048576, + "max_tokens": 1048576, + "max_output_tokens": 1048576, + "input_cost_per_token": 1.5e-07, + "output_cost_per_token": 5e-07, + "cache_read_input_token_cost": 3e-08, + "supports_prompt_caching": true, + "reasoning_effort_levels": [ + "low", + "high", + "max" + ], + "supports_parallel_function_calling": true, + "supports_response_schema": true, + "supports_native_structured_output": true, + "supports_system_messages": true, + "supports_tool_choice": true, + "mode": "chat", + "comment": "Native multimodal GLM model for efficient coding and long-horizon agent tasks", + "source": "https://api.friendli.ai/serverless/v1/models", + "supports_vision": true, + "supports_image_input": true, + "supports_video_input": true + }, "ft:babbage-002": { "deprecation_date": "2026-10-23", "input_cost_per_token": 1.6e-06, From 40edeaaecb0c7b0d7c0fac06a74fa67d0b57400c Mon Sep 17 00:00:00 2001 From: "devin-ai-integration[bot]" <158243242+devin-ai-integration[bot]@users.noreply.github.com> Date: Mon, 31 Aug 2026 13:59:59 -0700 Subject: [PATCH 49/49] fix(otel): emit cache token counts on OTel v2 LLM spans (#38716) * fix(otel): emit cache token counts on OTel v2 LLM spans Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(otel): trim comment in LLMUsage adapter Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(otel): drop casts in LLMUsage cache token adapter Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> * fix(deps): bump restrictedpython to 8.3 for GHSA-ffg3-p8fm-mjx2 Co-Authored-By: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --------- Co-authored-by: Devin AI <158243242+devin-ai-integration[bot]@users.noreply.github.com> --- litellm/integrations/otel/mappers/genai.py | 2 ++ litellm/integrations/otel/model/payloads.py | 22 ++++++++++++++----- litellm/integrations/otel/model/semconv.py | 2 ++ .../otel/test_otel_v2_components.py | 22 +++++++++++++++++++ .../otel/test_otel_v2_sources_of_truth.py | 22 +++++++++++++++++++ 5 files changed, 65 insertions(+), 5 deletions(-) diff --git a/litellm/integrations/otel/mappers/genai.py b/litellm/integrations/otel/mappers/genai.py index b09498f9292..3ac92b04c27 100644 --- a/litellm/integrations/otel/mappers/genai.py +++ b/litellm/integrations/otel/mappers/genai.py @@ -62,6 +62,8 @@ class GenAIMapper: GenAI.RESPONSE_TIME_TO_FIRST_CHUNK: lambda d: d.time_to_first_chunk_seconds, GenAI.USAGE_INPUT_TOKENS: lambda d: d.usage.input_tokens, GenAI.USAGE_OUTPUT_TOKENS: lambda d: d.usage.output_tokens, + GenAI.USAGE_CACHE_CREATION_INPUT_TOKENS: lambda d: d.usage.cache_creation_input_tokens, + GenAI.USAGE_CACHE_READ_INPUT_TOKENS: lambda d: d.usage.cache_read_input_tokens, Error.TYPE: lambda d: d.error.error_type if d.error else None, Server.ADDRESS: lambda d: d.server.address if d.server else None, Server.PORT: lambda d: d.server.port if d.server else None, diff --git a/litellm/integrations/otel/model/payloads.py b/litellm/integrations/otel/model/payloads.py index f70c777e1a7..d35405538f6 100644 --- a/litellm/integrations/otel/model/payloads.py +++ b/litellm/integrations/otel/model/payloads.py @@ -95,6 +95,22 @@ class LLMUsage: input_tokens: int | None = None output_tokens: int | None = None total_tokens: int | None = None + cache_creation_input_tokens: int | None = None + cache_read_input_tokens: int | None = None + + @classmethod + def from_standard_logging_payload(cls, payload: StandardLoggingPayload) -> LLMUsage: + # Cache token counts only exist on the raw provider usage object under metadata + metadata: Final[Mapping[str, object]] = payload.get("metadata") or {} + raw_usage: Final = metadata.get("usage_object") + usage_object: Final[Mapping[str, object]] = raw_usage if isinstance(raw_usage, Mapping) else {} + return cls( + input_tokens=as_int(payload.get("prompt_tokens")), + output_tokens=as_int(payload.get("completion_tokens")), + total_tokens=as_int(payload.get("total_tokens")), + cache_creation_input_tokens=as_int(usage_object.get("cache_creation_input_tokens")), + cache_read_input_tokens=as_int(usage_object.get("cache_read_input_tokens")), + ) @dataclass(frozen=True) @@ -363,11 +379,7 @@ class LLMCallSpanData: response_model=context.response_model, response_id=as_str(response.get("id")), request_params=LLMRequestParams.from_model_parameters(params), - usage=LLMUsage( - input_tokens=as_int(payload.get("prompt_tokens")), - output_tokens=as_int(payload.get("completion_tokens")), - total_tokens=as_int(payload.get("total_tokens")), - ), + usage=LLMUsage.from_standard_logging_payload(payload), finish_reasons=finish_reasons, error=_parse_error(payload), response_cost=as_float(payload.get("response_cost")), diff --git a/litellm/integrations/otel/model/semconv.py b/litellm/integrations/otel/model/semconv.py index 4ad0cb5d1b4..f7a6280f95b 100644 --- a/litellm/integrations/otel/model/semconv.py +++ b/litellm/integrations/otel/model/semconv.py @@ -110,6 +110,8 @@ class GenAI: # usage USAGE_INPUT_TOKENS: Final = "gen_ai.usage.input_tokens" USAGE_OUTPUT_TOKENS: Final = "gen_ai.usage.output_tokens" + USAGE_CACHE_CREATION_INPUT_TOKENS: Final = "gen_ai.usage.cache_creation.input_tokens" + USAGE_CACHE_READ_INPUT_TOKENS: Final = "gen_ai.usage.cache_read.input_tokens" # content (opt-in, gated by capture mode) INPUT_MESSAGES: Final = "gen_ai.input.messages" OUTPUT_MESSAGES: Final = "gen_ai.output.messages" diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_components.py b/tests/test_litellm/integrations/otel/test_otel_v2_components.py index 115e385eda4..4aa28b5abfd 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_components.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_components.py @@ -3,6 +3,7 @@ baggage helpers, metrics, the typed coercion helpers, mapper branches, span-name builders, and the registry validator's failure paths. Needs the OTel SDK.""" import json +from dataclasses import replace import pytest @@ -215,6 +216,27 @@ def test_genai_mapper_all_request_params(): assert attrs["server.port"] == 443 +def test_genai_mapper_cache_token_attrs(): + cached = replace( + _full_llm_call(), + usage=LLMUsage( + input_tokens=10, + output_tokens=5, + total_tokens=15, + cache_creation_input_tokens=7, + cache_read_input_tokens=3, + ), + ) + attrs = GenAIMapper().map(cached) + assert attrs[GenAI.USAGE_CACHE_CREATION_INPUT_TOKENS] == 7 + assert attrs[GenAI.USAGE_CACHE_READ_INPUT_TOKENS] == 3 + + # No cache usage keeps the span sparse: neither key present. + uncached = GenAIMapper().map(_full_llm_call()) + assert GenAI.USAGE_CACHE_CREATION_INPUT_TOKENS not in uncached + assert GenAI.USAGE_CACHE_READ_INPUT_TOKENS not in uncached + + def test_genai_mapper_stamps_input_output_messages(): data = LLMCallSpanData( operation=GenAIOperation.CHAT, diff --git a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py index baa72b5a7fe..ca628aa3405 100644 --- a/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py +++ b/tests/test_litellm/integrations/otel/test_otel_v2_sources_of_truth.py @@ -525,6 +525,28 @@ def test_llm_call_adapter_extracts_all_fields(): assert data.identity.key_hash == "hsh" +def test_llm_call_adapter_extracts_cache_tokens_from_usage_object(): + payload = _sample_payload() + payload["metadata"] = { + **payload["metadata"], + "usage_object": { + "prompt_tokens": 10, + "completion_tokens": 5, + "cache_creation_input_tokens": 7, + "cache_read_input_tokens": 3, + }, + } + data = LLMCallSpanData.from_standard_logging_payload(payload) + assert data.usage.cache_creation_input_tokens == 7 + assert data.usage.cache_read_input_tokens == 3 + + +def test_llm_call_adapter_cache_tokens_none_without_usage_object(): + data = LLMCallSpanData.from_standard_logging_payload(_sample_payload()) + assert data.usage.cache_creation_input_tokens is None + assert data.usage.cache_read_input_tokens is None + + def test_llm_call_adapter_failure_path(): payload = _sample_payload( status="failure",