diff --git a/litellm/constants.py b/litellm/constants.py index 9f29baac562..665ab771726 100644 --- a/litellm/constants.py +++ b/litellm/constants.py @@ -1705,6 +1705,7 @@ AZURE_SPEECH_FAST_TRANSCRIPTION_PATH: Final = "/speechtotext/transcriptions:tran AZURE_SPEECH_STT_DOMAIN: Final = "stt.speech.microsoft.com" AZURE_SPEECH_COGNITIVE_SERVICES_DOMAIN: Final = "api.cognitive.microsoft.com" AZURE_SPEECH_SUBSCRIPTION_KEY_HEADER: Final = "Ocp-Apim-Subscription-Key" +AWS_ERROR_TYPE_HEADER: Final = "x-amzn-errortype" AZURE_SPEECH_SHORT_AUDIO_MODEL: Final = "short-audio" AZURE_SPEECH_BATCH_MODEL: Final = "batch-transcription" AZURE_SPEECH_FAST_TRANSCRIPTION_MODEL: Final = "fast-transcription" diff --git a/litellm/integrations/custom_guardrail.py b/litellm/integrations/custom_guardrail.py index 2a2a2a0326d..93dd3d0f5ee 100644 --- a/litellm/integrations/custom_guardrail.py +++ b/litellm/integrations/custom_guardrail.py @@ -695,9 +695,7 @@ class CustomGuardrail(CustomLogger): supported_event_hooks: list[GuardrailEventHooks], ) -> None: allowed_hooks: Final = frozenset(supported_event_hooks) | ( - frozenset((GuardrailEventHooks.logging_only,)) - if self.uses_apply_guardrail_interface() and not self.use_native_lifecycle_hooks - else frozenset() + frozenset((GuardrailEventHooks.logging_only,)) if self.accepts_logging_only_mode() else frozenset() ) def _validate_event_hook_list_is_in_supported_event_hooks( @@ -869,8 +867,13 @@ class CustomGuardrail(CustomLogger): return True return False - def uses_apply_guardrail_interface(self) -> bool: - return type(self).apply_guardrail is not CustomGuardrail.apply_guardrail + @classmethod + def uses_apply_guardrail_interface(cls) -> bool: + return cls.apply_guardrail is not CustomGuardrail.apply_guardrail + + @classmethod + def accepts_logging_only_mode(cls) -> bool: + return cls.uses_apply_guardrail_interface() and not cls.use_native_lifecycle_hooks @classmethod def supports_logging_only_scope(cls) -> bool: @@ -1029,7 +1032,7 @@ class CustomGuardrail(CustomLogger): ) return kwargs, result litellm_params: Final = kwargs.get("litellm_params") or {} - scratch_metadata: Final = { + scratch_metadata: Final[dict[str, object]] = { # mutable-ok: apply_guardrail records its verdict here key: value for key, value in (litellm_params.get("metadata") or {}).items() if key != "standard_logging_guardrail_information" @@ -1083,11 +1086,11 @@ class CustomGuardrail(CustomLogger): async def _scan_logged_call( self, - kwargs: dict, # mutable-ok: CustomLogger.async_logging_hook contract + kwargs: dict[str, object], # mutable-ok: CustomLogger.async_logging_hook contract response: object | None, translation: "BaseTranslation", output_translation: "BaseTranslation", - scratch_metadata: dict, # mutable-ok: apply_guardrail records its verdict into request metadata + scratch_metadata: dict[str, object], # mutable-ok: apply_guardrail records its verdict into request metadata ) -> None: scratch_fields: Final = self._copy_scratch_request_fields(kwargs) scratch_input, scratch_tools = scratch_fields or (None, None) diff --git a/litellm/integrations/gitlab/gitlab_prompt_manager.py b/litellm/integrations/gitlab/gitlab_prompt_manager.py index 85b91821f6c..adaacd615a8 100644 --- a/litellm/integrations/gitlab/gitlab_prompt_manager.py +++ b/litellm/integrations/gitlab/gitlab_prompt_manager.py @@ -219,7 +219,7 @@ class GitLabTemplateManager: elif value.replace(".", "").isdigit(): try: result[key] = float(value) - except Exception: + except ValueError: result[key] = value else: result[key] = value.strip("\"'") diff --git a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py index 6c7760d7da8..ce5dceb2e5c 100644 --- a/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py +++ b/litellm/litellm_core_utils/llm_cost_calc/tool_call_cost_tracking.py @@ -2,12 +2,13 @@ Helper utilities for tracking the cost of built-in tools. """ -from collections.abc import Mapping +from collections.abc import Mapping, Sequence from typing import Final, Literal, cast from pydantic import ValidationError import litellm +from litellm._logging import verbose_logger from litellm.constants import OPENAI_FILE_SEARCH_COST_PER_1K_CALLS from litellm.litellm_core_utils.llm_cost_calc.utils import ( get_web_search_requests_from_usage, @@ -30,20 +31,36 @@ from litellm.types.utils import ( ) +def _output_item_field(output_item: object, field: str) -> object: + if isinstance(output_item, dict): + fields: Final[Mapping[str, object]] = cast(Mapping[str, object], output_item) # cast-ok: narrowed by isinstance + return fields.get(field) + return getattr(output_item, field, None) + + def _output_item_type(output_item: object) -> str | None: - item_type: Final = output_item.get("type") if isinstance(output_item, dict) else getattr(output_item, "type", None) + item_type: Final = _output_item_field(output_item, "type") return item_type if isinstance(item_type, str) else None -def _reported_web_search_requests(response_object: ResponsesAPIResponse) -> int | None: - tool_usage: Final = getattr(response_object, "tool_usage", None) - if tool_usage is None: +def _responses_tool_usage(response_object: ResponsesAPIResponse) -> ResponsesToolUsage | None: + top_level: Final = getattr(response_object, "tool_usage", None) + raw: Final = ( + top_level if top_level is not None else getattr(getattr(response_object, "usage", None), "tool_usage", None) + ) + if raw is None: return None try: - web_search: Final = ResponsesToolUsage.model_validate(tool_usage).web_search + return ResponsesToolUsage.model_validate(raw) except ValidationError: return None - return None if web_search is None else web_search.num_requests + + +def _reported_web_search_requests(response_object: ResponsesAPIResponse) -> int | None: + tool_usage: Final = _responses_tool_usage(response_object) + if tool_usage is None: + return None + return None if tool_usage.web_search is None else tool_usage.web_search.num_requests def _usage_reports_server_side_web_search_calls(usage: Usage) -> bool: @@ -87,31 +104,48 @@ class StandardBuiltInToolCostTracking: usage=usage, ) + image_generation_cost: Final = StandardBuiltInToolCostTracking._handle_image_generation_cost( + response_object=response_object, + custom_llm_provider=custom_llm_provider, + ) + # Handle web search if StandardBuiltInToolCostTracking.response_object_includes_web_search_call( response_object=response_object, usage=usage ): - return google_maps_grounding_cost + StandardBuiltInToolCostTracking._handle_web_search_cost( - model=model, - custom_llm_provider=custom_llm_provider, - usage=usage, - standard_built_in_tools_params=standard_built_in_tools_params, - response_object=response_object, + return ( + google_maps_grounding_cost + + image_generation_cost + + StandardBuiltInToolCostTracking._handle_web_search_cost( + model=model, + custom_llm_provider=custom_llm_provider, + usage=usage, + standard_built_in_tools_params=standard_built_in_tools_params, + response_object=response_object, + ) ) # Handle file search if StandardBuiltInToolCostTracking.response_object_includes_file_search_call(response_object=response_object): - return google_maps_grounding_cost + StandardBuiltInToolCostTracking._handle_file_search_cost( + return ( + google_maps_grounding_cost + + image_generation_cost + + StandardBuiltInToolCostTracking._handle_file_search_cost( + model=model, + custom_llm_provider=custom_llm_provider, + standard_built_in_tools_params=standard_built_in_tools_params, + ) + ) + + # Handle Azure assistant features + return ( + google_maps_grounding_cost + + image_generation_cost + + StandardBuiltInToolCostTracking._handle_azure_assistant_costs( model=model, custom_llm_provider=custom_llm_provider, standard_built_in_tools_params=standard_built_in_tools_params, ) - - # Handle Azure assistant features - return google_maps_grounding_cost + StandardBuiltInToolCostTracking._handle_azure_assistant_costs( - model=model, - custom_llm_provider=custom_llm_provider, - standard_built_in_tools_params=standard_built_in_tools_params, ) @staticmethod @@ -210,6 +244,87 @@ class StandardBuiltInToolCostTracking: ) return max(count, 1) + @staticmethod + def _image_generation_tool_model(response_object: ResponsesAPIResponse) -> str | None: + tools: Final[Sequence[object]] = cast(Sequence[object], response_object.tools or ()) # cast-ok: dict tool union + for tool in tools: + if _output_item_field(tool, "type") != "image_generation": + continue + if isinstance(model := _output_item_field(tool, "model"), str) and model: + return model + return None + + @staticmethod + def _image_generation_call_cost(output_item: object, tool_model: str, custom_llm_provider: str | None) -> float: + from litellm.cost_calculator import ( + default_image_cost_calculator, # pyright: ignore[reportUnknownVariableType] # optional_params param is untyped + ) + + quality: Final = _output_item_field(output_item, "quality") + size: Final = _output_item_field(output_item, "size") + try: + return default_image_cost_calculator( + model=tool_model, + custom_llm_provider=custom_llm_provider, + quality=quality if isinstance(quality, str) and quality != "auto" else None, + n=1, + size=size if isinstance(size, str) and size != "auto" else None, + ) + except Exception as e: # noqa: BLE001 # pricing helpers raise bare Exception for unmapped models; bill 0.0 + verbose_logger.debug("Could not price Responses API image_generation_call item: %s", e) + return 0.0 + + @staticmethod + def _image_generation_token_cost( + response_object: ResponsesAPIResponse, tool_model: str, custom_llm_provider: str | None + ) -> float | None: + tool_usage: Final = _responses_tool_usage(response_object) + if tool_usage is None or tool_usage.image_gen is None or tool_usage.image_gen.total_tokens <= 0: + return None + try: + model_info: Final = litellm.get_model_info(model=tool_model, custom_llm_provider=custom_llm_provider) + except Exception as e: # noqa: BLE001 # get_model_info raises bare Exception for unmapped models; fall back + verbose_logger.debug("Could not resolve pricing for image tool model %s: %s", tool_model, e) + return None + image_gen: Final = tool_usage.image_gen + input_details: Final = image_gen.input_tokens_details + output_details: Final = image_gen.output_tokens_details + token_cost: Final = ( + (input_details.text_tokens if input_details else 0) * (model_info.get("input_cost_per_token") or 0) + + (input_details.image_tokens if input_details else 0) * (model_info.get("input_cost_per_image_token") or 0) + + (output_details.image_tokens if output_details else 0) + * (model_info.get("output_cost_per_image_token") or 0) + + (output_details.text_tokens if output_details else 0) * (model_info.get("output_cost_per_token") or 0) + ) + return token_cost if token_cost > 0 else None + + @staticmethod + def _handle_image_generation_cost(response_object: object, custom_llm_provider: str | None) -> float: + if not isinstance(response_object, ResponsesAPIResponse): + return 0.0 + output: Final[Sequence[object]] = cast(Sequence[object], response_object.output) # cast-ok: isinstance-narrowed + completed_items: Final = tuple( + output_item + for output_item in output + if _output_item_type(output_item) == "image_generation_call" + and _output_item_field(output_item, "status") == "completed" + ) + if not completed_items: + return 0.0 + tool_model: Final = StandardBuiltInToolCostTracking._image_generation_tool_model(response_object) + if tool_model is None: + verbose_logger.debug("Responses API image_generation_call has no tools[].model to price, billing 0.0") + return 0.0 + token_cost: Final = StandardBuiltInToolCostTracking._image_generation_token_cost( + response_object, tool_model, custom_llm_provider + ) + if token_cost is not None: + return token_cost + return sum( + StandardBuiltInToolCostTracking._image_generation_call_cost(output_item, tool_model, custom_llm_provider) + for output_item in completed_items + ) + @staticmethod def _handle_file_search_cost( model: str, @@ -445,9 +560,12 @@ class StandardBuiltInToolCostTracking: return False elif isinstance(response_object, ResponsesAPIResponse): # response api explicitly includes web_search_call in the output - return StandardBuiltInToolCostTracking.response_includes_output_type( + if StandardBuiltInToolCostTracking.response_includes_output_type( response_object=response_object, output_type="web_search_call" - ) + ): + return True + reported: Final = _reported_web_search_requests(response_object) + return reported is not None and reported > 0 elif usage is not None: if get_web_search_requests_from_usage(usage) is not None or ( hasattr(usage, "prompt_tokens_details") diff --git a/litellm/litellm_core_utils/llm_response_utils/get_headers.py b/litellm/litellm_core_utils/llm_response_utils/get_headers.py index b86cc954b0e..3c08f16fa80 100644 --- a/litellm/litellm_core_utils/llm_response_utils/get_headers.py +++ b/litellm/litellm_core_utils/llm_response_utils/get_headers.py @@ -1,8 +1,11 @@ from collections.abc import Mapping from typing import Final +import httpx from pydantic import TypeAdapter, ValidationError +from litellm.constants import PROVIDER_REQUEST_ID_HEADERS + _HEADERS: Final = TypeAdapter(Mapping[str, str]) _REQUEST_ID_HEADERS: Final = ("request-id", "x-request-id", "llm_provider-request-id", "llm_provider-x-request-id") @@ -16,6 +19,16 @@ def get_provider_request_id(headers: object) -> str | None: return next((normalized[name] for name in _REQUEST_ID_HEADERS if normalized.get(name)), None) +def provider_request_id_headers(headers: Mapping[str, object] | None) -> httpx.Headers | None: + """The provider request id headers in `headers` whose values are ASCII strings, or None when there are none.""" + request_id_headers: Final = tuple( + (name, value) + for name, value in (headers or {}).items() + if name.lower() in PROVIDER_REQUEST_ID_HEADERS and isinstance(value, str) and value.isascii() + ) + return httpx.Headers(request_id_headers) if request_id_headers else None + + def get_response_headers(_response_headers: Mapping[str, str] | None = None) -> dict: """ diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 0c18ce38b13..76adadb2e1f 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -2582,6 +2582,48 @@ def _with_native_redacted_data( _AnthropicMessageList: TypeAlias = list[AllAnthropicPassThroughMessageValues] +def _client_tool_result(message: AllMessageValues, replaced_server_tool_answers: frozenset[str]) -> bool: + return ( + message_field(message, "role") in ("tool", "function") + and message_field(message, "tool_call_id") not in replaced_server_tool_answers + ) + + +def _tool_result_later_in_turn( + messages: Sequence[AllMessageValues], replaced_server_tool_answers: frozenset[str] +) -> tuple[bool, ...]: + from_the_end: Final = itertools.accumulate( + reversed(messages), + lambda later, message: ( + message_field(message, "role") in ("user", "tool", "function") + and (later or _client_tool_result(message, replaced_server_tool_answers)) + ), + initial=False, + ) + return tuple(from_the_end)[::-1][1:] + + +def _empty_text_placeholder(message: AllMessageValues) -> AnthropicMessagesTextParam: + placeholder: Final = AnthropicMessagesTextParam(type="text", text=_EMPTY_TEXT_PLACEHOLDER) + add_cache_control_to_content(anthropic_content_element=placeholder, original_content_element=message) + return placeholder + + +def _cache_breakpoint_placeholder( + message: AllMessageValues, tool_result_follows: bool +) -> tuple[AnthropicMessagesTextParam, ...]: + if tool_result_follows or not message_field(message, "cache_control"): + return () + return (_empty_text_placeholder(message),) + + +def _empty_user_turn_placeholder(span: Sequence[AllMessageValues]) -> AnthropicMessagesTextParam | None: + empty_user_message: Final = next( + (message for message in reversed(span) if message["role"] == "user" and not message.get("content")), None + ) + return None if empty_user_message is None else _empty_text_placeholder(empty_user_message) + + def anthropic_messages_pt( messages: list[AllMessageValues], model: str, @@ -2636,6 +2678,7 @@ def anthropic_messages_pt( is_vertex_ai = llm_provider.startswith("vertex_ai") if llm_provider else False force_base64 = is_bedrock_invoke or is_vertex_ai + tool_result_later_in_turn: Final = _tool_result_later_in_turn(messages, replaced_server_tool_answers) msg_i = 0 while msg_i < len(messages): user_content: list[AnthropicMessagesUserMessageValues] = [] @@ -2648,8 +2691,12 @@ def anthropic_messages_pt( ChatCompletionToolMessage | ChatCompletionUserMessage | ChatCompletionFunctionMessage ) = messages[msg_i] if user_message_types_block["role"] == "user": - if isinstance(user_message_types_block["content"], list): - for m in user_message_types_block["content"]: + if not (user_block_content := user_message_types_block.get("content")): + user_content.extend( + _cache_breakpoint_placeholder(user_message_types_block, tool_result_later_in_turn[msg_i]) + ) + elif isinstance(user_block_content, list): + for m in user_block_content: if m.get("type", "") == "image_url": m = cast(ChatCompletionImageObject, m) format = m["image_url"].get("format") if isinstance(m["image_url"], dict) else None @@ -2720,10 +2767,10 @@ def anthropic_messages_pt( _file_content_element, ) ) - elif isinstance(user_message_types_block["content"], str): + elif isinstance(user_block_content, str): _anthropic_content_text_element: AnthropicMessagesTextParam = { "type": "text", - "text": user_message_types_block["content"], + "text": user_block_content, } _content_element = add_cache_control_to_content( anthropic_content_element=_anthropic_content_text_element, @@ -2747,6 +2794,8 @@ def anthropic_messages_pt( if user_content: new_messages.append({"role": "user", "content": user_content}) + elif (placeholder := _empty_user_turn_placeholder(messages[init_msg_i:msg_i])) is not None: + new_messages.append({"role": "user", "content": [placeholder]}) # Track unique tool IDs in this merge block to avoid duplication unique_tool_ids: set[str] = set() diff --git a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py index 4d1e77126a9..fd0096f9f4d 100644 --- a/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py +++ b/litellm/litellm_core_utils/prompt_templates/mid_conversation_system.py @@ -414,10 +414,11 @@ def _assistant_renders(message: object, answered: frozenset[str]) -> bool: def _renders(message: object, answered: frozenset[str]) -> bool: """Whether ``anthropic_messages_pt`` puts a block on the wire for this message. - A tool message always becomes a ``tool_result`` and a user message with string - content always becomes a text block (empty text gets a placeholder). A user list - renders only through parts of a type the converter emits; ``None``, an empty list, - and a list of other parts vanish. Assistant messages follow ``_assistant_renders``. + A tool message always becomes a ``tool_result``, and a user message with string or + missing content (``None``, an empty list) always puts its turn on the wire, an empty + turn getting a placeholder. A user list renders only through parts of a type the + converter emits, so a list of only other parts vanishes. Assistant messages follow + ``_assistant_renders``. """ role: Final = message_field(message, "role") if role in _TOOL_ROLES: @@ -425,8 +426,10 @@ def _renders(message: object, answered: frozenset[str]) -> bool: if role == "assistant": return _assistant_renders(message, answered) content: Final = message_field(message, "content") - return isinstance(content, str) or any( - message_field(part, "type") in _RENDERED_PART_TYPES for part in parts_of(content) + return ( + isinstance(content, str) + or not content + or any(message_field(part, "type") in _RENDERED_PART_TYPES for part in parts_of(content)) ) diff --git a/litellm/llms/anthropic/count_tokens/token_counter.py b/litellm/llms/anthropic/count_tokens/token_counter.py index d78de0037a0..55eaa0d2ac6 100644 --- a/litellm/llms/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/anthropic/count_tokens/token_counter.py @@ -6,6 +6,7 @@ from typing import Any, Final from pydantic import JsonValue +import litellm from litellm._logging import verbose_logger from litellm.exceptions import AuthenticationError from litellm.llms.anthropic.count_tokens.handler import AnthropicCountTokensHandler @@ -61,14 +62,24 @@ class AnthropicTokenCounter(BaseTokenCounter): try: auth_header: Final = await AnthropicModelInfo.aget_auth_header( - api_key=litellm_params.get("api_key"), + api_key=litellm_params.get("api_key") or litellm.anthropic_key or litellm.api_key, api_base=api_base, litellm_params=litellm_params, allow_workload_identity=True, ) if auth_header is None: - verbose_logger.warning("No Anthropic credential found for token counting") - return None + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="anthropic_api", + error=True, + error_message=( + f"No Anthropic API key is set for model {request_model or model_to_use}, so tokens can't be " + "counted. Set api_key on the deployment or ANTHROPIC_API_KEY in the proxy environment." + ), + status_code=401, + ) result: Final = await anthropic_count_tokens_handler.handle_count_tokens_request( model=model_to_use, diff --git a/litellm/llms/anthropic/pass_through/convertible_blocks.py b/litellm/llms/anthropic/pass_through/convertible_blocks.py new file mode 100644 index 00000000000..543a6efd4c4 --- /dev/null +++ b/litellm/llms/anthropic/pass_through/convertible_blocks.py @@ -0,0 +1,106 @@ +from collections.abc import Iterator, Mapping, Sequence +from itertools import chain +from types import MappingProxyType +from typing import Final + +from pydantic import JsonValue, TypeAdapter, ValidationError +from typing_extensions import assert_never + +from litellm.llms.anthropic.pass_through.messages.utils import normalize_native_messages +from litellm.types.utils import TokenCountMessagesFormat + +_CONVERTIBLE_BLOCK_TYPES: Final = MappingProxyType( + { + "user": frozenset({"text", "image", "document", "tool_result"}), + "assistant": frozenset({"text", "tool_use", "thinking", "redacted_thinking"}), + } +) +_SOURCED_BLOCK_TYPES: Final = frozenset({"image", "document"}) +_TOOL_RESULT_BLOCK_TYPES: Final = frozenset({"text", "image", "document"}) + + +_JSON_MESSAGES: Final = TypeAdapter(tuple[Mapping[str, JsonValue], ...]) + + +def _json_messages(messages: Sequence[Mapping[str, object]]) -> tuple[Mapping[str, JsonValue], ...]: + try: + return _JSON_MESSAGES.validate_python(messages) + except ValidationError: + return () + + +def _source_label(block_type: str, source: JsonValue) -> str | None: + if not isinstance(source, dict): + return f"{block_type} without a source" + match source.get("type"): + case "base64" if source.get("data"): + return None + case "url" if source.get("url"): + return None + case "base64" | "url" as source_type: + return f"{block_type} with an empty {source_type} source" + case str() as source_type: + return f"{block_type} with a {source_type} source" + case _: + return f"{block_type} with an untyped source" + + +def _block_label(block: JsonValue, convertible: frozenset[str]) -> str | None: + if not isinstance(block, dict): + return None + block_type: Final = block.get("type") + if not isinstance(block_type, str): + return "a block without a type" + if block_type not in convertible: + return block_type + return _source_label(block_type, block.get("source")) if block_type in _SOURCED_BLOCK_TYPES else None + + +def _tool_result_labels(block: JsonValue) -> Iterator[str]: + if not isinstance(block, dict) or block.get("type") != "tool_result": + return iter(()) + inner: Final = block.get("content") + if not isinstance(inner, list): + return iter(()) + return ( + f"{label} inside a tool_result" + for item in inner + if (label := _block_label(item, _TOOL_RESULT_BLOCK_TYPES)) is not None + ) + + +def _content_labels(block: JsonValue, convertible: frozenset[str]) -> Iterator[str]: + label: Final = _block_label(block, convertible) + return chain(() if label is None else (label,), _tool_result_labels(block)) + + +def _message_labels(message: Mapping[str, JsonValue]) -> Iterator[str]: + role: Final = message.get("role") + convertible: Final = _CONVERTIBLE_BLOCK_TYPES.get(role) if isinstance(role, str) else None + content: Final = message.get("content") + if convertible is None or not isinstance(content, list): + return iter(()) + return chain.from_iterable(_content_labels(block, convertible) for block in content) + + +def unconvertible_anthropic_blocks(messages: Sequence[Mapping[str, object]]) -> tuple[str, ...]: + normalized: Final = _json_messages(normalize_native_messages(messages)) + return tuple(sorted(set(chain.from_iterable(_message_labels(message) for message in normalized)))) + + +def unconvertible_blocks_error( + messages: Sequence[Mapping[str, object]] | None, messages_format: TokenCountMessagesFormat +) -> str | None: + match messages_format: + case "openai": + return None + case "anthropic": + labels: Final = unconvertible_anthropic_blocks(messages or ()) + return ( + "count_tokens can't convert these content blocks for this model, so the count would leave them out: " + + ", ".join(labels) + if labels + else None + ) + case _: + assert_never(messages_format) diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py index 0e0b68df449..53734a77f3c 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/handler.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/handler.py @@ -30,7 +30,7 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): self, model: str, messages: list[dict[str, JsonValue]], - api_key: str, + api_key: str | None, api_base: str, litellm_params: dict[str, object] | None = None, timeout: float | httpx.Timeout | None = None, @@ -80,6 +80,15 @@ class AzureAIAnthropicCountTokensHandler(AzureAIAnthropicCountTokensConfig): api_key=api_key, litellm_params=litellm_params, ) + if not {"x-api-key", "Authorization"} & headers.keys(): + raise AnthropicError( + status_code=401, + message=( + f"No Azure AI API key or Entra ID credential is set for model {model}, so tokens can't be " + "counted. Set api_key or an Entra ID credential on the deployment, or AZURE_AI_API_KEY or " + "AZURE_API_KEY in the proxy environment." + ), + ) # Use LiteLLM's async httpx client async_client: Final = get_async_httpx_client(llm_provider=litellm.LlmProviders.AZURE_AI) diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py index 2707d24ee07..0ed432cc394 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/token_counter.py @@ -7,6 +7,7 @@ from typing import Any, Final from pydantic import JsonValue +import litellm from litellm._logging import verbose_logger from litellm.llms.azure_ai.anthropic.count_tokens.handler import ( AzureAIAnthropicCountTokensHandler, @@ -18,6 +19,21 @@ from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCou azure_ai_anthropic_count_tokens_handler: Final = AzureAIAnthropicCountTokensHandler() +def _missing_api_base_response(request_model: str, model_to_use: str) -> TokenCountResponse: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="azure_ai_anthropic_api", + error=True, + error_message=( + f"No Azure AI API base is set for model {request_model or model_to_use}, so tokens can't be counted. " + "Set api_base on the deployment or AZURE_AI_API_BASE or AZURE_API_BASE in the proxy environment." + ), + status_code=500, + ) + + class AzureAIAnthropicTokenCounter(BaseTokenCounter): """Token counter implementation for Azure AI Anthropic provider using the CountTokens API.""" @@ -65,18 +81,15 @@ class AzureAIAnthropicTokenCounter(BaseTokenCounter): if not api_key: api_key = os.getenv("AZURE_AI_API_KEY") - # Get API base from deployment config or environment - api_base = litellm_params.get("api_base") - if not api_base: - api_base = os.getenv("AZURE_AI_API_BASE") - - if not api_key: - verbose_logger.warning("No Azure AI API key found for token counting") - return None + api_base: Final = ( + litellm_params.get("api_base") + or os.getenv("AZURE_AI_API_BASE") + or os.getenv("AZURE_API_BASE") + or litellm.api_base + ) if not api_base: - verbose_logger.warning("No Azure AI API base found for token counting") - return None + return _missing_api_base_response(request_model, model_to_use) try: result: Final = await azure_ai_anthropic_count_tokens_handler.handle_count_tokens_request( diff --git a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py index d7455439151..5946a1f4480 100644 --- a/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py +++ b/litellm/llms/azure_ai/anthropic/count_tokens/transformation.py @@ -24,7 +24,7 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): def get_required_headers( self, - api_key: str, + api_key: str | None, litellm_params: dict[str, object] | None = None, ) -> dict[str, str]: """ @@ -45,12 +45,12 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): "Content-Type": "application/json", "anthropic-version": "2023-06-01", "anthropic-beta": ANTHROPIC_TOKEN_COUNTING_BETA_VERSION, - "x-api-key": api_key, # Azure AI Anthropic requires this header + **({"x-api-key": api_key} if api_key else {}), } # Also set up Azure auth headers for flexibility litellm_params = litellm_params or {} - if "api_key" not in litellm_params: + if api_key and "api_key" not in litellm_params: litellm_params["api_key"] = api_key litellm_params_obj: Final = GenericLiteLLMParams.model_validate(litellm_params) @@ -58,8 +58,9 @@ class AzureAIAnthropicCountTokensConfig(AnthropicCountTokensConfig): # Get Azure auth headers (api-key or Authorization) azure_headers = BaseAzureLLM.base_validate_azure_environment(headers={}, litellm_params=litellm_params_obj) - # Merge Azure auth headers headers.update(azure_headers) + if "api-key" in azure_headers: + headers.setdefault("x-api-key", azure_headers["api-key"]) return headers diff --git a/litellm/llms/bedrock/base_aws_llm.py b/litellm/llms/bedrock/base_aws_llm.py index 848120150cd..f1426aa9375 100644 --- a/litellm/llms/bedrock/base_aws_llm.py +++ b/litellm/llms/bedrock/base_aws_llm.py @@ -99,11 +99,16 @@ def _assume_role_params( ) -_SecureTransportBool = TypedDict("_SecureTransportBool", {"aws:SecureTransport": ReadOnly[Literal["true"]]}) +_SecureTransportBool: Final = TypedDict("_SecureTransportBool", {"aws:SecureTransport": ReadOnly[Literal["true"]]}) + +_CalledViaLastEquals: Final = TypedDict( + "_CalledViaLastEquals", {"aws:CalledViaLast": ReadOnly[Literal["aws-external-anthropic.amazonaws.com"]]} +) -class _SecureTransportCondition(TypedDict): +class _SessionPolicyCondition(TypedDict): Bool: ReadOnly[_SecureTransportBool] + StringEquals: NotRequired[ReadOnly[_CalledViaLastEquals]] class _SessionPolicyStatement(TypedDict): @@ -111,7 +116,7 @@ class _SessionPolicyStatement(TypedDict): Effect: ReadOnly[Literal["Allow"]] Action: ReadOnly[tuple[str, ...]] Resource: ReadOnly[Literal["*"]] - Condition: ReadOnly[_SecureTransportCondition] + Condition: ReadOnly[_SessionPolicyCondition] class WebIdentitySessionPolicy(TypedDict): @@ -147,11 +152,24 @@ _WEB_IDENTITY_SESSION_POLICY_ACTIONS: Final[Mapping[str, tuple[str, ...]]] = Map "aws-external-anthropic:Get*", "aws-external-anthropic:List*", ), + "ClaudePlatformOutboundIdentityLiteLLM": ( + "sts:GetWebIdentityToken", + "sts:TagGetWebIdentityToken", + ), "BedrockMantleLiteLLM": ("bedrock-mantle:CreateInference",), } ) -_SECURE_TRANSPORT_ONLY: Final = _SecureTransportCondition(Bool=_SecureTransportBool({"aws:SecureTransport": "true"})) +_SECURE_TRANSPORT_ONLY: Final = _SessionPolicyCondition(Bool=_SecureTransportBool({"aws:SecureTransport": "true"})) + +_WEB_IDENTITY_SESSION_POLICY_CONDITIONS: Final[Mapping[str, _SessionPolicyCondition]] = MappingProxyType( + { + "ClaudePlatformOutboundIdentityLiteLLM": _SessionPolicyCondition( + Bool=_SecureTransportBool({"aws:SecureTransport": "true"}), + StringEquals=_CalledViaLastEquals({"aws:CalledViaLast": "aws-external-anthropic.amazonaws.com"}), + ), + } +) def build_web_identity_session_policy() -> WebIdentitySessionPolicy: @@ -163,7 +181,7 @@ def build_web_identity_session_policy() -> WebIdentitySessionPolicy: Effect="Allow", Action=actions, Resource="*", - Condition=_SECURE_TRANSPORT_ONLY, + Condition=_WEB_IDENTITY_SESSION_POLICY_CONDITIONS.get(sid, _SECURE_TRANSPORT_ONLY), ) for sid, actions in _WEB_IDENTITY_SESSION_POLICY_ACTIONS.items() ), diff --git a/litellm/llms/bedrock/chat/invoke_handler.py b/litellm/llms/bedrock/chat/invoke_handler.py index a207bceb418..6ad400661be 100644 --- a/litellm/llms/bedrock/chat/invoke_handler.py +++ b/litellm/llms/bedrock/chat/invoke_handler.py @@ -13,6 +13,7 @@ from litellm.caching.caching import InMemoryCache from litellm.constants import RESPONSE_FORMAT_TOOL_NAME from litellm.litellm_core_utils.core_helpers import map_finish_reason from litellm.litellm_core_utils.litellm_logging import Logging +from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers from litellm.litellm_core_utils.optional_imports import ensure_optional_import from litellm.litellm_core_utils.prompt_templates.common_utils import native_redacted_thinking_data from litellm.llms.anthropic.chat.handler import ( @@ -448,6 +449,7 @@ class _EventStreamTally: f"x-amzn-requestid={_response_header(response_headers, 'x-amzn-requestid')!r}, " f"{self.bytes_received} bytes received" ) + request_id_headers: Final = provider_request_id_headers(response_headers) if not self.events: return BedrockError( status_code=502, @@ -455,6 +457,7 @@ class _EventStreamTally: "Bedrock answered the stream with HTTP 200 but its body decoded to no events " f"({detail}, first bytes={self.head!r})" ), + headers=request_id_headers, ) if not self.recognized_events: return BedrockError( @@ -464,10 +467,12 @@ class _EventStreamTally: f"event type (event types={sorted(self.unrecognized_event_types)}, {detail}, " f"first payload={self.unrecognized_head!r})" ), + headers=request_id_headers, ) return BedrockError( status_code=502, message=f"Bedrock stream ended with {undecoded} undecoded bytes after {self.events} events ({detail})", + headers=request_id_headers, ) @@ -837,11 +842,12 @@ class AWSEventStreamDecoder: event_stream_buffer: Final = EventStreamBuffer() tally: Final = _EventStreamTally(get_bedrock_stream_event_statuses()) + request_id_headers: Final = provider_request_id_headers(response_headers) for chunk in iterator: event_stream_buffer.add_data(chunk) tally.add_chunk(chunk) for event in event_stream_buffer: - message = self._decode_event(event, tally) + message = self._decode_event(event, tally, request_id_headers) if message: # sse_event = ServerSentEvent(data=message, event="completion") _data = json.loads(message) @@ -858,11 +864,12 @@ class AWSEventStreamDecoder: event_stream_buffer: Final = EventStreamBuffer() tally: Final = _EventStreamTally(get_bedrock_stream_event_statuses()) + request_id_headers: Final = provider_request_id_headers(response_headers) async for chunk in iterator: event_stream_buffer.add_data(chunk) tally.add_chunk(chunk) for event in event_stream_buffer: - message = self._decode_event(event, tally) + message = self._decode_event(event, tally, request_id_headers) if message: _data = json.loads(message) yield self.chunk_parser(chunk_data=_data) @@ -882,27 +889,34 @@ class AWSEventStreamDecoder: ) return response_stream_shape - def _decode_event(self, event: "EventStreamMessage", tally: _EventStreamTally) -> str | None: + def _decode_event( + self, event: "EventStreamMessage", tally: _EventStreamTally, request_id_headers: httpx.Headers | None + ) -> str | None: response_stream_shape: Final = self._response_stream_shape() response_dict: Final = bedrock_event_stream_response(event) tally.add_event(event, response_dict["headers"]) - return self._parse_message_from_response(response_dict, response_stream_shape) + return self._parse_message_from_response(response_dict, response_stream_shape, request_id_headers) def _parse_message_from_event(self, event: "EventStreamMessage") -> str | None: response_stream_shape: Final = self._response_stream_shape() return self._parse_message_from_response(bedrock_event_stream_response(event), response_stream_shape) def _parse_message_from_response( - self, response_dict: BedrockEventStreamResponseDict, response_stream_shape: "Shape" + self, + response_dict: BedrockEventStreamResponseDict, + response_stream_shape: "Shape", + request_id_headers: httpx.Headers | None = None, ) -> str | None: parsed_response: Final = self.parser.parse(response_dict, response_stream_shape) if response_dict["status_code"] != 200: - raise build_bedrock_stream_error(response_dict, response_stream_shape) + raise build_bedrock_stream_error(response_dict, response_stream_shape, headers=request_id_headers) event_type: Final = bedrock_event_stream_header(response_dict["headers"], ":event-type") event_error_status: Final = bedrock_stream_event_error_status(event_type) if event_type is not None and event_error_status is not None: - raise build_bedrock_stream_event_error(event_type, event_error_status, response_dict["body"]) + raise build_bedrock_stream_event_error( + event_type, event_error_status, response_dict["body"], headers=request_id_headers + ) if "chunk" in parsed_response: chunk = parsed_response.get("chunk") if not chunk: diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index e79865ef6bf..5862b05a619 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -1913,13 +1913,18 @@ def bedrock_event_stream_header(headers: Mapping[str, object], name: str) -> str return value if isinstance(value, str) else None -def build_bedrock_stream_event_error(event_type: str, status_code: int, body: bytes) -> BedrockError: - return BedrockError(status_code=status_code, message=f"{event_type} {body.decode(errors='replace')}") +def build_bedrock_stream_event_error( + event_type: str, status_code: int, body: bytes, headers: httpx.Headers | None = None +) -> BedrockError: + return BedrockError( + status_code=status_code, message=f"{event_type} {body.decode(errors='replace')}", headers=headers + ) def build_bedrock_stream_error( response_dict: BedrockEventStreamResponseDict, response_stream_shape: Shape | None, + headers: httpx.Headers | None = None, ) -> BedrockError: """Build a BedrockError for a non-200 event-stream error event. @@ -1929,12 +1934,14 @@ def build_bedrock_stream_error( """ exception_type: Final = bedrock_event_stream_header(response_dict["headers"], ":exception-type") if exception_type is None: - return BedrockError(status_code=response_dict["status_code"], message=response_dict["body"].decode()) + return BedrockError( + status_code=response_dict["status_code"], message=response_dict["body"].decode(), headers=headers + ) member: Final = _structure_members(response_stream_shape).get(exception_type) modeled_status: Final = None if member is None else _modeled_error_status(member) status_code: Final = response_dict["status_code"] if modeled_status is None else modeled_status - return build_bedrock_stream_event_error(exception_type, status_code, response_dict["body"]) + return build_bedrock_stream_event_error(exception_type, status_code, response_dict["body"], headers=headers) class BedrockEventStreamDecoderBase: diff --git a/litellm/llms/gemini/common_utils.py b/litellm/llms/gemini/common_utils.py index bac5824b981..3d04954dd41 100644 --- a/litellm/llms/gemini/common_utils.py +++ b/litellm/llms/gemini/common_utils.py @@ -507,6 +507,7 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): gemini_count_contents, inference_request_view, provider_error_response, + unconvertible_messages_error, vertex_gemini_count_context, ) @@ -514,6 +515,9 @@ class GoogleAIStudioTokenCounter(BaseTokenCounter): count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools) if isinstance(count_context, VertexGeminiCountContextError): return count_context_error_response(count_context, request_model, model_to_use) + blocks_error: Final = unconvertible_messages_error(request_view.messages, messages_format) + if blocks_error is not None: + return count_context_error_response(blocks_error, request_model, model_to_use) deployment = deployment or {} count_tokens_params_request: Final = copy.deepcopy(deployment.get("litellm_params", {})) diff --git a/litellm/llms/openai/responses/count_tokens/token_counter.py b/litellm/llms/openai/responses/count_tokens/token_counter.py index b97e1329675..7c0697626ca 100644 --- a/litellm/llms/openai/responses/count_tokens/token_counter.py +++ b/litellm/llms/openai/responses/count_tokens/token_counter.py @@ -2,19 +2,21 @@ OpenAI Token Counter implementation using the Responses API /input_tokens endpoint. """ -import os from collections.abc import Mapping, Sequence from typing import Any, Final +import openai from pydantic import JsonValue from typing_extensions import assert_never from litellm._logging import verbose_logger +from litellm.litellm_core_utils.asyncify import asyncify +from litellm.llms.anthropic.pass_through.convertible_blocks import unconvertible_blocks_error from litellm.llms.anthropic.pass_through.responses_adapters.transformation import ( LiteLLMAnthropicToResponsesAPIAdapter, ) from litellm.llms.base_llm.base_utils import BaseTokenCounter -from litellm.llms.openai.common_utils import OpenAIError +from litellm.llms.openai.common_utils import OpenAIError, get_openai_credentials from litellm.llms.openai.responses.count_tokens.handler import ( OpenAICountTokensHandler, ) @@ -22,12 +24,26 @@ from litellm.llms.openai.responses.count_tokens.transformation import ( OpenAICountTokensConfig, ResponsesCountRequestParts, ) +from litellm.llms.openai.workload_identity import ( + OpenAIWorkloadIdentityConfig, + get_workload_identity_bearer_token, + resolve_openai_workload_identity_config, +) from litellm.types.utils import LlmProviders, TokenCountMessagesFormat, TokenCountResponse # Global handler instance - reuse across all token counting requests openai_count_tokens_handler: Final = OpenAICountTokensHandler() +def _workload_identity_config(litellm_params: Mapping[str, object]) -> OpenAIWorkloadIdentityConfig | None: + api_base: Final = litellm_params.get("api_base") + return resolve_openai_workload_identity_config( + api_key=None, + api_base=api_base if isinstance(api_base, str) else None, + litellm_params=litellm_params, + ) + + class OpenAITokenCounter(BaseTokenCounter): """Token counter implementation for OpenAI provider using the Responses API.""" @@ -91,30 +107,42 @@ class OpenAITokenCounter(BaseTokenCounter): if messages_format == "anthropic" else None ) - if malformed_error is not None: + request_error: Final = malformed_error or unconvertible_blocks_error(messages, messages_format) + if request_error is not None: return TokenCountResponse( total_tokens=0, request_model=request_model, model_used=model_to_use, tokenizer_type="openai_api", error=True, - error_message=malformed_error, + error_message=request_error, status_code=400, ) deployment = deployment or {} litellm_params: Final = deployment.get("litellm_params", {}) - # Get OpenAI API key from deployment config or environment - api_key = litellm_params.get("api_key") - if not api_key: - api_key = os.getenv("OPENAI_API_KEY") + credentials: Final = get_openai_credentials( + api_base=litellm_params.get("api_base"), api_key=litellm_params.get("api_key") + ) + credential: Final = credentials.api_key or _workload_identity_config(litellm_params) - if not api_key: - verbose_logger.warning("No OpenAI API key found for token counting") - return None + if credential is None: + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="openai_api", + error=True, + error_message=( + f"No OpenAI API key is set for model {request_model or model_to_use}, so tokens can't be " + "counted. Set api_key or workload identity on the deployment, or OPENAI_API_KEY in the proxy " + "environment." + ), + status_code=401, + ) - api_base: Final = litellm_params.get("api_base") + api_base: Final = credentials.api_base input_items, instructions, counted_tools, text = self._responses_request_parts( model_to_use=model_to_use, @@ -130,6 +158,11 @@ class OpenAITokenCounter(BaseTokenCounter): return None try: + api_key: Final = ( + credential + if isinstance(credential, str) + else await asyncify(get_workload_identity_bearer_token)(credential) + ) result: Final = await openai_count_tokens_handler.handle_count_tokens_request( model=model_to_use, input=input_items, @@ -161,6 +194,17 @@ class OpenAITokenCounter(BaseTokenCounter): error_message=e.message, status_code=e.status_code, ) + except openai.AuthenticationError as e: + verbose_logger.warning("OpenAI workload identity token exchange failed: %s", e.message) + return TokenCountResponse( + total_tokens=0, + request_model=request_model, + model_used=model_to_use, + tokenizer_type="openai_api", + error=True, + error_message=e.message, + status_code=401, + ) except Exception as e: verbose_logger.warning("Error calling OpenAI CountTokens API: %s", e) return TokenCountResponse( diff --git a/litellm/llms/soniox/audio_transcription/handler.py b/litellm/llms/soniox/audio_transcription/handler.py index b53955ae75d..f86bc142898 100644 --- a/litellm/llms/soniox/audio_transcription/handler.py +++ b/litellm/llms/soniox/audio_transcription/handler.py @@ -19,11 +19,14 @@ import asyncio import math import time from collections.abc import Coroutine, Mapping, Sequence +from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final import httpx +from pydantic import TypeAdapter from typing_extensions import ReadOnly, TypedDict +import litellm from litellm.litellm_core_utils.audio_utils.utils import ( get_audio_file_name, process_audio_file, @@ -35,7 +38,12 @@ from litellm.llms.custom_httpx.http_handler import ( get_httpx_client, ) from litellm.llms.soniox.audio_transcription.transformation import ( + SONIOX_HANDLER_ONLY_PARAMS, SonioxAudioTranscriptionConfig, + SonioxInvalidFormParam, + decode_soniox_form_params, + raise_soniox_form_error, + soniox_language_hints, ) from litellm.llms.soniox.common_utils import ( SONIOX_DEFAULT_CLEANUP, @@ -58,6 +66,39 @@ else: LiteLLMLoggingObj = Any +_CLEANUP_ENTRIES: Final = TypeAdapter(tuple[object, ...]) +_NOT_FORWARDED_TO_SONIOX: Final[frozenset[str]] = frozenset( + (*SONIOX_HANDLER_ONLY_PARAMS, "audio_url", "file_id", "response_format", "language") +) + + +def _optional_str(value: object) -> str | None: + return None if value is None else str(value) + + +def _float_or_default(value: object, default: float) -> float: + try: + parsed: Final = float(str(value)) + except ValueError: + return default + return parsed if math.isfinite(parsed) else default + + +def _int_or_default(value: object, default: int) -> int: + try: + return int(value) if isinstance(value, (int, float)) else int(str(value)) + except (ValueError, OverflowError): + return default + + +def _cleanup_targets(raw: object) -> tuple[str, ...]: + if raw is None: + return () + if isinstance(raw, str): + return (raw,) + return tuple(target for target in _CLEANUP_ENTRIES.validate_python(raw) if isinstance(target, str)) + + class _TranscriptionMeta(TypedDict, total=False): """Fields the handler reads from a Soniox transcription object.""" @@ -182,7 +223,7 @@ class SonioxAudioTranscriptionHandler: ) -> tuple[ dict[str, str], # auth headers str, # api_base (no trailing slash) - dict[str, object], # body for POST /v1/transcriptions (without file_id/audio_url) + Mapping[str, object], # body for POST /v1/transcriptions (without file_id/audio_url) _HandlerOptions, # handler-only options (poll interval, cleanup, ...) ]: # Validate env -> auth headers. @@ -198,52 +239,42 @@ class SonioxAudioTranscriptionHandler: base_url: Final = get_soniox_api_base(api_base) - # Operate on a local copy so we don't mutate the caller's dict - # (the caller may reuse `optional_params` for retries or logging). - params: Final = dict(optional_params) - - # Pull handler-only kwargs out of params so they aren't sent - # to Soniox. - poll_interval = float(params.pop("soniox_polling_interval", SONIOX_DEFAULT_POLL_INTERVAL)) - try: - max_attempts = int(params.pop("soniox_max_polling_attempts", SONIOX_DEFAULT_MAX_POLL_ATTEMPTS)) - except (ValueError, OverflowError): - max_attempts = SONIOX_DEFAULT_MAX_POLL_ATTEMPTS - cleanup_raw: Final = params.pop("soniox_cleanup", SONIOX_DEFAULT_CLEANUP) - if cleanup_raw is None: - cleanup: list[str] = [] - elif isinstance(cleanup_raw, str): - cleanup = [cleanup_raw] - else: - cleanup = list(cleanup_raw) - filename_override: Final = params.pop("filename", None) + decoded: Final = decode_soniox_form_params( + optional_params, drop_params=litellm_params.get("drop_params") is True or litellm.drop_params is True + ) + if isinstance(decoded, SonioxInvalidFormParam): + raise_soniox_form_error(decoded) # Server-side clamps. Caller-supplied poll settings (from request kwargs) # are bounded so an authenticated caller cannot force a worker into a # tight poll loop (zero interval) or pin it indefinitely (huge attempt # count). Total polling time is bounded by # SONIOX_MAX_POLL_ATTEMPTS * SONIOX_MAX_POLL_INTERVAL. - if not math.isfinite(poll_interval): - poll_interval = SONIOX_DEFAULT_POLL_INTERVAL - clamped_poll_interval: Final = max(SONIOX_MIN_POLL_INTERVAL, min(poll_interval, SONIOX_MAX_POLL_INTERVAL)) - clamped_max_attempts: Final = max(1, min(max_attempts, SONIOX_MAX_POLL_ATTEMPTS)) + poll_interval: Final = _float_or_default(decoded.get("soniox_polling_interval"), SONIOX_DEFAULT_POLL_INTERVAL) + max_attempts: Final = _int_or_default( + decoded.get("soniox_max_polling_attempts"), SONIOX_DEFAULT_MAX_POLL_ATTEMPTS + ) + cleanup: Final = _cleanup_targets(decoded.get("soniox_cleanup", SONIOX_DEFAULT_CLEANUP)) + language_hints: Final = soniox_language_hints(decoded.get("language"), decoded.get("language_hints")) # response_format is handled by LiteLLM post-processing, not Soniox. handler_opts: Final[_HandlerOptions] = { - "poll_interval": clamped_poll_interval, - "max_attempts": clamped_max_attempts, + "poll_interval": max(SONIOX_MIN_POLL_INTERVAL, min(poll_interval, SONIOX_MAX_POLL_INTERVAL)), + "max_attempts": max(1, min(max_attempts, SONIOX_MAX_POLL_ATTEMPTS)), "cleanup": cleanup, - "filename_override": filename_override, - "audio_url": params.pop("audio_url", None), - "file_id": params.pop("file_id", None), - "response_format": params.pop("response_format", None), + "filename_override": _optional_str(decoded.get("filename")), + "audio_url": _optional_str(decoded.get("audio_url")), + "file_id": _optional_str(decoded.get("file_id")), + "response_format": _optional_str(decoded.get("response_format")), } - # Soniox does not accept `language` directly; map_openai_params should - # already have translated it, but drop any leftover to be safe. - params.pop("language", None) - - return auth_headers, base_url, params, handler_opts + provider_params: Final = MappingProxyType( + { + **{key: value for key, value in decoded.items() if key not in _NOT_FORWARDED_TO_SONIOX}, + **({} if language_hints is None else {"language_hints": language_hints}), + } + ) + return auth_headers, base_url, provider_params, handler_opts def _build_create_body( self, diff --git a/litellm/llms/soniox/audio_transcription/transformation.py b/litellm/llms/soniox/audio_transcription/transformation.py index aa1ff239320..30d872f248f 100644 --- a/litellm/llms/soniox/audio_transcription/transformation.py +++ b/litellm/llms/soniox/audio_transcription/transformation.py @@ -9,9 +9,14 @@ async API requires multiple HTTP calls and does not fit the single-request contract of `base_llm_http_handler.audio_transcriptions`. """ -from typing import Any, Final +import re +from collections.abc import Mapping +from dataclasses import dataclass +from types import MappingProxyType +from typing import Any, Final, NoReturn from httpx import Headers, Response +from pydantic import TypeAdapter, ValidationError from litellm.llms.base_llm.audio_transcription.transformation import ( AudioTranscriptionRequestData, @@ -26,6 +31,7 @@ from litellm.llms.soniox.common_utils import ( render_soniox_tokens_as_srt, render_soniox_tokens_as_vtt, ) +from litellm.llms.soniox.types import SonioxContext from litellm.types.llms.openai import ( AllMessageValues, OpenAIAudioTranscriptionOptionalParams, @@ -58,6 +64,152 @@ SONIOX_HANDLER_ONLY_PARAMS: Final[list[str]] = [ ] +SONIOX_BOOL_PARAMS: Final[frozenset[str]] = frozenset( + {"enable_speaker_diarization", "enable_language_identification", "language_hints_strict"} +) +_JSON_OBJECT: Final = TypeAdapter(dict[str, object]) +_STRUCTURED_CONTEXT: Final = TypeAdapter(SonioxContext) +_FORM_BOOL: Final = TypeAdapter(bool) +_ITEMS: Final = TypeAdapter(tuple[object, ...]) +_BRACKETED_FORM_FIELD: Final = re.compile( + r"(?Pcontext|translation)\[(?P\w+)\](?:\[\]\[(?P\w+)\])?" +) +_BracketedField = tuple[str, str, str | None, object] + + +@dataclass(frozen=True, slots=True) +class SonioxInvalidFormParam: + key: str + value: str + expected: str + + +def raise_soniox_form_error(error: SonioxInvalidFormParam) -> NoReturn: + raise SonioxException( + message=f"`{error.key}` must be {error.expected}, got {error.value!r}", status_code=400, headers=None + ) + + +def soniox_language_hints(language: object, hints: object) -> object: + """Soniox has no `language` field, so the OpenAI `language` becomes `language_hints` when none are set.""" + if hints is not None or not isinstance(language, str) or not language: + return hints + return (language,) + + +def _decode_bool(key: str, value: str) -> bool | SonioxInvalidFormParam: + try: + return _FORM_BOOL.validate_python(value.strip()) + except ValidationError: + return SonioxInvalidFormParam(key=key, value=value, expected="a boolean") + + +def _decode_polling_interval(value: str) -> str | SonioxInvalidFormParam: + try: + float(value) + except ValueError: + return SonioxInvalidFormParam(key="soniox_polling_interval", value=value, expected="a number") + return value + + +def _decode_context(value: str) -> SonioxContext | str: + if not value.lstrip().startswith("{"): + return value + try: + return _STRUCTURED_CONTEXT.validate_json(value) + except ValidationError: + return value + + +def _decode_translation(value: str) -> Mapping[str, object] | SonioxInvalidFormParam: + try: + return _JSON_OBJECT.validate_json(value) + except ValidationError: + return SonioxInvalidFormParam(key="translation", value=value, expected="a JSON object") + + +def _decode_language_hints(value: str) -> tuple[object, ...] | SonioxInvalidFormParam: + if not value.lstrip().startswith("["): + return tuple(hint.strip() for hint in value.split(",") if hint.strip()) + try: + return _ITEMS.validate_json(value) + except ValidationError: + return SonioxInvalidFormParam(key="language_hints", value=value, expected="a JSON list") + + +def _decode_form_value(key: str, value: object) -> object | SonioxInvalidFormParam: + if not isinstance(value, str): + return value + if key in SONIOX_BOOL_PARAMS: + return _decode_bool(key, value) + if key == "soniox_polling_interval": + return _decode_polling_interval(value) + if key == "context": + return _decode_context(value) + if key == "translation": + return _decode_translation(value) + if key == "language_hints": + return _decode_language_hints(value) + return value + + +def _nested_form_param(param: str, fields: tuple[_BracketedField, ...]) -> Mapping[str, object]: + owned: Final = tuple(field for field in fields if field[0] == param) + flat: Final = {name: value for _, name, entry_field, value in owned if entry_field is None} + listed: Final = {name for _, name, entry_field, _ in owned if entry_field is not None} + return { + **flat, + **{ + name: ( + { + entry_field: value + for _, owner, entry_field, value in owned + if owner == name and entry_field is not None + }, + ) + for name in listed + }, + } + + +def _is_soniox_context(value: object) -> bool: + try: + _STRUCTURED_CONTEXT.validate_python(value) + except ValidationError: + return False + return True + + +def _nest_bracketed_form_fields(params: Mapping[str, object]) -> Mapping[str, object]: + """The OpenAI SDK sends a dict in `extra_body` as `context[terms]` and `context[general][][key]` form keys.""" + matches: Final = {key: _BRACKETED_FORM_FIELD.fullmatch(key) for key in params} + fields: Final = tuple( + (match["param"], match["field"], match["entry_field"], params[key]) for key, match in matches.items() if match + ) + nested: Final = {param: _nested_form_param(param, fields) for param in {field[0] for field in fields}} + kept: Final = {param: value for param, value in nested.items() if param != "context" or _is_soniox_context(value)} + unbracketed: Final = { + key: value for key, value in params.items() if (match := matches[key]) is None or match["param"] not in kept + } + return {**kept, **unbracketed} + + +def decode_soniox_form_params( + optional_params: Mapping[str, object], + drop_params: bool = False, +) -> Mapping[str, object] | SonioxInvalidFormParam: + """Multipart form fields reach the proxy as strings; restore the JSON types Soniox expects.""" + decoded: Final = { + key: _decode_form_value(key, value) for key, value in _nest_bracketed_form_fields(optional_params).items() + } + invalid: Final = next((value for value in decoded.values() if isinstance(value, SonioxInvalidFormParam)), None) + if invalid is not None and not drop_params: + return invalid + return MappingProxyType( + {key: value for key, value in decoded.items() if not isinstance(value, SonioxInvalidFormParam)} + ) + + class SonioxAudioTranscriptionConfig(BaseAudioTranscriptionConfig): """Configuration for Soniox async speech-to-text transcription.""" @@ -74,25 +226,26 @@ class SonioxAudioTranscriptionConfig(BaseAudioTranscriptionConfig): model: str, drop_params: bool, ) -> dict: - # Translate the OpenAI `language` param into Soniox `language_hints`. - if "language" in non_default_params and non_default_params["language"]: - language: Final = non_default_params["language"] - existing_hints: Final = optional_params.get("language_hints") - if not existing_hints: - optional_params["language_hints"] = [language] - elif language not in existing_hints: - optional_params["language_hints"] = [language] + list(existing_hints) - - # Capture response_format for post-processing (not sent to Soniox API). - if "response_format" in non_default_params: - optional_params["response_format"] = non_default_params["response_format"] - - # Pass through Soniox-native kwargs unchanged. - for key in SONIOX_PASSTHROUGH_PARAMS + SONIOX_HANDLER_ONLY_PARAMS: - if key in non_default_params and non_default_params[key] is not None: - optional_params[key] = non_default_params[key] - - return optional_params + language: Final[object] = non_default_params.get("language") + language_params: Final[Mapping[str, object]] = ( + { + "language": language, + "language_hints": soniox_language_hints(language, optional_params.get("language_hints")), + } + if language + else {} + ) + response_format_params: Final[Mapping[str, object]] = ( + {"response_format": non_default_params["response_format"]} + if "response_format" in non_default_params + else {} + ) + passthrough_params: Final[Mapping[str, object]] = { + key: non_default_params[key] + for key in SONIOX_PASSTHROUGH_PARAMS + SONIOX_HANDLER_ONLY_PARAMS + if non_default_params.get(key) is not None + } + return {**optional_params, **language_params, **response_format_params, **passthrough_params} def get_error_class(self, error_message: str, status_code: int, headers: dict | Headers) -> BaseLLMException: return SonioxException(message=error_message, status_code=status_code, headers=headers) diff --git a/litellm/llms/soniox/types.py b/litellm/llms/soniox/types.py new file mode 100644 index 00000000000..806d2a8fb0e --- /dev/null +++ b/litellm/llms/soniox/types.py @@ -0,0 +1,24 @@ +from collections.abc import Sequence + +from pydantic import ConfigDict, with_config +from typing_extensions import ReadOnly, TypedDict + + +class SonioxContextGeneralEntry(TypedDict): + key: ReadOnly[str] + value: ReadOnly[str] + + +class SonioxTranslationTerm(TypedDict): + source: ReadOnly[str] + target: ReadOnly[str] + + +@with_config(ConfigDict(extra="forbid")) +class SonioxContext(TypedDict, total=False): + """Soniox `context` request field: https://soniox.com/docs/stt/concepts/context""" + + general: ReadOnly[Sequence[SonioxContextGeneralEntry]] + text: ReadOnly[str] + terms: ReadOnly[Sequence[str]] + translation_terms: ReadOnly[Sequence[SonioxTranslationTerm]] diff --git a/litellm/llms/vertex_ai/common_utils.py b/litellm/llms/vertex_ai/common_utils.py index 4c0026783f7..98bc679a78e 100644 --- a/litellm/llms/vertex_ai/common_utils.py +++ b/litellm/llms/vertex_ai/common_utils.py @@ -1362,6 +1362,7 @@ class VertexAITokenCounter(BaseTokenCounter): gemini_count_contents, inference_request_view, provider_error_response, + unconvertible_messages_error, vertex_gemini_count_context, ) from litellm.llms.vertex_ai.count_tokens.handler import VertexAITokenCounter @@ -1370,6 +1371,9 @@ class VertexAITokenCounter(BaseTokenCounter): count_context: Final = vertex_gemini_count_context(system=request_view.system, tools=tools) if isinstance(count_context, VertexGeminiCountContextError): return count_context_error_response(count_context, request_model, model_to_use) + blocks_error: Final = unconvertible_messages_error(request_view.messages, messages_format) + if blocks_error is not None: + return count_context_error_response(blocks_error, request_model, model_to_use) resolved_contents: Final = ( contents diff --git a/litellm/llms/vertex_ai/count_tokens/context.py b/litellm/llms/vertex_ai/count_tokens/context.py index 3f9289c8244..e094f4c5c0c 100644 --- a/litellm/llms/vertex_ai/count_tokens/context.py +++ b/litellm/llms/vertex_ai/count_tokens/context.py @@ -138,6 +138,15 @@ def vertex_gemini_count_context( ) +def unconvertible_messages_error( + messages: Sequence[Mapping[str, object]] | None, messages_format: TokenCountMessagesFormat +) -> VertexGeminiCountContextError | None: + from litellm.llms.anthropic.pass_through.convertible_blocks import unconvertible_blocks_error + + error_message: Final = unconvertible_blocks_error(messages, messages_format) + return None if error_message is None else VertexGeminiCountContextError(error_message=error_message) + + def count_context_error_response( error: VertexGeminiCountContextError, request_model: str, model_used: str ) -> TokenCountResponse: diff --git a/litellm/model_prices_and_context_window_backup.json b/litellm/model_prices_and_context_window_backup.json index a204c49967e..b744032fb6c 100644 --- a/litellm/model_prices_and_context_window_backup.json +++ b/litellm/model_prices_and_context_window_backup.json @@ -29256,22 +29256,32 @@ "output_cost_per_token_priority": 1.8e-05 }, "gemini-3.1-pro-preview": { - "prompt_cache_min_tokens": 4096, + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_flex": 1e-06, + "input_cost_per_token_priority": 3.6e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", + "output_cost_per_image": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, - "output_cost_per_image": 0.00012, + "output_cost_per_token_flex": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "prompt_cache_min_tokens": 4096.0, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -29289,33 +29299,18 @@ ], "supports_audio_input": true, "supports_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "supports_url_context": true, - "supports_native_streaming": true, - "input_cost_per_token_priority": 3.6e-06, - "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, - "output_cost_per_token_priority": 2.16e-05, - "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, - "cache_read_input_token_cost_priority": 3.6e-07, - "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, - "cache_read_input_token_cost_flex": 2e-07, - "search_context_cost_per_query": { - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014, - "search_context_size_high": 0.014 - }, - "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_token_flex": 1e-06, - "output_cost_per_token_flex": 6e-06 + "web_search_billing_unit": "per_query" }, "gemini-3.1-pro-preview-customtools": { "prompt_cache_min_tokens": 4096, @@ -29428,15 +29423,25 @@ "vertex_ai/gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, "cache_read_input_token_cost": 5e-08, - "cache_read_input_token_cost_flex": 5e-08, - "input_cost_per_token": 5e-07, + "cache_read_input_token_cost_priority": 9e-08, + "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token_batches": 5e-07, + "input_cost_per_audio_token_priority": 1.8e-06, + "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, + "input_cost_per_token_flex": 2.5e-07, + "input_cost_per_token_priority": 9e-07, "litellm_provider": "vertex_ai", "max_input_tokens": 1048576, "max_output_tokens": 65535, "max_tokens": 65535, "mode": "chat", "output_cost_per_token": 3e-06, + "output_cost_per_token_batches": 1.5e-06, + "output_cost_per_token_flex": 1.5e-06, + "output_cost_per_token_priority": 5.4e-06, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -29454,6 +29459,7 @@ ], "supports_audio_input": true, "supports_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -29463,23 +29469,7 @@ "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "supports_native_streaming": true, - "input_cost_per_token_priority": 9e-07, - "input_cost_per_audio_token_priority": 1.8e-06, - "output_cost_per_token_priority": 5.4e-06, - "cache_read_input_token_cost_priority": 9e-08, - "search_context_cost_per_query": { - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014, - "search_context_size_high": 0.014 - }, - "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_audio_token_batches": 5e-07, - "input_cost_per_token_batches": 2.5e-07, - "input_cost_per_token_flex": 2.5e-07, - "output_cost_per_token_batches": 1.5e-06, - "output_cost_per_token_flex": 1.5e-06 + "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.5-flash": { "prompt_cache_min_tokens": 4096, @@ -29807,22 +29797,32 @@ "output_cost_per_video_token": 1e-06 }, "vertex_ai/gemini-3.1-pro-preview": { - "prompt_cache_min_tokens": 4096, + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_flex": 1e-06, + "input_cost_per_token_priority": 3.6e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 1048576, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", + "output_cost_per_image": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, - "output_cost_per_image": 0.00012, + "output_cost_per_token_flex": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "prompt_cache_min_tokens": 4096.0, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -29840,33 +29840,18 @@ ], "supports_audio_input": true, "supports_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "supports_url_context": true, - "supports_native_streaming": true, - "input_cost_per_token_priority": 3.6e-06, - "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, - "output_cost_per_token_priority": 2.16e-05, - "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, - "cache_read_input_token_cost_priority": 3.6e-07, - "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, - "cache_read_input_token_cost_flex": 2e-07, - "search_context_cost_per_query": { - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014, - "search_context_size_high": 0.014 - }, - "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_token_flex": 1e-06, - "output_cost_per_token_flex": 6e-06 + "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.1-pro-preview-customtools": { "prompt_cache_min_tokens": 4096, @@ -31666,9 +31651,15 @@ "gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, "cache_read_input_token_cost": 5e-08, - "cache_read_input_token_cost_flex": 5e-08, + "cache_read_input_token_cost_priority": 9e-08, + "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token_batches": 5e-07, + "input_cost_per_audio_token_priority": 1.8e-06, "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, + "input_cost_per_token_flex": 2.5e-07, + "input_cost_per_token_priority": 9e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 65535, @@ -31676,6 +31667,10 @@ "mode": "chat", "output_cost_per_reasoning_token": 3e-06, "output_cost_per_token": 3e-06, + "output_cost_per_token_batches": 1.5e-06, + "output_cost_per_token_flex": 1.5e-06, + "output_cost_per_token_priority": 5.4e-06, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -31693,6 +31688,7 @@ ], "supports_audio_output": false, "supports_function_calling": true, + "supports_native_streaming": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -31703,23 +31699,7 @@ "supports_url_context": true, "supports_vision": true, "supports_web_search": true, - "supports_native_streaming": true, - "input_cost_per_token_priority": 9e-07, - "input_cost_per_audio_token_priority": 1.8e-06, - "output_cost_per_token_priority": 5.4e-06, - "cache_read_input_token_cost_priority": 9e-08, - "search_context_cost_per_query": { - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014, - "search_context_size_high": 0.014 - }, - "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_audio_token_batches": 5e-07, - "input_cost_per_token_batches": 2.5e-07, - "input_cost_per_token_flex": 2.5e-07, - "output_cost_per_token_batches": 1.5e-06, - "output_cost_per_token_flex": 1.5e-06 + "web_search_billing_unit": "per_query" }, "gemini-omni-flash-preview": { "input_cost_per_audio_token": 1.5e-06, diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 97188217f63..420f42c690c 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -6,6 +6,7 @@ from __future__ import annotations import asyncio import contextvars +import weakref from collections.abc import AsyncGenerator, AsyncIterator, Awaitable, Coroutine, Generator, Iterator from functools import partial from types import TracebackType @@ -26,6 +27,23 @@ from litellm.utils import client base_llm_http_handler = BaseLLMHTTPHandler() from .utils import BasePassthroughUtils +_provider_errors: Final[weakref.WeakSet[httpx.HTTPStatusError]] = weakref.WeakSet() + + +def raise_for_provider_status(response: httpx.Response) -> None: + """``response.raise_for_status()`` for the provider's response to a passthrough request, recording the raised + error as the provider's own.""" + try: + response.raise_for_status() + except httpx.HTTPStatusError as error: + _provider_errors.add(error) + raise + + +def is_provider_error(error: httpx.HTTPStatusError) -> bool: + """Whether ``error`` is the provider's own error response to a passthrough request.""" + return error in _provider_errors + async def _as_async_generator(iterable: AsyncIterator[bytes]) -> AsyncGenerator[bytes, bytes]: async for chunk in iterable: @@ -112,7 +130,7 @@ class AsyncPassthroughStreamingResponse(AsyncGenerator[bytes, bytes]): self.status_code = self._response.status_code self._initialized = True try: - self._response.raise_for_status() + raise_for_provider_status(self._response) self._iterator = _as_async_generator(self._response.aiter_bytes()) except Exception: # noqa: BLE001 # Safe catch-all for cleanup logic try: @@ -626,7 +644,7 @@ async def _async_passthrough_request( else: response: Final = await response_result await response.aread() - response.raise_for_status() + raise_for_provider_status(response) return response else: # Fallback for sync-like behavior (shouldn't happen in async path) diff --git a/litellm/proxy/_types.py b/litellm/proxy/_types.py index e4a0a5242b4..d62c6b0d51f 100644 --- a/litellm/proxy/_types.py +++ b/litellm/proxy/_types.py @@ -3814,6 +3814,16 @@ class ModelAccessDeniedProxyException(ProxyException): return self.internal_message.replace("\r", "").replace("\n", "") +class AwsErrorTypeProxyException(Exception): + """A ProxyException answered to an AWS SDK, which reads the error code from the ``x-amzn-errortype`` header and + the message from a top-level ``message`` field.""" + + def __init__(self, error: ProxyException, aws_error_type: str) -> None: + super().__init__(error.message) + self.error: Final = error + self.aws_error_type: Final = aws_error_type + + from litellm.types.proxy.common_proxy_errors import ( # noqa: E402 # public re-export CommonProxyErrors as CommonProxyErrors, # noqa: PLC0414 # public re-export ) diff --git a/litellm/proxy/anthropic_endpoints/endpoints.py b/litellm/proxy/anthropic_endpoints/endpoints.py index 1f17f0a42d2..ee1ac79c820 100644 --- a/litellm/proxy/anthropic_endpoints/endpoints.py +++ b/litellm/proxy/anthropic_endpoints/endpoints.py @@ -29,6 +29,7 @@ from litellm.proxy.common_request_processing import ( create_response, log_llm_api_exception, proxy_exception_from_http_exception, + proxy_exception_from_route_error, resolve_litellm_call_id, ) from litellm.proxy.common_utils.error_body_call_id import error_body_call_id @@ -43,7 +44,6 @@ from litellm.proxy.common_utils.openai_error_payload import ( openai_error_type, with_litellm_call_id, ) -from litellm.types.utils import TokenCountResponse router: Final = APIRouter() @@ -340,10 +340,10 @@ async def count_tokens( messages: Final = data.get("messages", []) if not model_name: - raise HTTPException(status_code=400, detail={"error": "model parameter is required"}) + raise HTTPException(status_code=400, detail="model parameter is required") if not messages: - raise HTTPException(status_code=400, detail={"error": "messages parameter is required"}) + raise HTTPException(status_code=400, detail="messages parameter is required") # Create TokenCountRequest for the internal endpoint from litellm.proxy._types import TokenCountRequest @@ -361,30 +361,15 @@ async def count_tokens( messages_format="anthropic", output_format=_requested_output_format(data), ) - _token_response_dict: dict = {} - if isinstance(token_response, TokenCountResponse): - _token_response_dict = token_response.model_dump() - elif isinstance(token_response, dict): - _token_response_dict = token_response - - # Convert the internal response to Anthropic API format - return {"input_tokens": _token_response_dict.get("total_tokens", 0)} - - except HTTPException: - raise + except HTTPException as e: + return _anthropic_error_json_response(proxy_exception_from_http_exception(e, {}), request) except ProxyException as e: - status_code: Final = int(e.code) if e.code and e.code.isdigit() else 500 - detail: Final = AnthropicExceptionMapping.transform_to_anthropic_error( - status_code=status_code, - raw_message=e.message, - ) - raise HTTPException( - status_code=status_code, - detail=detail, - ) + return _anthropic_error_json_response(e, request) except Exception as e: log_llm_api_exception(e, litellm_call_id) - raise HTTPException(status_code=500, detail={"error": f"Internal server error: {e}"}) + return _anthropic_error_json_response(proxy_exception_from_route_error(e), request) + else: + return {"input_tokens": token_response.total_tokens} @router.post( diff --git a/litellm/proxy/google_endpoints/endpoints.py b/litellm/proxy/google_endpoints/endpoints.py index 784e3d4719f..d382ff53ffb 100644 --- a/litellm/proxy/google_endpoints/endpoints.py +++ b/litellm/proxy/google_endpoints/endpoints.py @@ -1,16 +1,22 @@ from collections.abc import Mapping from itertools import chain +from types import MappingProxyType from typing import Final from fastapi import APIRouter, Depends, HTTPException, Request, Response -from fastapi.responses import ORJSONResponse +from fastapi.responses import JSONResponse, ORJSONResponse from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter, ValidationError from litellm.litellm_core_utils.json_validation_rule import normalize_json_schema_types from litellm.proxy._types import * -from litellm.proxy._types import TokenCountRequest +from litellm.proxy._types import ProxyException, TokenCountRequest from litellm.proxy.auth.user_api_key_auth import UserAPIKeyAuth, user_api_key_auth -from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + log_llm_api_exception, + proxy_exception_from_http_exception, + resolve_litellm_call_id, +) from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports _read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export read_request_body, @@ -63,6 +69,37 @@ def _input_schema(declaration: _GoogleCountTokensFunctionDeclaration) -> object: return {"type": "object", "properties": {}} +_GOOGLE_RPC_STATUS: Final = MappingProxyType( + { + 400: "INVALID_ARGUMENT", + 401: "UNAUTHENTICATED", + 403: "PERMISSION_DENIED", + 404: "NOT_FOUND", + 409: "ABORTED", + 429: "RESOURCE_EXHAUSTED", + 499: "CANCELLED", + 501: "UNIMPLEMENTED", + 503: "UNAVAILABLE", + 504: "DEADLINE_EXCEEDED", + } +) + + +def _google_error_response(request: Request, error: ProxyException) -> JSONResponse: + from litellm.proxy.proxy_server import ( + _close_dangling_otel_server_span, # pyright: ignore[reportPrivateUsage] # proxy_server keeps the span-close helper private; error JSONResponses returned by the route must stamp the OTel server span like the global ProxyException handler does + ) + + status_code: Final = int(error.code) if error.code is not None and error.code.isdigit() else 500 + _close_dangling_otel_server_span(request, status_code, exc=error) + status: Final = _GOOGLE_RPC_STATUS.get(status_code, "INTERNAL" if status_code >= 500 else "FAILED_PRECONDITION") + return JSONResponse( + status_code=status_code, + content={"error": {"code": status_code, "message": error.message, "status": status}}, + headers=error.headers, + ) + + def _google_count_tokens_request(model_name: str, data: object) -> TokenCountRequest: from litellm.google_genai.adapters.transformation import GoogleGenAIAdapter @@ -251,16 +288,26 @@ async def google_count_tokens(request: Request, model_name: str): data: Final = await read_request_body(request=request) try: token_request: Final = _google_count_tokens_request(model_name, data) - except ValidationError as error: - raise HTTPException(status_code=400, detail="Invalid Google countTokens request body") from error + except ValidationError: + return _google_error_response( + request, + ProxyException( + message="Invalid Google countTokens request body", type="invalid_request_error", param=None, code=400 + ), + ) - # Call the internal token counter function with direct request flag set to False - token_response: Final = await internal_token_counter( - request=token_request, - call_endpoint=True, - ) - if token_response is None: - return TokenCountDetailsResponse(totalTokens=0, promptTokensDetails=[]) + try: + token_response: Final = await internal_token_counter(request=token_request, call_endpoint=True) + except ProxyException as error: + return _google_error_response(request, error) + except HTTPException as error: + return _google_error_response(request, proxy_exception_from_http_exception(error, {})) + except Exception as error: + log_llm_api_exception(error, resolve_litellm_call_id(request.headers.get("x-litellm-call-id"))) + return _google_error_response( + request, + ProxyException(message="Internal server error", type="internal_server_error", param=None, code=500), + ) original_response: Final[dict] = token_response.original_response or {} return TokenCountDetailsResponse( totalTokens=original_response.get("totalTokens") or token_response.total_tokens or 0, diff --git a/litellm/proxy/guardrails/guardrail_endpoints.py b/litellm/proxy/guardrails/guardrail_endpoints.py index d713a18f470..1d5c530a616 100644 --- a/litellm/proxy/guardrails/guardrail_endpoints.py +++ b/litellm/proxy/guardrails/guardrail_endpoints.py @@ -1483,6 +1483,17 @@ async def get_guardrail_info(guardrail_id: str): raise HTTPException(status_code=500, detail=str(e)) +def _undeclared_logging_only( + guardrail_class: type[CustomGuardrail], hooks: Sequence[GuardrailEventHooks] +) -> tuple[GuardrailEventHooks, ...]: + """logging_only when the class accepts it at save time without declaring it in its hooks.""" + return ( + (GuardrailEventHooks.logging_only,) + if guardrail_class.accepts_logging_only_mode() and GuardrailEventHooks.logging_only not in hooks + else () + ) + + @router.get( "/guardrails/ui/add_guardrail_settings", tags=["Guardrails"], @@ -1515,7 +1526,7 @@ async def get_guardrail_ui_settings() -> GuardrailUIAddGuardrailSettings: ] supported_modes_by_provider: Final = { - provider: [hook.value for hook in hooks] + provider: [hook.value for hook in (*hooks, *_undeclared_logging_only(guardrail_class, hooks))] for provider, guardrail_class in guardrail_class_registry.items() if (hooks := guardrail_class.get_supported_event_hooks()) is not None } | MappingProxyType( diff --git a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py index f959292d567..40e44151e65 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py +++ b/litellm/proxy/guardrails/guardrail_hooks/litellm_content_filter/content_filter.py @@ -6,6 +6,7 @@ to detect and block/mask sensitive content. """ import asyncio +import contextvars import itertools import json import os @@ -50,6 +51,7 @@ from litellm.types.utils import ( if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj + from litellm.llms.base_llm.guardrail_translation.base_translation import BaseTranslation from litellm.types.guardrails import ( BlockedWord, @@ -101,6 +103,13 @@ WORD_NUMBER_SEQUENCE_PATTERN: Final = re.compile( ) WORD_NUMBER_TOKEN_FINDER: Final = re.compile(rf"(?:{WORD_NUMBER_TOKEN_REGEX})", re.IGNORECASE) +_logging_only_scan_guardrail: Final[contextvars.ContextVar["ContentFilterGuardrail | None"]] = contextvars.ContextVar( + "litellm_content_filter_logging_only_scan", default=None +) +_logging_only_scan_would_reject: Final[contextvars.ContextVar[bool]] = contextvars.ContextVar( + "litellm_content_filter_logging_only_scan_would_reject", default=False +) + class ConditionalCategoryConfig(TypedDict): identifier_words: Sequence[str] @@ -1209,6 +1218,18 @@ class ContentFilterGuardrail(CustomGuardrail): return (keyword, action, description) return None + def _is_logging_only_scan(self) -> bool: + return _logging_only_scan_guardrail.get() is self + + def _enforces_block(self, action: ContentFilterAction) -> bool: + return action == ContentFilterAction.BLOCK and not self._is_logging_only_scan() + + @staticmethod + def _asks_to_block(detection: ContentFilterDetection) -> bool: + if detection["type"] == "competitor_intent": + return detection["action_hint"] in ("refuse", "reframe") + return detection["action"] == ContentFilterAction.BLOCK.value + def _handle_conditional_match( self, matched_phrase: str, @@ -1228,7 +1249,7 @@ class ContentFilterGuardrail(CustomGuardrail): } detections.append(category_detection) - if action == ContentFilterAction.BLOCK: + if self._enforces_block(action): error_msg: Final = ( f"Content blocked: {category_name} conditional match '{matched_phrase}' detected (severity: {severity})" ) @@ -1269,7 +1290,7 @@ class ContentFilterGuardrail(CustomGuardrail): } detections.append(category_detection) - if action == ContentFilterAction.BLOCK: + if self._enforces_block(action): error_msg = f"Content blocked: {category_name} category keyword '{keyword}' detected (severity: {severity})" verbose_proxy_logger.warning(error_msg) raise HTTPException( @@ -1312,7 +1333,7 @@ class ContentFilterGuardrail(CustomGuardrail): } detections.append(pattern_detection) - if action == ContentFilterAction.BLOCK: + if self._enforces_block(action): error_msg: Final = f"Content blocked: {pattern_name} pattern detected" verbose_proxy_logger.warning(error_msg) raise HTTPException( @@ -1346,7 +1367,7 @@ class ContentFilterGuardrail(CustomGuardrail): } detections.append(blocked_word_detection) - if action == ContentFilterAction.BLOCK: + if self._enforces_block(action): error_msg = f"Content blocked: keyword '{keyword}' detected" if description: error_msg += f" ({description})" @@ -1619,7 +1640,7 @@ class ContentFilterGuardrail(CustomGuardrail): - Number of entities masked - Whether the guardrail blocked the request (max risk) """ - if status == "guardrail_intervened": + if status in ("guardrail_intervened", "guardrail_flagged"): return 10.0 total_masked: Final = sum(masked_entity_count.values()) if masked_entity_count else 0 @@ -1665,6 +1686,8 @@ class ContentFilterGuardrail(CustomGuardrail): } detections.append(detection) + if self._is_logging_only_scan(): + return if action_hint_val == "refuse": msg = "Content blocked: competitor comparison or ranking intent detected." if self._competitor_intent_checker and getattr( @@ -1763,18 +1786,18 @@ class ContentFilterGuardrail(CustomGuardrail): return raw_name return None + def _reject(self, error: str) -> None: + """Raise the 400 an enforcing mode returns, or note it on the leg a logging_only scan is observing.""" + if self._is_logging_only_scan(): + _logging_only_scan_would_reject.set(True) + return + raise HTTPException(status_code=400, detail={"error": error}) + def _assert_argument_label_clean( self, text: str, detections: list[ContentFilterDetection], context_label: str ) -> None: if self._filter_single_text(text, detections=detections) != text: - raise HTTPException( - status_code=400, - detail={ - "error": ( - f"Content blocked: {context_label} argument matched a masking rule on a non-rewritable field" - ) - }, - ) + self._reject(f"Content blocked: {context_label} argument matched a masking rule on a non-rewritable field") def _filter_argument_value( self, @@ -1784,10 +1807,8 @@ class ContentFilterGuardrail(CustomGuardrail): depth: int = 0, ) -> object: if depth > DEFAULT_MAX_RECURSE_DEPTH: - raise HTTPException( - status_code=400, - detail={"error": f"Content blocked: {context_label} arguments exceed the maximum nesting depth"}, - ) + self._reject(f"Content blocked: {context_label} arguments exceed the maximum nesting depth") + return value if isinstance(value, str): return self._filter_single_text(value, detections=detections) if isinstance(value, (int, float)) and not isinstance(value, bool): @@ -1874,6 +1895,21 @@ class ContentFilterGuardrail(CustomGuardrail): if filtered_arguments != arguments: self._set_tool_call_arguments(tool_call, filtered_arguments) + async def _scan_logged_call( + self, + kwargs: dict[str, object], # mutable-ok: CustomLogger.async_logging_hook contract + response: object | None, + translation: "BaseTranslation", + output_translation: "BaseTranslation", + scratch_metadata: dict[str, object], # mutable-ok: apply_guardrail records its verdict into request metadata + ) -> None: + """Scan the logged call with every BLOCK match recorded instead of raised.""" + token: Final = _logging_only_scan_guardrail.set(self) + try: + await super()._scan_logged_call(kwargs, response, translation, output_translation, scratch_metadata) + finally: + _logging_only_scan_guardrail.reset(token) + async def apply_guardrail( self, inputs: "GenericGuardrailAPIInputs", @@ -1904,6 +1940,7 @@ class ContentFilterGuardrail(CustomGuardrail): masked_entity_count: Final[dict[str, int]] = {} status: GuardrailStatus = "success" exception_str: str = "" + would_reject_token: Final = _logging_only_scan_would_reject.set(False) try: texts: Final = inputs.get("texts", []) @@ -1947,11 +1984,13 @@ class ContentFilterGuardrail(CustomGuardrail): exception_str = str(e) raise e finally: + would_reject: Final = _logging_only_scan_would_reject.get() or any(map(self._asks_to_block, detections)) + _logging_only_scan_would_reject.reset(would_reject_token) # Log guardrail information self._log_guardrail_information( request_data=request_data, detections=detections, - status=status, + status="guardrail_flagged" if status == "success" and would_reject else status, start_time=start_time, masked_entity_count=masked_entity_count, exception_str=exception_str, diff --git a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py index 5d825099ceb..464cfee270a 100644 --- a/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py +++ b/litellm/proxy/pass_through_endpoints/llm_passthrough_endpoints.py @@ -34,6 +34,7 @@ from litellm import get_llm_provider from litellm._logging import verbose_proxy_logger from litellm.constants import ( ALLOWED_VERTEX_AI_PASSTHROUGH_HEADERS, + AWS_ERROR_TYPE_HEADER, AZURE_SPEECH_BATCH_PATH_PREFIX, AZURE_SPEECH_COGNITIVE_SERVICES_DOMAIN, AZURE_SPEECH_CUSTOM_LLM_PROVIDER, @@ -45,7 +46,8 @@ from litellm.constants import ( BEDROCK_AGENT_RUNTIME_PASS_THROUGH_ROUTES, TYPESAFE_DANGER_ZONE_HEADER_PREFIX, ) -from litellm.integrations.custom_guardrail import guardrail_request_data_with_streaming +from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException +from litellm.integrations.custom_guardrail import guardrail_request_data_with_streaming, is_guardrail_intervention from litellm.litellm_core_utils.aws_partition import get_aws_dns_suffix from litellm.llms.anthropic.common_utils import AnthropicModelInfo, merge_anthropic_beta_headers from litellm.llms.azure.passthrough.transformation import ( @@ -67,7 +69,7 @@ from litellm.llms.openai.common_utils import OpenAIError as LiteLLMOpenAIError from litellm.llms.openai.workload_identity import get_workload_identity_bearer_token_for_api_base from litellm.llms.oss_decision import OssDecisionProvider, oss_connection, validate_oss_request from litellm.llms.vertex_ai.vertex_llm_base import VertexBase -from litellm.passthrough.main import AsyncPassthroughStreamingResponse +from litellm.passthrough.main import AsyncPassthroughStreamingResponse, is_provider_error from litellm.proxy._types import * from litellm.proxy.auth.auth_checks import enforced_model_allowlists from litellm.proxy.auth.handle_jwt import JWTHandler @@ -80,7 +82,11 @@ from litellm.proxy.auth.user_api_key_auth import ( # noqa: F401 # legacy modul user_api_key_auth_websocket, user_api_key_auth_websocket_for_model, ) -from litellm.proxy.common_request_processing import model_deployments_support_stream_options, open_sse_before_first_byte +from litellm.proxy.common_request_processing import ( + ProxyBaseLLMRequestProcessing, + model_deployments_support_stream_options, + open_sse_before_first_byte, +) from litellm.proxy.common_utils.http_parsing_utils import ( # noqa: F401 # legacy module exports _read_request_body, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export _safe_get_request_headers, # pyright: ignore[reportPrivateUsage,reportUnusedImport] # backwards-compatible package export @@ -1064,6 +1070,68 @@ def extract_model_from_bedrock_endpoint(endpoint: str) -> str: _extract_model_from_bedrock_endpoint: Final = extract_model_from_bedrock_endpoint +def _guardrail_refusal_status(error: Exception) -> int | None: + """The status a guardrail blocked the request with. None for a request check that answers with the same 400, 403 + or 422, for a guardrail that failed closed because it could not get a verdict, and for a guardrail block raised + with a plain-string detail. + + Guardrail hooks stamp ``guardrail_name`` onto the dict detail of an HTTPException they raise.""" + match error: + case HTTPException(): + detail: Final[object] = error.detail + stamped: Final = isinstance(detail, dict) and "guardrail_name" in detail + return error.status_code if stamped and is_guardrail_intervention(error) else None + case GuardrailRaisedException(blocked_content=True) | BlockedPiiEntityError(): + return error.status_code + case _: + return None + + +def _bedrock_error_response(error: Exception, handled: HTTPException) -> Response | None: + """Bedrock's own error response, for an error Bedrock answered the passthrough request with an AWS error + type.""" + if not isinstance(error, httpx.HTTPStatusError) or not is_provider_error(error): + return None + upstream: Final = error.response + if AWS_ERROR_TYPE_HEADER not in upstream.headers: + return None + passed_through: Final = { + name: upstream.headers[name] for name in (AWS_ERROR_TYPE_HEADER, "content-type") if name in upstream.headers + } + return Response( + content=upstream.content, + status_code=upstream.status_code, + headers={**(handled.headers or {}), **passed_through}, + ) + + +async def bedrock_passthrough_error_response( + processor: ProxyBaseLLMRequestProcessing, + error: Exception, + user_api_key_dict: UserAPIKeyAuth, + proxy_logging_obj: ProxyLoggingType, +) -> Response: + """Run the shared error handling for a Bedrock passthrough request, then answer a guardrail block or an error + Bedrock returned the way Bedrock does, so an AWS SDK reads its code and message. Anything else is raised as the + shared handler raises it.""" + try: + await processor.handle_llm_api_exception( + e=error, + user_api_key_dict=user_api_key_dict, + proxy_logging_obj=proxy_logging_obj, + ) + except HTTPException as handled: + response: Final = _bedrock_error_response(error, handled) + if response is None: + raise + return response + except ProxyException as handled: + refusal_status: Final = _guardrail_refusal_status(error) + if refusal_status is not None and str(refusal_status) == handled.code: + raise AwsErrorTypeProxyException(handled, "ValidationException") from handled + raise + + async def handle_bedrock_passthrough_router_model( model: str, endpoint: str, @@ -1148,11 +1216,8 @@ async def handle_bedrock_passthrough_router_model( ) return result except Exception as e: - # Use common exception handling - raise await base_llm_response_processor.handle_llm_api_exception( - e=e, - user_api_key_dict=user_api_key_dict, - proxy_logging_obj=proxy_logging_obj, + return await bedrock_passthrough_error_response( + base_llm_response_processor, e, user_api_key_dict, proxy_logging_obj ) @@ -1351,10 +1416,8 @@ async def bedrock_llm_proxy_route( return result except Exception as e: - raise await base_llm_response_processor.handle_llm_api_exception( - e=e, - user_api_key_dict=user_api_key_dict, - proxy_logging_obj=proxy_logging_obj, + return await bedrock_passthrough_error_response( + base_llm_response_processor, e, user_api_key_dict, proxy_logging_obj ) diff --git a/litellm/proxy/proxy_server.py b/litellm/proxy/proxy_server.py index 4b4c6a13b1d..97d8eda1fad 100644 --- a/litellm/proxy/proxy_server.py +++ b/litellm/proxy/proxy_server.py @@ -67,6 +67,7 @@ from litellm.constants import ( AIOHTTP_NEEDS_CLEANUP_CLOSED, AIOHTTP_TTL_DNS_CACHE, AUDIO_SPEECH_CHUNK_SIZE, + AWS_ERROR_TYPE_HEADER, BACKGROUND_HEALTH_CHECK_DB_SAVE_JOB_NAME, BASE_MCP_ROUTE, DAILY_TAG_SPEND_BATCH_MULTIPLIER, @@ -97,6 +98,7 @@ from litellm.litellm_core_utils.safe_json_loads import safe_json_loads from litellm.litellm_core_utils.token_counter import offload_token_count from litellm.proxy._types import ( UI_TEAM_ID, + AwsErrorTypeProxyException, CallbackDelete, CallInfo, CommonProxyErrors, @@ -2061,6 +2063,17 @@ async def openai_exception_handler(request: Request, exc: ProxyException): ) +@app.exception_handler(AwsErrorTypeProxyException) +async def aws_error_type_exception_handler(request: Request, exc: AwsErrorTypeProxyException) -> Response: + """The ProxyException response plus the ``x-amzn-errortype`` header and top-level ``message`` an AWS SDK reads.""" + response: Final = await openai_exception_handler(request, exc.error) + return JSONResponse( + status_code=response.status_code, + content={"message": exc.error.message, **JSON_OBJECT.validate_json(bytes(response.body))}, + headers={**exc.error.headers, AWS_ERROR_TYPE_HEADER: exc.aws_error_type}, + ) + + @app.exception_handler(StarletteHTTPException) async def otlp_http_exception_handler(connection: Request | WebSocket, exc: StarletteHTTPException) -> Response | None: if isinstance(connection, WebSocket): diff --git a/litellm/proxy/response_api_endpoints/endpoints.py b/litellm/proxy/response_api_endpoints/endpoints.py index cdaf54e10ce..d9f5dd9c39d 100644 --- a/litellm/proxy/response_api_endpoints/endpoints.py +++ b/litellm/proxy/response_api_endpoints/endpoints.py @@ -1908,7 +1908,7 @@ async def _responses_websocket_session( ) await websocket.close(code=1011, reason="Internal server error") finally: - await proxy_logging_obj.arelease_max_parallel_requests_on_disconnect(user_api_key_dict) + await proxy_logging_obj.arelease_responses_websocket_slot(user_api_key_dict) @router.websocket("/v1/responses") diff --git a/litellm/proxy/utils.py b/litellm/proxy/utils.py index 5d2ebe73cac..be6532ba36b 100644 --- a/litellm/proxy/utils.py +++ b/litellm/proxy/utils.py @@ -4405,12 +4405,20 @@ class ProxyLogging: double-decrement under the limiter's in-memory fallback. """ limiter: Final = self.get_proxy_hook("parallel_request_limiter") - if not isinstance(limiter, (PROXY_MaxParallelRequestsHandler_v3, _PROXY_MaxParallelRequestsHandler)): + if not isinstance(limiter, PROXY_MaxParallelRequestsHandler_v3): return await limiter.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) _arelease_max_parallel_requests_on_disconnect = arelease_max_parallel_requests_on_disconnect + async def arelease_responses_websocket_slot(self, user_api_key_dict: UserAPIKeyAuth) -> None: + """Release the slot a Responses WebSocket took at admission. Only call this once that admission succeeded: + the legacy limiter keeps no record of its acquisitions and decrements unconditionally.""" + limiter: Final = self.get_proxy_hook("parallel_request_limiter") + if not isinstance(limiter, (PROXY_MaxParallelRequestsHandler_v3, _PROXY_MaxParallelRequestsHandler)): + return + await limiter.async_release_max_parallel_requests_on_disconnect(user_api_key_dict) + async def enforce_mcp_server_rate_limits( self, user_api_key_dict: UserAPIKeyAuth | None, diff --git a/litellm/responses/streaming_iterator.py b/litellm/responses/streaming_iterator.py index 33adb663021..fd9a7889487 100644 --- a/litellm/responses/streaming_iterator.py +++ b/litellm/responses/streaming_iterator.py @@ -31,6 +31,7 @@ from litellm.litellm_core_utils.core_helpers import process_response_headers from litellm.litellm_core_utils.internal_call_metadata import without_budget_reservation from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.llm_response_utils.get_api_base import get_api_base +from litellm.litellm_core_utils.llm_response_utils.get_headers import provider_request_id_headers from litellm.litellm_core_utils.llm_response_utils.response_metadata import ( update_response_metadata, ) @@ -262,7 +263,9 @@ def stream_error_status_and_message(error_obj: object) -> tuple[int, str]: return _status_code_for_error_fields(error_type, error_code), message -def _map_stream_error_to_exception(error_obj: object, model: str, custom_llm_provider: str) -> Exception: +def _map_stream_error_to_exception( + error_obj: object, model: str, custom_llm_provider: str, headers: httpx.Headers | None = None +) -> Exception: from litellm.llms.base_llm.chat.transformation import BaseLLMException error_message, error_type, error_code = _error_event_fields(error_obj) @@ -272,6 +275,7 @@ def _map_stream_error_to_exception(error_obj: object, model: str, custom_llm_pro status_code=status_code, message=f"Error code: {status_code} - {{'error': {error_body}}}", body=error_body, + headers=headers, ) try: return litellm.exception_type( @@ -667,7 +671,12 @@ class BaseResponsesAPIStreamingIterator: ) def _map_error_event_exception(self, error_obj: object) -> Exception: - return _map_stream_error_to_exception(error_obj, self.model or "", self.custom_llm_provider or "") + return _map_stream_error_to_exception( + error_obj, + self.model or "", + self.custom_llm_provider or "", + headers=provider_request_id_headers(self._raw_response_headers), + ) def _maybe_raise_for_error_event(self, result: object) -> None: chunk_type: Final = getattr(result, "type", None) @@ -683,11 +692,17 @@ class BaseResponsesAPIStreamingIterator: mapped_exception: Final = self._map_error_event_exception(error_obj) if not _mid_stream_fallback_eligible(mapped_exception): raise mapped_exception + mapped_status: Final = getattr(mapped_exception, "status_code", None) raise MidStreamFallbackError( message=str(mapped_exception), model=self.model or "", llm_provider=self.custom_llm_provider or "", original_exception=mapped_exception, + response=httpx.Response( + status_code=int(mapped_status) if mapped_status is not None else 503, + headers=provider_request_id_headers(self._raw_response_headers), + request=httpx.Request(method="POST", url=f"https://{self.custom_llm_provider or ''}.com/v1/"), + ), generated_content=self._generated_content, is_pre_first_chunk=not self._yielded_first_chunk, ) diff --git a/litellm/router.py b/litellm/router.py index 94debd5b03d..682fb71421d 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -128,6 +128,7 @@ from litellm.llms.base_llm.vector_store.transformation import ( RouterVectorStoreEmbeddingExecutor, vector_store_request_metadata, ) +from litellm.llms.bedrock.common_utils import get_bedrock_base_model from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler, get_async_httpx_client from litellm.llms.openai_like.json_loader import JSONProviderRegistry from litellm.llms.openai_like.model_info import ( @@ -374,6 +375,7 @@ from litellm.utils import ( get_non_default_completion_params, get_secret, get_utc_datetime, + is_generalized_model_info, is_region_allowed, provider_rejectable_params, set_live_deployment_replay, @@ -991,6 +993,49 @@ def as_output_cap(value: object) -> int | None: return cap if cap >= 0 else None +def _exact_base_model_max_input_tokens(base_model: str) -> int | None: + try: + model_info: Final = litellm.get_model_info(model=base_model) + except Exception: # noqa: BLE001 # get_model_info raises ModelNotMappedError or a bare Exception for an unmapped model + return None + if is_generalized_model_info(model_info): + return None + return model_info.get("max_input_tokens") + + +def _provider_model_id(model_name: str, custom_llm_provider: str) -> str: + model_id: Final = model_name.rsplit("/", 1)[-1] + return get_bedrock_base_model(model_id) if custom_llm_provider.startswith("bedrock") else model_id + + +def _names_route_model(base_model: str, route_key: str, custom_llm_provider: str) -> bool: + route_model_id: Final = _provider_model_id(route_key, custom_llm_provider) + base_model_id: Final = _provider_model_id(base_model, custom_llm_provider) + return route_model_id in (base_model_id, base_model_id.split("@", 1)[0]) + + +def _declared_base_model(user_model_info: Mapping[str, object]) -> str | None: + match user_model_info: + case {"base_model": str() as base_model}: + return base_model + case _: + return None + + +def _base_model_widened_window( + route_model_info: ModelMapInfo, base_model: str | None, custom_llm_provider: str +) -> Mapping[str, int]: + route_max_input_tokens: Final = route_model_info.get("max_input_tokens") + if base_model is None or route_max_input_tokens is None: + return MappingProxyType({}) + if not _names_route_model(base_model, route_model_info["key"], custom_llm_provider): + return MappingProxyType({}) + base_max_input_tokens: Final = _exact_base_model_max_input_tokens(base_model) + if base_max_input_tokens is None or base_max_input_tokens <= route_max_input_tokens: + return MappingProxyType({}) + return MappingProxyType({"max_input_tokens": base_max_input_tokens}) + + class Router: @property def _routing_groups(self) -> dict[str, RoutingGroup]: @@ -11371,6 +11416,7 @@ class Router: Augment litellm info with additional params set in `model_info`. For azure models, ignore the `model:`. Only set max tokens, cost values if base_model is set. + For every other provider, a base_model naming the same model as `model:` can only widen its context window. Returns - ModelInfo - If found -> typed dict with max tokens, input cost, etc. @@ -11467,20 +11513,23 @@ class Router: ## CHECK USER SET MODEL INFO raw_user_model_info: Final = deployment.get("model_info") - user_model_info: Final = ( + user_model_info: Final = _MODEL_INFO_ADAPTER.validate_python( raw_user_model_info.model_dump(exclude_none=True) if isinstance(raw_user_model_info, BaseModel) - else raw_user_model_info + else raw_user_model_info or {} ) # get_model_info() hands back an lru_cache'd dict, so merge into a copy; unset # values are skipped or Deployment's None pricing defaults would erase the map's merged_model_info: Final[ModelMapInfo] = { **copy.deepcopy(model_info), - **self.get_discovered_model_info((deployment.get("model_info") or {}).get("id")), - **MappingProxyType( - {key: value for key, value in (user_model_info or MappingProxyType({})).items() if value is not None} + **_base_model_widened_window( + model_info, + _declared_base_model(user_model_info) or _declared_base_model(litellm_params.model_extra or {}), + custom_llm_provider, ), + **self.get_discovered_model_info((deployment.get("model_info") or {}).get("id")), + **MappingProxyType({key: value for key, value in user_model_info.items() if value is not None}), } return merged_model_info diff --git a/litellm/types/llms/openai.py b/litellm/types/llms/openai.py index 8731bc2905a..8cc9e912e80 100644 --- a/litellm/types/llms/openai.py +++ b/litellm/types/llms/openai.py @@ -1421,10 +1421,26 @@ class WebSearchToolUsage(LiteLLMBaseModel): num_requests: NonNegativeInt +class ImageGenTokenDetails(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + image_tokens: NonNegativeInt = 0 + text_tokens: NonNegativeInt = 0 + + +class ImageGenToolUsage(LiteLLMBaseModel): + model_config = ConfigDict(frozen=True) + + total_tokens: NonNegativeInt = 0 + input_tokens_details: ImageGenTokenDetails | None = None + output_tokens_details: ImageGenTokenDetails | None = None + + class ResponsesToolUsage(LiteLLMBaseModel): model_config = ConfigDict(frozen=True) web_search: WebSearchToolUsage | None = None + image_gen: ImageGenToolUsage | None = None ResponsesAPIStatus = Literal["completed", "failed", "in_progress", "cancelled", "queued", "incomplete"] diff --git a/model_prices_and_context_window.json b/model_prices_and_context_window.json index a204c49967e..b744032fb6c 100644 --- a/model_prices_and_context_window.json +++ b/model_prices_and_context_window.json @@ -29256,22 +29256,32 @@ "output_cost_per_token_priority": 1.8e-05 }, "gemini-3.1-pro-preview": { - "prompt_cache_min_tokens": 4096, + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_flex": 1e-06, + "input_cost_per_token_priority": 3.6e-06, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", + "output_cost_per_image": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, - "output_cost_per_image": 0.00012, + "output_cost_per_token_flex": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "prompt_cache_min_tokens": 4096.0, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -29289,33 +29299,18 @@ ], "supports_audio_input": true, "supports_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "supports_url_context": true, - "supports_native_streaming": true, - "input_cost_per_token_priority": 3.6e-06, - "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, - "output_cost_per_token_priority": 2.16e-05, - "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, - "cache_read_input_token_cost_priority": 3.6e-07, - "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, - "cache_read_input_token_cost_flex": 2e-07, - "search_context_cost_per_query": { - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014, - "search_context_size_high": 0.014 - }, - "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_token_flex": 1e-06, - "output_cost_per_token_flex": 6e-06 + "web_search_billing_unit": "per_query" }, "gemini-3.1-pro-preview-customtools": { "prompt_cache_min_tokens": 4096, @@ -29428,15 +29423,25 @@ "vertex_ai/gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, "cache_read_input_token_cost": 5e-08, - "cache_read_input_token_cost_flex": 5e-08, - "input_cost_per_token": 5e-07, + "cache_read_input_token_cost_priority": 9e-08, + "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token_batches": 5e-07, + "input_cost_per_audio_token_priority": 1.8e-06, + "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, + "input_cost_per_token_flex": 2.5e-07, + "input_cost_per_token_priority": 9e-07, "litellm_provider": "vertex_ai", "max_input_tokens": 1048576, "max_output_tokens": 65535, "max_tokens": 65535, "mode": "chat", "output_cost_per_token": 3e-06, + "output_cost_per_token_batches": 1.5e-06, + "output_cost_per_token_flex": 1.5e-06, + "output_cost_per_token_priority": 5.4e-06, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -29454,6 +29459,7 @@ ], "supports_audio_input": true, "supports_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, @@ -29463,23 +29469,7 @@ "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "supports_native_streaming": true, - "input_cost_per_token_priority": 9e-07, - "input_cost_per_audio_token_priority": 1.8e-06, - "output_cost_per_token_priority": 5.4e-06, - "cache_read_input_token_cost_priority": 9e-08, - "search_context_cost_per_query": { - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014, - "search_context_size_high": 0.014 - }, - "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_audio_token_batches": 5e-07, - "input_cost_per_token_batches": 2.5e-07, - "input_cost_per_token_flex": 2.5e-07, - "output_cost_per_token_batches": 1.5e-06, - "output_cost_per_token_flex": 1.5e-06 + "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.5-flash": { "prompt_cache_min_tokens": 4096, @@ -29807,22 +29797,32 @@ "output_cost_per_video_token": 1e-06 }, "vertex_ai/gemini-3.1-pro-preview": { - "prompt_cache_min_tokens": 4096, + "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, "cache_read_input_token_cost": 2e-07, "cache_read_input_token_cost_above_200k_tokens": 4e-07, - "cache_creation_input_token_cost_above_200k_tokens": 2.5e-07, + "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, + "cache_read_input_token_cost_priority": 3.6e-07, + "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_token": 2e-06, "input_cost_per_token_above_200k_tokens": 4e-06, + "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, "input_cost_per_token_batches": 1e-06, + "input_cost_per_token_flex": 1e-06, + "input_cost_per_token_priority": 3.6e-06, "litellm_provider": "vertex_ai", "max_input_tokens": 1048576, "max_output_tokens": 65536, "max_tokens": 65536, "mode": "chat", + "output_cost_per_image": 0.00012, "output_cost_per_token": 1.2e-05, "output_cost_per_token_above_200k_tokens": 1.8e-05, + "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, "output_cost_per_token_batches": 6e-06, - "output_cost_per_image": 0.00012, + "output_cost_per_token_flex": 6e-06, + "output_cost_per_token_priority": 2.16e-05, + "prompt_cache_min_tokens": 4096.0, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -29840,33 +29840,18 @@ ], "supports_audio_input": true, "supports_function_calling": true, + "supports_native_streaming": true, "supports_pdf_input": true, "supports_prompt_caching": true, "supports_reasoning": true, "supports_response_schema": true, "supports_system_messages": true, "supports_tool_choice": true, + "supports_url_context": true, "supports_video_input": true, "supports_vision": true, "supports_web_search": true, - "supports_url_context": true, - "supports_native_streaming": true, - "input_cost_per_token_priority": 3.6e-06, - "input_cost_per_token_above_200k_tokens_priority": 7.2e-06, - "output_cost_per_token_priority": 2.16e-05, - "output_cost_per_token_above_200k_tokens_priority": 3.24e-05, - "cache_read_input_token_cost_priority": 3.6e-07, - "cache_read_input_token_cost_above_200k_tokens_priority": 7.2e-07, - "cache_read_input_token_cost_flex": 2e-07, - "search_context_cost_per_query": { - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014, - "search_context_size_high": 0.014 - }, - "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_token_flex": 1e-06, - "output_cost_per_token_flex": 6e-06 + "web_search_billing_unit": "per_query" }, "vertex_ai/gemini-3.1-pro-preview-customtools": { "prompt_cache_min_tokens": 4096, @@ -31666,9 +31651,15 @@ "gemini-3-flash-preview": { "cache_read_input_audio_token_cost": 1e-07, "cache_read_input_token_cost": 5e-08, - "cache_read_input_token_cost_flex": 5e-08, + "cache_read_input_token_cost_priority": 9e-08, + "google_maps_grounding_cost_per_query": 0.014, "input_cost_per_audio_token": 1e-06, + "input_cost_per_audio_token_batches": 5e-07, + "input_cost_per_audio_token_priority": 1.8e-06, "input_cost_per_token": 5e-07, + "input_cost_per_token_batches": 2.5e-07, + "input_cost_per_token_flex": 2.5e-07, + "input_cost_per_token_priority": 9e-07, "litellm_provider": "vertex_ai-language-models", "max_input_tokens": 1048576, "max_output_tokens": 65535, @@ -31676,6 +31667,10 @@ "mode": "chat", "output_cost_per_reasoning_token": 3e-06, "output_cost_per_token": 3e-06, + "output_cost_per_token_batches": 1.5e-06, + "output_cost_per_token_flex": 1.5e-06, + "output_cost_per_token_priority": 5.4e-06, + "search_context_cost_per_query": {"search_context_size_low":0.014,"search_context_size_medium":0.014,"search_context_size_high":0.014}, "source": "https://cloud.google.com/gemini-enterprise-agent-platform/generative-ai/pricing", "supported_endpoints": [ "/v1/chat/completions", @@ -31693,6 +31688,7 @@ ], "supports_audio_output": false, "supports_function_calling": true, + "supports_native_streaming": true, "supports_parallel_function_calling": true, "supports_pdf_input": true, "supports_prompt_caching": true, @@ -31703,23 +31699,7 @@ "supports_url_context": true, "supports_vision": true, "supports_web_search": true, - "supports_native_streaming": true, - "input_cost_per_token_priority": 9e-07, - "input_cost_per_audio_token_priority": 1.8e-06, - "output_cost_per_token_priority": 5.4e-06, - "cache_read_input_token_cost_priority": 9e-08, - "search_context_cost_per_query": { - "search_context_size_low": 0.014, - "search_context_size_medium": 0.014, - "search_context_size_high": 0.014 - }, - "web_search_billing_unit": "per_query", - "google_maps_grounding_cost_per_query": 0.014, - "input_cost_per_audio_token_batches": 5e-07, - "input_cost_per_token_batches": 2.5e-07, - "input_cost_per_token_flex": 2.5e-07, - "output_cost_per_token_batches": 1.5e-06, - "output_cost_per_token_flex": 1.5e-06 + "web_search_billing_unit": "per_query" }, "gemini-omni-flash-preview": { "input_cost_per_audio_token": 1.5e-06, diff --git a/tests/benchmarks/test_prompt_template_benchmarks.py b/tests/benchmarks/test_prompt_template_benchmarks.py new file mode 100644 index 00000000000..c7858771d35 --- /dev/null +++ b/tests/benchmarks/test_prompt_template_benchmarks.py @@ -0,0 +1,21 @@ +from collections.abc import Mapping +from typing import Final + +import pytest + +from litellm.litellm_core_utils.prompt_templates.factory import anthropic_messages_pt + +LONG_RUN_OF_EMPTY_CACHED_USER_MESSAGES: Final[tuple[Mapping[str, object], ...]] = ( + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + *({"role": "user", "content": [], "cache_control": {"type": "ephemeral"}} for _ in range(1_000)), +) + + +@pytest.mark.benchmark +def test_anthropic_messages_pt_long_run_of_empty_cached_user_messages() -> None: + result: Final = anthropic_messages_pt( + messages=list(LONG_RUN_OF_EMPTY_CACHED_USER_MESSAGES), model="claude-opus-5", llm_provider="vertex_ai" + ) + + assert len(result[-1]["content"]) == 1_000 diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 17a10e98092..c1e1ebb87b9 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -114,6 +114,11 @@ ignored_function_names = [ "_deployment_that_just_failed", # Tested through the same-boundary hop test in test_router_order_fallback.py "_collect_ptu_share_teams", # Cached behind team_ptu_ceiling, tested through the live model list test in test_router.py "_resolve_team_ptu_ceiling", # Cached behind team_ptu_ceiling, tested through the live model list test in test_router.py + "_exact_base_model_max_input_tokens", # Tested through get_router_model_info in test_router_model_cost_isolation.py + "_base_model_widened_window", # Tested through get_router_model_info in test_router_model_cost_isolation.py + "_provider_model_id", # Tested through get_router_model_info in test_router_model_cost_isolation.py + "_names_route_model", # Tested through get_router_model_info in test_router_model_cost_isolation.py + "_declared_base_model", # Tested through get_router_model_info in test_router_model_cost_isolation.py "async_admit_deployment_turn", # Bound as the WebSocket turn_admission partial, tested through it in test_router_helper_utils.py ] diff --git a/tests/e2e/AGENTS.md b/tests/e2e/AGENTS.md index e7a26c12db8..d267059cdc7 100644 --- a/tests/e2e/AGENTS.md +++ b/tests/e2e/AGENTS.md @@ -44,7 +44,7 @@ Each subdirectory under `tests/e2e/` is one suite, scoped to an endpoint family - `logging/` - logging-integration delivery (datadog and friends) - `security/` - secret handling and log-leak protection - `router/` - routing and reliability behavior (fallbacks, cooldowns) plus the memory tests (`test_reliability_memory_e2e.py`: every worker's RSS as read at collection time, before any test traffic, must sit under a fixed idle budget, the release-gate check for a DB-backed boot that idles near the pod limit the way v1.100.x did, and a few hundred failing requests with retries and fallbacks must not grow proxy RSS past a fixed budget nor store a request snapshot past a fixed size, the release-gate check for the v1.100.0 retry-breadcrumb leak) -- `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in. The perf cases (`test_perf_cases.py`, one `PerformanceTestCase` per module in `perf/cases/`, no `_e2e` suffix) are the request-performance runs. A case names what the gateway was deployed with (model groups and guardrails, written out in the gateway's config: the `proxy_config` block of project-releaser's perf helm values in the lane, `gateway/perf_ci_config.yml` locally, with each deployment's mock knobs in an `x-mock-knobs` extra header so the model ids are real ones like `gpt-6-luna` and `anthropic.claude-sonnet-5`) and lists what the test creates and sends: teams, Keycloak JWT users, virtual keys, `RequestShape`s (a coding-agent session from `perf/setup/corpus/` whose input-token size is drawn per request from the shape's range, with or without tools, and the output cap), a `Traffic` plan (seed, rps, request count, endpoint, shape and stream shares) and `Limit`s (one metric at one percentile under one ceiling). The infra side lives in `perf/setup/`: `provision.py` checks the deployed config matches the case and builds the callers, `plan.py` draws the request list from the seed, `run.py` replays it from a separate Locust process (`locustfile.py`) that builds each body right before sending and records status, header, first-byte, first-token and total time per request, and `observe.py` joins all of that with what the mocks received and the workers' RSS to score every `Limit`. Limits are about what the caller sees and what the gateway costs: every non-blocked request is a 200 (`unexpected_status_count`), the seconds the gateway added per request at p50 and p99 (`gateway_added_seconds`, to the first token on streams), the client's time to the first token of a stream at p50 and p99 (`time_to_first_token_seconds`, mock delay included), worker memory growth (`memory_growth_megabytes`), and the Postgres transactions and Redis commands each request cost (`postgres_calls_per_request`, `redis_calls_per_request`, from `pg_stat_database` and Redis `INFO` read by `datastore.py` before and after the run). Whether a body, a guardrail call or a spend row is right is printed as a diagnostic, never a limit, that is the functional suites' job. The providers and the guardrail are out-of-process servers in `perf/setup/mocks/` (`server.py openai|anthropic|bedrock|guardrail`, each answering with a reply generated for the request in that provider's wire format, sized from the prompt and output cap, a tool call when tools were offered, waiting the case's `mock_latency` first, request logs at `/mock/requests` and `/mock/calls`), reached through `E2E_PERF_MOCK_*_URL` by both the gateway config and the test. Adding a perf test is adding a module with a `CASE` to `perf/cases/`; a new deployment or guardrail is a config change; a new metric is a name in `Metric` plus a branch in `observe.observe` +- `perf/` - performance-category tests, kept OUT of the main suite: throughput/load SLO tests are a different testing category from functional e2e (variance-driven, historically flaky) and live outside this suite until re-implemented as their own pipeline (LIT-5163). Do not add a live load test that runs in the default collection. What lives here: the weekly session-anomaly test (`test_weekly_session_anomaly_e2e.py`, Claude Code-shaped multi-turn sessions against real providers with ceilings on error rate, cache read/write, turn time, and spend, marked `weekly` and deselected unless `E2E_WEEKLY_ANOMALY` is set, driven by `.github/workflows/weekly_load_anomaly.yml`), the Redis chaos test (`test_redis_chaos_e2e.py`, locust load against mock deployments split round robin over `/chat/completions` and `/v1/messages`, one endpoint per simulated user, with `CLIENT PAUSE ALL` on the proxy's Redis mid-run to simulate it being down outright, asserting zero failed requests on every endpoint, budgeting RSS and CPU-per-request as ratios against the same run's healthy phase, and holding p50/p90/p99 latency and log-bytes-per-request to flat ceilings (a ratio cannot bound those two: an open breaker skips Redis instead of waiting on it, so the chaos phase can measure cheaper than baseline while still being far slower than a user should see), needs a proxy booted from `gateway/redis_chaos_ci_config.yml` on the same host with `E2E_PROXY_PID` and `E2E_PROXY_LOG` set, marked `redis_chaos`, deselected unless `E2E_PERF_TESTS` is set and excluded from the per-PR selector like the rest of `perf/`, driven by `.github/workflows/test-e2e-redis-chaos.yml` and by the Buildkite `e2e-redis-chaos` step in project-releaser, which runs the proxy, Postgres and Valkey co-located with pytest in one pod and sets the opt-in). New perf tests go here too, marked `perf`, behind the same `E2E_PERF_TESTS` opt-in. The perf cases (`test_perf_cases.py`, one `PerformanceTestCase` per module in `perf/cases/`, no `_e2e` suffix) are the request-performance runs. A case names what the gateway was deployed with (model groups and guardrails, written out in the gateway's config: the `proxy_config` block of project-releaser's perf helm values in the lane, `gateway/perf_ci_config.yml` locally, with each deployment's mock knobs in an `x-mock-knobs` extra header so the model ids are real ones like `gpt-6-luna` and `anthropic.claude-sonnet-5`) and lists what the test creates and sends: teams, Keycloak JWT users, virtual keys, `RequestShape`s (a coding-agent session from `perf/setup/corpus/` whose input-token size is drawn per request from the shape's range, with or without tools, and the output cap), a `Traffic` plan (seed, rps, request count, endpoint, shape and stream shares) and `Limit`s (one metric at one percentile under one ceiling). The infra side lives in `perf/setup/`: `provision.py` checks the deployed config matches the case and builds the callers, `plan.py` draws the request list from the seed, `run.py` replays it from a separate Locust process (`locustfile.py`) that builds each body right before sending and records status, header, first-byte, first-token and total time per request, and `observe.py` joins all of that with what the mocks received and the workers' RSS to score every `Limit`. Limits are about what the caller sees and what the gateway costs: every non-blocked request is a 200 (`unexpected_status_count`), the seconds the gateway added per request at p50 and p99 (`gateway_added_seconds`, to the first token on streams), the client's time to the first token of a stream at p50 and p99 (`time_to_first_token_seconds`, mock delay included), worker memory growth (`memory_growth_megabytes`), and the Postgres statements and Redis commands a request waits on before it is forwarded (`postgres_calls_per_request`, `redis_calls_per_request`: auth and routing reads, the rate-limiter script), with the work the gateway does after the response or on a timer (spend writes, daily aggregates, config reloads) reported separately as `postgres_background_calls_per_request` and `redis_background_calls_per_request`; `datastore.py` reads `pg_stat_statements` before and after the run and holds a Redis `MONITOR` connection during it, sorts each statement by verb and table and each command by key, and the run output lists the top statements and commands per bucket so a red run names what grew. Whether a body, a guardrail call or a spend row is right is printed as a diagnostic, never a limit, that is the functional suites' job. The providers and the guardrail are out-of-process servers in `perf/setup/mocks/` (`server.py openai|anthropic|bedrock|guardrail`, each answering with a reply generated for the request in that provider's wire format, sized from the prompt and output cap, a tool call when tools were offered, waiting the case's `mock_latency` first, request logs at `/mock/requests` and `/mock/calls`), reached through `E2E_PERF_MOCK_*_URL` by both the gateway config and the test. Adding a perf test is adding a module with a `CASE` to `perf/cases/`; a new deployment or guardrail is a config change; a new metric is a name in `Metric` plus a branch in `observe.observe` - `other/` - the holding-pen suite for the `other.*` registry cluster with no home of its own yet: the master-key auth gate, JWT auth (access tokens issued by a real Keycloak realm, `idp.py` plus `idp_realm.json`, whose JWKS the proxy's `JWT_PUBLIC_KEY_URL` points at, see CONTRIBUTING.md for the start command and config block), and the process-lifecycle health probes (liveness, public readiness, authenticated readiness diagnostics). Promote a cluster out once it is large/stable enough for its own suite - `secret_manager/` - the gateway's `key_management_system` against a real secret manager: deployment keys resolved from it (`os.environ/` where the name exists only in the manager) and virtual keys written to and deleted from it. The tests are backend-agnostic and each backend is its own lane, because the setting is global to the proxy: `E2E_SECRET_MANAGER=` opts in and picks the backend from `secret_backends.BACKENDS`, the proxy is booted from `gateway/secret_manager__ci_config.yml` against the live manager, and the tests reach that manager through the backend's `SecretStore` (`secret_store_.py`). A test needing something not every backend does carries `requires_capability(...)` and is deselected on lanes that lack it. `secret_manager/backend.sh up ` runs a backend in Docker and writes the proxy's and the tests' env. Marked `secret_manager`, deselected unless `E2E_SECRET_MANAGER` is set, and kept out of the per-PR selector. Backends today: `hashicorp_vault` and `cyberark` (CyberArk Conjur, which cannot delete, so the delete test is Vault-only) - `gateway/` - proxy configuration only (`litellm-config.yml`). No tests diff --git a/tests/e2e/perf/cases/mixed_size_requests_2rps.py b/tests/e2e/perf/cases/mixed_size_requests_2rps.py new file mode 100644 index 00000000000..41b61cdc987 --- /dev/null +++ b/tests/e2e/perf/cases/mixed_size_requests_2rps.py @@ -0,0 +1,74 @@ +"""Mixed-size requests at 2 rps on the small gateway: three real model groups over the mock providers, +two teams plus users in no team, JWT and key callers, a pre and +a post guardrail, half the requests streaming, short chats, long contexts and tool calls. Budgets +are real money at cost-map prices (a 335k-token `claude-sonnet-5` request is about $0.70 and the +run spends roughly $600), so they are sized so no caller runs dry mid-run.""" + +from __future__ import annotations + +from typing import Final + +from setup.performance_test_case import ( + LONG, + SHORT, + Guardrail, + Limit, + MockLatency, + PerformanceTestCase, + Team, + Traffic, + User, + VirtualKey, +) + +CASE: Final = PerformanceTestCase( + name="mixed_size_requests_2rps", + teams=( + Team("team-a", models=("gpt-6-luna", "claude-sonnet-5"), max_budget=5000.0, rpm_limit=6000), + Team("team-b", models=("claude-sonnet-5-bedrock", "claude-sonnet-5"), max_budget=5000.0), + ), + users=( + User("alice", team="team-a"), + User("bob", team="team-a"), + User("carol", team="team-a"), + User("dave", team="team-b"), + User("erin", team="team-b"), + User("frank", team="team-b"), + User("grace"), + User("heidi"), + User("ivan"), + User("judy", max_budget=1000.0), + ), + virtual_keys=( + VirtualKey("key-a1", team="team-a", models=("gpt-6-luna",), rpm_limit=1200), + VirtualKey("key-a2", team="team-a", user="alice"), + VirtualKey("key-b1", team="team-b"), + VirtualKey("key-b2", team="team-b", user="dave"), + VirtualKey("key-plain"), + ), + model_groups=("gpt-6-luna", "claude-sonnet-5", "claude-sonnet-5-bedrock"), + mock_latency=MockLatency(seconds=1.0, jitter_seconds=0.25), + guardrails=(Guardrail("perf-pre", mode="pre_call"), Guardrail("perf-post", mode="post_call")), + traffic=Traffic( + seed=7, + rps=2.0, + total_requests=1200, + endpoints={"chat_completions": 0.5, "responses": 0.25, "messages": 0.25}, + shapes={SHORT: 0.75, LONG: 0.25}, + stream_share=0.5, + intervene_share=0.0, + block_share=0.0, + concurrency=100, + ), + limits=( + Limit("unexpected_status_count", max=3), + Limit("gateway_added_seconds", percentile=50, max=0.58), + Limit("gateway_added_seconds", percentile=99, max=6.1), + Limit("time_to_first_token_seconds", percentile=50, max=2.7), + Limit("time_to_first_token_seconds", percentile=99, max=8.8), + Limit("memory_growth_megabytes", max=600), + Limit("cpu_seconds_per_request", max=1.26), + Limit("postgres_calls_per_request", max=3.5), + Limit("redis_calls_per_request", max=21.0), + ), +) diff --git a/tests/e2e/perf/cases/small_bodies_5rps.py b/tests/e2e/perf/cases/small_bodies_5rps.py index 52fe41623e3..00a9fdefdb5 100644 --- a/tests/e2e/perf/cases/small_bodies_5rps.py +++ b/tests/e2e/perf/cases/small_bodies_5rps.py @@ -64,11 +64,11 @@ CASE: Final = PerformanceTestCase( ), limits=( Limit("unexpected_status_count", max=3), - Limit("gateway_added_seconds", percentile=50, max=0.38), - Limit("time_to_first_token_seconds", percentile=99, max=3.0), + Limit("gateway_added_seconds", percentile=50, max=0.29), + Limit("time_to_first_token_seconds", percentile=99, max=3.7), Limit("memory_growth_megabytes", max=100), - Limit("postgres_calls_per_request", max=13), - Limit("redis_calls_per_request", max=35), - Limit("cpu_seconds_per_request", max=0.19), + Limit("postgres_calls_per_request", max=1.8), + Limit("redis_calls_per_request", max=18), + Limit("cpu_seconds_per_request", max=0.29), ), ) diff --git a/tests/e2e/perf/setup/datastore.py b/tests/e2e/perf/setup/datastore.py index e123d129cc5..b4ad10b5a5a 100644 --- a/tests/e2e/perf/setup/datastore.py +++ b/tests/e2e/perf/setup/datastore.py @@ -1,26 +1,90 @@ -"""Postgres and Redis work counters, read before and after the traffic phase so a case can put a ceiling -on how much database and cache work the gateway does per request. Both come from the servers' own -statistics (`pg_stat_database`, Redis `INFO`), so nothing here depends on how the gateway talks to them. -`DATABASE_URL` and `REDIS_HOST` are the runner's own env, as in the other suites; without them the -counters are None and any Limit on them reports as missing.""" +"""Postgres and Redis work the gateway did during the traffic phase, split into what a request waits on and +what runs after the response. Both come from the servers' own accounting, so nothing here depends on how the +gateway talks to them: Postgres statement counts from `pg_stat_statements` (read before and after the run) and +Redis commands from a `MONITOR` connection held open for the run. `DATABASE_URL` and `REDIS_HOST` are the +runner's own env, as in the other suites; without them the counters are None and any Limit on them reports as +missing. `E2E_DATABASE_ADMIN_URL`, when set, is used once to create the `pg_stat_statements` extension. Only the +runtime user's own statements are read: other users' rows come back with a NULL queryid and hidden text.""" from __future__ import annotations import os +import re +import threading +from collections import Counter +from collections.abc import Generator +from contextlib import contextmanager from dataclasses import dataclass -from typing import Final +from typing import TYPE_CHECKING, Final, Literal from pydantic import TypeAdapter -_DB_QUERY: Final = "SELECT xact_commit FROM pg_stat_database WHERE datname = current_database()" -_DB_ROW: Final = TypeAdapter(tuple[int]) +if TYPE_CHECKING: + import redis + +Bucket = Literal["awaited", "background", "plumbing", "unclassified"] + +_COMMITS_QUERY: Final = "SELECT xact_commit FROM pg_stat_database WHERE datname = current_database()" +_STATEMENTS_QUERY: Final = ( + "SELECT queryid, calls, query FROM pg_stat_statements " + "WHERE dbid = (SELECT oid FROM pg_database WHERE datname = current_database()) " + "AND userid = (SELECT usesysid FROM pg_user WHERE usename = current_user) AND queryid IS NOT NULL" +) +_CREATE_EXTENSION: Final = "CREATE EXTENSION IF NOT EXISTS pg_stat_statements" +_COMMITS_ROW: Final = TypeAdapter(tuple[int]) +_STATEMENT_ROWS: Final = TypeAdapter(list[tuple[int, int, str]]) _INFO: Final = TypeAdapter(dict[str, object]) _INT: Final = TypeAdapter(int) +_MONITOR_LINE: Final = TypeAdapter(dict[str, object]) + +_TABLE: Final = re.compile(r'"(LiteLLM_\w+)"') +_AWAITED_TABLES: Final = frozenset( + { + "LiteLLM_VerificationToken", + "LiteLLM_UserTable", + "LiteLLM_TeamTable", + "LiteLLM_TeamMembership", + "LiteLLM_OrganizationTable", + "LiteLLM_OrganizationMembership", + "LiteLLM_EndUserTable", + "LiteLLM_BudgetTable", + "LiteLLM_ObjectPermissionTable", + } +) +_PLUMBING_VERBS: Final = frozenset({"BEGIN", "COMMIT", "ROLLBACK", "DEALLOCATE", "SET", "SHOW"}) +_REDIS_READS: Final = frozenset({"GET", "MGET", "HGET", "HGETALL", "EXISTS", "TTL", "PTTL"}) +_REDIS_NAMESPACE: Final = "litellm.caching:" +_REDIS_WRITES: Final = frozenset( + {"SET", "SETEX", "DEL", "EXPIRE", "INCRBYFLOAT", "INCR", "RPUSH", "LPUSH", "LPOP", "RPOP", "LTRIM", "HSET", "SADD"} +) + + +@dataclass(frozen=True, slots=True) +class StatementCalls: + queryid: int + query: str + calls: int + + @property + def bucket(self) -> Bucket: + """A SELECT on the key, user, team, org, end-user or budget tables is what a request waits on at auth + and routing. Transaction control and session statements are plumbing. Everything else (spend and + audit writes, daily aggregates, config and model reloads) runs after the response or on a timer.""" + verb: Final = self.query.lstrip("(").split(maxsplit=1)[0].upper() if self.query.strip() else "" + tables: Final = _TABLE.findall(self.query) + if verb in _PLUMBING_VERBS or not tables: + return "plumbing" + return "awaited" if verb == "SELECT" and tables[0] in _AWAITED_TABLES else "background" + + @property + def summary(self) -> str: + return re.sub(r"\s+", " ", self.query)[:110] @dataclass(frozen=True, slots=True) class DbCounters: commits: int + statements: tuple[StatementCalls, ...] | None @dataclass(frozen=True, slots=True) @@ -29,36 +93,169 @@ class RedisCounters: connected_clients: int +@dataclass(frozen=True, slots=True) +class RedisCommand: + command: str + key: str + inside_script: bool + count: int + + @property + def bucket(self) -> Bucket: + """Budget reads, cooldown reads and the rate-limiter script (with every command it issues) happen + before the request is forwarded. Spend increments, token counters, config rewrites, cache + invalidations (PUBLISH, DEL), the spend buffers the pods flush and any other write happen after the + response or on a timer. Reads on keys these rules do not know stay unclassified and count in neither.""" + key: Final = self.key.removeprefix(_REDIS_NAMESPACE) + if key.startswith(("litellm_config:", "daily_", "cronjob_lock:", "litellm_spend", "litellm_daily")): + return "background" + if self.inside_script or self.command in {"EVALSHA", "EVAL"}: + return "awaited" + if self.command in {"PUBLISH", "DEL", "SCRIPT", "INFO", "MONITOR"}: + return "background" + if key.startswith("spend:"): + return "awaited" if self.command in _REDIS_READS else "background" + if key.startswith("deployment:") or key.endswith(("}:requests", "}:window")): + return "awaited" + if key.endswith("}:tokens") or self.command in _REDIS_WRITES: + return "background" + return "unclassified" + + @property + def summary(self) -> str: + key: Final = re.sub(r"[0-9a-f]{16,}|sk-[A-Za-z0-9_-]+|[0-9a-f-]{36}", "*", self.key)[:70] + return f"{self.command} {key}{' (in script)' if self.inside_script else ''}".rstrip() + + +def statement_deltas(before: DbCounters | None, after: DbCounters | None) -> tuple[StatementCalls, ...] | None: + """Calls per statement during the run: `after` minus `before` by queryid, new statements whole.""" + if before is None or after is None or before.statements is None or after.statements is None: + return None + earlier: Final = {statement.queryid: statement.calls for statement in before.statements} + return tuple( + StatementCalls( + queryid=statement.queryid, query=statement.query, calls=statement.calls - earlier.get(statement.queryid, 0) + ) + for statement in after.statements + if statement.calls > earlier.get(statement.queryid, 0) + ) + + +def calls_in(bucket: Bucket, items: tuple[StatementCalls, ...] | tuple[RedisCommand, ...] | None) -> int | None: + if items is None: + return None + return sum(_count(item) for item in items if item.bucket == bucket) + + +def _count(item: StatementCalls | RedisCommand) -> int: + return item.calls if isinstance(item, StatementCalls) else item.count + + def read_db_counters() -> DbCounters | None: url: Final = os.environ.get("DATABASE_URL") if url is None: return None import psycopg + _ensure_statements_extension(url) with psycopg.connect(url) as conn: - row: Final = conn.execute(_DB_QUERY).fetchone() - (commits,) = _DB_ROW.validate_python(row) - return DbCounters(commits=commits) + (commits,) = _COMMITS_ROW.validate_python(conn.execute(_COMMITS_QUERY).fetchone()) + try: + rows = _STATEMENT_ROWS.validate_python(conn.execute(_STATEMENTS_QUERY).fetchall()) + except psycopg.errors.UndefinedTable: + return DbCounters(commits=commits, statements=None) + return DbCounters( + commits=commits, + statements=tuple(StatementCalls(queryid=queryid, query=query, calls=calls) for queryid, calls, query in rows), + ) -def read_redis_counters() -> RedisCounters | None: - """Standalone Redis only, which is what the perf lane's chart-bundled Redis and a local run are.""" - host: Final = os.environ.get("REDIS_HOST") - if host is None: - return None +def _ensure_statements_extension(url: str) -> None: + """Creating the extension needs a superuser, so prefer the admin URL; the runtime user's attempt is + allowed to fail (a database without `shared_preload_libraries = pg_stat_statements` cannot have it).""" + import psycopg + + try: + with psycopg.connect(os.environ.get("E2E_DATABASE_ADMIN_URL", url), autocommit=True) as conn: + _ = conn.execute(_CREATE_EXTENSION) + except psycopg.Error: + return + + +def _redis_client() -> redis.Redis[str]: import redis - client: Final = redis.Redis( - host=host, + return redis.Redis( + host=os.environ["REDIS_HOST"], port=int(os.environ.get("REDIS_PORT", "6379")), ssl=os.environ.get("REDIS_SSL", "false").lower() == "true", password=os.environ.get("REDIS_PASSWORD") or None, decode_responses=True, socket_connect_timeout=2, ) + + +def read_redis_counters() -> RedisCounters | None: + """Standalone Redis only, which is what the perf lane's chart-bundled Redis and a local run are.""" + if os.environ.get("REDIS_HOST") is None: + return None + client: Final = _redis_client() stats: Final = _INFO.validate_python(client.info("stats")) clients: Final = _INFO.validate_python(client.info("clients")) return RedisCounters( commands=_INT.validate_python(stats["total_commands_processed"]), connected_clients=_INT.validate_python(clients["connected_clients"]), ) + + +class RedisMonitor: + """Every command Redis executed while the monitor was open, counted by command, key and whether a Lua + script issued it. `MONITOR` echoes each command to this connection, so it costs Redis one write per + command; fine at the rates these cases run.""" + + def __init__(self) -> None: + self._counts: Counter[tuple[str, str, bool]] = Counter() + self._thread: threading.Thread | None = None + self._pool: redis.ConnectionPool | None = None + self.commands: tuple[RedisCommand, ...] | None = None + + def _run(self) -> None: + import redis + + client: Final = _redis_client() + self._pool = client.connection_pool # pyright: ignore[reportAny] # redis.Redis leaves connection_pool untyped + with client.monitor() as monitor: + try: + while True: + line = _MONITOR_LINE.validate_python(monitor.next_command()) + words = str(line["command"]).split(" ", 2) + key = words[1].strip('"') if len(words) > 1 else "" + self._counts[(words[0].upper(), key, line["client_type"] == "lua")] += 1 + except (redis.ConnectionError, OSError, ValueError): + return + + def start(self) -> None: + self._thread = threading.Thread(target=self._run, name="redis-monitor", daemon=True) + self._thread.start() + + def stop(self) -> tuple[RedisCommand, ...]: + if self._pool is not None: + self._pool.disconnect() + if self._thread is not None: + self._thread.join(timeout=5) + return tuple( + RedisCommand(command=command, key=key, inside_script=inside_script, count=count) + for (command, key, inside_script), count in self._counts.items() + ) + + +@contextmanager +def redis_monitor() -> Generator[RedisMonitor, None, None]: + """Counts Redis commands for the duration of the block; `.commands` is set once the block ends.""" + monitor: Final = RedisMonitor() + if os.environ.get("REDIS_HOST") is not None: + monitor.start() + try: + yield monitor + finally: + monitor.commands = monitor.stop() if os.environ.get("REDIS_HOST") is not None else None diff --git a/tests/e2e/perf/setup/observe.py b/tests/e2e/perf/setup/observe.py index a322e5a7443..98391a02418 100644 --- a/tests/e2e/perf/setup/observe.py +++ b/tests/e2e/perf/setup/observe.py @@ -15,7 +15,14 @@ from e2e_http import NoBody, get_external, post_external, unwrap from memory_readings import RssCapture, RssReading, WorkerKey from pydantic import BaseModel -from setup.datastore import DbCounters, RedisCounters +from setup.datastore import ( + DbCounters, + RedisCommand, + RedisCounters, + StatementCalls, + calls_in, + statement_deltas, +) from setup.mocks.reply_common import Provider from setup.performance_test_case import Limit, Metric, MockLatency, PerformanceTestCase from setup.wire import WireRecord @@ -106,6 +113,7 @@ class Observation: db_after: DbCounters | None = None redis_before: RedisCounters | None = None redis_after: RedisCounters | None = None + redis_commands: tuple[RedisCommand, ...] | None = None gateway_cpu_seconds: float | None = None def per_request(self, before: int | None, after: int | None) -> float | None: @@ -113,6 +121,13 @@ class Observation: return None return (after - before) / len(self.records) + def count_per_request(self, count: int | None) -> float | None: + return None if count is None or not self.records else count / len(self.records) + + @property + def statements(self) -> tuple[StatementCalls, ...] | None: + return statement_deltas(self.db_before, self.db_after) + @property def retried_markers(self) -> frozenset[str]: counts: Final = Counter(seen.marker for seen in self.mock_requests if seen.marker is not None) @@ -191,11 +206,13 @@ def observe(metric: Metric, percentile_rank: int | None, observation: Observatio if metric == "unexpected_status_count": return float(observation.unexpected_status_count()) if metric == "postgres_calls_per_request": - db_before, db_after = observation.db_before, observation.db_after - return observation.per_request(db_before and db_before.commits, db_after and db_after.commits) + return observation.count_per_request(calls_in("awaited", observation.statements)) + if metric == "postgres_background_calls_per_request": + return observation.count_per_request(calls_in("background", observation.statements)) if metric == "redis_calls_per_request": - redis_before, redis_after = observation.redis_before, observation.redis_after - return observation.per_request(redis_before and redis_before.commands, redis_after and redis_after.commands) + return observation.count_per_request(calls_in("awaited", observation.redis_commands)) + if metric == "redis_background_calls_per_request": + return observation.count_per_request(calls_in("background", observation.redis_commands)) if metric == "cpu_seconds_per_request": cpu: Final = observation.gateway_cpu_seconds return None if cpu is None or not observation.records else cpu / len(observation.records) @@ -234,7 +251,9 @@ REPORTED: Final[tuple[tuple[Metric, int | None], ...]] = ( ("time_to_first_token_seconds", 99), ("memory_growth_megabytes", None), ("postgres_calls_per_request", None), + ("postgres_background_calls_per_request", None), ("redis_calls_per_request", None), + ("redis_background_calls_per_request", None), ("cpu_seconds_per_request", None), ) @@ -298,6 +317,7 @@ def render(case: PerformanceTestCase, observation: Observation, rows: list[Row]) flag = " OVER" if row.over else " MISSING" if row.missing else "" lines.append(f" {row.label:32} {observed:>12} {limit:>12}{flag}") lines.extend(_diagnostics(observation)) + lines.extend(_datastore_lines(observation)) for record in failures: lines.append( f" wrong: seq {record.seq} {record.caller} {record.endpoint}{' stream' if record.stream else ''} status {record.status}: {record.error}" @@ -316,9 +336,49 @@ def _diagnostics(observation: Observation) -> list[str]: f" guardrail missed {observation.guardrail_missed() if observation.guardrail_calls else 'n/a'}", f" retried share {len(observation.retried_markers) / count:.4f}", f" spend logged share {'n/a' if share is None else f'{share:.3f}'}", + f" postgres transactions per request {_fmt(observation.per_request(_commits(observation.db_before), _commits(observation.db_after)))}", + f" redis commands per request {_fmt(observation.per_request(_commands(observation.redis_before), _commands(observation.redis_after)))}", f" redis connected clients {_clients(observation.redis_before)} -> {_clients(observation.redis_after)}", ] +def _datastore_lines(observation: Observation) -> list[str]: + """What the per-request Postgres and Redis counts are made of, so a red run names the statement or + command that grew. `plumbing` (BEGIN, COMMIT, session settings) and `unclassified` count in neither limit.""" + count: Final = max(1, len(observation.records)) + lines: list[str] = [] + for title, items in ( + ("postgres statements", observation.statements), + ("redis commands", observation.redis_commands), + ): + if items is None: + lines.append(f" {title}: n/a") + continue + for bucket in ("awaited", "background", "plumbing", "unclassified"): + in_bucket = sorted((item for item in items if item.bucket == bucket), key=_item_count, reverse=True) + total = sum(_item_count(item) for item in in_bucket) + if not in_bucket: + continue + lines.append(f" {title} {bucket}: {total / count:.3f} per request, {len(in_bucket)} distinct") + lines.extend(f" {_item_count(item) / count:8.3f} {item.summary}" for item in in_bucket[:8]) + return lines + + +def _item_count(item: StatementCalls | RedisCommand) -> int: + return item.calls if isinstance(item, StatementCalls) else item.count + + +def _commits(counters: DbCounters | None) -> int | None: + return None if counters is None else counters.commits + + +def _commands(counters: RedisCounters | None) -> int | None: + return None if counters is None else counters.commands + + +def _fmt(value: float | None) -> str: + return "n/a" if value is None else f"{value:.3f}" + + def _clients(counters: RedisCounters | None) -> str: return "n/a" if counters is None else str(counters.connected_clients) diff --git a/tests/e2e/perf/setup/performance_test_case.py b/tests/e2e/perf/setup/performance_test_case.py index d90b262329c..76e6efa3203 100644 --- a/tests/e2e/perf/setup/performance_test_case.py +++ b/tests/e2e/perf/setup/performance_test_case.py @@ -22,8 +22,10 @@ Metric = Literal[ "time_to_first_token_seconds", "memory_growth_megabytes", "postgres_calls_per_request", + "postgres_background_calls_per_request", "cpu_seconds_per_request", "redis_calls_per_request", + "redis_background_calls_per_request", ] diff --git a/tests/e2e/perf/test_perf_cases.py b/tests/e2e/perf/test_perf_cases.py index d077cb9f8c7..d89a8634ab3 100644 --- a/tests/e2e/perf/test_perf_cases.py +++ b/tests/e2e/perf/test_perf_cases.py @@ -9,11 +9,13 @@ import time from typing import Final import pytest +from cases import CASES from idp import Keycloak from lifecycle import ResourceManager from memory_readings import read_rss_everywhere from models import SpendLogsParams -from cases import CASES +from proxy_client import ProxyClient +from setup.datastore import read_db_counters, read_redis_counters, redis_monitor from setup.observe import ( Observation, evaluate, @@ -26,11 +28,9 @@ from setup.observe import ( ) from setup.performance_test_case import PerformanceTestCase from setup.plan import build_plan -from setup.datastore import read_db_counters, read_redis_counters from setup.podcpu import gateway_cpu_sampler from setup.provision import MOCK_PROVIDERS, build_setup, guardrail_url, mock_url from setup.run import run_plan -from proxy_client import ProxyClient pytestmark = [pytest.mark.perf, pytest.mark.quiet_stack] _SPEND_SAMPLE: Final = 40 @@ -68,9 +68,11 @@ def test_perf_case( db_before: Final = read_db_counters() redis_before: Final = read_redis_counters() run_started: Final = time.monotonic() - with gateway_cpu_sampler() as cpu: + with gateway_cpu_sampler() as cpu, redis_monitor() as monitor: outcome: Final = run_plan(case, plan, setup) run_finished: Final = time.monotonic() + db_after: Final = read_db_counters() + redis_after: Final = read_redis_counters() rss_after: Final = read_rss_everywhere(proxy) rss_read: Final = time.monotonic() assert outcome.records, f"locust wrote no records (exit {outcome.locust_exit}):\n{outcome.locust_output}" @@ -81,8 +83,6 @@ def test_perf_case( proxy, served_call_ids[:: max(1, len(served_call_ids) // _SPEND_SAMPLE)][:_SPEND_SAMPLE] ) spend_read: Final = time.monotonic() - db_after: Final = read_db_counters() - redis_after: Final = read_redis_counters() mock_requests: Final = tuple(seen for url in mock_urls.values() for seen in read_mock_requests(url)) guardrail_calls: Final = tuple(read_guardrail_calls(guardrail_url())) if case.guardrails else () mocks_read: Final = time.monotonic() @@ -98,6 +98,7 @@ def test_perf_case( db_after=db_after, redis_before=redis_before, redis_after=redis_after, + redis_commands=monitor.commands, gateway_cpu_seconds=cpu.seconds, ) rows, exceeded_limits = evaluate(case, observation) diff --git a/tests/e2e_harness/perf/setup/test_observe.py b/tests/e2e_harness/perf/setup/test_observe.py index 998d7d9b87e..f38e7927fb7 100644 --- a/tests/e2e_harness/perf/setup/test_observe.py +++ b/tests/e2e_harness/perf/setup/test_observe.py @@ -5,7 +5,7 @@ from typing import Final from cases.small_bodies_5rps import CASE from memory_readings import RssCapture -from setup.datastore import DbCounters +from setup.datastore import DbCounters, StatementCalls from setup.observe import Observation, evaluate, junit_properties from setup.performance_test_case import Limit from setup.wire import WireRecord @@ -35,6 +35,10 @@ def served(seq: int, total_s: float) -> WireRecord: ) +def _key_lookup(calls: int) -> StatementCalls: + return StatementCalls(queryid=1, query='SELECT * FROM "LiteLLM_VerificationToken" WHERE token = $1', calls=calls) + + def test_junit_properties_carry_each_observed_metric_and_each_limit_including_unobserved_ones() -> None: case: Final = replace( CASE, @@ -50,8 +54,8 @@ def test_junit_properties_carry_each_observed_metric_and_each_limit_including_un rss_before=NO_RSS, rss_after=NO_RSS, wall_seconds=3.0, - db_before=DbCounters(commits=100), - db_after=DbCounters(commits=130), + db_before=DbCounters(commits=100, statements=(_key_lookup(calls=100),)), + db_after=DbCounters(commits=130, statements=(_key_lookup(calls=130),)), ) rows, _ = evaluate(case, observation) @@ -64,5 +68,6 @@ def test_junit_properties_carry_each_observed_metric_and_each_limit_including_un "perf_gateway_added_seconds_p50_limit": 0.625, "perf_gateway_added_seconds_p99": 0.75, "perf_postgres_calls_per_request": 10.0, + "perf_postgres_background_calls_per_request": 0.0, "perf_redis_calls_per_request_limit": 35.0, } diff --git a/tests/integration/cost_calculation/cases/openai.py b/tests/integration/cost_calculation/cases/openai.py index 9a3885af259..cfb7fda4d71 100644 --- a/tests/integration/cost_calculation/cases/openai.py +++ b/tests/integration/cost_calculation/cases/openai.py @@ -11192,6 +11192,87 @@ GPT_5_6_SLASH_ALIAS_STREAM: Final = GPT_5_6_STREAM.model_copy( update={"name": "gpt-5.6-slash_alias_stream", "alias": "openai/custom/gpt-5.6-$MARKER"} ) +GPT_5_6_RESPONSES_IMAGE_GENERATION_TOOL_TOKENS: Final = CostTrackingTestCase( + name="gpt-5.6-responses_image_generation_tool_tokens", + covers="quota_management.spend_tracking.cost_matrix.logs_cost", + model="gpt-5.6", + endpoint="/v1/responses", + request={ + "model": "$MODEL", + "input": "draw a duck", + "tools": [{"type": "image_generation", "model": "gpt-image-next", "quality": "low", "size": "1024x1024"}], + }, + response=JsonResponse( + content_type="application/json", + body={ + "id": "resp_$REQUEST_ID", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-5.6", + "tools": [ + { + "type": "image_generation", + "background": "auto", + "model": "gpt-image-next", + "moderation": "auto", + "n": 1, + "output_compression": 100, + "output_format": "png", + "quality": "low", + "size": "1024x1024", + } + ], + "output": [ + { + "type": "image_generation_call", + "id": "ig_$REQUEST_ID", + "status": "completed", + "action": "generate", + "background": "opaque", + "output_format": "png", + "quality": "low", + "revised_prompt": "a duck", + "size": "1024x1024", + "result": "AAAA", + }, + { + "type": "message", + "id": "msg_$REQUEST_ID", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "scripted response", "annotations": []}], + }, + ], + "usage": { + "input_tokens": 1840, + "output_tokens": 412, + "total_tokens": 2252, + "input_tokens_details": {"cached_tokens": 0}, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + "tool_usage": { + "image_gen": { + "input_tokens": 32, + "input_tokens_details": {"image_tokens": 0, "text_tokens": 32}, + "output_tokens": 196, + "output_tokens_details": {"image_tokens": 196, "text_tokens": 0}, + "total_tokens": 228, + }, + "web_search": {"num_requests": 0}, + }, + }, + ), + expected=ExactExpected( + spend=0.01004232, + input_cost=0.00322, + output_cost=0.005768, + prompt_tokens=1840, + completion_tokens=412, + tool_usage_cost=0.00105432, + ), +) + CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GPT_5_6_SLASH_ALIAS, GPT_5_6_SLASH_ALIAS_STREAM, @@ -11336,6 +11417,7 @@ CASES: Final[tuple[CostTrackingTestCase, ...]] = ( GPT_5_6_RESPONSES_WEB_SEARCH_REPORTED_COUNT_MIXED_ACTIONS, GPT_5_6_RESPONSES_WEB_SEARCH_REPORTED_COUNT_OPEN_PAGE_ONLY, GPT_5_6_RESPONSES_WEB_SEARCH_REPORTED_COUNT_MIXED_ACTIONS_STREAM, + GPT_5_6_RESPONSES_IMAGE_GENERATION_TOOL_TOKENS, GPT_5_3_CODEX_RESPONSES_FILE_SEARCH, GPT_5_6_RESPONSES_SERVICE_TIER_FLEX, GPT_5_6_RESPONSES_SERVICE_TIER_PRIORITY, diff --git a/tests/integration/observability/test_bedrock_error_request_id.py b/tests/integration/observability/test_bedrock_error_request_id.py index 230619a2863..436dcc572d5 100644 --- a/tests/integration/observability/test_bedrock_error_request_id.py +++ b/tests/integration/observability/test_bedrock_error_request_id.py @@ -1,13 +1,208 @@ +import asyncio +import base64 import json import uuid -from typing import Final +from collections.abc import Mapping +from concurrent.futures import ThreadPoolExecutor +from types import MappingProxyType +from typing import Final, Literal, NamedTuple +from urllib.parse import unquote -from integration._support.client import Gateway, eventually, object_value +import anthropic +import httpx +import openai +import pytest +from integration._support.client import Gateway, Scenario, eventually, object_value from integration._support.database import read_rows +from integration._support.responses_stream import ( + created, + delta, + error_event, + failed, + frame, + healthy_stream, + openai_error_event, + spend_row, +) +from integration._support.upstream import _aws_event_frame, aws_event_stream_frame from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue + +import litellm +from litellm.exceptions import MidStreamFallbackError +from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id _MODEL: Final = "bedrock/converse/anthropic.claude-sonnet-4-5-20250929-v1:0" +_BEDROCK_RESPONSES_MODEL: Final = "bedrock/us.openai.gpt-5.6-sol" _TOKEN: Final = "synthetic-bedrock-bearer" +_SHAPES: Final = ("nested-error", "top-level-error", "response.failed", "error-before-output") +_ERROR_CODES: Final = ("server_error", "rate_limit_exceeded", "invalid_request_error", "content_policy_violation") +_INVOKE_MODEL_ID: Final = "anthropic.claude-3-haiku-20240307-v1:0" +_BEDROCK_AWS: Final[Mapping[str, JsonValue]] = MappingProxyType( + { + "aws_access_key_id": "AKIASCRIPTEDPROVIDER", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": "us-east-1", + } +) +_EVENT_STREAM: Final = "application/vnd.amazon.eventstream" +_BEDROCK_STREAM_TARGETS: Final = frozenset( + ( + "/model/anthropic.claude-sonnet-4-5-20250929-v1:0/converse-stream", + f"/model/{_INVOKE_MODEL_ID}/invoke-with-response-stream", + ) +) + +BedrockFailure = Literal["exception-event", "exception-message", "no-known-event", "truncated-frame"] + + +class _FailureShape(NamedTuple): + bedrock_status: int + converse_class: str + after_output: tuple[bool, ...] + + +BedrockEndpoint = Literal["chat", "responses", "messages", "messages-invoke"] +_BEDROCK_ENDPOINTS: Final[tuple[BedrockEndpoint, ...]] = ("chat", "responses", "messages", "messages-invoke") +_BEDROCK_FAILURES: Final[Mapping[BedrockFailure, _FailureShape]] = MappingProxyType( + { + "exception-event": _FailureShape(500, "ServiceUnavailableError", (False, True)), + "exception-message": _FailureShape(429, "RateLimitError", (False, True)), + "no-known-event": _FailureShape(502, "BadGatewayError", (False,)), + "truncated-frame": _FailureShape(502, "BadGatewayError", (True,)), + } +) +_BEDROCK_CELLS: Final = tuple( + pytest.param( + endpoint, failure, after_output, id=f"{endpoint}-{failure}-{'after' if after_output else 'before'}-output" + ) + for endpoint in _BEDROCK_ENDPOINTS + for failure, shape in _BEDROCK_FAILURES.items() + for after_output in shape.after_output +) + + +def _stream(identity: str, shape: str, error_code: str, message: str) -> tuple[bytes, ...]: + opening: Final = (frame(created(identity)), frame(delta(identity, "partial output"))) + nested: Final = frame(error_event({"type": error_code, "code": error_code, "message": message})) + match shape: + case "nested-error": + return (*opening, nested) + case "top-level-error": + return (*opening, frame(openai_error_event(error_code, message))) + case "response.failed": + return (*opening, frame(failed(identity, error_code, message))) + case "error-before-output": + return (nested,) + raise AssertionError(shape) + + +def _bedrock_frame(event_type: str, payload: dict[str, JsonValue]) -> bytes: + return _aws_event_frame(event_type, payload, "sc", "u") + + +def _invoke_chunk(event: dict[str, JsonValue]) -> bytes: + return _bedrock_frame("chunk", {"bytes": base64.b64encode(json.dumps(event).encode()).decode()}) + + +def _bedrock_output(invoke: bool, text: str = "partial output") -> tuple[bytes, ...]: + if not invoke: + return ( + _bedrock_frame("messageStart", {"role": "assistant"}), + _bedrock_frame("contentBlockDelta", {"delta": {"text": text}, "contentBlockIndex": 0}), + ) + started: Final[dict[str, JsonValue]] = { + "id": f"msg_{uuid.uuid4().hex}", + "type": "message", + "role": "assistant", + "content": [], + "model": _INVOKE_MODEL_ID, + "stop_reason": None, + "usage": {"input_tokens": 3, "output_tokens": 1}, + } + return ( + _invoke_chunk({"type": "message_start", "message": started}), + _invoke_chunk({"type": "content_block_start", "index": 0, "content_block": {"type": "text", "text": ""}}), + _invoke_chunk({"type": "content_block_delta", "index": 0, "delta": {"type": "text_delta", "text": text}}), + ) + + +def _bedrock_failure_frame(failure: BedrockFailure, message: str) -> bytes: + match failure: + case "exception-event": + return _bedrock_frame("internalServerException", {"message": message}) + case "exception-message": + headers: Final = { + ":exception-type": "throttlingException", + ":content-type": "application/json", + ":message-type": "exception", + } + return aws_event_stream_frame(headers, json.dumps({"message": message}).encode()) + case "no-known-event": + return _bedrock_frame("somethingBedrockAddedLater", {"message": message}) + case "truncated-frame": + return _bedrock_frame("contentBlockDelta", {"delta": {"text": message}, "contentBlockIndex": 0})[:-7] + + +def _bedrock_failure_text(failure: BedrockFailure, message: str) -> str: + return "undecoded bytes after" if failure == "truncated-frame" else message + + +def _bedrock_stream_request( + endpoint: BedrockEndpoint, model: str, prompt: str, *, stream: bool = True +) -> tuple[str, dict[str, JsonValue]]: + extra: Final[dict[str, JsonValue]] = {"stream": stream, "num_retries": 0, "cache": {"no-cache": True}} + turn: Final[list[JsonValue]] = [{"role": "user", "content": prompt}] + match endpoint: + case "chat": + return "/v1/chat/completions", {"model": model, "messages": turn, **extra} + case "responses": + return "/v1/responses", {"model": model, "input": prompt, **extra} + case "messages" | "messages-invoke": + return "/v1/messages", {"model": model, "messages": turn, "max_tokens": 16, **extra} + + +def _failure_information(call_id: str) -> dict[str, JsonValue]: + row: Final = spend_row(call_id) + assert row["status"] == "failure", row + metadata: Final = row["metadata"] + parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata) + return object_value(parsed["error_information"]) + + +def _streamed_response(gateway: Gateway, model: str, prompt: str) -> httpx.Response: + response: Final = gateway.request( + "POST", + "/v1/responses", + {"model": model, "input": prompt, "stream": True, "num_retries": 0, "cache": {"no-cache": True}}, + ) + assert response.status_code == 200, response.text + return response + + +async def _consume_sdk_stream(asynchronous: bool, api_base: str, identity: str) -> None: + if asynchronous: + async_stream: Final = await litellm.aresponses( + model="openai/gpt-6", + input=identity, + stream=True, + api_base=api_base, + api_key="synthetic-provider-key", + num_retries=0, + ) + async for _ in async_stream: + pass + return + tuple( + litellm.responses( + model="openai/gpt-6", + input=identity, + stream=True, + api_base=api_base, + api_key="synthetic-provider-key", + num_retries=0, + ) + ) def test_bedrock_500_keeps_amzn_request_id_on_error_headers_and_failure_log(gateway: Gateway) -> None: @@ -54,3 +249,556 @@ def test_bedrock_500_keeps_amzn_request_id_on_error_headers_and_failure_log(gate parsed: Final = json.loads(metadata) if isinstance(metadata, str) else object_value(metadata) error_information: Final = object_value(parsed["error_information"]) assert error_information["error_provider_request_id"] == amzn_request_id, error_information + + +@pytest.mark.parametrize("shape", _SHAPES) +@pytest.mark.parametrize("error_code", _ERROR_CODES) +def test_responses_stream_failure_log_retains_provider_request_id( + gateway: Gateway, shape: str, error_code: str +) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + provider_request_id: Final = str(uuid.uuid4()) + message: Final = f"stream failure {identity}" + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/responses", request.target + assert json.loads(request.body)["stream"] is True + return Reply( + headers={"x-amzn-RequestId": provider_request_id}, + content_type="text/event-stream", + chunks=_stream(identity, shape, error_code, message), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model="openai/gpt-6", api_base=wire.url, num_retries=0) + response: Final = _streamed_response(gateway, model, identity) + assert message in response.text, response.text + assert ("partial output" in response.text) is (shape != "error-before-output"), response.text + assert response.headers.get("llm_provider-x-amzn-requestid") == provider_request_id, dict(response.headers) + assert len(wire.drain()) == 1 + information: Final = _failure_information(response.headers["x-litellm-call-id"]) + assert information["error_provider_request_id"] == provider_request_id, information + + +@pytest.mark.parametrize("shape", _SHAPES) +def test_responses_stream_failure_without_provider_request_id_logs_none(gateway: Gateway, shape: str) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + message: Final = f"stream failure {identity}" + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/responses", request.target + return Reply(content_type="text/event-stream", chunks=_stream(identity, shape, "server_error", message)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model="openai/gpt-6", api_base=wire.url, num_retries=0) + response: Final = _streamed_response(gateway, model, identity) + assert message in response.text, response.text + assert "llm_provider-x-amzn-requestid" not in response.headers, dict(response.headers) + assert len(wire.drain()) == 1 + information: Final = _failure_information(response.headers["x-litellm-call-id"]) + assert information["error_class"] == "InternalServerError", information + assert information["error_provider_request_id"] is None, information + + +@pytest.mark.parametrize("shape", ("nested-error", "response.failed")) +def test_bedrock_responses_stream_server_error_logs_amzn_request_id(gateway: Gateway, shape: str) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + amzn_request_id: Final = str(uuid.uuid4()) + message: Final = f"bedrock stream failure {identity}" + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/openai/v1/responses", request.target + assert request.headers["authorization"] == f"Bearer {_TOKEN}", request.headers + assert json.loads(request.body)["stream"] is True + return Reply( + headers={"x-amzn-RequestId": amzn_request_id}, + content_type="text/event-stream", + chunks=_stream(identity, shape, "server_error", message), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=_BEDROCK_RESPONSES_MODEL, + api_key=_TOKEN, + api_base=wire.url, + aws_region_name="us-east-1", + num_retries=0, + ) + response: Final = _streamed_response(gateway, model, identity) + assert "partial output" in response.text and message in response.text, response.text + assert response.headers.get("llm_provider-x-amzn-requestid") == amzn_request_id, dict(response.headers) + assert len(wire.drain()) == 1 + information: Final = _failure_information(response.headers["x-litellm-call-id"]) + assert information["llm_provider"] == "bedrock", information + assert information["error_provider_request_id"] == amzn_request_id, information + + +def test_responses_stream_failure_log_reads_openai_x_request_id(gateway: Gateway) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + provider_request_id: Final = f"req_{uuid.uuid4().hex}" + message: Final = f"stream failure {identity}" + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/responses", request.target + return Reply( + headers={"x-request-id": provider_request_id}, + content_type="text/event-stream", + chunks=_stream(identity, "top-level-error", "server_error", message), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model="openai/gpt-6", api_base=wire.url, num_retries=0) + response: Final = _streamed_response(gateway, model, identity) + assert message in response.text, response.text + assert len(wire.drain()) == 1 + information: Final = _failure_information(response.headers["x-litellm-call-id"]) + assert information["error_provider_request_id"] == provider_request_id, information + + +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +@pytest.mark.parametrize("event_type", ("error", "response.failed")) +async def test_responses_sdk_stream_error_retains_provider_request_id(asynchronous: bool, event_type: str) -> None: + identity: Final = f"resp_{uuid.uuid4().hex}" + provider_request_id: Final = str(uuid.uuid4()) + shape: Final = "nested-error" if event_type == "error" else "response.failed" + + def respond(request: Request) -> Reply: + assert request.method == "POST" and request.target == "/responses", request.target + return Reply( + headers={"x-amzn-RequestId": provider_request_id, "retry-after": "7"}, + content_type="text/event-stream", + chunks=_stream(identity, shape, "server_error", identity), + ) + + with wire_server(respond) as wire: + with pytest.raises(MidStreamFallbackError) as raised: + await _consume_sdk_stream(asynchronous, wire.url, identity) + assert raised.value.status_code == 500 + assert raised.value.response.status_code == 500 + assert str(raised.value.response.request.url) == "https://openai.com/v1/" + assert raised.value.generated_content == "partial output" + assert dict(raised.value.response.headers) == {"x-amzn-requestid": provider_request_id} + assert _get_provider_request_id(raised.value) == provider_request_id + assert raised.value.original_exception is not None + assert _get_provider_request_id(raised.value.original_exception) == provider_request_id + assert len(wire.drain()) == 1 + + +def _concurrent_reply(request: Request, failing: tuple[str, ...], recovery: str) -> Reply: + prompt: Final = json.loads(request.body)["input"] + assert prompt in (*failing, recovery) + chunks: Final = ( + healthy_stream(f"resp_{prompt}", "recovered") + if prompt == recovery + else _stream(f"resp_{prompt}", "nested-error", "server_error", prompt) + ) + return Reply(content_type="text/event-stream", headers={"x-amzn-requestid": prompt}, chunks=chunks) + + +def test_concurrent_responses_failures_keep_distinct_request_ids_and_recover(gateway: Gateway) -> None: + identity: Final = uuid.uuid4().hex + prompts: Final = tuple(f"{identity}-{index}" for index in range(4)) + recovery: Final = f"{identity}-recovery" + + with ( + wire_server(lambda request: _concurrent_reply(request, prompts, recovery)) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model(model="openai/gpt-6", api_base=wire.url, num_retries=0) + + def call(prompt: str) -> str: + response: Final = _streamed_response(gateway, model, prompt) + assert (prompt if prompt != recovery else "recovered") in response.text, response.text + return response.headers["x-litellm-call-id"] + + with ThreadPoolExecutor(max_workers=4) as workers: + call_ids: Final = tuple(workers.map(call, prompts)) + recovered_call_id: Final = call(recovery) + assert len(wire.drain()) == 5 + for prompt, call_id in zip(prompts, call_ids, strict=True): + assert _failure_information(call_id)["error_provider_request_id"] == prompt + assert spend_row(recovered_call_id)["status"] == "success" + + +def _bedrock_logged_error(endpoint: BedrockEndpoint, failure: BedrockFailure) -> tuple[str, str]: + shape: Final = _BEDROCK_FAILURES[failure] + if endpoint == "messages-invoke": + return "BedrockError", str(shape.bedrock_status) + return shape.converse_class, "503" if failure == "exception-event" else str(shape.bedrock_status) + + +def _bedrock_client_status(endpoint: BedrockEndpoint, failure: BedrockFailure, after_output: bool) -> int: + logged_code: Final = int(_bedrock_logged_error(endpoint, failure)[1]) + if after_output or endpoint in ("responses", "messages"): + return 200 + if endpoint == "messages-invoke" and failure == "no-known-event": + return 200 + return logged_code + + +@pytest.mark.parametrize(("endpoint", "failure", "after_output"), _BEDROCK_CELLS) +def test_bedrock_event_stream_failure_log_retains_amzn_request_id( + gateway: Gateway, endpoint: BedrockEndpoint, failure: BedrockFailure, after_output: bool +) -> None: + amzn_request_id: Final = str(uuid.uuid4()) + message: Final = f"bedrock stream failure {amzn_request_id}" + invoke: Final = endpoint == "messages-invoke" + frames: Final = (*(_bedrock_output(invoke) if after_output else ()), _bedrock_failure_frame(failure, message)) + + def respond(request: Request) -> Reply: + assert request.method == "POST" and unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target + return Reply( + headers={"x-amzn-RequestId": amzn_request_id, "retry-after": "7"}, + content_type=_EVENT_STREAM, + body=b"".join(frames), + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model=f"bedrock/invoke/{_INVOKE_MODEL_ID}" if invoke else _MODEL, + api_base=wire.url, + num_retries=0, + **_BEDROCK_AWS, + ) + path, body = _bedrock_stream_request(endpoint, model, message) + with gateway.client.stream( + "POST", path, json=body, headers={"Authorization": f"Bearer {gateway.key}"} + ) as response: + text: Final = "\n".join(line for line in response.iter_lines() if line) + status: Final = response.status_code + response_headers: Final = dict(response.headers) + assert status == _bedrock_client_status(endpoint, failure, after_output), text + assert _bedrock_failure_text(failure, message) in text and ("partial output" in text) is after_output, text + assert len(wire.drain()) == 1 + information: Final = _failure_information(response_headers["x-litellm-call-id"]) + assert (information["error_class"], information["error_code"]) == _bedrock_logged_error(endpoint, failure) + assert information["error_provider_request_id"] == amzn_request_id, information + if status != 200 and not invoke: + assert response_headers.get("llm_provider-x-amzn-requestid") == amzn_request_id, response_headers + assert "llm_provider-retry-after" not in response_headers, response_headers + + +def _bedrock_model(scenario: Scenario, endpoint: BedrockEndpoint, api_base: str) -> str: + return scenario.model( + model=f"bedrock/invoke/{_INVOKE_MODEL_ID}" if endpoint == "messages-invoke" else _MODEL, + api_base=api_base, + num_retries=0, + **_BEDROCK_AWS, + ) + + +def _bedrock_stream( + gateway: Gateway, endpoint: BedrockEndpoint, model: str, prompt: str +) -> tuple[int, str, dict[str, str]]: + path, body = _bedrock_stream_request(endpoint, model, prompt) + with gateway.client.stream("POST", path, json=body, headers={"Authorization": f"Bearer {gateway.key}"}) as response: + text: Final = "\n".join(line for line in response.iter_lines() if line) + return response.status_code, text, dict(response.headers) + + +@pytest.mark.parametrize("after_output", (False, True), ids=("before-output", "after-output")) +@pytest.mark.parametrize("endpoint", _BEDROCK_ENDPOINTS) +def test_bedrock_event_stream_failure_without_amzn_request_id_logs_none( + gateway: Gateway, endpoint: BedrockEndpoint, after_output: bool +) -> None: + message: Final = f"bedrock stream failure {uuid.uuid4()}" + frames: Final = ( + *(_bedrock_output(endpoint == "messages-invoke") if after_output else ()), + _bedrock_failure_frame("exception-event", message), + ) + + def respond(request: Request) -> Reply: + assert unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target + return Reply(headers={"retry-after": "7"}, content_type=_EVENT_STREAM, body=b"".join(frames)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + status, text, response_headers = _bedrock_stream( + gateway, endpoint, _bedrock_model(scenario, endpoint, wire.url), message + ) + assert status == _bedrock_client_status(endpoint, "exception-event", after_output), text + assert message in text, text + assert not any("amzn" in name for name in response_headers), response_headers + assert len(wire.drain()) == 1 + information: Final = _failure_information(response_headers["x-litellm-call-id"]) + assert (information["error_class"], information["error_code"]) == _bedrock_logged_error( + endpoint, "exception-event" + ) + assert information["error_provider_request_id"] is None, information + + +@pytest.mark.parametrize("endpoint", ("responses", "messages", "messages-invoke")) +def test_bedrock_non_stream_500_logs_amzn_request_id_on_every_endpoint( + gateway: Gateway, endpoint: BedrockEndpoint +) -> None: + amzn_request_id: Final = str(uuid.uuid4()) + targets: Final = frozenset( + ( + "/model/anthropic.claude-sonnet-4-5-20250929-v1:0/converse", + f"/model/{_INVOKE_MODEL_ID}/invoke", + ) + ) + + def respond(request: Request) -> Reply: + assert request.method == "POST" and unquote(request.target) in targets, request.target + return Reply( + status=500, + headers={"x-amzn-RequestId": amzn_request_id}, + body=b'{"message":"synthetic bedrock failure"}', + ) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + path, body = _bedrock_stream_request( + endpoint, _bedrock_model(scenario, endpoint, wire.url), amzn_request_id, stream=False + ) + response: Final = gateway.request("POST", path, body) + assert response.status_code == 503, response.text + if endpoint == "responses": + assert response.headers.get("llm_provider-x-amzn-requestid") == amzn_request_id, dict(response.headers) + assert len(wire.drain()) == 1 + information: Final = _failure_information(response.headers["x-litellm-call-id"]) + assert information["error_provider_request_id"] == amzn_request_id, information + + +@pytest.mark.parametrize("amzn_request_id", ("", "a" * 5000, "dup-1, dup-1"), ids=("empty", "5kb", "joined-duplicate")) +def test_bedrock_stream_failure_logs_odd_request_ids_like_the_non_stream_error( + gateway: Gateway, amzn_request_id: str +) -> None: + stream_body: Final = b"".join((*_bedrock_output(False), _bedrock_failure_frame("exception-event", "odd id"))) + + def respond(request: Request) -> Reply: + if unquote(request.target) in _BEDROCK_STREAM_TARGETS: + return Reply(headers={"x-amzn-RequestId": amzn_request_id}, content_type=_EVENT_STREAM, body=stream_body) + return Reply(status=500, headers={"x-amzn-RequestId": amzn_request_id}, body=b'{"message":"odd id"}') + + with wire_server(respond) as wire, gateway.scenario() as scenario: + model: Final = _bedrock_model(scenario, "chat", wire.url) + status, text, stream_headers = _bedrock_stream(gateway, "chat", model, "odd id") + path, body = _bedrock_stream_request("chat", model, "odd id", stream=False) + non_stream: Final = gateway.request("POST", path, body) + assert status == 200 and "odd id" in text, text + assert non_stream.status_code == 503, non_stream.text + assert len(wire.drain()) == 2 + streamed: Final = _failure_information(stream_headers["x-litellm-call-id"]) + unstreamed: Final = _failure_information(non_stream.headers["x-litellm-call-id"]) + assert streamed["error_provider_request_id"] == unstreamed["error_provider_request_id"], (streamed, unstreamed) + + +_BURST_CALLS: Final = tuple( + (endpoint, failure, after_output) + for endpoint in _BEDROCK_ENDPOINTS + for failure, shape in _BEDROCK_FAILURES.items() + for after_output in shape.after_output +) + + +def _healthy_bedrock_stream(invoke: bool) -> bytes: + closing: Final = ( + ( + _invoke_chunk({"type": "content_block_stop", "index": 0}), + _invoke_chunk( + {"type": "message_delta", "delta": {"stop_reason": "end_turn"}, "usage": {"output_tokens": 1}} + ), + _invoke_chunk({"type": "message_stop"}), + ) + if invoke + else ( + _bedrock_frame("contentBlockStop", {"contentBlockIndex": 0}), + _bedrock_frame("messageStop", {"stopReason": "end_turn"}), + _bedrock_frame("metadata", {"usage": {"inputTokens": 3, "outputTokens": 1, "totalTokens": 4}}), + ) + ) + return b"".join((*_bedrock_output(invoke, "recovered"), *closing)) + + +def test_concurrent_bedrock_stream_failures_on_every_endpoint_keep_distinct_request_ids(gateway: Gateway) -> None: + identity: Final = uuid.uuid4().hex + failing: Final = {f"{identity}-{index:02d}": call for index, call in enumerate(_BURST_CALLS)} + recoveries: Final = {f"{identity}-r{index}": endpoint for index, endpoint in enumerate(_BEDROCK_ENDPOINTS)} + + def respond(request: Request) -> Reply: + (token,) = (token for token in (*failing, *recoveries) if token in request.body.decode()) + if token in recoveries: + body = _healthy_bedrock_stream(recoveries[token] == "messages-invoke") + else: + endpoint, failure, after_output = failing[token] + opening = _bedrock_output(endpoint == "messages-invoke") if after_output else () + body = b"".join((*opening, _bedrock_failure_frame(failure, token))) + return Reply(headers={"x-amzn-RequestId": token, "retry-after": "7"}, content_type=_EVENT_STREAM, body=body) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + models: Final = {endpoint: _bedrock_model(scenario, endpoint, wire.url) for endpoint in _BEDROCK_ENDPOINTS} + + def call(token: str, endpoint: BedrockEndpoint) -> tuple[int, str, dict[str, str]]: + return _bedrock_stream(gateway, endpoint, models[endpoint], token) + + with ThreadPoolExecutor(max_workers=8) as workers: + outcomes: Final = dict( + zip(failing, workers.map(call, failing, (endpoint for endpoint, _, _ in failing.values())), strict=True) + ) + recovered: Final = {token: call(token, endpoint) for token, endpoint in recoveries.items()} + assert len(wire.drain()) == len(failing) + len(recoveries) + for token, (endpoint, failure, after_output) in failing.items(): + status, text, response_headers = outcomes[token] + assert status == _bedrock_client_status(endpoint, failure, after_output), (token, text) + assert _bedrock_failure_text(failure, token) in text, (token, text) + information = _failure_information(response_headers["x-litellm-call-id"]) + assert information["error_provider_request_id"] == token, (token, information) + for token, (status, text, response_headers) in recovered.items(): + assert status == 200 and "recovered" in text, (token, text) + assert spend_row(response_headers["x-litellm-call-id"])["status"] == "success" + + +def _sdk_stream_text( + client_kind: str, asynchronous: bool, gateway: Gateway, model: str, prompt: str +) -> tuple[str, str]: + extra_body: Final[dict[str, JsonValue]] = {"num_retries": 0, "cache": {"no-cache": True}} + turn: Final = [{"role": "user", "content": prompt}] + + def open_stream(client: openai.OpenAI | openai.AsyncOpenAI | anthropic.Anthropic | anthropic.AsyncAnthropic): + match client_kind: + case "openai-chat": + return client.chat.completions.with_raw_response.create( + model=model, messages=turn, stream=True, extra_body=extra_body + ) + case "openai-responses": + return client.responses.with_raw_response.create( + model=model, input=prompt, stream=True, extra_body=extra_body + ) + case "anthropic-messages": + return client.messages.with_raw_response.create( + model=model, messages=turn, max_tokens=16, stream=True, extra_body=extra_body + ) + raise AssertionError(client_kind) + + seen: Final[list[str]] = [] + base_url: Final = str(gateway.client.base_url).rstrip("/") + if not asynchronous: + with httpx.Client(timeout=15, trust_env=False) as transport: + client = ( + anthropic.Anthropic(api_key=gateway.key, base_url=base_url, max_retries=0, http_client=transport) + if client_kind == "anthropic-messages" + else openai.OpenAI(api_key=gateway.key, base_url=f"{base_url}/v1", max_retries=0, http_client=transport) + ) + raw = open_stream(client) + try: + seen.extend(repr(event) for event in raw.parse()) + except (openai.APIError, anthropic.APIError) as error: + seen.append(str(error)) + return raw.headers["x-litellm-call-id"], "\n".join(seen) + + async def run() -> str: + async with httpx.AsyncClient(timeout=15, trust_env=False) as transport: + client = ( + anthropic.AsyncAnthropic(api_key=gateway.key, base_url=base_url, max_retries=0, http_client=transport) + if client_kind == "anthropic-messages" + else openai.AsyncOpenAI( + api_key=gateway.key, base_url=f"{base_url}/v1", max_retries=0, http_client=transport + ) + ) + raw = await open_stream(client) + try: + async for event in raw.parse(): + seen.append(repr(event)) + except (openai.APIError, anthropic.APIError) as error: + seen.append(str(error)) + return raw.headers["x-litellm-call-id"] + + return asyncio.run(run()), "\n".join(seen) + + +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +@pytest.mark.parametrize("client_kind", ("openai-chat", "openai-responses", "anthropic-messages")) +def test_sdk_clients_streaming_bedrock_failure_through_proxy_log_amzn_request_id( + gateway: Gateway, client_kind: str, asynchronous: bool +) -> None: + amzn_request_id: Final = str(uuid.uuid4()) + message: Final = f"bedrock stream failure {amzn_request_id}" + frames: Final = (*_bedrock_output(False), _bedrock_failure_frame("exception-event", message)) + + def respond(request: Request) -> Reply: + assert unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target + return Reply(headers={"x-amzn-RequestId": amzn_request_id}, content_type=_EVENT_STREAM, body=b"".join(frames)) + + with wire_server(respond) as wire, gateway.scenario() as scenario: + call_id, seen = _sdk_stream_text( + client_kind, asynchronous, gateway, _bedrock_model(scenario, "chat", wire.url), message + ) + assert "partial output" in seen and message in seen, seen + assert len(wire.drain()) == 1 + information: Final = _failure_information(call_id) + assert information["error_provider_request_id"] == amzn_request_id, information + + +async def _consume_bedrock_sdk_stream(api: Literal["chat", "responses"], asynchronous: bool, api_base: str) -> str: + common: Final[dict[str, JsonValue]] = { + "model": _MODEL, + "stream": True, + "api_base": api_base, + "num_retries": 0, + **_BEDROCK_AWS, + } + turn: Final[list[JsonValue]] = [{"role": "user", "content": "probe"}] + match api, asynchronous: + case "chat", True: + return "".join([repr(chunk) async for chunk in await litellm.acompletion(messages=turn, **common)]) + case "chat", False: + return "".join(repr(chunk) for chunk in litellm.completion(messages=turn, **common)) + case "responses", True: + return "".join([repr(event) async for event in await litellm.aresponses(input="probe", **common)]) + case "responses", False: + return "".join(repr(event) for event in litellm.responses(input="probe", **common)) + + +@pytest.mark.parametrize("asynchronous", (False, True), ids=("sync", "async")) +@pytest.mark.parametrize("api", ("chat", "responses")) +async def test_bedrock_sdk_stream_exception_carries_amzn_request_id( + api: Literal["chat", "responses"], asynchronous: bool +) -> None: + amzn_request_id: Final = str(uuid.uuid4()) + message: Final = f"bedrock stream failure {uuid.uuid4()}" + frames: Final = (*_bedrock_output(False), _bedrock_failure_frame("exception-event", message)) + + def respond(request: Request) -> Reply: + assert unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target + return Reply( + headers={"x-amzn-RequestId": amzn_request_id, "retry-after": "7"}, + content_type=_EVENT_STREAM, + body=b"".join(frames), + ) + + with wire_server(respond) as wire: + with pytest.raises(MidStreamFallbackError) as raised: + await _consume_bedrock_sdk_stream(api, asynchronous, wire.url) + assert len(wire.drain()) == 1 + provider_error: Final = raised.value.original_exception + assert isinstance(provider_error, litellm.ServiceUnavailableError), provider_error + assert raised.value.generated_content == "partial output" + assert message in str(provider_error) + assert _get_provider_request_id(provider_error) == amzn_request_id + assert "retry-after" not in provider_error.response.headers, provider_error.response.headers + + +@pytest.mark.parametrize("api", ("chat", "responses")) +@pytest.mark.parametrize("healthy", (True, False), ids=("healthy", "failing")) +async def test_sync_bedrock_sdk_stream_with_non_ascii_request_id_header_keeps_streaming( + api: Literal["chat", "responses"], healthy: bool +) -> None: + body: Final = ( + _healthy_bedrock_stream(False) + if healthy + else b"".join((*_bedrock_output(False), _bedrock_failure_frame("exception-event", "non-ascii id failure"))) + ) + + def respond(request: Request) -> Reply: + assert unquote(request.target) in _BEDROCK_STREAM_TARGETS, request.target + return Reply(headers={"x-amzn-RequestId": "req-\u00fc-1"}, content_type=_EVENT_STREAM, body=body) + + with wire_server(respond) as wire: + if healthy: + assert "recovered" in await _consume_bedrock_sdk_stream(api, False, wire.url) + else: + with pytest.raises(MidStreamFallbackError) as raised: + await _consume_bedrock_sdk_stream(api, False, wire.url) + assert "non-ascii id failure" in str(raised.value.original_exception) + assert _get_provider_request_id(raised.value.original_exception) is None + assert len(wire.drain()) == 1 diff --git a/tests/integration/observability/test_content_filter_logging_only.py b/tests/integration/observability/test_content_filter_logging_only.py new file mode 100644 index 00000000000..ab1779c497d --- /dev/null +++ b/tests/integration/observability/test_content_filter_logging_only.py @@ -0,0 +1,712 @@ +"""`litellm_content_filter` in `mode: logging_only`, driven through a real proxy. + +The provider is the only double, an owned wire peer. The proxy, its guardrail registry, Postgres and Redis run +for real with two workers. Every check reads the spend row, which is written after the logging-only scans finish +""" + +from __future__ import annotations + +import json +import uuid +from collections.abc import Iterator +from concurrent.futures import ThreadPoolExecutor +from dataclasses import dataclass +from datetime import datetime, timezone +from hashlib import sha256 +from pathlib import Path +from typing import Final + +import anthropic +import openai +import pytest +import yaml +from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment, object_value +from integration._support.database import read_rows +from integration._support.mcp import EntryPoint, McpCaller, echo_tool, register_mcp, scripted_peer, tool_calls +from integration._support.otlp_sink import AuditConfigWriter, SpanSinks, recorded_spans +from integration._support.process import owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server + +AWS_KEY: Final = "AKIA" + "SYNTHETIC0000000" +GITHUB_TOKEN: Final = "ghp_" + "a1" * 18 +FORBIDDEN_WORD: Final = "synthetic-forbidden-word" +LEAK_MARK: Final = "SYNTHETIC-LEAK-TOKEN" +OUTAGE_MARK: Final = "SYNTHETIC-PROVIDER-OUTAGE" +OBSERVE_BLOCK: Final = "cf-observe-block" +OBSERVE_MASK: Final = "cf-observe-mask" +ENFORCE_BLOCK: Final = "cf-enforce-block" +ENFORCE_MASK: Final = "cf-enforce-mask" +OBSERVE_CATEGORY: Final = "cf-observe-category" +OBSERVE_COMPETITOR: Final = "cf-observe-competitor" +SECRETS: Final = f"AWS_ACCESS_KEY_ID={AWS_KEY}\nGITHUB_TOKEN={GITHUB_TOKEN}" + + +def _answer(marker: str) -> str: + return "synthetic answer " + marker + + +def _provider_answer(body: bytes) -> str: + text: Final = body.decode() + start: Final = text.find("mark-") + assert start >= 0, text + answer: Final = _answer(text[start : start + 37]) + return answer + " " + GITHUB_TOKEN if LEAK_MARK in text else answer + + +def _sse(events: tuple[dict[str, object], ...]) -> tuple[bytes, ...]: + return tuple(("data: " + json.dumps(event) + "\n\n").encode() for event in events) + + +def _chat_reply(answer: str, streaming: bool) -> Reply: + if not streaming: + return Reply( + body=json.dumps( + { + "id": "chatcmpl-" + uuid.uuid4().hex, + "object": "chat.completion", + "created": 1, + "model": "gpt-5.4-mini", + "choices": [ + {"index": 0, "message": {"role": "assistant", "content": answer}, "finish_reason": "stop"} + ], + "usage": {"prompt_tokens": 10, "completion_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + identity: Final = "chatcmpl-" + uuid.uuid4().hex + + def chunk(delta: dict[str, object], finish: str | None) -> dict[str, object]: + return { + "id": identity, + "object": "chat.completion.chunk", + "created": 1, + "model": "gpt-5.4-mini", + "choices": [{"index": 0, "delta": delta, "finish_reason": finish}], + } + + events: Final = (chunk({"role": "assistant", "content": answer[:9]}, None), chunk({"content": answer[9:]}, "stop")) + return Reply(chunks=(*_sse(events), b"data: [DONE]\n\n"), content_type="text/event-stream") + + +def _messages_reply(answer: str) -> Reply: + return Reply( + body=json.dumps( + { + "id": "msg_" + uuid.uuid4().hex, + "type": "message", + "role": "assistant", + "model": "claude-haiku-5-5", + "content": [{"type": "text", "text": answer}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 10, "output_tokens": 5}, + } + ).encode() + ) + + +def _responses_reply(answer: str) -> Reply: + return Reply( + body=json.dumps( + { + "id": "resp_" + uuid.uuid4().hex, + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.4-mini", + "output": [ + { + "type": "message", + "id": "msgo_" + uuid.uuid4().hex, + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": answer, "annotations": []}], + } + ], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + } + ).encode() + ) + + +def _provider(request: Request) -> Reply: + if OUTAGE_MARK.encode() in request.body: + return Reply(status=503, body=b'{"error": {"message": "synthetic provider outage", "type": "server_error"}}') + answer: Final = _provider_answer(request.body) + if request.target.endswith("/v1/messages"): + return _messages_reply(answer) + if request.target.endswith("/v1/responses"): + return _responses_reply(answer) + assert request.target.endswith("/v1/chat/completions"), request.target + return _chat_reply(answer, bool(json.loads(request.body).get("stream"))) + + +def _content_filter(name: str, mode: str, action: str, *, default_on: bool) -> dict[str, object]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "litellm_content_filter", + "mode": mode, + "default_on": default_on, + "patterns": [ + {"pattern_type": "regex", "pattern": "AKIA[0-9A-Z]{16}", "name": "aws_access_key", "action": action}, + {"pattern_type": "regex", "pattern": "ghp_[A-Za-z0-9]{36}", "name": "github_pat", "action": action}, + ], + "blocked_words": [{"keyword": FORBIDDEN_WORD, "action": action}], + }, + } + + +def _observe_only(name: str, **rules: object) -> dict[str, object]: + return { + "guardrail_name": name, + "litellm_params": {"guardrail": "litellm_content_filter", "mode": "logging_only", "default_on": False, **rules}, + } + + +def _rig_config(root: Path) -> Path: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["litellm_settings"]["cache"] = False + config["general_settings"]["store_prompts_in_spend_logs"] = True + config["guardrails"] = [ + _content_filter(OBSERVE_BLOCK, "logging_only", "BLOCK", default_on=True), + _content_filter(OBSERVE_MASK, "logging_only", "MASK", default_on=True), + _content_filter(ENFORCE_BLOCK, "pre_call", "BLOCK", default_on=False), + _content_filter(ENFORCE_MASK, "pre_call", "MASK", default_on=False), + _observe_only( + OBSERVE_CATEGORY, + categories=[ + {"category": "harmful_violence", "enabled": True, "action": "BLOCK", "severity_threshold": "medium"} + ], + blocked_words=[{"keyword": FORBIDDEN_WORD, "action": "BLOCK"}], + ), + _observe_only( + OBSERVE_COMPETITOR, + competitor_intent_config={ + "brand_self": ["emirates"], + "competitors": ["qatar airways"], + "domain_words": ["airline", "flight"], + "policy": {"competitor_comparison": "refuse"}, + "threshold_high": 0.5, + }, + ), + ] + path: Final = root / "content-filter-logging-only.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + provider: Wire + chat_model: str + claude_model: str + responses_model: str + scenario: Scenario + + def base(self) -> str: + return str(self.proxy.client.base_url).rstrip("/") + + def provider_bodies(self, marker: str) -> tuple[str, ...]: + return tuple(request.body.decode() for request in self.provider.drain() if marker.encode() in request.body) + + def entries(self, response_id: str, name: str) -> tuple[dict[str, object], ...]: + return _named_entries( + 'SELECT request_id, status, metadata FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (response_id,), name + ) + + +def _named_entries(query: str, parameters: tuple[str, ...], name: str) -> tuple[dict[str, object], ...]: + rows: Final = eventually(lambda: read_rows(query, parameters), lambda values: len(values) >= 1, seconds=70) + assert len(rows) == 1, rows + recorded: Final = object_value(rows[0]["metadata"]).get("guardrail_information") + assert isinstance(recorded, list), rows[0] + return tuple(entry for entry in (object_value(item) for item in recorded) if entry.get("guardrail_name") == name) + + +def _findings(entry: dict[str, object]) -> tuple[str, ...]: + response: Final = entry.get("guardrail_response") + assert isinstance(response, list), entry + return tuple( + str(detection.get("pattern_name") or detection.get("keyword") or detection.get("type")) + for detection in (object_value(item) for item in response) + ) + + +def _verdicts(entries: tuple[dict[str, object], ...]) -> list[tuple[object, object, tuple[str, ...]]]: + return [(entry["guardrail_mode"], entry["guardrail_status"], _findings(entry)) for entry in entries] + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + root: Final = tmp_path_factory.mktemp("content-filter-logging-only") + with ( + gateway_from_environment() as gateway, + wire_server(_provider) as provider, + owned_proxy_process(gateway, root, {}, config=_rig_config(root), workers=2) as owned, + owned.gateway.scenario() as scenario, + ): + chat: Final = scenario.model( + model="openai/gpt-5.4-mini", api_base=provider.url + "/v1", api_key="synthetic-openai-key" + ) + claude: Final = scenario.model( + model="anthropic/claude-haiku-5-5", api_base=provider.url, api_key="synthetic-anthropic-key" + ) + responses: Final = scenario.model( + model="openai/gpt-5.4-mini", api_base=provider.url + "/v1", api_key="synthetic-openai-key" + ) + yield Rig(owned.gateway, provider, chat, claude, responses, scenario) + + +def _marker() -> str: + return "mark-" + uuid.uuid4().hex + + +def _sync_chat(rig: Rig, text: str) -> tuple[str, str]: + client: Final = openai.OpenAI(base_url=rig.base() + "/v1", api_key=rig.proxy.key, max_retries=0) + completion: Final = client.chat.completions.create( + model=rig.chat_model, messages=[{"role": "user", "content": text}] + ) + return completion.id, completion.choices[0].message.content or "" + + +async def _async_chat(rig: Rig, text: str) -> tuple[str, str]: + client: Final = openai.AsyncOpenAI(base_url=rig.base() + "/v1", api_key=rig.proxy.key, max_retries=0) + completion: Final = await client.chat.completions.create( + model=rig.chat_model, messages=[{"role": "user", "content": text}] + ) + return completion.id, completion.choices[0].message.content or "" + + +def _streamed_chat(rig: Rig, text: str) -> tuple[str, str]: + client: Final = openai.OpenAI(base_url=rig.base() + "/v1", api_key=rig.proxy.key, max_retries=0) + chunks: Final = tuple( + client.chat.completions.create(model=rig.chat_model, messages=[{"role": "user", "content": text}], stream=True) + ) + return chunks[0].id, "".join(chunk.choices[0].delta.content or "" for chunk in chunks if chunk.choices) + + +def _messages(rig: Rig, text: str) -> tuple[str, str]: + client: Final = anthropic.Anthropic(base_url=rig.base(), api_key=rig.proxy.key, max_retries=0) + message: Final = client.messages.create( + model=rig.claude_model, max_tokens=64, messages=[{"role": "user", "content": text}] + ) + assert message.content[0].type == "text" + return message.id, message.content[0].text + + +def _responses(rig: Rig, text: str) -> tuple[str, str]: + client: Final = openai.OpenAI(base_url=rig.base() + "/v1", api_key=rig.proxy.key, max_retries=0) + result: Final = client.responses.create(model=rig.responses_model, input=text) + return result.id, result.output_text + + +SYNC_CALLS: Final = {"chat": _sync_chat, "chat_stream": _streamed_chat, "messages": _messages, "responses": _responses} +EVERY_FINDING: Final = ("aws_access_key", "github_pat", FORBIDDEN_WORD) + + +@pytest.mark.parametrize("endpoint", list(SYNC_CALLS)) +def test_logging_only_block_records_every_match_without_touching_the_call(rig: Rig, endpoint: str) -> None: + marker: Final = _marker() + prompt: Final = f"{marker} what did the file contain?\n{SECRETS}\n{FORBIDDEN_WORD}" + response_id, answer = SYNC_CALLS[endpoint](rig, prompt) + assert answer == _answer(marker) + sent: Final = rig.provider_bodies(marker) + assert len(sent) == 1 and AWS_KEY in sent[0] and GITHUB_TOKEN in sent[0] and FORBIDDEN_WORD in sent[0], sent + + assert _verdicts(rig.entries(response_id, OBSERVE_BLOCK)) == [ + ("logging_only", "guardrail_flagged", EVERY_FINDING), + ("logging_only", "success", ()), + ] + + +@pytest.mark.asyncio +async def test_logging_only_block_records_every_match_for_an_async_client(rig: Rig) -> None: + marker: Final = _marker() + response_id, answer = await _async_chat(rig, f"{marker} {SECRETS}") + assert answer == _answer(marker) + + assert _verdicts(rig.entries(response_id, OBSERVE_BLOCK)) == [ + ("logging_only", "guardrail_flagged", ("aws_access_key", "github_pat")), + ("logging_only", "success", ()), + ] + + +def test_logging_only_block_still_scans_the_response_after_a_request_match(rig: Rig) -> None: + marker: Final = _marker() + response_id, answer = _sync_chat(rig, f"{marker} {LEAK_MARK} key {AWS_KEY}") + assert answer == _answer(marker) + " " + GITHUB_TOKEN + + assert _verdicts(rig.entries(response_id, OBSERVE_BLOCK)) == [ + ("logging_only", "guardrail_flagged", ("aws_access_key",)), + ("logging_only", "guardrail_flagged", ("github_pat",)), + ] + + +def _tool_call_history(text: str, arguments: str) -> list[dict[str, object]]: + return [ + {"role": "user", "content": text}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": arguments}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "done"}, + ] + + +@pytest.mark.parametrize( + ("arguments", "request_findings"), + ( + pytest.param(json.dumps({AWS_KEY: 1}), ("aws_access_key",), id="masked-argument-key"), + pytest.param("[" * 150 + "1" + "]" * 150, (), id="too-deep"), + ), +) +def test_logging_only_flags_tool_arguments_an_enforcing_mode_rejects_and_keeps_scanning( + rig: Rig, arguments: str, request_findings: tuple[str, ...] +) -> None: + marker: Final = _marker() + messages: Final = _tool_call_history(f"{marker} {LEAK_MARK}", arguments) + headers: Final = {"Authorization": f"Bearer {rig.proxy.key}"} + enforced: Final = rig.proxy.client.post( + "/v1/chat/completions", + json={"model": rig.chat_model, "messages": messages, "guardrails": [ENFORCE_MASK]}, + headers=headers, + ) + assert enforced.status_code == 400, enforced.text + assert rig.provider_bodies(marker) == () + + observed: Final = rig.proxy.client.post( + "/v1/chat/completions", json={"model": rig.chat_model, "messages": messages}, headers=headers + ) + assert observed.status_code == 200, observed.text + completion: Final = object_value(observed.json()) + assert len(rig.provider_bodies(marker)) == 1 + + assert _verdicts(rig.entries(str(completion["id"]), OBSERVE_MASK)) == [ + ("logging_only", "guardrail_flagged", request_findings), + ("logging_only", "success", ("github_pat",)), + ] + + +@pytest.mark.parametrize("entry", ("mcp", "rest")) +def test_logging_only_block_records_every_match_on_an_mcp_tool_call(rig: Rig, entry: EntryPoint) -> None: + alias: Final = "cfmcp" + uuid.uuid4().hex[:8] + arguments: Final = {"note": f"{_marker()} {SECRETS}"} + with scripted_peer(echo_tool("echo")) as peer: + server_id: Final = register_mcp(rig.scenario, peer, alias) + key: Final = rig.scenario.key(object_permission={"mcp_servers": [server_id]}) + outcome: Final = McpCaller(rig.proxy, key, entry, alias).call(f"{alias}-echo", arguments, server_id=server_id) + assert (outcome.error, outcome.text) == (None, json.dumps(arguments, sort_keys=True)), outcome.raw + assert [call["body"]["params"]["arguments"] for call in tool_calls(peer.drain())] == [arguments] + + entries: Final = _named_entries( + 'SELECT metadata FROM "LiteLLM_SpendLogs" WHERE api_key = %s AND call_type = %s', + (sha256(key.encode()).hexdigest(), "call_mcp_tool"), + OBSERVE_BLOCK, + ) + assert _verdicts(entries) == [("logging_only", "guardrail_flagged", ("aws_access_key", "github_pat"))], entries + + +@pytest.mark.parametrize( + ("guardrail", "prompt", "findings"), + ( + pytest.param( + OBSERVE_CATEGORY, + f"how to kill someone {FORBIDDEN_WORD}", + ("how to kill someone", FORBIDDEN_WORD), + id="category-keyword", + ), + pytest.param( + OBSERVE_COMPETITOR, "Is Qatar Airways better than Emirates?", ("competitor_intent",), id="competitor-intent" + ), + ), +) +def test_logging_only_block_flags_category_and_competitor_matches_without_blocking( + rig: Rig, guardrail: str, prompt: str, findings: tuple[str, ...] +) -> None: + marker: Final = _marker() + response: Final = rig.proxy.client.post( + "/v1/chat/completions", + json={ + "model": rig.chat_model, + "messages": [{"role": "user", "content": f"{marker} {prompt}"}], + "guardrails": [guardrail], + }, + headers={"Authorization": f"Bearer {rig.proxy.key}"}, + ) + assert response.status_code == 200, response.text + completion: Final = object_value(response.json()) + assert len(rig.provider_bodies(marker)) == 1 + + assert _verdicts(rig.entries(str(completion["id"]), guardrail)) == [ + ("logging_only", "guardrail_flagged", findings), + ("logging_only", "success", ()), + ] + + +def test_logging_only_mask_records_every_match(rig: Rig) -> None: + marker: Final = _marker() + response_id, answer = _sync_chat(rig, f"{marker} {SECRETS}\n{FORBIDDEN_WORD}") + assert answer == _answer(marker) + + entries: Final = rig.entries(response_id, OBSERVE_MASK) + assert _verdicts(entries) == [ + ("logging_only", "success", EVERY_FINDING), + ("logging_only", "success", ()), + ] + assert entries[0]["masked_entity_count"] == {"aws_access_key": 1, "github_pat": 1, "blocked_word": 1}, entries + + +def test_logging_only_records_nothing_for_a_clean_call(rig: Rig) -> None: + marker: Final = _marker() + response_id, answer = _sync_chat(rig, f"{marker} nothing secret here") + assert answer == _answer(marker) + + assert _verdicts(rig.entries(response_id, OBSERVE_BLOCK)) == [ + ("logging_only", "success", ()), + ("logging_only", "success", ()), + ] + + +def test_pre_call_block_still_rejects_the_request(rig: Rig) -> None: + marker: Final = _marker() + response: Final = rig.proxy.client.post( + "/v1/chat/completions", + json={ + "model": rig.chat_model, + "messages": [{"role": "user", "content": f"{marker} {SECRETS}"}], + "guardrails": [ENFORCE_BLOCK], + }, + headers={"Authorization": f"Bearer {rig.proxy.key}"}, + ) + assert response.status_code == 400, response.text + assert "aws_access_key" in response.text + assert rig.provider_bodies(marker) == () + + +def test_apply_guardrail_endpoint_still_blocks_with_a_logging_only_content_filter(rig: Rig) -> None: + response: Final = rig.proxy.client.post( + "/guardrails/apply_guardrail", + json={"guardrail_name": OBSERVE_BLOCK, "text": f"key {AWS_KEY}"}, + headers={"Authorization": f"Bearer {rig.proxy.key}"}, + ) + assert response.status_code == 400, response.text + assert "aws_access_key" in response.text + + +def _advertised_modes(rig: Rig) -> dict[str, object]: + response: Final = rig.proxy.client.get( + "/guardrails/ui/add_guardrail_settings", headers={"Authorization": f"Bearer {rig.proxy.key}"} + ) + assert response.status_code == 200, response.text + return object_value(response.json()["supported_modes_by_provider"]) + + +@pytest.mark.parametrize("provider", ["litellm_content_filter", "generic_guardrail_api", "bedrock", "akto"]) +def test_dashboard_offers_logging_only_where_the_config_accepts_it(rig: Rig, provider: str) -> None: + modes: Final = _advertised_modes(rig) + offered: Final = modes[provider] + assert isinstance(offered, list) and "logging_only" in offered, modes + assert offered.count("logging_only") == 1, offered + + +@pytest.mark.parametrize("provider", ["azure/prompt_shield", "lakera_v2"]) +def test_dashboard_does_not_offer_logging_only_where_the_config_rejects_it(rig: Rig, provider: str) -> None: + offered: Final = _advertised_modes(rig)[provider] + assert isinstance(offered, list) and "logging_only" not in offered, offered + + +def test_a_logging_only_content_filter_created_through_the_api_records_every_match(rig: Rig) -> None: + name: Final = "cf-observe-api-" + uuid.uuid4().hex[:8] + created: Final = rig.proxy.client.post( + "/guardrails", + json={"guardrail": {"guardrail_name": name, **_content_filter(name, "logging_only", "BLOCK", default_on=True)}}, + headers={"Authorization": f"Bearer {rig.proxy.key}"}, + ) + assert created.status_code == 200, created.text + guardrail_id: Final = str(object_value(created.json())["guardrail_id"]) + try: + marker: Final = _marker() + response_id: Final = eventually( + lambda: _sync_chat(rig, f"{marker} {SECRETS}")[0], + lambda identity: bool(rig.entries(identity, name)), + seconds=70, + ) + assert _verdicts(rig.entries(response_id, name))[0] == ( + "logging_only", + "guardrail_flagged", + ("aws_access_key", "github_pat"), + ) + finally: + rig.proxy.client.delete(f"/guardrails/{guardrail_id}", headers={"Authorization": f"Bearer {rig.proxy.key}"}) + + +def test_guardrails_monitor_counts_observe_only_matches_as_flagged_not_blocked(rig: Rig) -> None: + name: Final = "cf-observe-usage-" + uuid.uuid4().hex[:8] + headers: Final = {"Authorization": f"Bearer {rig.proxy.key}"} + created: Final = rig.proxy.client.post( + "/guardrails", + json={"guardrail": {"guardrail_name": name, **_content_filter(name, "logging_only", "BLOCK", default_on=True)}}, + headers=headers, + ) + assert created.status_code == 200, created.text + guardrail_id: Final = str(object_value(created.json())["guardrail_id"]) + try: + eventually( + lambda: _sync_chat(rig, f"{_marker()} {SECRETS}")[0], + lambda identity: bool(rig.entries(identity, name)), + seconds=70, + ) + today: Final = datetime.now(timezone.utc).date().isoformat() + detail: Final = eventually( + lambda: object_value( + rig.proxy.client.get( + f"/guardrails/usage/detail/{guardrail_id}", + params={"start_date": today, "end_date": today}, + headers=headers, + ).json() + ), + lambda body: int(str(body.get("requestsEvaluated", 0))) >= 1, + seconds=60, + ) + assert (detail["failRate"], detail["status"]) == (0.0, "healthy"), detail + finally: + rig.proxy.client.delete(f"/guardrails/{guardrail_id}", headers=headers) + + +@pytest.mark.parametrize( + ("environment", "name_key", "status_key"), + ( + pytest.param({}, "guardrail_name", "guardrail_status", id="otel-v1"), + pytest.param({"LITELLM_OTEL_V2": "1"}, "litellm.guardrail.name", "litellm.guardrail.status", id="otel-v2"), + ), +) +def test_otel_guardrail_spans_carry_the_flagged_status_without_an_error( + gateway: Gateway, + audit_sinks: SpanSinks, + otel_audit_config: AuditConfigWriter, + tmp_path: Path, + environment: dict[str, str], + name_key: str, + status_key: str, +) -> None: + config_path: Final = otel_audit_config(tmp_path, {"cache": False}) + config: Final = yaml.safe_load(config_path.read_text()) + config["guardrails"] = [_content_filter(OBSERVE_BLOCK, "logging_only", "BLOCK", default_on=True)] + config_path.write_text(yaml.safe_dump(config)) + overrides: Final = {"OTEL_EXPORTER": "http/json", "OTEL_ENDPOINT": audit_sinks.operator, **environment} + with ( + wire_server(_provider) as provider, + owned_proxy_process(gateway, tmp_path, overrides, config=config_path, workers=2) as owned, + owned.gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model="openai/gpt-5.4-mini", api_base=provider.url + "/v1", api_key="synthetic-openai-key" + ) + start, _ = recorded_spans(audit_sinks.operator) + marker: Final = _marker() + client: Final = openai.OpenAI( + base_url=str(owned.gateway.client.base_url).rstrip("/") + "/v1", api_key=owned.gateway.key, max_retries=0 + ) + completion: Final = client.chat.completions.create( + model=model, messages=[{"role": "user", "content": f"{marker} {SECRETS}"}] + ) + assert completion.choices[0].message.content == _answer(marker) + + spans: Final = eventually( + lambda: tuple( + span + for span in recorded_spans(audit_sinks.operator, start)[1] + if span["attributes"].get(name_key) == OBSERVE_BLOCK + ), + lambda found: len(found) >= 2, + seconds=60, + return_last_on_timeout=True, + ) + assert sorted( + (str(span["attributes"].get(status_key)), "error.type" in span["attributes"]) for span in spans + ) == [("guardrail_flagged", False), ("success", False)], spans + + +def _burst_call(rig: Rig, index: int) -> tuple[str, int, str, str]: + marker: Final = _marker() + if index % 4 == 0: + enforced: Final = rig.proxy.client.post( + "/v1/chat/completions", + json={ + "model": rig.chat_model, + "messages": [{"role": "user", "content": f"{marker} {SECRETS}"}], + "guardrails": [ENFORCE_BLOCK], + }, + headers={"Authorization": f"Bearer {rig.proxy.key}"}, + ) + return "enforce", enforced.status_code, marker, "" + if index % 4 == 1: + failed: Final = rig.proxy.client.post( + "/v1/chat/completions", + json={ + "model": rig.chat_model, + "messages": [{"role": "user", "content": f"{marker} {OUTAGE_MARK} {SECRETS}"}], + }, + headers={"Authorization": f"Bearer {rig.proxy.key}"}, + ) + return "outage", failed.status_code, marker, failed.headers.get("x-litellm-call-id", "") + endpoint: Final = list(SYNC_CALLS)[index % len(SYNC_CALLS)] + response_id, answer = SYNC_CALLS[endpoint](rig, f"{marker} {SECRETS}\n{FORBIDDEN_WORD}") + assert answer == _answer(marker), (endpoint, answer) + return "observe:" + endpoint, 200, marker, response_id + + +def _failure_row(call_id: str) -> tuple[object, tuple[dict[str, object], ...]]: + rows: Final = eventually( + lambda: read_rows('SELECT status, metadata FROM "LiteLLM_SpendLogs" WHERE request_id = %s', (call_id,)), + lambda values: len(values) >= 1, + seconds=70, + ) + assert len(rows) == 1, rows + recorded: Final = object_value(rows[0]["metadata"]).get("guardrail_information") or [] + assert isinstance(recorded, list), rows[0] + return rows[0]["status"], tuple( + entry for entry in (object_value(item) for item in recorded) if entry.get("guardrail_name") == OBSERVE_BLOCK + ) + + +def test_concurrent_burst_with_a_provider_outage_records_every_observed_call_once(rig: Rig) -> None: + def call(index: int) -> tuple[str, int, str, str]: + return _burst_call(rig, index) + + with ThreadPoolExecutor(max_workers=12) as pool: + results: Final = tuple(pool.map(call, range(32))) + + assert sorted({(kind.split(":")[0], status) for kind, status, _, _ in results}) == [ + ("enforce", 400), + ("observe", 200), + ("outage", 503), + ], results + sent: Final = rig.provider_bodies("mark-") + for kind, _, marker, _ in results: + if kind == "enforce": + assert not any(marker in body for body in sent), marker + outages: Final = tuple(call_id for kind, _, _, call_id in results if kind == "outage") + assert len(outages) == 8 and all(outages), results + assert [_failure_row(call_id) for call_id in outages] == [("failure", ())] * 8 + observed: Final = tuple((kind, response_id) for kind, _, _, response_id in results if kind.startswith("observe")) + assert len(observed) == 16 + for kind, response_id in observed: + assert _verdicts(rig.entries(response_id, OBSERVE_BLOCK)) == [ + ("logging_only", "guardrail_flagged", EVERY_FINDING), + ("logging_only", "success", ()), + ], (kind, response_id) + + marker: Final = _marker() + response_id, answer = _sync_chat(rig, f"{marker} {SECRETS}") + assert answer == _answer(marker) + assert _verdicts(rig.entries(response_id, OBSERVE_BLOCK))[0] == ( + "logging_only", + "guardrail_flagged", + ("aws_access_key", "github_pat"), + ) diff --git a/tests/integration/observability/test_logging_only_scope_config.py b/tests/integration/observability/test_logging_only_scope_config.py index c8c765a68a2..354e3b464fe 100644 --- a/tests/integration/observability/test_logging_only_scope_config.py +++ b/tests/integration/observability/test_logging_only_scope_config.py @@ -179,16 +179,16 @@ def test_native_content_filter_scope_logs_without_blocking( ( identity, "logging_only", - "guardrail_intervened" if direction == blocked_side else "success", + "guardrail_flagged" if direction == blocked_side else "success", ) for direction in expected_directions ), entries - intervened_entries: Final = tuple( - entry for entry in entries if entry["guardrail_status"] == "guardrail_intervened" + flagged_entries: Final = tuple( + entry for entry in entries if entry["guardrail_status"] == "guardrail_flagged" ) - assert len(intervened_entries) == 1, entries - assert intervened_entries[0]["guardrail_response"] is not None, entries - assert intervened_entries[0]["guardrail_response"] == "REDACTED_BY_LITELM", entries + assert len(flagged_entries) == 1, entries + assert flagged_entries[0]["guardrail_response"] is not None, entries + assert flagged_entries[0]["guardrail_response"] == "REDACTED_BY_LITELM", entries finally: delete_scenario(upstream_handle) diff --git a/tests/integration/observability/test_otel_excluded_services_matrix.py b/tests/integration/observability/test_otel_excluded_services_matrix.py index 07646a0b189..8a7236f8626 100644 --- a/tests/integration/observability/test_otel_excluded_services_matrix.py +++ b/tests/integration/observability/test_otel_excluded_services_matrix.py @@ -358,6 +358,10 @@ def _has_root(spans: tuple[Span, ...]) -> bool: return any(span["kind"] == SERVER for span in spans) +def _has_model_span(spans: tuple[Span, ...]) -> bool: + return any("gen_ai.operation.name" in span["attributes"] for span in spans) + + def _trace_of_call(sink: str, call_id: str, since: int) -> tuple[Span, ...]: _, spans = recorded_spans(sink, since) traces: Final = {span["trace_id"] for span in spans if span["attributes"].get("litellm.call_id") == call_id} @@ -367,7 +371,7 @@ def _trace_of_call(sink: str, call_id: str, since: int) -> tuple[Span, ...]: def _operator_trace(rig: Rig, sent: Sent, cursors: Cursors) -> tuple[Span, ...]: trace: Final = eventually( lambda: _trace_of_call(rig.sinks.operator, sent.call_id, cursors.operator), - lambda spans: _has_root(spans) and "redis" in _db_systems(spans), + lambda spans: _has_root(spans) and "redis" in _db_systems(spans) and _has_model_span(spans), seconds=40, ) assert len({span["trace_id"] for span in trace}) == 1, _names(trace) @@ -405,7 +409,7 @@ def _assert_tenant_mirrors(rig: Rig, operator: tuple[Span, ...], cursors: Cursor def _assert_withheld(rig: Rig, sent: Sent, cursors: Cursors) -> tuple[Span, ...]: tenant: Final = _assert_tenant_mirrors(rig, _operator_trace(rig, sent, cursors), cursors) - assert any("gen_ai.operation.name" in span["attributes"] for span in tenant), _names(tenant) + assert _has_model_span(tenant), _names(tenant) return tenant @@ -719,7 +723,7 @@ def _assert_tenant_kept( return ( sum(1 for span in spans if span["kind"] == SERVER) == 1 and "redis" in _db_systems(spans) - and (not needs_model_span or any("gen_ai.operation.name" in span["attributes"] for span in spans)) + and (not needs_model_span or _has_model_span(spans)) ) tenant: Final = eventually( @@ -730,7 +734,7 @@ def _assert_tenant_kept( ) assert sum(1 for span in tenant if span["kind"] == SERVER) == 1, _names(tenant) assert "redis" in _db_systems(tenant), f"redis spans missing at the tenant: {_names(tenant)}" - assert not needs_model_span or any("gen_ai.operation.name" in span["attributes"] for span in tenant), _names(tenant) + assert not needs_model_span or _has_model_span(tenant), _names(tenant) return tenant diff --git a/tests/integration/providers/test_anthropic_empty_user_turn_wire.py b/tests/integration/providers/test_anthropic_empty_user_turn_wire.py new file mode 100644 index 00000000000..654a88c60bb --- /dev/null +++ b/tests/integration/providers/test_anthropic_empty_user_turn_wire.py @@ -0,0 +1,216 @@ +import json +from collections.abc import Iterator, Mapping, Sequence +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import pytest +import yaml +from integration._support.client import Gateway, gateway_from_environment, object_value +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._cache_control_marks_support import ( + anthropic_deployment, + anthropic_peer, + final_text, + new_marker, +) +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MODEL: Final = "empty-user-turn" +_PLACEHOLDER: Final = "[System: Empty message content sanitised to satisfy protocol]" +_REPLY: Final = "Noted." +_EMPTY_ROWS: Final[Mapping[str, Mapping[str, JsonValue]]] = MappingProxyType( + { + "list": MappingProxyType({"role": "user", "content": []}), + "null": MappingProxyType({"role": "user", "content": None}), + "missing": MappingProxyType({"role": "user"}), + "empty-string": MappingProxyType({"role": "user", "content": ""}), + } +) +_CHAT_EMPTY_KINDS: Final = ("list", "null", "missing") + +Turn = tuple[str, tuple[str, ...]] + + +@dataclass(frozen=True, slots=True) +class _Rig: + gateway: Gateway + provider: Wire + + +def _empty_row(kind: str) -> dict[str, JsonValue]: + return dict(_EMPTY_ROWS[kind]) + + +def _guardrail_peer(request: Request) -> Reply: + kind: Final = request.target.split("/")[1] + rows: Final = _JSON_OBJECT.validate_json(request.body)["structured_messages"] + assert isinstance(rows, list), rows + verdict: Final = {"action": "GUARDRAIL_INTERVENED", "structured_messages": [*rows, _empty_row(kind)]} + return Reply(body=json.dumps(verdict).encode()) + + +def _guardrail(kind: str, api_base: str) -> dict[str, JsonValue]: + return { + "guardrail_name": f"append-{kind}", + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": "pre_call", + "default_on": False, + "api_base": f"{api_base}/{kind}", + "api_key": "synthetic-guardrail-key", + }, + } + + +def _config(directory: Path, provider_url: str, guardrail_url: str) -> Path: + stock: Final = _JSON_OBJECT.validate_python(yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text())) + config: Final[dict[str, JsonValue]] = { + **stock, + "model_list": [anthropic_deployment(_MODEL, provider_url)], + "router_settings": {**object_value(stock["router_settings"]), "num_retries": 0}, + "guardrails": [_guardrail(kind, guardrail_url) for kind in _EMPTY_ROWS], + } + path: Final = directory / "empty-user-turn.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("empty-user-turn") + with ( + wire_server(anthropic_peer) as provider, + wire_server(_guardrail_peer) as guardrail, + gateway_from_environment() as gateway, + ): + config: Final = _config(directory, provider.url, guardrail.url) + with owned_proxy(gateway, directory, {}, config=config, workers=2) as owned: + yield _Rig(owned, provider) + + +def _block(block: JsonValue) -> str: + part: Final = object_value(block) + return str(part["text"]) if part.get("type") == "text" else str(part.get("type")) + + +def _turn(message: JsonValue) -> Turn: + row: Final = object_value(message) + content: Final = row["content"] + if isinstance(content, str): + return (str(row["role"]), (content,)) + assert isinstance(content, list), row + return (str(row["role"]), tuple(_block(block) for block in content)) + + +def _send(rig: _Rig, path: str, body: Mapping[str, JsonValue]) -> None: + response: Final = rig.gateway.client.post( + path, json=dict(body), headers={"Authorization": f"Bearer {rig.gateway.key}"} + ) + assert response.status_code == 200, response.text + + +def _provider_turns(rig: _Rig, marker: str) -> tuple[Turn, ...]: + needle: Final = marker.encode() + requests: Final = [request for request in rig.provider.drain() if needle in request.body] + assert len(requests) == 1, requests + messages: Final = _JSON_OBJECT.validate_json(requests[0].body)["messages"] + assert isinstance(messages, list), messages + return tuple(_turn(message) for message in messages) + + +def _ask_row(ask: str) -> dict[str, JsonValue]: + return {"role": "user", "content": ask} + + +def _chat(messages: Sequence[JsonValue]) -> dict[str, JsonValue]: + return {"model": _MODEL, "max_tokens": 64, "messages": list(messages)} + + +@pytest.mark.parametrize("kind", _CHAT_EMPTY_KINDS) +def test_chat_trailing_empty_user_message_keeps_a_user_turn(rig: _Rig, kind: str) -> None: + marker: Final = new_marker() + ask: Final = final_text(marker) + _send( + rig, + "/v1/chat/completions", + _chat([_ask_row(ask), {"role": "assistant", "content": _REPLY}, _empty_row(kind)]), + ) + assert _provider_turns(rig, marker) == (("user", (ask,)), ("assistant", (_REPLY,)), ("user", (_PLACEHOLDER,))) + + +@pytest.mark.parametrize("kind", _CHAT_EMPTY_KINDS) +def test_chat_empty_user_message_beside_text_sends_only_the_text(rig: _Rig, kind: str) -> None: + marker: Final = new_marker() + ask: Final = final_text(marker) + _send(rig, "/v1/chat/completions", _chat([_ask_row(ask), _empty_row(kind)])) + assert _provider_turns(rig, marker) == (("user", (ask,)),) + + +def _system_then_empty_turns(rig: _Rig, kind: str) -> tuple[Turn, ...]: + marker: Final = new_marker() + ask: Final = final_text(marker) + _send( + rig, + "/v1/chat/completions", + _chat( + [ + _ask_row(ask), + {"role": "assistant", "content": _REPLY}, + {"role": "system", "content": "Answer in French."}, + _empty_row(kind), + ] + ), + ) + turns: Final = _provider_turns(rig, marker) + assert turns[0] == ("user", (ask,)), turns + return turns[1:] + + +@pytest.mark.parametrize("kind", _CHAT_EMPTY_KINDS) +def test_chat_system_message_before_an_empty_user_message_is_placed_like_an_empty_string(rig: _Rig, kind: str) -> None: + assert _system_then_empty_turns(rig, kind) == _system_then_empty_turns(rig, "empty-string") + + +def test_responses_trailing_empty_input_message_keeps_a_user_turn(rig: _Rig) -> None: + marker: Final = new_marker() + ask: Final = final_text(marker) + body: Final[dict[str, JsonValue]] = { + "model": _MODEL, + "max_output_tokens": 64, + "input": [ + _ask_row(ask), + {"role": "assistant", "content": _REPLY}, + {"role": "user", "content": []}, + ], + } + _send(rig, "/v1/responses", body) + assert _provider_turns(rig, marker) == (("user", (ask,)), ("assistant", (_REPLY,)), ("user", (_PLACEHOLDER,))) + + +@pytest.mark.parametrize("kind", tuple(_EMPTY_ROWS)) +def test_messages_guardrail_empty_user_row_is_sent_like_an_empty_string(rig: _Rig, kind: str) -> None: + marker: Final = new_marker() + ask: Final = final_text(marker) + body: Final[dict[str, JsonValue]] = { + "model": _MODEL, + "max_tokens": 64, + "messages": [_ask_row(ask)], + "guardrails": [f"append-{kind}"], + } + _send(rig, "/v1/messages", body) + assert _provider_turns(rig, marker) == (("user", (ask,)), ("user", (_PLACEHOLDER,))) + + +def test_a_long_run_of_empty_user_messages_is_sent_as_one_text_turn(rig: _Rig) -> None: + marker: Final = new_marker() + ask: Final = final_text(marker) + _send( + rig, "/v1/chat/completions", _chat([*({"role": "user", "content": []} for _ in range(10_000)), _ask_row(ask)]) + ) + assert _provider_turns(rig, marker) == (("user", (ask,)),) diff --git a/tests/integration/providers/test_bedrock_passthrough_aws_errors.py b/tests/integration/providers/test_bedrock_passthrough_aws_errors.py new file mode 100644 index 00000000000..bb87b15b363 --- /dev/null +++ b/tests/integration/providers/test_bedrock_passthrough_aws_errors.py @@ -0,0 +1,443 @@ +"""Errors on the Bedrock passthrough reach the AWS SDK with their message (LIT-9617). + +boto3 reads an error's code from the ``x-amzn-errortype`` header and its message from the body's ``message``. A +guardrail block on ``/bedrock/model//...`` that happens before any frame was sent reaches the caller as a +``ValidationException`` carrying the guardrail's message, and an error Bedrock itself returned, such as a 429 +``ThrottlingException``, reaches the caller with Bedrock's status, ``x-amzn-errortype`` and body.""" + +from __future__ import annotations + +import json +import os +import uuid +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from pathlib import Path +from types import MappingProxyType +from typing import Final + +import boto3 +import httpx +import pytest +import yaml +from botocore.config import Config +from botocore.exceptions import ClientError +from integration._support.budget import case_timeout_seconds +from integration._support.client import Gateway, JsonValue +from integration._support.process import owned_proxy_process +from integration._support.upstream import aws_event_stream_frame +from integration._support.wire import Reply, Request, Wire, wire_server + +MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0" +ROUTER_MODEL: Final = "aws-errors" +EVENT_STREAM: Final = "application/vnd.amazon.eventstream" +POST_CALL: Final = "aws-errors-post-call" +PRE_CALL: Final = "aws-errors-pre-call" +BEDROCK: Final = "aws-errors-bedrock" +VENDOR_DOWN: Final = "aws-errors-vendor-down" +CATO: Final = "aws-errors-cato" +UNGUARDED: Final = "unguarded" +OVER_BUDGET: Final = "over-budget" +ONE_RPM: Final = "one-rpm" +BLOCKED_WORD: Final = "FORBIDDEN" +BLOCK_REASON: Final = "synthetic policy says no" +BEDROCK_BLOCK_MESSAGE: Final = "Violated guardrail policy" +THROTTLE_MESSAGE: Final = "Too many requests, please wait before trying again." +THROTTLE_TYPE: Final = "ThrottlingException:http://internal.amazon.com/coral/com.amazon.bedrock/" +INVALID_MESSAGE: Final = "The provided model identifier is invalid." +INVALID_TYPE: Final = "ValidationException:http://internal.amazon.com/coral/com.amazon.bedrock/" +ANTHROPIC_THROTTLE: Final = b'{"type":"error","error":{"type":"rate_limit_error","message":"slow down"}}' + + +@dataclass(frozen=True, slots=True) +class UpstreamError: + status: int + errortype: str + message: str + + @property + def body(self) -> bytes: + return json.dumps({"message": self.message}).encode() + + @property + def code(self) -> str: + return self.errortype.split(":", 1)[0] + + +THROTTLED: Final = UpstreamError(429, THROTTLE_TYPE, THROTTLE_MESSAGE) +INVALID: Final = UpstreamError(400, INVALID_TYPE, INVALID_MESSAGE) +UPSTREAM_ERRORS: Final = MappingProxyType({"throttled": THROTTLED, "invalid": INVALID}) +ROUTES: Final = (pytest.param(ROUTER_MODEL, id="router"), pytest.param(MODEL_ID, id="direct")) + + +def upstream_marker(kind: str) -> str: + """A prompt the scripted Bedrock answers with the ``UPSTREAM_ERRORS`` entry named ``kind``.""" + return f"upstream-{kind}-{uuid.uuid4().hex}" + + +def upstream_error_of(marker: str) -> UpstreamError | None: + prefix, _, rest = marker.partition("-") + return UPSTREAM_ERRORS.get(rest.partition("-")[0]) if prefix == "upstream" else None + + +def event_frame(event_type: str, payload: Mapping[str, JsonValue]) -> bytes: + return aws_event_stream_frame( + {":event-type": event_type, ":content-type": "application/json", ":message-type": "event"}, + json.dumps(payload, separators=(",", ":")).encode(), + ) + + +def converse_frames(text: str) -> tuple[bytes, ...]: + return ( + event_frame("messageStart", {"role": "assistant"}), + event_frame("contentBlockDelta", {"contentBlockIndex": 0, "delta": {"text": text}}), + event_frame("contentBlockStop", {"contentBlockIndex": 0}), + event_frame("messageStop", {"stopReason": "end_turn"}), + event_frame("metadata", {"usage": {"inputTokens": 9, "outputTokens": 4, "totalTokens": 13}, "metrics": {}}), + ) + + +def converse_reply(text: str) -> bytes: + return json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": text}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 9, "outputTokens": 4, "totalTokens": 13}, + "metrics": {"latencyMs": 1}, + } + ).encode() + + +def invoke_body(marker: str) -> str: + return json.dumps( + { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 16, + "messages": [{"role": "user", "content": [{"type": "text", "text": marker}]}], + } + ) + + +def marker_of(request: Request) -> str: + return str(json.loads(request.body)["messages"][0]["content"][0]["text"]) + + +def bedrock_texts(request: Request) -> tuple[str, ...]: + return tuple(str(item["text"]["text"]) for item in json.loads(request.body)["content"]) + + +@dataclass(frozen=True, slots=True) +class Rig: + gateway: Gateway + provider: Wire + vendor: Wire + keys: Mapping[str, str] + + def provider_calls(self, marker: str) -> int: + return sum(1 for request in self.provider.drain() if marker.encode() in request.body) + + def client(self, guardrail: str, attempts: int = 1): # noqa: ANN201 # boto3 clients are untyped + client: Final = boto3.client( + "bedrock-runtime", + region_name="us-east-1", + endpoint_url=str(self.gateway.client.base_url).rstrip("/") + "/bedrock", + aws_access_key_id="AKIASYNTHETICCALLER", + aws_secret_access_key="synthetic-caller-secret", + config=Config( + read_timeout=case_timeout_seconds(), + retries={"total_max_attempts": attempts, "mode": "standard"}, + ), + ) + key: Final = self.keys[guardrail] + + def authenticate(request: object, **_: object) -> None: + getattr(request, "headers")["x-litellm-api-key"] = key + + client.meta.events.register("before-send.bedrock-runtime.*", authenticate) + return client + + def call( + self, guardrail: str, action: str, marker: str, attempts: int = 1, model: str = ROUTER_MODEL + ) -> ClientError | None: + client: Final = self.client(guardrail, attempts) + messages: Final = [{"role": "user", "content": [{"text": marker}]}] + try: + if action == "converse-stream": + for _ in client.converse_stream(modelId=model, messages=messages)["stream"]: + pass + elif action == "invoke": + client.invoke_model(modelId=model, body=invoke_body(marker)) + else: + client.converse(modelId=model, messages=messages) + except ClientError as error: + return error + return None + + def raw(self, guardrail: str, action: str, marker: str) -> httpx.Response: + return self.gateway.client.post( + f"/bedrock/model/{ROUTER_MODEL}/{action}", + json={"messages": [{"role": "user", "content": [{"text": marker}]}]}, + headers={"x-litellm-api-key": self.keys[guardrail]}, + timeout=case_timeout_seconds(), + ) + + +def _generic_guardrail(name: str, mode: str, vendor: Wire, **settings: JsonValue) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "generic_guardrail_api", + "mode": mode, + "default_on": False, + "api_base": f"{vendor.url}/{name}", + "api_key": "synthetic-guardrail-key", + **settings, + }, + } + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + def provider(request: Request) -> Reply: + if request.target == "/v1/messages": + return Reply(status=429, body=ANTHROPIC_THROTTLE) + assert request.target in tuple( + f"/model/{MODEL_ID}/{action}" for action in ("converse-stream", "converse", "invoke") + ), request.target + marker: Final = marker_of(request) + error: Final = upstream_error_of(marker) + if error is not None: + return Reply(status=error.status, body=error.body, headers={"x-amzn-errortype": error.errortype}) + text: Final = f"The answer to {marker} is {BLOCKED_WORD}" + if request.target.endswith("/converse-stream"): + return Reply(content_type=EVENT_STREAM, chunks=converse_frames(text)) + return Reply(body=converse_reply(text)) + + def vendor(request: Request) -> Reply: + if request.target.startswith(f"/{BEDROCK}/"): + if not any(BLOCKED_WORD in text for text in bedrock_texts(request)): + return Reply(body=json.dumps({"action": "NONE", "outputs": [], "assessments": []}).encode()) + return Reply( + body=json.dumps( + { + "action": "GUARDRAIL_INTERVENED", + "outputs": [{"text": "Sorry, the model cannot answer this."}], + "assessments": [ + {"wordPolicy": {"customWords": [{"match": BLOCKED_WORD, "action": "BLOCKED"}]}} + ], + } + ).encode() + ) + if request.target.startswith(f"/{VENDOR_DOWN}"): + return Reply(status=500, body=b'{"detail":"vendor is down"}') + if request.target.startswith(f"/{CATO}/"): + return Reply( + status=THROTTLED.status, body=THROTTLED.body, headers={"x-amzn-errortype": THROTTLED.errortype} + ) + texts: Final = [str(text) for text in json.loads(request.body)["texts"]] + if any(BLOCKED_WORD in text for text in texts): + return Reply(body=json.dumps({"action": "BLOCKED", "blocked_reason": BLOCK_REASON}).encode()) + return Reply(body=json.dumps({"action": "NONE"}).encode()) + + with httpx.Client(base_url=os.environ["INTEGRATION_PROXY_URL"], timeout=30, trust_env=False) as root_client: + root: Final = Gateway( + root_client, + os.environ.get("INTEGRATION_MASTER_KEY", "sk-integration-master"), + os.environ["INTEGRATION_UPSTREAM_URL"], + ) + directory: Final = tmp_path_factory.mktemp("bedrock-aws-errors") + with wire_server(provider) as upstream, wire_server(vendor) as sink: + config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + config["guardrails"] = [ + _generic_guardrail(POST_CALL, "post_call", sink, streaming_buffer_until_moderated=True), + _generic_guardrail(PRE_CALL, "pre_call", sink), + _generic_guardrail(VENDOR_DOWN, "pre_call", sink), + { + "guardrail_name": CATO, + "litellm_params": { + "guardrail": "cato_networks", + "mode": "pre_call", + "default_on": False, + "api_base": f"{sink.url}/{CATO}", + "api_key": "synthetic-cato-key", + }, + }, + { + "guardrail_name": BEDROCK, + "litellm_params": { + "guardrail": "bedrock", + "mode": "post_call", + "default_on": False, + "guardrailIdentifier": "aws-errors", + "guardrailVersion": "1", + "aws_access_key_id": "AKIASYNTHETICGUARD", + "aws_secret_access_key": "synthetic-guardrail-secret", + "aws_region_name": "us-east-1", + "aws_bedrock_runtime_endpoint": f"{sink.url}/{BEDROCK}", + }, + }, + ] + config["model_list"] = [ + { + "model_name": name, + "litellm_params": { + "model": f"bedrock/{model}", + "api_base": upstream.url, + "aws_access_key_id": "AKIASYNTHETICPROVIDER", + "aws_secret_access_key": "synthetic-provider-secret", + "aws_region_name": "us-east-1", + }, + } + for name, model in ((ROUTER_MODEL, MODEL_ID), ("anthropic.*", "anthropic.*")) + ] + config["router_settings"] = {**config.get("router_settings", {}), "num_retries": 0} + config_path: Final = directory / "bedrock-aws-errors.yaml" + config_path.write_text(yaml.safe_dump(config)) + overrides: Final = {"ANTHROPIC_API_BASE": upstream.url, "ANTHROPIC_API_KEY": "sk-ant-synthetic"} + with owned_proxy_process(root, directory, overrides, config=config_path, workers=2) as owned: + with owned.gateway.scenario() as scenario: + yield Rig( + owned.gateway, + upstream, + sink, + { + UNGUARDED: scenario.key(), + OVER_BUDGET: scenario.key(max_budget=0), + ONE_RPM: scenario.key(rpm_limit=1), + **{ + name: scenario.key(guardrails=[name]) + for name in (POST_CALL, PRE_CALL, BEDROCK, VENDOR_DOWN, CATO) + }, + }, + ) + + +BLOCKS: Final = ( + pytest.param(POST_CALL, "converse-stream", BLOCK_REASON, 1, id="post-call-held-converse-stream"), + pytest.param(POST_CALL, "converse", BLOCK_REASON, 1, id="post-call-converse"), + pytest.param(PRE_CALL, "converse-stream", BLOCK_REASON, 0, id="pre-call-converse-stream"), + pytest.param(PRE_CALL, "converse", BLOCK_REASON, 0, id="pre-call-converse"), + pytest.param(BEDROCK, "converse-stream", BEDROCK_BLOCK_MESSAGE, 1, id="bedrock-guardrail-converse-stream"), + pytest.param(BEDROCK, "converse", BEDROCK_BLOCK_MESSAGE, 1, id="bedrock-guardrail-converse"), + pytest.param(PRE_CALL, "invoke", BLOCK_REASON, 0, id="pre-call-invoke"), +) + + +@pytest.mark.parametrize("model", ROUTES) +@pytest.mark.parametrize(("guardrail", "action", "message", "provider_calls"), BLOCKS) +def test_guardrail_block_reaches_boto3_as_a_validation_exception_with_the_guardrail_message( + rig: Rig, guardrail: str, action: str, message: str, provider_calls: int, model: str +) -> None: + marker: Final = f"{BLOCKED_WORD}-{uuid.uuid4().hex}" + error: Final = rig.call(guardrail, action, marker, model=model) + assert error is not None, f"{guardrail} {action} was not blocked" + assert error.response["ResponseMetadata"]["HTTPStatusCode"] == 400, error.response + assert (error.response["Error"]["Code"], error.response["Error"]["Message"]) == ( + "ValidationException", + message, + ), error.response + assert rig.provider_calls(marker) == provider_calls + + +@pytest.mark.parametrize("action", ("converse-stream", "converse")) +def test_guardrail_block_body_adds_the_aws_message_and_keeps_the_litellm_error(rig: Rig, action: str) -> None: + marker: Final = f"{BLOCKED_WORD}-{uuid.uuid4().hex}" + response: Final = rig.raw(POST_CALL, action, marker) + assert response.status_code == 400, response.text + assert response.headers.get("x-amzn-errortype") == "ValidationException", (dict(response.headers), response.text) + assert response.headers.get("content-type", "").startswith("application/json"), dict(response.headers) + body: Final = response.json() + assert body["message"] == BLOCK_REASON, body + assert body["error"]["message"] == BLOCK_REASON, body + assert response.headers.get("x-litellm-call-id"), dict(response.headers) + + +@pytest.mark.parametrize("kind", tuple(UPSTREAM_ERRORS)) +@pytest.mark.parametrize("action", ("converse-stream", "converse")) +@pytest.mark.parametrize("guardrail", (UNGUARDED, POST_CALL)) +def test_bedrock_error_reaches_the_caller_as_bedrock_sent_it(rig: Rig, kind: str, action: str, guardrail: str) -> None: + upstream: Final = UPSTREAM_ERRORS[kind] + marker: Final = upstream_marker(kind) + response: Final = rig.raw(guardrail, action, marker) + assert response.status_code == upstream.status, response.text + assert response.headers.get("x-amzn-errortype") == upstream.errortype, (dict(response.headers), response.text) + assert response.content == upstream.body, response.text + assert response.headers.get("content-type") == "application/json", dict(response.headers) + assert response.headers.get("x-litellm-call-id"), dict(response.headers) + assert rig.provider_calls(marker) == 1 + + +@pytest.mark.parametrize("model", ROUTES) +@pytest.mark.parametrize("action", ("converse-stream", "converse", "invoke")) +@pytest.mark.parametrize("guardrail", (UNGUARDED, POST_CALL)) +def test_boto3_sees_bedrock_throttling_and_retries_it_as_against_bedrock( + rig: Rig, action: str, guardrail: str, model: str +) -> None: + marker: Final = upstream_marker("throttled") + error: Final = rig.call(guardrail, action, marker, attempts=3, model=model) + assert error is not None, f"{action} was not throttled" + assert (error.response["Error"]["Code"], error.response["Error"]["Message"]) == ( + THROTTLED.code, + THROTTLE_MESSAGE, + ), error.response + assert error.response["ResponseMetadata"]["HTTPStatusCode"] == 429, error.response + assert error.response["ResponseMetadata"]["RetryAttempts"] == 2, error.response + assert rig.provider_calls(marker) == 3 + + +def converse_with_key(rig: Rig, model: str, key: str) -> httpx.Response: + return rig.gateway.client.post( + f"/bedrock/model/{model}/converse", + json={"messages": [{"role": "user", "content": [{"text": f"plain-{uuid.uuid4().hex}"}]}]}, + headers={"x-litellm-api-key": key}, + timeout=case_timeout_seconds(), + ) + + +def assert_litellm_error(response: httpx.Response, status: int) -> None: + assert response.status_code == status, response.text + assert "x-amzn-errortype" not in response.headers, dict(response.headers) + assert set(response.json()) == {"error"}, response.text + + +@pytest.mark.parametrize( + ("guardrail", "model", "status"), + ( + pytest.param(VENDOR_DOWN, ROUTER_MODEL, 500, id="guardrail-vendor-down"), + pytest.param(UNGUARDED, "not-a-configured-model", 400, id="unknown-model"), + pytest.param(OVER_BUDGET, ROUTER_MODEL, 422, id="over-budget"), + ), +) +def test_errors_that_are_not_guardrail_blocks_keep_the_litellm_error( + rig: Rig, guardrail: str, model: str, status: int +) -> None: + assert_litellm_error(converse_with_key(rig, model, rig.keys[guardrail]), status) + + +def test_guardrail_service_error_shaped_like_aws_keeps_the_litellm_error(rig: Rig) -> None: + response: Final = converse_with_key(rig, ROUTER_MODEL, rig.keys[CATO]) + assert response.status_code == THROTTLED.status, response.text + assert "x-amzn-errortype" not in response.headers, dict(response.headers) + assert response.content != THROTTLED.body, response.text + + +def test_unknown_key_keeps_the_litellm_error(rig: Rig) -> None: + assert_litellm_error(converse_with_key(rig, ROUTER_MODEL, f"sk-unknown-{uuid.uuid4().hex}"), 401) + + +def test_litellm_rate_limit_keeps_the_litellm_error(rig: Rig) -> None: + admitted: Final = converse_with_key(rig, ROUTER_MODEL, rig.keys[ONE_RPM]) + assert admitted.status_code == 200, admitted.text + assert_litellm_error(converse_with_key(rig, ROUTER_MODEL, rig.keys[ONE_RPM]), 429) + + +def test_other_provider_passthrough_errors_carry_no_aws_error_type(rig: Rig) -> None: + response: Final = rig.gateway.client.post( + "/anthropic/v1/messages", + json={"model": "claude-sonnet-5", "max_tokens": 8, "messages": [{"role": "user", "content": "hi"}]}, + headers={"x-litellm-api-key": rig.keys[UNGUARDED]}, + timeout=case_timeout_seconds(), + ) + assert response.status_code == 429, response.text + assert "x-amzn-errortype" not in response.headers, dict(response.headers) + assert response.content == ANTHROPIC_THROTTLE, response.text diff --git a/tests/integration/providers/test_count_tokens_error_shapes_wire.py b/tests/integration/providers/test_count_tokens_error_shapes_wire.py new file mode 100644 index 00000000000..200f5e81673 --- /dev/null +++ b/tests/integration/providers/test_count_tokens_error_shapes_wire.py @@ -0,0 +1,205 @@ +import json +from collections.abc import Iterator +from typing import Final + +import anthropic +import openai +import pytest +from integration._support.client import Gateway, gateway_from_environment +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, wire_server +from pydantic import JsonValue, TypeAdapter + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_JSON_OBJECT: Final = TypeAdapter(dict[str, JsonValue]) +_MESSAGES: Final[list[JsonValue]] = [{"role": "user", "content": "Count this synthetic request."}] +_PROVIDER_KEY_VARIABLES: Final = ( + "ANTHROPIC_API_KEY", + "ANTHROPIC_AUTH_TOKEN", + "AZURE_AI_API_KEY", + "AZURE_API_KEY", + "AZURE_OPENAI_API_KEY", + "AZURE_AD_TOKEN", + "AZURE_CLIENT_ID", + "AZURE_CLIENT_SECRET", + "AZURE_TENANT_ID", + "OPENAI_API_KEY", +) + + +@pytest.fixture(scope="module") +def keyless_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]: + with ( + gateway_from_environment() as gateway, + owned_proxy_process( + gateway, + tmp_path_factory.mktemp("count-tokens-error-shapes"), + {}, + remove_environment=_PROVIDER_KEY_VARIABLES, + workers=1, + ) as owned, + ): + yield owned + + +@pytest.fixture(scope="module") +def azure_environment_key_proxy(tmp_path_factory: pytest.TempPathFactory) -> Iterator[OwnedProxy]: + with ( + gateway_from_environment() as gateway, + owned_proxy_process( + gateway, + tmp_path_factory.mktemp("count-tokens-azure-environment-key"), + {"AZURE_API_KEY": "synthetic-azure-env-key"}, + remove_environment=tuple(name for name in _PROVIDER_KEY_VARIABLES if name != "AZURE_API_KEY"), + workers=1, + ) as owned, + ): + yield owned + + +def _anthropic(gateway: Gateway) -> anthropic.Anthropic: + return anthropic.Anthropic( + api_key=gateway.key, base_url=str(gateway.client.base_url), max_retries=0, http_client=gateway.client + ) + + +def _never_called(request: Request) -> Reply: + raise AssertionError(f"count reached the provider without a key: {request.target}") + + +@pytest.mark.parametrize( + ("deployment", "missing"), + ( + ({"model": "anthropic/claude-sonnet-4-6"}, "No Anthropic API key is set for model "), + ({"model": "azure_ai/claude-sonnet-4-6"}, "No Azure AI API key or Entra ID credential is set for model "), + ), + ids=("anthropic", "azure_ai_anthropic"), +) +def test_anthropic_count_without_a_provider_credential_is_an_anthropic_401( + keyless_proxy: OwnedProxy, deployment: dict[str, JsonValue], missing: str +) -> None: + gateway: Final = keyless_proxy.gateway + with wire_server(_never_called) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(**deployment, api_base=wire.url, api_key=None, num_retries=0) + with pytest.raises(anthropic.AuthenticationError) as raised: + _anthropic(gateway).messages.count_tokens(model=model, messages=_MESSAGES) + assert raised.value.body == { + "type": "error", + "error": {"type": "authentication_error", "message": raised.value.body["error"]["message"]}, + } + assert missing in raised.value.message, raised.value.message + assert wire.drain() == () + + +def test_azure_ai_anthropic_count_sends_the_azure_environment_key_as_x_api_key( + azure_environment_key_proxy: OwnedProxy, +) -> None: + gateway: Final = azure_environment_key_proxy.gateway + + def count(request: Request) -> Reply: + assert request.headers.get("x-api-key") == "synthetic-azure-env-key", request.headers + return Reply(body=json.dumps({"input_tokens": 58}).encode()) + + with wire_server(count) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="azure_ai/claude-sonnet-4-6", + api_base=wire.url, + api_key=None, + num_retries=0, + ) + result: Final = _anthropic(gateway).messages.count_tokens(model=model, messages=_MESSAGES) + assert result.input_tokens == 58 + + +def test_azure_ai_anthropic_count_authenticates_with_an_entra_token(keyless_proxy: OwnedProxy) -> None: + gateway: Final = keyless_proxy.gateway + + def count(request: Request) -> Reply: + assert request.headers.get("authorization") == "Bearer synthetic-entra-token", request.headers + assert "x-api-key" not in request.headers, request.headers + return Reply(body=json.dumps({"input_tokens": 61}).encode()) + + with wire_server(count) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="azure_ai/claude-sonnet-4-6", + api_base=wire.url, + api_key=None, + azure_ad_token="synthetic-entra-token", + num_retries=0, + ) + result: Final = _anthropic(gateway).messages.count_tokens(model=model, messages=_MESSAGES) + assert result.input_tokens == 61 + assert len(wire.drain()) == 1 + + +def test_anthropic_count_request_errors_use_the_anthropic_envelope(keyless_proxy: OwnedProxy) -> None: + gateway: Final = keyless_proxy.gateway + upstream_error: Final = {"type": "error", "error": {"type": "invalid_request_error", "message": "scripted reject"}} + with ( + wire_server(lambda _request: Reply(status=400, body=json.dumps(upstream_error).encode())) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model="anthropic/claude-sonnet-4-6", api_base=wire.url, api_key="synthetic-key", num_retries=0 + ) + missing_messages: Final = gateway.request("POST", "/v1/messages/count_tokens", {"model": model}) + assert missing_messages.status_code == 400, missing_messages.text + assert _JSON_OBJECT.validate_json(missing_messages.content) == { + "type": "error", + "error": {"type": "invalid_request_error", "message": "messages parameter is required"}, + } + + rejected: Final = gateway.request( + "POST", "/v1/messages/count_tokens", {"model": model, "messages": _MESSAGES} + ) + assert rejected.status_code == 400, rejected.text + body: Final = _JSON_OBJECT.validate_json(rejected.content) + assert set(body) <= {"type", "error", "request_id"}, body + assert body["error"] == upstream_error["error"], body + assert len(wire.drain()) == 1 + + +def test_openai_input_tokens_without_a_provider_key_is_an_openai_401(keyless_proxy: OwnedProxy) -> None: + gateway: Final = keyless_proxy.gateway + with wire_server(_never_called) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model="openai/gpt-4o-mini", api_base=f"{wire.url}/v1", api_key=None, num_retries=0 + ) + client: Final = openai.OpenAI( + api_key=gateway.key, base_url=str(gateway.client.base_url), max_retries=0, http_client=gateway.client + ) + with pytest.raises(openai.AuthenticationError) as raised: + client.responses.input_tokens.count(model=model, input="Count this synthetic request.") + assert f"No OpenAI API key is set for model {model}" in raised.value.message + assert wire.drain() == () + + +def test_google_count_tokens_errors_use_the_google_envelope(keyless_proxy: OwnedProxy) -> None: + gateway: Final = keyless_proxy.gateway + upstream_error: Final = { + "error": {"code": 400, "message": "API key not valid. Please pass a valid API key.", "status": "INVALID_ARGUMENT"} + } + with ( + wire_server(lambda _request: Reply(status=400, body=json.dumps(upstream_error).encode())) as wire, + gateway.scenario() as scenario, + ): + model: Final = scenario.model( + model="gemini/gemini-2.5-flash", api_base=wire.url, api_key="synthetic-key", num_retries=0 + ) + rejected: Final = gateway.request( + "POST", + f"/v1beta/models/{model}:countTokens", + {"contents": [{"role": "user", "parts": [{"text": "Count this synthetic request."}]}]}, + ) + assert rejected.status_code == 400, rejected.text + error: Final = _JSON_OBJECT.validate_json(rejected.content)["error"] + assert isinstance(error, dict) and error["code"] == 400 and error["status"] == "INVALID_ARGUMENT", error + assert "API key not valid" in str(error["message"]), error + assert len(wire.drain()) == 1 + + malformed: Final = gateway.request("POST", f"/v1beta/models/{model}:countTokens", {"contents": "not-a-list"}) + assert malformed.status_code == 400, malformed.text + assert _JSON_OBJECT.validate_json(malformed.content) == { + "error": {"code": 400, "message": "Invalid Google countTokens request body", "status": "INVALID_ARGUMENT"} + } diff --git a/tests/integration/providers/test_openai_count_tokens_tools_wire.py b/tests/integration/providers/test_openai_count_tokens_tools_wire.py index 3ab5826006f..2b23ccdf4f2 100644 --- a/tests/integration/providers/test_openai_count_tokens_tools_wire.py +++ b/tests/integration/providers/test_openai_count_tokens_tools_wire.py @@ -321,6 +321,48 @@ def test_messages_count_tokens_sends_the_input_v1_messages_inference_sends( assert count_body["model"] == inference_body["model"] +_UNCONVERTIBLE_BODY: Final[Mapping[str, JsonValue]] = { + "messages": [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the notes."}, + {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "notes"}}, + {"type": "search_result", "source": "https://example.com/notes", "title": "Notes", "content": []}, + ], + }, + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "search", "input": {}}]}, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_01", + "content": [ + { + "type": "search_result", + "source": "https://example.com/hits", + "title": "Hits", + "content": [{"type": "text", "text": "hit"}], + } + ], + } + ], + }, + ], +} + + +def test_messages_count_tokens_rejects_blocks_it_cannot_convert(counting_proxy: OwnedProxy) -> None: + gateway: Final = counting_proxy.gateway + with wire_server(_responses_peer) as wire, gateway.scenario() as scenario: + model: Final = _deployment(scenario, wire.url) + response: Final = gateway.request("POST", "/v1/messages/count_tokens", {"model": model, **_UNCONVERTIBLE_BODY}) + assert response.status_code == 400, response.text + assert "document with a text source, search_result, search_result inside a tool_result" in response.text + assert wire.drain() == () + + def test_messages_count_tokens_without_a_role_answers_400_before_calling_openai(counting_proxy: OwnedProxy) -> None: gateway: Final = counting_proxy.gateway with wire_server(_count_response) as wire, gateway.scenario() as scenario: diff --git a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py index 7fef79845a1..5fe82d37e9e 100644 --- a/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py +++ b/tests/integration/providers/test_vertex_gemini_count_tokens_wire.py @@ -254,6 +254,53 @@ def test_count_tokens_sends_the_contents_messages_inference_sends(gateway: Gatew assert _counted_contents(count_body) == generated_body["contents"] +_UNCONVERTIBLE_MESSAGES: Final[JsonValue] = [ + { + "role": "user", + "content": [ + {"type": "text", "text": "Summarize the notes."}, + {"type": "document", "source": {"type": "text", "media_type": "text/plain", "data": "notes"}}, + {"type": "search_result", "source": "https://example.com/notes", "title": "Notes", "content": []}, + ], + }, + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "search", "input": {}}]}, + { + "role": "user", + "content": [ + { + "type": "tool_result", + "tool_use_id": "toolu_01", + "content": [ + { + "type": "search_result", + "source": "https://example.com/hits", + "title": "Hits", + "content": [{"type": "text", "text": "hit"}], + } + ], + } + ], + }, +] + + +@pytest.mark.parametrize("provider", ("vertex_ai", "gemini")) +def test_count_tokens_rejects_blocks_it_cannot_convert(gateway: Gateway, provider: str) -> None: + peer: Final = _GeminiPeer() + with wire_server(peer) as wire, gateway.scenario() as scenario: + model: Final = ( + _deployment(gateway, scenario, wire.url) + if provider == "vertex_ai" + else _ai_studio_deployment(scenario, wire.url) + ) + response: Final = gateway.request( + "POST", "/v1/messages/count_tokens", {"model": model, "messages": _UNCONVERTIBLE_MESSAGES} + ) + assert response.status_code == 400, response.text + assert "document with a text source, search_result, search_result inside a tool_result" in response.text + assert peer.drain() == () + + def test_count_tokens_keeps_each_text_block_in_system_instruction(gateway: Gateway) -> None: peer: Final = _GeminiPeer() with wire_server(peer) as wire, gateway.scenario() as scenario: diff --git a/tests/integration/routing/test_base_model_context_window_owned_proxy.py b/tests/integration/routing/test_base_model_context_window_owned_proxy.py new file mode 100644 index 00000000000..d09ec868c78 --- /dev/null +++ b/tests/integration/routing/test_base_model_context_window_owned_proxy.py @@ -0,0 +1,171 @@ +import json +import uuid +from collections.abc import Iterator, Mapping +from dataclasses import dataclass +from typing import Final + +import pytest +from integration._support.client import gateway_from_environment +from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process +from integration._support.wire import Reply, Request, Wire, wire_server +from integration.providers._cache_control_marks_support import marker_of, owned_config +from pydantic import JsonValue + +_TOKEN: Final = "synthetic-bedrock-bearer" +_KIMI: Final = "moonshotai.kimi-k2.5" +_SONNET: Final = "global.anthropic.claude-sonnet-4-6" +_KIMI_BASE: Final = f"bedrock/us-east-1/{_KIMI}" +_KIMI_TARGET: Final = f"/model/{_KIMI}/converse" +_SONNET_TARGET: Final = f"/model/{_SONNET}/converse" +_PAST_ROUTE_ROW_LINES: Final = 16250 +_PAST_ESCALATION_BUFFER_LINES: Final = 15480 +_SMALL_LINES: Final = 1300 +_REPLY: Final = json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": "8115"}]}}, + "stopReason": "end_turn", + "usage": {"inputTokens": 9, "outputTokens": 2, "totalTokens": 11}, + "metrics": {"latencyMs": 1}, + } +).encode() + + +@dataclass(frozen=True, slots=True) +class _Rig: + owned: OwnedProxy + wire: Wire + + +def _converse_peer(request: Request) -> Reply: + assert request.headers["authorization"] == f"Bearer {_TOKEN}", request.headers + return Reply(body=_REPLY) + + +def _deployment( + name: str, + model: str, + api_base: str, + *, + model_info: Mapping[str, JsonValue] | None = None, + **params: JsonValue, +) -> dict[str, JsonValue]: + return { + "model_name": name, + "litellm_params": { + "model": f"bedrock/converse/{model}", + "api_key": _TOKEN, + "aws_region_name": "us-east-1", + "aws_bedrock_runtime_endpoint": api_base, + **params, + }, + "model_info": {"id": f"{name}-{uuid.uuid4().hex}", **(model_info or {})}, + } + + +@pytest.fixture(scope="module") +def rig(tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Rig]: + directory: Final = tmp_path_factory.mktemp("base-model-window") + with gateway_from_environment() as gateway, wire_server(_converse_peer) as wire: + config: Final = owned_config( + directory, + [ + _deployment("base-in-model-info", _KIMI, wire.url, model_info={"base_model": _KIMI_BASE}), + _deployment("base-in-litellm-params", _KIMI, wire.url, base_model=_KIMI_BASE), + _deployment("no-base", _KIMI, wire.url), + _deployment("other-model-base", _KIMI, wire.url, model_info={"base_model": _SONNET}), + _deployment( + "declared-window", + _KIMI, + wire.url, + model_info={"base_model": _KIMI_BASE, "max_input_tokens": 257000}, + ), + _deployment( + "unmapped-base", _KIMI, wire.url, model_info={"base_model": f"{_KIMI_BASE}-not-in-cost-map"} + ), + _deployment( + "smaller-other-model-base", + _KIMI, + wire.url, + model_info={"base_model": "bedrock/us-east-1/meta.llama3-8b-instruct-v1:0"}, + ), + _deployment("one-million", _SONNET, wire.url), + { + "model_name": "complexity", + "litellm_params": { + "model": "auto_router/complexity_router", + "complexity_router_config": { + "tiers": { + "SIMPLE": "base-in-model-info", + "MEDIUM": "base-in-model-info", + "COMPLEX": "one-million", + "REASONING": "one-million", + }, + "enable_context_window_escalation": True, + }, + }, + }, + ], + litellm_settings={"cache": False}, + router_settings={"enable_pre_call_checks": True}, + ) + with owned_proxy_process(gateway, directory, {}, config=config, workers=2, client_timeout=120) as owned: + yield _Rig(owned, wire) + + +def _prompt(lines: int, marker: str) -> str: + listing: Final = "\n".join(f"Line {i}: the value of item {i} is {(i * 7) % 1000}." for i in range(1, lines + 1)) + return f"{listing}\n\nWhat is the value of item 12345? Answer with the number only. marker-{marker}" + + +def _request(surface: str, model: str, text: str) -> tuple[str, dict[str, JsonValue]]: + if surface == "messages": + return "/v1/messages", {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": text}]} + if surface == "responses": + return "/v1/responses", {"model": model, "max_output_tokens": 16, "input": text} + return "/v1/chat/completions", {"model": model, "max_tokens": 16, "messages": [{"role": "user", "content": text}]} + + +def _send(rig: _Rig, surface: str, model: str, lines: int) -> tuple[int, str, tuple[str, ...]]: + marker: Final = uuid.uuid4().hex + path, body = _request(surface, model, _prompt(lines, marker)) + response: Final = rig.owned.gateway.request("POST", path, body) + targets: Final = tuple(request.target for request in rig.wire.drain() if marker_of(request) == marker) + return response.status_code, response.text, targets + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +@pytest.mark.parametrize("surface", ["chat", "messages", "responses"]) +def test_a_same_model_base_model_row_widens_the_window_pre_call_checks_enforce(rig: _Rig, surface: str) -> None: + status, text, targets = _send(rig, surface, "base-in-model-info", _PAST_ROUTE_ROW_LINES) + assert status == 200, text + assert targets == (_KIMI_TARGET,) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_a_same_model_base_model_in_litellm_params_widens_the_window_too(rig: _Rig) -> None: + status, text, targets = _send(rig, "chat", "base-in-litellm-params", _PAST_ROUTE_ROW_LINES) + assert status == 200, text + assert targets == (_KIMI_TARGET,) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +@pytest.mark.parametrize("model", ["no-base", "other-model-base", "declared-window", "unmapped-base"]) +def test_a_base_model_that_cannot_widen_keeps_the_route_window_and_calls_no_upstream(rig: _Rig, model: str) -> None: + status, text, targets = _send(rig, "chat", model, _PAST_ROUTE_ROW_LINES) + assert status == 400, text + assert "ContextWindowExceededError" in text, text + assert targets == () + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_a_smaller_other_model_base_model_never_narrows_the_route_window(rig: _Rig) -> None: + status, text, targets = _send(rig, "chat", "smaller-other-model-base", _SMALL_LINES) + assert status == 200, text + assert targets == (_KIMI_TARGET,) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +def test_complexity_router_keeps_a_prompt_on_the_tier_whose_base_model_row_holds_it(rig: _Rig) -> None: + status, text, targets = _send(rig, "chat", "complexity", _PAST_ESCALATION_BUFFER_LINES) + assert status == 200, text + assert targets == (_KIMI_TARGET,) diff --git a/tests/unit/integrations/test_custom_guardrail.py b/tests/unit/integrations/test_custom_guardrail.py index 9bd41f4a602..614ff97553a 100644 --- a/tests/unit/integrations/test_custom_guardrail.py +++ b/tests/unit/integrations/test_custom_guardrail.py @@ -2790,7 +2790,7 @@ class TestLoggingOnlyApplyGuardrail: assert out_kwargs["messages"] == kwargs["messages"] assert ( out_kwargs["standard_logging_object"]["guardrail_information"][0]["guardrail_status"] - == "guardrail_intervened" + == "guardrail_flagged" ) @pytest.mark.asyncio diff --git a/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py index 320a4eff371..c8dbbd00d66 100644 --- a/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py +++ b/tests/unit/litellm_core_utils/llm_cost_calc/test_tool_call_cost_tracking.py @@ -593,6 +593,422 @@ _BEDROCK_MANTLE_WEB_SEARCH_MODELS = ( _BEDROCK_MANTLE_WEB_SEARCH_RATE = 0.012 + + +def _openai_responses_response(model, output, usage=None, tools=None): + return ResponsesAPIResponse.model_validate( + { + "id": "resp_1", + "created_at": 1754900000, + "model": model, + "object": "response", + "status": "completed", + "output": output, + "usage": usage or {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, + **({"tools": tools} if tools is not None else {}), + } + ) + + +_ASSISTANT_MESSAGE_OUTPUT_ITEM = { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "done", "annotations": []}], +} + +_GPT_IMAGE_1_HIGH_1024_COST_KEY = "high/1024-x-1024/gpt-image-1" +_GPT_IMAGE_1_TOOLS = [{"type": "image_generation", "model": "gpt-image-1"}] + + +@pytest.mark.parametrize("custom_llm_provider", ["openai", None]) +def test_responses_image_generation_call_billed_as_tool_usage_cost(local_model_cost_map, custom_llm_provider): + """A completed image_generation_call bills at the tools[].model rate for its quality/size, provider given or inferred.""" + expected_image_cost = litellm.model_cost[_GPT_IMAGE_1_HIGH_1024_COST_KEY]["input_cost_per_image"] + response = _openai_responses_response( + "gpt-5", + [ + { + "type": "image_generation_call", + "id": "ig_1", + "status": "completed", + "quality": "high", + "size": "1024x1024", + "result": "AAAA", + }, + dict(_ASSISTANT_MESSAGE_OUTPUT_ITEM), + ], + tools=_GPT_IMAGE_1_TOOLS, + ) + + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model="gpt-5", + response_object=response, + usage=None, + custom_llm_provider=custom_llm_provider, + standard_built_in_tools_params=None, + ) + + assert cost > 0 + assert cost == pytest.approx(expected_image_cost) + + +def test_responses_image_generation_call_without_tool_model_bills_zero(local_model_cost_map): + """With no tools[].model to price against, the image tool part bills $0 instead of guessing a model or raising.""" + response = _openai_responses_response( + "gpt-5", + [ + { + "type": "image_generation_call", + "id": "ig_1", + "status": "completed", + "quality": "high", + "size": "1024x1024", + "result": "AAAA", + } + ], + usage=_image_gen_token_usage(10, 0, 50, 0), + ) + + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model="gpt-5", + response_object=response, + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + + assert cost == 0.0 + + +def test_responses_web_search_and_image_generation_costs_are_additive(local_model_cost_map): + """A response billed for web search must still also bill its image_generation_call items.""" + model = "gpt-4o-search-preview" + image_cost = litellm.model_cost[_GPT_IMAGE_1_HIGH_1024_COST_KEY]["input_cost_per_image"] + image_item = { + "type": "image_generation_call", + "id": "ig_1", + "status": "completed", + "quality": "high", + "size": "1024x1024", + "result": "AAAA", + } + web_search_item = {"type": "web_search_call", "id": "ws_1", "status": "completed"} + + combined = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model=model, + response_object=_openai_responses_response(model, [web_search_item, image_item], tools=_GPT_IMAGE_1_TOOLS), + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + web_search_only = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model=model, + response_object=_openai_responses_response(model, [web_search_item]), + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + image_only = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model=model, + response_object=_openai_responses_response(model, [image_item], tools=_GPT_IMAGE_1_TOOLS), + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + + assert image_only == pytest.approx(image_cost) + assert web_search_only > 0 + assert combined == pytest.approx(web_search_only + image_only) + + +def test_responses_incomplete_image_generation_call_not_billed(local_model_cost_map): + """A failed image_generation_call produced no billable image, so it must cost $0.""" + response = _openai_responses_response( + "gpt-5", + [ + { + "type": "image_generation_call", + "id": "ig_1", + "status": "failed", + "quality": "high", + "size": "1024x1024", + "result": None, + }, + dict(_ASSISTANT_MESSAGE_OUTPUT_ITEM), + ], + tools=_GPT_IMAGE_1_TOOLS, + ) + + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model="gpt-5", + response_object=response, + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + + assert cost == 0.0 + + +def test_completion_cost_includes_responses_image_generation_tool_cost(local_model_cost_map): + """The image tool fee must flow through completion_cost on top of the token-only baseline.""" + image_item = { + "type": "image_generation_call", + "id": "ig_1", + "status": "completed", + "quality": "high", + "size": "1024x1024", + "result": "AAAA", + } + response_with_image = _openai_responses_response( + "gpt-5", [image_item, dict(_ASSISTANT_MESSAGE_OUTPUT_ITEM)], tools=_GPT_IMAGE_1_TOOLS + ) + response_without_image = _openai_responses_response("gpt-5", [dict(_ASSISTANT_MESSAGE_OUTPUT_ITEM)]) + + cost_with_image = litellm.completion_cost( + completion_response=response_with_image, + model="gpt-5", + custom_llm_provider="openai", + call_type="aresponses", + ) + cost_without_image = litellm.completion_cost( + completion_response=response_without_image, + model="gpt-5", + custom_llm_provider="openai", + call_type="aresponses", + ) + + assert cost_with_image > cost_without_image + assert cost_with_image - cost_without_image == pytest.approx( + litellm.model_cost[_GPT_IMAGE_1_HIGH_1024_COST_KEY]["input_cost_per_image"] + ) + + +def test_responses_usage_tool_usage_web_search_billed_without_output_item(local_model_cost_map): + """usage.tool_usage.web_search.num_requests bills web search even when no web_search_call item is present.""" + model = "gpt-5.4-mini" + per_call = litellm.get_model_info(model)["search_context_cost_per_query"]["search_context_size_medium"] + + for num_requests in (1, 2): + response = _openai_responses_response( + model, + [dict(_ASSISTANT_MESSAGE_OUTPUT_ITEM)], + usage={ + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + "tool_usage": {"web_search": {"num_requests": num_requests}}, + }, + ) + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model=model, + response_object=response, + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + assert cost == pytest.approx(num_requests * per_call) + + +def _image_gen_token_usage(input_text, input_image, output_image, output_text): + return { + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + "tool_usage": { + "image_gen": { + "input_tokens": input_text + input_image, + "output_tokens": output_image + output_text, + "total_tokens": input_text + input_image + output_image + output_text, + "input_tokens_details": {"image_tokens": input_image, "text_tokens": input_text}, + "output_tokens_details": {"image_tokens": output_image, "text_tokens": output_text}, + } + }, + } + + +def _expected_image_gen_token_cost(model_info, input_text, input_image, output_image, output_text): + return ( + input_text * (model_info.get("input_cost_per_token") or 0) + + input_image * (model_info.get("input_cost_per_image_token") or 0) + + output_image * (model_info.get("output_cost_per_image_token") or 0) + + output_text * (model_info.get("output_cost_per_token") or 0) + ) + + +def test_responses_image_tool_model_from_tools_bills_token_usage(local_model_cost_map): + """The image tool's tools[].model is used and usage.tool_usage.image_gen tokens bill at that model's rates.""" + tool_model = "gpt-image-2" + model_info = litellm.get_model_info(tool_model, custom_llm_provider="openai") + tools = [{"type": "image_generation", "model": tool_model, "quality": "low", "size": "1024x1024"}] + output = [ + { + "type": "image_generation_call", + "id": "ig_1", + "status": "completed", + "quality": "low", + "size": "1024x1024", + "result": "AAAA", + }, + dict(_ASSISTANT_MESSAGE_OUTPUT_ITEM), + ] + + response = _openai_responses_response("gpt-5", output, usage=_image_gen_token_usage(10, 0, 50, 0), tools=tools) + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model="gpt-5", + response_object=response, + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + assert cost > 0 + assert cost == pytest.approx(_expected_image_gen_token_cost(model_info, 10, 0, 50, 0)) + + response_more_tokens = _openai_responses_response( + "gpt-5", output, usage=_image_gen_token_usage(10, 0, 80, 0), tools=tools + ) + cost_more = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model="gpt-5", + response_object=response_more_tokens, + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + assert cost_more != cost + assert cost_more == pytest.approx(_expected_image_gen_token_cost(model_info, 10, 0, 80, 0)) + + +def test_responses_zero_image_gen_tokens_fall_back_to_per_image_pricing(local_model_cost_map): + """An all-zero image_gen usage block keeps the per-image path for the tool's model/quality/size.""" + tool_model = "gpt-image-1" + quality = "low" + size = "1024x1024" + tools = [{"type": "image_generation", "model": tool_model, "quality": quality, "size": size}] + response = _openai_responses_response( + "gpt-5", + [ + { + "type": "image_generation_call", + "id": "ig_1", + "status": "completed", + "quality": quality, + "size": size, + "result": "AAAA", + } + ], + usage=_image_gen_token_usage(0, 0, 0, 0), + tools=tools, + ) + + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model="gpt-5", + response_object=response, + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + + from litellm.cost_calculator import default_image_cost_calculator + + assert cost == pytest.approx( + default_image_cost_calculator( + model=tool_model, + custom_llm_provider="openai", + quality=quality, + n=1, + size=size, + ) + ) + assert cost > 0 + + +def test_responses_auto_size_image_generation_call_billed_at_default_size(local_model_cost_map): + """An image_generation_call with size "auto" bills at the default size instead of erroring to $0.""" + from litellm.cost_calculator import default_image_cost_calculator + + response = _openai_responses_response( + "gpt-5", + [ + { + "type": "image_generation_call", + "id": "ig_1", + "status": "completed", + "quality": "high", + "size": "auto", + "result": "AAAA", + } + ], + tools=_GPT_IMAGE_1_TOOLS, + ) + + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model="gpt-5", + response_object=response, + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + + assert cost == pytest.approx( + default_image_cost_calculator( + model="gpt-image-1", + custom_llm_provider="openai", + quality="high", + n=1, + size=None, + ) + ) + assert cost > 0 + + +def test_responses_image_gen_total_without_token_details_falls_back_to_per_image(local_model_cost_map): + """A positive image_gen total with no token details falls back to per-image pricing, not $0.""" + tool_model = "gpt-image-1" + quality = "low" + size = "1024x1024" + response = _openai_responses_response( + "gpt-5", + [ + { + "type": "image_generation_call", + "id": "ig_1", + "status": "completed", + "quality": quality, + "size": size, + "result": "AAAA", + } + ], + usage={ + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + "tool_usage": {"image_gen": {"input_tokens": 5, "output_tokens": 5, "total_tokens": 10}}, + }, + tools=[{"type": "image_generation", "model": tool_model, "quality": quality, "size": size}], + ) + + cost = StandardBuiltInToolCostTracking.get_cost_for_built_in_tools( + model="gpt-5", + response_object=response, + usage=None, + custom_llm_provider="openai", + standard_built_in_tools_params=None, + ) + + from litellm.cost_calculator import default_image_cost_calculator + + assert cost == pytest.approx( + default_image_cost_calculator( + model=tool_model, + custom_llm_provider="openai", + quality=quality, + n=1, + size=size, + ) + ) + assert cost > 0 @pytest.mark.parametrize( ("model_key", "web_search_requests", "medium_cost_per_query", "expected_cost"), [ diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_get_headers.py b/tests/unit/litellm_core_utils/llm_response_utils/test_get_headers.py index f52e837136a..1eccf5010ec 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_get_headers.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_get_headers.py @@ -8,6 +8,7 @@ from litellm.litellm_core_utils.llm_response_utils.get_headers import ( _get_llm_provider_headers, get_response_headers, get_provider_request_id, + provider_request_id_headers, ) from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome @@ -89,6 +90,31 @@ def test_invalid_provider_request_ids_remain_absent(headers: object) -> None: assert get_provider_request_id(headers) is None +@pytest.mark.parametrize( + ("headers", "expected"), + ( + ( + {"X-Amzn-RequestId": "amzn-1", "retry-after": "7", "content-type": "text/event-stream"}, + {"x-amzn-requestid": "amzn-1"}, + ), + ( + {"x-request-id": "req_1", "cf-ray": "ray-1", "x-ratelimit-remaining-requests": "0"}, + {"x-request-id": "req_1", "cf-ray": "ray-1"}, + ), + ({"retry-after": "7"}, None), + ({"x-request-id": 7, "request-id": None, "x-amzn-requestid": "amzn-2"}, {"x-amzn-requestid": "amzn-2"}), + ({"x-amzn-requestid": "amzn-\u00fc", "x-request-id": "req_2"}, {"x-request-id": "req_2"}), + ({}, None), + (None, None), + ), +) +def test_provider_request_id_headers_keeps_only_request_id_headers( + headers: dict[str, object] | None, expected: dict[str, str] | None +) -> None: + result = provider_request_id_headers(headers) + assert (None if result is None else dict(result)) == expected + + @pytest.fixture(autouse=True) def _vcr_outcome_gate(request, vcr): install_live_call_probe(request, vcr) diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py index dba1730208d..ad216548d08 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_factory.py @@ -4,9 +4,9 @@ import logging import os import re import time -from collections.abc import Callable +from collections.abc import Callable, Mapping, Sequence from copy import deepcopy -from typing import Final, List +from typing import Final, List, TypeAlias from unittest.mock import MagicMock, patch import httpx @@ -16,6 +16,7 @@ import respx import litellm from litellm.litellm_core_utils.prompt_templates.factory import ( BEDROCK_DOCUMENT_PLACEHOLDER_TEXT, + _EMPTY_TEXT_PLACEHOLDER, BedrockConverseMessagesProcessor, BedrockImageProcessor, _bedrock_converse_messages_pt, @@ -795,6 +796,138 @@ def test_anthropic_messages_pt_keeps_signed_thinking_block(): assert thinking_blocks[0]["thinking"] == "genuine anthropic reasoning" +EMPTY_USER_MESSAGES: Final = pytest.mark.parametrize( + "empty_user_message", + [{"role": "user", "content": None}, {"role": "user", "content": []}, {"role": "user"}], + ids=["null", "empty-list", "content-key-dropped"], +) + +Message: TypeAlias = Mapping[str, object] + + +def _convert(messages: Sequence[Message]) -> Sequence[Message]: + return anthropic_messages_pt(messages=list(messages), model="claude-opus-5", llm_provider="vertex_ai") + + +@EMPTY_USER_MESSAGES +def test_anthropic_messages_pt_preserves_trailing_user_message_with_empty_content(empty_user_message: Message) -> None: + result: Final = _convert( + [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + empty_user_message, + ] + ) + + assert [message["role"] for message in result] == ["user", "assistant", "user"] + assert result[-1]["content"] == [{"type": "text", "text": _EMPTY_TEXT_PLACEHOLDER}] + + +@EMPTY_USER_MESSAGES +@pytest.mark.parametrize( + "cache_control", [{}, {"cache_control": {"type": "ephemeral"}}], ids=["no-cache-control", "with-cache-control"] +) +def test_anthropic_messages_pt_empty_user_message_beside_tool_result_keeps_tool_result_only( + empty_user_message: Message, cache_control: Message +) -> None: + result: Final = _convert( + [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "toolu_1", "type": "function", "function": {"name": "f", "arguments": "{}"}}], + }, + {**empty_user_message, **cache_control}, + {"role": "tool", "tool_call_id": "toolu_1", "content": "ok"}, + ] + ) + + assert [message["role"] for message in result] == ["user", "assistant", "user"] + assert [block["type"] for block in result[-1]["content"]] == ["tool_result"] + + +@EMPTY_USER_MESSAGES +def test_anthropic_messages_pt_empty_user_message_keeps_cache_control_like_an_empty_string( + empty_user_message: Message, +) -> None: + def last_turn(last_user_message: Message) -> Message: + return _convert( + [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + {**last_user_message, "cache_control": {"type": "ephemeral"}}, + ] + )[-1] + + assert last_turn(empty_user_message)["content"][0]["cache_control"] == {"type": "ephemeral"} + assert last_turn(empty_user_message) == last_turn({"role": "user", "content": ""}) + + +TOOL_CALL_AND_RESULT: Final[tuple[Message, ...]] = ( + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": None, + "tool_calls": [{"id": "toolu_1", "type": "function", "function": {"name": "f", "arguments": "{}"}}], + }, + {"role": "tool", "tool_call_id": "toolu_1", "content": "ok"}, +) + + +@EMPTY_USER_MESSAGES +@pytest.mark.parametrize( + "messages_before", + [ + [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + {"role": "user", "content": "q"}, + ], + [{"role": "user", "content": "hi"}, {"role": "assistant", "content": "hello"}, {"role": "user", "content": []}], + TOOL_CALL_AND_RESULT, + ], + ids=["after-user-text", "after-another-empty-message", "after-tool-result"], +) +def test_anthropic_messages_pt_empty_user_message_beside_other_content_sends_nothing_extra( + empty_user_message: Message, messages_before: Sequence[Message] +) -> None: + assert _convert([*messages_before, empty_user_message]) == _convert(messages_before) + + +@EMPTY_USER_MESSAGES +def test_anthropic_messages_pt_empty_user_message_with_cache_control_after_text_keeps_its_breakpoint( + empty_user_message: Message, +) -> None: + def last_turn(last_user_message: Message) -> Message: + return _convert( + [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + {"role": "user", "content": "long document"}, + {**last_user_message, "cache_control": {"type": "ephemeral"}}, + ] + )[-1] + + assert last_turn(empty_user_message)["content"][-1]["cache_control"] == {"type": "ephemeral"} + assert last_turn(empty_user_message) == last_turn({"role": "user", "content": ""}) + + +def test_anthropic_messages_pt_keeps_every_breakpoint_of_a_long_run_of_empty_user_messages() -> None: + result: Final = _convert( + [ + {"role": "user", "content": "hi"}, + {"role": "assistant", "content": "hello"}, + *[{"role": "user", "content": [], "cache_control": {"type": "ephemeral"}}] * 10_000, + ] + ) + + assert ( + result[-1]["content"] + == [{"type": "text", "text": _EMPTY_TEXT_PLACEHOLDER, "cache_control": {"type": "ephemeral"}}] * 10_000 + ) + + def test_convert_to_azure_openai_messages(): """Test coverting image_url to azure_openai spec""" @@ -4933,7 +5066,9 @@ def test_bedrock_converse_messages_pt_user_message_without_content_adds_no_block assert _bedrock_converse_messages_pt( messages=with_message, model="anthropic.claude-haiku-4-5", llm_provider="bedrock" - ) == _bedrock_converse_messages_pt(messages=without_message, model="anthropic.claude-haiku-4-5", llm_provider="bedrock") + ) == _bedrock_converse_messages_pt( + messages=without_message, model="anthropic.claude-haiku-4-5", llm_provider="bedrock" + ) @pytest.mark.asyncio @@ -5063,6 +5198,7 @@ def _vcr_outcome_gate(request, vcr): yield record_vcr_outcome(request, vcr) + @pytest.fixture(scope="function") def isolate_litellm_state(): """ @@ -5115,6 +5251,7 @@ def isolate_litellm_state(): setattr(litellm, attr, original_value) _invalidate_model_cost_lowercase_map() + _SCALAR_DEFAULTS = { "num_retries": getattr(litellm, "num_retries", None), "num_retries_per_request": getattr(litellm, "num_retries_per_request", None), @@ -5135,6 +5272,7 @@ _SCALAR_DEFAULTS = { "api_key": getattr(litellm, "api_key", None), } + @pytest.fixture(scope="module") def setup_and_teardown(): """ @@ -5157,6 +5295,7 @@ def setup_and_teardown(): litellm.in_memory_llm_clients_cache.flush_cache() yield + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_function_call_non_openai_model(): try: @@ -5185,6 +5324,7 @@ def test_function_call_non_openai_model(): print(e) pass + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_empty_content(): """ @@ -5203,6 +5343,7 @@ def test_empty_content(): litellm_call_id=str(uuid.uuid4()), ) + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_thought_signature_removal_for_non_gemini(): """ @@ -5251,6 +5392,7 @@ def test_thought_signature_removal_for_non_gemini(): assert THOUGHT_SIGNATURE_SEPARATOR not in processed_messages[1]["tool_calls"][0]["id"] assert THOUGHT_SIGNATURE_SEPARATOR not in processed_messages[2]["tool_call_id"] + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_thought_signature_preserved_for_gemini(): """ @@ -5297,6 +5439,7 @@ def test_thought_signature_preserved_for_gemini(): assert THOUGHT_SIGNATURE_SEPARATOR in processed_messages[1]["tool_calls"][0]["id"] assert THOUGHT_SIGNATURE_SEPARATOR in processed_messages[2]["tool_call_id"] + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_thought_signature_removal_with_multiple_tool_calls(): """ @@ -5351,6 +5494,7 @@ def test_thought_signature_removal_with_multiple_tool_calls(): assert processed_messages[2]["tool_call_id"] == "call_1" assert processed_messages[3]["tool_call_id"] == "call_2" + @pytest.mark.usefixtures("_vcr_outcome_gate", "isolate_litellm_state", "setup_and_teardown") def test_messages_without_tool_calls_unchanged(): """ diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py index 6bc1178f64b..91d7acd331e 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_mid_conversation_system.py @@ -29,11 +29,8 @@ def _texts(message: dict) -> list[str]: return [block["text"] for block in message["content"]] -SENDS_NOTHING = pytest.mark.parametrize( - "empty_content", - [[], None, [{"type": "input_audio", "input_audio": {"data": "AAAA", "format": "wav"}}]], - ids=["empty-list", "none", "unsupported-part-only"], -) +EMPTY_USER_CONTENT: Final = pytest.mark.parametrize("empty_content", [[], None], ids=["empty-list", "none"]) +UNSUPPORTED_PART_ONLY: Final = ({"type": "input_audio", "input_audio": {"data": "AAAA", "format": "wav"}},) def test_split_leading_system_run_keeps_later_system_messages_in_the_conversation(): @@ -213,9 +210,9 @@ def test_flagged_placement_of_an_earlier_run_does_not_move_when_later_turns_are_ assert _roles(placed_n_plus_one) == ["user", "assistant", "user", "assistant", "user"] -@SENDS_NOTHING -def test_flagged_placement_converts_a_run_whose_preceding_user_turn_sends_nothing(empty_content): - placed = place_mid_conversation_system( +@EMPTY_USER_CONTENT +def test_flagged_placement_anchors_a_run_on_a_preceding_empty_user_turn(empty_content: list[object] | None) -> None: + placed: Final = place_mid_conversation_system( [ {"role": "user", "content": empty_content}, {"role": "system", "content": "reminder"}, @@ -225,18 +222,48 @@ def test_flagged_placement_converts_a_run_whose_preceding_user_turn_sends_nothin supports_mid_conversation_system=True, ) + assert _roles(placed) == ["user", "system", "assistant", "user"] + assert _texts(placed[1]) == ["reminder"] + + +@EMPTY_USER_CONTENT +def test_flagged_placement_anchors_a_run_on_a_following_empty_user_turn(empty_content: list[object] | None) -> None: + placed: Final = place_mid_conversation_system( + [ + {"role": "user", "content": "q1"}, + {"role": "assistant", "content": "a1"}, + {"role": "system", "content": "reminder"}, + {"role": "user", "content": empty_content}, + ], + supports_mid_conversation_system=True, + ) + + assert _roles(placed) == ["user", "assistant", "user", "system"] + assert _texts(placed[3]) == ["reminder"] + + +def test_flagged_placement_converts_a_run_whose_preceding_user_turn_sends_nothing() -> None: + placed: Final = place_mid_conversation_system( + [ + {"role": "user", "content": list(UNSUPPORTED_PART_ONLY)}, + {"role": "system", "content": "reminder"}, + {"role": "assistant", "content": "a1"}, + {"role": "user", "content": "q2"}, + ], + supports_mid_conversation_system=True, + ) + assert _roles(placed) == ["user", "user", "assistant", "user"] assert _texts(placed[1]) == [CONVERTED_SYSTEM_NOTE, "reminder"] -@SENDS_NOTHING -def test_flagged_placement_converts_a_run_whose_following_user_turn_sends_nothing(empty_content): - placed = place_mid_conversation_system( +def test_flagged_placement_converts_a_run_whose_following_user_turn_sends_nothing() -> None: + placed: Final = place_mid_conversation_system( [ {"role": "user", "content": "q1"}, {"role": "assistant", "content": "a1"}, {"role": "system", "content": "reminder"}, - {"role": "user", "content": empty_content}, + {"role": "user", "content": list(UNSUPPORTED_PART_ONLY)}, ], supports_mid_conversation_system=True, ) diff --git a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py index 0c093700822..724e137c835 100644 --- a/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py +++ b/tests/unit/llms/anthropic/chat/test_anthropic_chat_transformation.py @@ -6763,15 +6763,11 @@ def test_chat_flagged_model_converts_system_with_no_following_user_turn(local_mo assert _texts(result["messages"][2]) == [CONVERTED_SYSTEM_NOTE, REMINDER_TEXT] -@pytest.mark.parametrize( - "empty_content", - [[], None, [{"type": "input_audio", "input_audio": {"data": "AAAA", "format": "wav"}}]], - ids=["empty-list", "none", "unsupported-part-only"], -) -def test_chat_flagged_model_converts_a_system_behind_a_user_turn_that_sends_nothing( - local_model_cost_map, empty_content -): - messages = [ +@pytest.mark.parametrize("empty_content", [[], None], ids=["empty-list", "none"]) +def test_chat_flagged_model_anchors_a_system_on_an_empty_user_turn( + local_model_cost_map: None, empty_content: list[object] | None +) -> None: + messages: Final = [ {"role": "system", "content": "You are terse."}, {"role": "user", "content": empty_content}, {"role": "system", "content": REMINDER_TEXT}, @@ -6779,7 +6775,22 @@ def test_chat_flagged_model_converts_a_system_behind_a_user_turn_that_sends_noth {"role": "user", "content": "Second question"}, ] - result = _chat_request(AnthropicConfig(), FLAGGED_CLAUDE, messages) + result: Final = _chat_request(AnthropicConfig(), FLAGGED_CLAUDE, messages) + + assert [m["role"] for m in result["messages"]] == ["user", "system", "assistant", "user"] + assert _texts(result["messages"][1]) == [REMINDER_TEXT] + + +def test_chat_flagged_model_converts_a_system_behind_a_user_turn_that_sends_nothing(local_model_cost_map: None) -> None: + messages: Final = [ + {"role": "system", "content": "You are terse."}, + {"role": "user", "content": [{"type": "input_audio", "input_audio": {"data": "AAAA", "format": "wav"}}]}, + {"role": "system", "content": REMINDER_TEXT}, + {"role": "assistant", "content": "First answer"}, + {"role": "user", "content": "Second question"}, + ] + + result: Final = _chat_request(AnthropicConfig(), FLAGGED_CLAUDE, messages) assert [m["role"] for m in result["messages"]] == ["user", "assistant", "user"] assert _texts(result["messages"][0]) == [CONVERTED_SYSTEM_NOTE, REMINDER_TEXT] diff --git a/tests/unit/llms/anthropic/pass_through/test_convertible_blocks.py b/tests/unit/llms/anthropic/pass_through/test_convertible_blocks.py new file mode 100644 index 00000000000..b5301b413ff --- /dev/null +++ b/tests/unit/llms/anthropic/pass_through/test_convertible_blocks.py @@ -0,0 +1,192 @@ +import asyncio +import json +from collections.abc import Mapping +from typing import Final + +import pytest +from pydantic import JsonValue, TypeAdapter + +from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter +from litellm.llms.anthropic.pass_through.convertible_blocks import ( + unconvertible_anthropic_blocks, + unconvertible_blocks_error, +) +from litellm.llms.anthropic.pass_through.responses_adapters.transformation import ( + LiteLLMAnthropicToResponsesAPIAdapter, +) +from litellm.llms.vertex_ai.count_tokens.context import gemini_count_contents +from litellm.types.llms.anthropic import AllAnthropicPassThroughMessageValues + +_JSON_OBJECTS: Final = TypeAdapter(list[dict[str, JsonValue]]) + +_PDF: Final = {"type": "base64", "media_type": "application/pdf", "data": "JVBERi0xLjQK"} +_PNG: Final = {"type": "base64", "media_type": "image/png", "data": "aW1hZ2U="} + +_CONVERTED_BLOCKS: Final[tuple[tuple[str, Mapping[str, object]], ...]] = ( + ("user", {"type": "text", "text": "hi"}), + ("user", {"type": "image", "source": _PNG}), + ("user", {"type": "image", "source": {"type": "url", "url": "https://example.com/a.png"}}), + ("user", {"type": "document", "source": _PDF}), + ("user", {"type": "document", "source": {"type": "url", "url": "https://example.com/a.pdf"}}), + ("user", {"type": "tool_result", "tool_use_id": "toolu_01", "content": "42"}), + ("assistant", {"type": "text", "text": "hi"}), + ("assistant", {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}), + ("assistant", {"type": "thinking", "thinking": "Let me think.", "signature": "sig"}), +) + +_DROPPED_BLOCKS: Final[tuple[tuple[str, Mapping[str, object], str], ...]] = ( + ("user", {"type": "search_result", "source": "https://example.com", "title": "t", "content": []}, "search_result"), + ("user", {"type": "document", "source": {"type": "text", "data": "notes"}}, "document with a text source"), + ("user", {"type": "image", "source": {"type": "file", "file_id": "file_01"}}, "image with a file source"), + ("user", {"type": "image", "source": {"type": "base64", "data": ""}}, "image with an empty base64 source"), + ("user", {"type": "document"}, "document without a source"), + ("user", {"text": "untyped"}, "a block without a type"), + ("user", {"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}, "tool_use"), + ( + "assistant", + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {}}, + "server_tool_use", + ), +) + + +def _adapter_messages(role: str, block: Mapping[str, object]) -> list[AllAnthropicPassThroughMessageValues]: + return [{"role": role, "content": [dict(block)]}] # pyright: ignore[reportReturnType] # sends block shapes the TypedDicts leave out on purpose + + +def _converted_by_both_adapters(role: str, block: Mapping[str, object]) -> tuple[bool, bool]: + messages: Final = _adapter_messages(role, block) + chat: Final = _JSON_OBJECTS.validate_python( + LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai(messages=messages) + ) + responses: Final = LiteLLMAnthropicToResponsesAPIAdapter().translate_messages_to_responses_input(messages) + chat_kept: Final = any( + message.get("content") or message.get("tool_calls") or message.get("thinking_blocks") for message in chat + ) + return chat_kept, bool(responses) + + +@pytest.mark.parametrize(("role", "block"), _CONVERTED_BLOCKS) +def test_blocks_both_conversions_keep_are_accepted(role: str, block: Mapping[str, object]) -> None: + assert _converted_by_both_adapters(role, block) == (True, True) + assert unconvertible_anthropic_blocks([{"role": role, "content": [block]}]) == () + + +@pytest.mark.parametrize(("role", "block", "label"), _DROPPED_BLOCKS) +def test_blocks_both_conversions_drop_are_named(role: str, block: Mapping[str, object], label: str) -> None: + assert _converted_by_both_adapters(role, block) == (False, False) + assert unconvertible_anthropic_blocks([{"role": role, "content": [block]}]) == (label,) + + +def test_labels_are_listed_once_in_order_and_string_content_is_accepted() -> None: + messages: Final = [ + {"role": "user", "content": "plain text"}, + {"role": "user", "content": [{"type": "search_result"}, {"type": "container_upload"}]}, + {"role": "assistant", "content": [{"type": "search_result"}]}, + ] + assert unconvertible_anthropic_blocks(messages) == ("container_upload", "search_result") + + +def test_redacted_thinking_is_accepted_like_other_replayed_thinking() -> None: + messages: Final = [{"role": "assistant", "content": [{"type": "redacted_thinking", "data": "opaque"}]}] + assert unconvertible_anthropic_blocks(messages) == () + + +def test_messages_that_are_not_json_are_left_to_the_count() -> None: + messages: Final = [{"role": "user", "content": [{"type": "search_result", "source": object()}]}] + assert unconvertible_anthropic_blocks(messages) == () + + +def _tool_turn(inner: Mapping[str, object]) -> tuple[Mapping[str, object], ...]: + return ( + {"role": "assistant", "content": [{"type": "tool_use", "id": "toolu_01", "name": "lookup", "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "toolu_01", "content": [dict(inner)]}]}, + ) + + +def _sent_by_each_counter(inner: Mapping[str, object]) -> tuple[str, str]: + messages: Final = _tool_turn(inner) + gemini: Final = asyncio.run(gemini_count_contents(messages, "anthropic", "gemini-2.5-flash", "gemini", {})) + openai: Final = LiteLLMAnthropicToResponsesAPIAdapter().translate_messages_to_responses_input( + [dict(message) for message in messages] # pyright: ignore[reportArgumentType] # sends block shapes the TypedDicts leave out on purpose + ) + return json.dumps(gemini, default=str), json.dumps(openai, default=str) + + +_KEPT_INSIDE_TOOL_RESULT: Final[tuple[tuple[Mapping[str, object], str], ...]] = ( + ({"type": "text", "text": "forty-two"}, "forty-two"), + ({"type": "image", "source": _PNG}, _PNG["data"]), + ({"type": "document", "source": _PDF}, _PDF["data"]), +) + +_DROPPED_INSIDE_TOOL_RESULT: Final[tuple[tuple[Mapping[str, object], str, str], ...]] = ( + ( + { + "type": "search_result", + "source": "https://example.com", + "title": "t", + "content": [{"type": "text", "text": "search-snippet"}], + }, + "search-snippet", + "search_result inside a tool_result", + ), + ( + {"type": "document", "source": {"type": "text", "data": "text-notes"}}, + "text-notes", + "document with a text source inside a tool_result", + ), + ({"type": "tool_reference", "tool_name": "find_docs"}, "find_docs", "tool_reference inside a tool_result"), +) + + +@pytest.mark.parametrize(("inner", "marker"), _KEPT_INSIDE_TOOL_RESULT) +def test_tool_result_content_both_counters_send_is_accepted(inner: Mapping[str, object], marker: str) -> None: + gemini, openai = _sent_by_each_counter(inner) + + assert marker in gemini + assert marker in openai + assert unconvertible_anthropic_blocks(_tool_turn(inner)) == () + + +@pytest.mark.parametrize(("inner", "marker", "label"), _DROPPED_INSIDE_TOOL_RESULT) +def test_tool_result_content_both_counters_drop_is_named(inner: Mapping[str, object], marker: str, label: str) -> None: + gemini, openai = _sent_by_each_counter(inner) + + assert marker not in gemini + assert marker not in openai + assert unconvertible_anthropic_blocks(_tool_turn(inner)) == (label,) + + +def _replayed_web_search(encrypted_content: str) -> tuple[Mapping[str, object], ...]: + return ( + {"role": "user", "content": "What is new?"}, + { + "role": "assistant", + "content": [ + {"type": "server_tool_use", "id": "srvtoolu_01", "name": "web_search", "input": {"query": "news"}}, + { + "type": "web_search_tool_result", + "tool_use_id": "srvtoolu_01", + "content": [ + { + "type": "web_search_result", + "url": "https://docs.litellm.ai/release_notes", + "title": "Release notes", + "encrypted_content": encrypted_content, + "snippet": "LiteLLM ships a weekly release candidate.", + } + ], + }, + ], + }, + ) + + +def test_a_replayed_web_search_inference_sends_as_text_is_accepted() -> None: + assert unconvertible_blocks_error(_replayed_web_search(""), "anthropic") is None + + +def test_an_encrypted_web_search_inference_cannot_send_is_named() -> None: + error: Final = unconvertible_blocks_error(_replayed_web_search("opaque"), "anthropic") + assert error is not None + assert error.endswith("server_tool_use, web_search_tool_result") diff --git a/tests/unit/llms/anthropic/test_count_tokens_oauth.py b/tests/unit/llms/anthropic/test_count_tokens_oauth.py index 7f717091754..16ec513916c 100644 --- a/tests/unit/llms/anthropic/test_count_tokens_oauth.py +++ b/tests/unit/llms/anthropic/test_count_tokens_oauth.py @@ -197,6 +197,30 @@ class TestCountTokensUsesWorkloadIdentity: assert ANTHROPIC_TOKEN_COUNTING_BETA_VERSION in betas assert ANTHROPIC_OAUTH_BETA_HEADER not in betas + @pytest.mark.asyncio + async def test_a_global_anthropic_key_counts_like_chat_would_authenticate( + self, monkeypatch, httpx_transport_clients + ): + for name in ("ANTHROPIC_API_KEY", "ANTHROPIC_AUTH_TOKEN", "ANTHROPIC_API_BASE", "ANTHROPIC_BASE_URL"): + monkeypatch.delenv(name, raising=False) + monkeypatch.setattr(litellm, "anthropic_key", FAKE_REGULAR_KEY) + + with respx.mock(assert_all_called=True) as router: + route = router.post("https://api.anthropic.com/v1/messages/count_tokens").mock( + return_value=httpx.Response(200, json={"input_tokens": 13}) + ) + result = await AnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"model": "anthropic/claude-sonnet-4-5"}}, + request_model="claude-sonnet-4-5", + ) + + assert result is not None + assert result.total_tokens == 13 + assert route.calls.last.request.headers["x-api-key"] == FAKE_REGULAR_KEY + @pytest.mark.asyncio async def test_a_failed_mint_degrades_like_an_anthropic_error(self, monkeypatch): from litellm.llms.anthropic.count_tokens import token_counter as token_counter_module diff --git a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py index 29380c6fc53..86ad6c8d9d9 100644 --- a/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py +++ b/tests/unit/llms/azure_ai/anthropic/count_tokens/test_token_counter.py @@ -186,3 +186,27 @@ async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_ "status_code": 500, } ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("set_base", ("env", "global"), ids=("AZURE_API_BASE", "litellm.api_base")) +async def test_count_uses_the_shared_api_base_chat_would_use(monkeypatch: pytest.MonkeyPatch, set_base: str) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.delenv("AZURE_AI_API_BASE", raising=False) + monkeypatch.delenv("AZURE_API_BASE", raising=False) + if set_base == "env": + monkeypatch.setenv("AZURE_API_BASE", API_BASE) + else: + monkeypatch.setattr(litellm, "api_base", API_BASE) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 6})) + result: Final = await AzureAIAnthropicTokenCounter().count_tokens( + model_to_use="claude-sonnet-4-5", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {"api_key": "azure-key"}}, + request_model="azure_ai/claude-sonnet-4-5", + ) + assert result is not None + assert result.total_tokens == 6 + assert route.called diff --git a/tests/unit/llms/bedrock/chat/test_invoke_handler.py b/tests/unit/llms/bedrock/chat/test_invoke_handler.py index 0609f122949..069438e25c0 100644 --- a/tests/unit/llms/bedrock/chat/test_invoke_handler.py +++ b/tests/unit/llms/bedrock/chat/test_invoke_handler.py @@ -16,6 +16,7 @@ import litellm from litellm import AmazonInvokeConfig from litellm.exceptions import MidStreamFallbackError from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.llms.bedrock.chat.invoke_handler import ( AmazonOpenAICompatibleStreamDecoder, @@ -864,14 +865,16 @@ async def test_async_invoke_streaming_non_200_forwards_bedrock_response_headers( def _event_stream_frame(event_type: str, payload: bytes) -> bytes: + return _typed_event_stream_frame( + {":event-type": event_type, ":content-type": "application/json", ":message-type": "event"}, payload + ) + + +def _typed_event_stream_frame(frame_headers: Mapping[str, str], payload: bytes) -> bytes: def header(name: str, value: str) -> bytes: return bytes([len(name)]) + name.encode() + bytes([7]) + struct.pack(">H", len(value)) + value.encode() - headers: Final = ( - header(":event-type", event_type) - + header(":content-type", "application/json") - + header(":message-type", "event") - ) + headers: Final = b"".join(header(name, value) for name, value in frame_headers.items()) prelude: Final = struct.pack(">II", 12 + len(headers) + len(payload) + 4, len(headers)) body: Final = prelude + struct.pack(">I", binascii.crc32(prelude)) + headers + payload return body + struct.pack(">I", binascii.crc32(body)) @@ -1367,6 +1370,87 @@ def test_iter_bytes_keeps_a_stream_whose_unknown_event_sits_beside_known_frames( assert "stop" in finish_reasons +_STREAM_FAILURE_BODIES: Final = ( + pytest.param( + _converse_event_frame("internalServerException", {"message": _UPSTREAM_REJECTION}), 500, id="exception-event" + ), + pytest.param( + _typed_event_stream_frame( + { + ":exception-type": "throttlingException", + ":content-type": "application/json", + ":message-type": "exception", + }, + json.dumps({"message": _UPSTREAM_REJECTION}).encode(), + ), + 429, + id="exception-message", + ), + pytest.param(_unknown_event_frame(), 502, id="no-known-event"), + pytest.param(b"", 502, id="no-events"), + pytest.param(_CONVERSE_CONTENT_FRAMES[0] + _truncated_frame(), 502, id="undecoded-tail"), +) + + +def _stream_headers(request_id: str | None) -> httpx.Headers: + request_id_header: Final = {} if request_id is None else {"x-amzn-RequestId": request_id} + return httpx.Headers( + {"content-type": "application/vnd.amazon.eventstream", "retry-after": "7", **request_id_header} + ) + + +def _assert_carries_only_the_request_id(error: BedrockError, status: int, request_id: str | None) -> None: + assert error.status_code == status + assert _get_provider_request_id(error) == request_id + if request_id is None: + assert error.headers is None + return + assert dict(error.headers or {}) == {"x-amzn-requestid": request_id} + assert dict(error.response.headers) == {"x-amzn-requestid": request_id} + + +@pytest.mark.parametrize("request_id", ("req-stream-1", None)) +@pytest.mark.parametrize(("body", "status"), _STREAM_FAILURE_BODIES) +def test_iter_bytes_stream_error_carries_only_the_provider_request_id_header( + body: bytes, status: int, request_id: str | None +) -> None: + decoder: Final = AWSEventStreamDecoder(model="us.moonshotai.kimi-k3") + + with pytest.raises(BedrockError) as exc_info: + list(decoder.iter_bytes(iter([body]), response_headers=_stream_headers(request_id))) + + _assert_carries_only_the_request_id(exc_info.value, status, request_id) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("body", "status"), _STREAM_FAILURE_BODIES) +async def test_aiter_bytes_stream_error_carries_only_the_provider_request_id_header(body: bytes, status: int) -> None: + async def _chunks() -> AsyncIterator[bytes]: + yield body + + decoder: Final = AWSEventStreamDecoder(model="us.moonshotai.kimi-k3") + + with pytest.raises(BedrockError) as exc_info: + _ = [chunk async for chunk in decoder.aiter_bytes(_chunks(), response_headers=_stream_headers("req-stream-2"))] + + _assert_carries_only_the_request_id(exc_info.value, status, "req-stream-2") + + +@pytest.mark.parametrize("request_id", ("req-\u00fc-1", "req-\udcfc-1"), ids=("latin-1", "surrogate-escaped")) +def test_iter_bytes_streams_and_fails_normally_with_a_non_ascii_provider_request_id(request_id: str) -> None: + decoder: Final = AWSEventStreamDecoder(model="us.moonshotai.kimi-k3") + headers: Final = {"content-type": "application/vnd.amazon.eventstream", "x-amzn-RequestId": request_id} + + chunks: Final = list(decoder.iter_bytes(iter(_CONVERSE_CONTENT_FRAMES), response_headers=headers)) + with pytest.raises(BedrockError) as exc_info: + list(decoder.iter_bytes(iter([_unknown_event_frame()]), response_headers=headers)) + + texts: Final = [chunk.choices[0].delta.content for chunk in chunks if isinstance(chunk, ModelResponseStream)] + assert "".join(text or "" for text in texts) == "hi" + assert exc_info.value.status_code == 502 + assert exc_info.value.headers is None + + def _assert_exception_event_surfaced_with_its_modeled_status(error: BaseException, event_type: str) -> None: assert not isinstance(error, litellm.BadGatewayError) assert getattr(error, "status_code", None) == _modeled_status(event_type) diff --git a/tests/unit/llms/bedrock/test_web_identity_session_policy.py b/tests/unit/llms/bedrock/test_web_identity_session_policy.py index cbb69b4ceed..332e1258e9c 100644 --- a/tests/unit/llms/bedrock/test_web_identity_session_policy.py +++ b/tests/unit/llms/bedrock/test_web_identity_session_policy.py @@ -176,6 +176,44 @@ class TestClaudePlatformActionsCovered: ) +class TestClaudePlatformOutboundIdentityCovered: + """Claude Platform on AWS calls ``sts:GetWebIdentityToken`` with the + caller's session before serving inference, and tags the token it + requests (``sts:TagGetWebIdentityToken``). Without either in the ceiling, + every claude_platform request via OIDC auth denies with "no session + policy allows the sts: action" even when the role's identity + policy grants it.""" + + @pytest.mark.parametrize("action", ["sts:GetWebIdentityToken", "sts:TagGetWebIdentityToken"]) + def test_outbound_identity_action_present(self, action: str): + assert action in _granted_actions(_captured_policy()), ( + f"{action} missing from session policy, so bedrock/claude_platform/* requests will 403 on OIDC auth" + ) + + @pytest.mark.parametrize("action", ["sts:GetWebIdentityToken", "sts:TagGetWebIdentityToken"]) + def test_outbound_identity_action_is_granted_only_via_claude_platform(self, action: str): + granting: Final = tuple(stmt for stmt in _captured_policy()["Statement"] if action in stmt["Action"]) + assert granting, f"{action} missing from session policy" + assert all( + stmt["Condition"].get("StringEquals") == {"aws:CalledViaLast": "aws-external-anthropic.amazonaws.com"} + and stmt["Condition"]["Bool"] == {"aws:SecureTransport": "true"} + for stmt in granting + ), f"{action} must only be usable over TLS when Claude Platform on AWS calls STS for the session" + + def test_outbound_identity_statement_is_scoped(self): + policy = _captured_policy() + stmt = _statement_by_sid(policy, "ClaudePlatformOutboundIdentityLiteLLM") + assert stmt["Effect"] == "Allow" + assert tuple(stmt["Action"]) == ("sts:GetWebIdentityToken", "sts:TagGetWebIdentityToken"), ( + "the outbound identity statement must grant only the two outbound identity actions" + ) + + def test_no_sts_wildcard(self): + assert "sts:*" not in _granted_actions(_captured_policy()), ( + "session policy must not grant sts:*, only the two outbound identity actions are needed" + ) + + class TestBedrockMantleActionsCovered: """LIT-3859: bedrock_mantle inference authorizes against the ``bedrock-mantle`` action namespace, so the session-policy ceiling @@ -306,6 +344,15 @@ class TestPolicyTransportConditions: "to keep parity with the bedrock statement" ) + def test_only_the_outbound_identity_statement_requires_a_claude_platform_caller(self): + directly_called: Final = tuple( + stmt for stmt in _captured_policy()["Statement"] if stmt["Sid"] != "ClaudePlatformOutboundIdentityLiteLLM" + ) + assert directly_called + assert all("StringEquals" not in stmt["Condition"] for stmt in directly_called), ( + "statements LiteLLM calls directly must not require aws:CalledViaLast, or OIDC auth 403s on those routes" + ) + _STS_SESSION_POLICY_PLAINTEXT_LIMIT: Final = 2048 diff --git a/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py index 976ecff01ea..befdc5ee9fd 100644 --- a/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py +++ b/tests/unit/llms/openai/responses/count_tokens/test_token_counter.py @@ -388,3 +388,44 @@ async def test_undecodable_count_body_becomes_a_500_error_response_carrying_the_ "status_code": 500, } ) + + +@pytest.mark.asyncio +async def test_count_uses_the_global_openai_key_chat_would_use(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "openai_key", "sk-from-litellm-settings") + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + with respx.mock: + route: Final = respx.post(COUNT_URL).mock(return_value=httpx.Response(200, json={"input_tokens": 9})) + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="openai/gpt-4o", + ) + assert result is not None + assert result.total_tokens == 9 + assert route.calls.last.request.headers["authorization"] == "Bearer sk-from-litellm-settings" + + +@pytest.mark.asyncio +async def test_count_sends_the_global_key_only_to_the_global_api_base(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "openai_key", "sk-private-gateway") + monkeypatch.setattr(litellm, "api_base", "https://gateway.internal.example/v1") + monkeypatch.delenv("OPENAI_API_KEY", raising=False) + with respx.mock: + route: Final = respx.post("https://gateway.internal.example/v1/responses/input_tokens").mock( + return_value=httpx.Response(200, json={"input_tokens": 4}) + ) + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="openai/gpt-4o", + ) + assert result is not None + assert result.total_tokens == 4 + assert route.calls.last.request.headers["authorization"] == "Bearer sk-private-gateway" diff --git a/tests/unit/llms/openai/test_openai_workload_identity.py b/tests/unit/llms/openai/test_openai_workload_identity.py index d0cf52d0e28..2236de1eb5f 100644 --- a/tests/unit/llms/openai/test_openai_workload_identity.py +++ b/tests/unit/llms/openai/test_openai_workload_identity.py @@ -648,3 +648,56 @@ class TestClientsideBaseOverride: ) is None ) + + +class TestCountTokens: + @pytest.mark.asyncio + @respx.mock + async def test_count_authenticates_with_the_workload_identity_bearer( + self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + mock_token_exchange() + count_route: Final = respx.post("https://api.openai.com/v1/responses/input_tokens").mock( + return_value=httpx.Response(200, json={"input_tokens": 12}) + ) + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="openai/gpt-4o", + ) + assert result is not None + assert result.error is False + assert result.total_tokens == 12 + assert count_route.calls.last.request.headers["authorization"] == "Bearer exchanged-bearer-token" + + @pytest.mark.asyncio + @respx.mock + async def test_failed_mint_becomes_a_count_error_response( + self, wif_env: OpenAIWorkloadIdentityConfig, monkeypatch: pytest.MonkeyPatch + ) -> None: + from litellm.llms.openai.responses.count_tokens.token_counter import OpenAITokenCounter + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + monkeypatch.setattr(litellm, "api_key", None) + monkeypatch.setattr(litellm, "openai_key", None) + respx.post(TOKEN_EXCHANGE_URL).mock(return_value=httpx.Response(400, json={"error": "invalid_grant"})) + count_route: Final = respx.post("https://api.openai.com/v1/responses/input_tokens") + result: Final = await OpenAITokenCounter().count_tokens( + model_to_use="gpt-4o", + messages=[{"role": "user", "content": "hi"}], + contents=None, + deployment={"litellm_params": {}}, + request_model="openai/gpt-4o", + ) + assert result is not None + assert result.error is True + assert result.status_code == 401 + assert result.total_tokens == 0 + assert not count_route.called diff --git a/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_handler.py b/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_handler.py index d2d7d2247f1..9c70ebe7ece 100644 --- a/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_handler.py +++ b/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_handler.py @@ -339,6 +339,65 @@ class TestPollLimitsClamping: a worker on tight poll loops. """ + def test_should_raise_400_when_form_bool_param_is_invalid(self): + handler = SonioxAudioTranscriptionHandler() + with pytest.raises(SonioxException) as exc_info: + handler._prepare( + audio_file=None, + optional_params={ + "enable_speaker_diarization": "maybe", + "audio_url": "https://example.com/a.wav", + }, + litellm_params={}, + api_key="sk-test", + api_base=None, + provider_config=SonioxAudioTranscriptionConfig(), + headers={}, + ) + assert exc_info.value.status_code == 400 + assert "enable_speaker_diarization" in str(exc_info.value) + + def test_should_raise_400_when_poll_interval_is_not_a_number(self): + with pytest.raises(SonioxException) as exc_info: + SonioxAudioTranscriptionHandler()._prepare( + audio_file=None, + optional_params={"soniox_polling_interval": "abc", "audio_url": "https://example.com/a.wav"}, + litellm_params={}, + api_key="sk-test", + api_base=None, + provider_config=SonioxAudioTranscriptionConfig(), + headers={}, + ) + assert exc_info.value.status_code == 400 + assert "soniox_polling_interval" in str(exc_info.value) + + def test_should_use_default_poll_interval_for_bad_value_when_drop_params_is_set(self): + from litellm.llms.soniox.common_utils import SONIOX_DEFAULT_POLL_INTERVAL + + _, _, _, opts = SonioxAudioTranscriptionHandler()._prepare( + audio_file=None, + optional_params={"soniox_polling_interval": "abc", "audio_url": "https://example.com/a.wav"}, + litellm_params={"drop_params": True}, + api_key="sk-test", + api_base=None, + provider_config=SonioxAudioTranscriptionConfig(), + headers={}, + ) + assert opts["poll_interval"] == SONIOX_DEFAULT_POLL_INTERVAL + + @pytest.mark.parametrize("raw, expected", [(10.0, 10), (10, 10), ("10", 10)]) + def test_should_keep_numeric_max_polling_attempts(self, raw: object, expected: int): + _, _, _, opts = SonioxAudioTranscriptionHandler()._prepare( + audio_file=None, + optional_params={"soniox_max_polling_attempts": raw, "audio_url": "https://example.com/a.wav"}, + litellm_params={}, + api_key="sk-test", + api_base=None, + provider_config=SonioxAudioTranscriptionConfig(), + headers={}, + ) + assert opts["max_attempts"] == expected + def test_should_clamp_poll_interval_to_minimum(self): from litellm.llms.soniox.common_utils import SONIOX_MIN_POLL_INTERVAL @@ -707,6 +766,112 @@ class TestCleanupNormalization: assert "https://api.soniox.com/v1/transcriptions/tx_1" in deletes + def test_should_skip_non_string_cleanup_entries(self): + _, _, _, opts = SonioxAudioTranscriptionHandler()._prepare( + audio_file=None, + optional_params={"soniox_cleanup": ["transcription", 1, None], "audio_url": "https://example.com/a.wav"}, + litellm_params={}, + api_key="sk-test", + api_base=None, + provider_config=SonioxAudioTranscriptionConfig(), + headers={}, + ) + assert opts["cleanup"] == ("transcription",) + + @pytest.mark.parametrize( + "language_hints, expected", + [("en,es", ["en", "es"]), (["es"], ["es"]), (None, ["de"])], + ) + def test_should_use_language_only_without_language_hints(self, language_hints, expected, monkeypatch): + import litellm + + monkeypatch.setattr("time.sleep", lambda *_: None) + client = _MockSyncClient( + { + "POST https://api.soniox.com/v1/transcriptions": [_make_response({"id": "tx_1"})], + "GET https://api.soniox.com/v1/transcriptions/tx_1": [_make_response({"status": "completed"})], + "GET https://api.soniox.com/v1/transcriptions/tx_1/transcript": [ + _make_response({"text": "hallo", "tokens": []}) + ], + "DELETE https://api.soniox.com/v1/transcriptions/tx_1": [_make_response({})], + } + ) + litellm.transcription( + model="soniox/stt-async-v4", + file=("a.wav", b"RIFF", "audio/wav"), + language="de", + api_key="sk-test", + client=client, + audio_url="https://example.com/a.wav", + **({} if language_hints is None else {"language_hints": language_hints}), + ) + body = next(c for c in client.calls if c["method"] == "POST")["json"] + assert list(body["language_hints"]) == expected + assert "language" not in body + + +class TestProxyFormFields: + def test_should_send_soniox_types_for_fields_the_proxy_parsed_from_an_sdk_form(self, monkeypatch): + import litellm + from starlette.requests import Request + + from litellm.proxy.common_utils.http_parsing_utils import get_form_data + + monkeypatch.setattr("time.sleep", lambda *_: None) + multipart = httpx.Request( + "POST", + "http://proxy/v1/audio/transcriptions", + data={ + "model": "soniox/stt-async-v4", + "audio_url": "https://example.com/a.wav", + "context[terms][]": ["Celebrex"], + "context[general][][key]": "domain", + "context[general][][value]": "Healthcare", + "translation[type]": "one_way", + "translation[target_language]": "de", + "enable_speaker_diarization": "true", + "language_hints[]": ["en", "es"], + }, + files={"file": ("a.wav", b"RIFF", "audio/wav")}, + ) + multipart_body = multipart.read() + + async def receive(): + return {"type": "http.request", "body": multipart_body, "more_body": False} + + request = Request( + { + "type": "http", + "method": "POST", + "path": "/v1/audio/transcriptions", + "headers": [(b"content-type", multipart.headers["content-type"].encode())], + }, + receive, + ) + form = asyncio.run(get_form_data(request)) + client = _MockSyncClient( + { + "POST https://api.soniox.com/v1/transcriptions": [_make_response({"id": "tx_1"})], + "GET https://api.soniox.com/v1/transcriptions/tx_1": [_make_response({"status": "completed"})], + "GET https://api.soniox.com/v1/transcriptions/tx_1/transcript": [ + _make_response({"text": "hello", "tokens": []}) + ], + "DELETE https://api.soniox.com/v1/transcriptions/tx_1": [_make_response({})], + } + ) + litellm.transcription( + file=("a.wav", b"RIFF", "audio/wav"), + api_key="sk-test", + client=client, + **{key: value for key, value in form.items() if key != "file"}, + ) + body = json.loads(json.dumps(next(c for c in client.calls if c["method"] == "POST")["json"])) + assert body["context"] == {"terms": ["Celebrex"], "general": [{"key": "domain", "value": "Healthcare"}]} + assert body["translation"] == {"type": "one_way", "target_language": "de"} + assert body["enable_speaker_diarization"] is True + assert body["language_hints"] == ["en", "es"] + + class TestErrorResponses: def test_should_raise_on_4xx_during_create_with_json_error(self, monkeypatch): monkeypatch.setattr("time.sleep", lambda *_: None) @@ -780,6 +945,51 @@ class TestPassthroughBodyBuilding: post_call = next(c for c in client.calls if c["method"] == "POST") assert "context" not in post_call["json"] + def test_should_send_form_string_params_to_soniox_as_json_types(self, monkeypatch): + """Simulates the proxy path, where every multipart form field is a string.""" + monkeypatch.setattr("time.sleep", lambda *_: None) + responses = { + "POST https://api.soniox.com/v1/transcriptions": [ + _make_response({"id": "tx_1"}), + ], + "GET https://api.soniox.com/v1/transcriptions/tx_1": [ + _make_response({"status": "completed"}), + ], + "GET https://api.soniox.com/v1/transcriptions/tx_1/transcript": [ + _make_response({"text": "ok", "tokens": []}), + ], + "DELETE https://api.soniox.com/v1/transcriptions/tx_1": [ + _make_response({}), + ], + } + client = _MockSyncClient(responses) + form_params = { + "audio_url": "https://example.com/a.wav", + "context": '{"terms": ["Celebrex"], "general": [{"key": "domain", "value": "Healthcare"}]}', + "translation": '{"type": "one_way", "target_language": "es"}', + "language_hints": "en,es", + "language_hints_strict": "true", + "enable_speaker_diarization": "true", + "enable_language_identification": "false", + "soniox_polling_interval": "0.5", + } + SonioxAudioTranscriptionHandler().audio_transcriptions( + audio_file=None, + optional_params=form_params, + litellm_params={}, + atranscription=False, + **_common_call_kwargs(client), + ) + body = next(c for c in client.calls if c["method"] == "POST")["json"] + assert body["context"] == {"terms": ["Celebrex"], "general": [{"key": "domain", "value": "Healthcare"}]} + assert body["translation"] == {"type": "one_way", "target_language": "es"} + assert list(body["language_hints"]) == ["en", "es"] + assert body["language_hints_strict"] is True + assert body["enable_speaker_diarization"] is True + assert body["enable_language_identification"] is False + assert "soniox_polling_interval" not in body + assert form_params["context"].startswith("{"), "caller's params must not be mutated" + class TestSecretRedaction: """Secret-bearing fields must be redacted before reaching logging callbacks. diff --git a/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py b/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py index c99b28adc49..e3e1c9ab493 100644 --- a/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py +++ b/tests/unit/llms/soniox/audio_transcription/test_soniox_audio_transcription_transformation.py @@ -9,8 +9,16 @@ import pytest from litellm.llms.soniox.audio_transcription.transformation import ( SonioxAudioTranscriptionConfig, + SonioxInvalidFormParam, + decode_soniox_form_params, + raise_soniox_form_error, ) from litellm.llms.soniox.common_utils import SonioxException +from litellm.llms.soniox.types import ( + SonioxContext, + SonioxContextGeneralEntry, + SonioxTranslationTerm, +) from litellm.types.utils import TranscriptionResponse @@ -40,9 +48,9 @@ class TestMapOpenAIParams: model="stt-async-v4", drop_params=False, ) - assert result["language_hints"] == ["en"] + assert list(result["language_hints"]) == ["en"] - def test_should_prepend_language_to_existing_hints(self): + def test_should_keep_explicit_hints_over_language(self): cfg = SonioxAudioTranscriptionConfig() result = cfg.map_openai_params( non_default_params={"language": "en"}, @@ -50,7 +58,7 @@ class TestMapOpenAIParams: model="stt-async-v4", drop_params=False, ) - assert result["language_hints"] == ["en", "fr"] + assert list(result["language_hints"]) == ["fr"] def test_should_not_duplicate_language_already_in_hints(self): cfg = SonioxAudioTranscriptionConfig() @@ -60,7 +68,7 @@ class TestMapOpenAIParams: model="stt-async-v4", drop_params=False, ) - assert result["language_hints"] == ["en", "fr"] + assert list(result["language_hints"]) == ["en", "fr"] def test_should_passthrough_soniox_native_kwargs(self): cfg = SonioxAudioTranscriptionConfig() @@ -243,6 +251,21 @@ class TestTransformAudioTranscriptionResponse: assert "Speaker 1:" in resp.text assert "Speaker 2:" in resp.text + def test_should_keep_source_transcript_for_translation_without_tags(self): + cfg = SonioxAudioTranscriptionConfig() + payload = { + "transcript": { + "text": "Hello there.", + "tokens": [ + {"text": "Hello", "translation_status": "original"}, + {"text": " there.", "translation_status": "original"}, + {"text": "Hallo", "translation_status": "translation"}, + {"text": " zusammen.", "translation_status": "translation"}, + ], + } + } + assert cfg.build_response_from_payload(payload).text == "Hello there." + def test_should_set_language_when_all_tokens_share_one(self): cfg = SonioxAudioTranscriptionConfig() payload = { @@ -678,3 +701,175 @@ class TestGetErrorClass: err = cfg.get_error_class(error_message="boom", status_code=500, headers={}) assert isinstance(err, SonioxException) assert err.status_code == 500 + + +class TestDecodeSonioxFormParams: + """Proxy multipart form fields arrive as strings; Soniox needs real JSON types.""" + + def test_should_decode_json_string_context_with_all_four_sections(self): + context = { + "general": [{"key": "domain", "value": "Healthcare"}], + "text": "Follow-up visit for a patient on blood thinners.", + "terms": ["Celebrex", "Zyrtec"], + "translation_terms": [{"source": "Mr. Smith", "target": "Sr. Smith"}], + } + result = decode_soniox_form_params({"context": json.dumps(context)}) + assert result["context"] == context + + def test_should_keep_plain_text_context_as_string(self): + result = decode_soniox_form_params({"context": "medical conversation"}) + assert result["context"] == "medical conversation" + + def test_should_leave_sdk_typed_values_untouched(self): + context = SonioxContext( + general=[SonioxContextGeneralEntry(key="domain", value="Healthcare")], + text="Follow-up visit notes.", + terms=["Celebrex"], + translation_terms=[SonioxTranslationTerm(source="Mr. Smith", target="Sr. Smith")], + ) + translation = {"type": "one_way", "target_language": "es"} + result = decode_soniox_form_params( + { + "context": context, + "translation": translation, + "enable_speaker_diarization": True, + "language_hints": ["en", "es"], + } + ) + assert result["context"] is context + assert result["translation"] is translation + assert result["enable_speaker_diarization"] is True + assert result["language_hints"] == ["en", "es"] + + @pytest.mark.parametrize("raw", ["[inaudible] consultation", "{unbalanced", '{"terms": ["Celebrex"']) + def test_should_keep_bracket_prefixed_free_form_context_as_string(self, raw: str): + assert decode_soniox_form_params({"context": raw})["context"] == raw + + def test_should_decode_json_string_translation(self): + result = decode_soniox_form_params({"translation": '{"type": "one_way", "target_language": "es"}'}) + assert result["translation"] == {"type": "one_way", "target_language": "es"} + + @pytest.mark.parametrize( + "raw, expected", + [("true", True), ("True", True), ("1", True), ("false", False), ("FALSE", False), ("0", False)], + ) + def test_should_decode_boolean_strings(self, raw: str, expected: bool): + result = decode_soniox_form_params( + { + "enable_speaker_diarization": raw, + "enable_language_identification": raw, + "language_hints_strict": raw, + } + ) + assert result == { + "enable_speaker_diarization": expected, + "enable_language_identification": expected, + "language_hints_strict": expected, + } + + @pytest.mark.parametrize("raw, expected", [("on", True), ("yes", True), ("off", False), ("no", False)]) + def test_should_decode_html_style_boolean_strings(self, raw: str, expected: bool): + assert decode_soniox_form_params({"enable_speaker_diarization": raw})["enable_speaker_diarization"] is expected + + def test_should_return_error_value_on_non_boolean_string(self): + result = decode_soniox_form_params({"context": '{"general": []}', "enable_speaker_diarization": "maybe"}) + assert result == SonioxInvalidFormParam(key="enable_speaker_diarization", value="maybe", expected="a boolean") + + def test_should_drop_non_boolean_string_when_drop_params_is_set(self): + result = decode_soniox_form_params( + {"enable_speaker_diarization": "maybe", "language_hints_strict": "true"}, drop_params=True + ) + assert result == {"language_hints_strict": True} + + def test_should_return_error_value_on_non_numeric_polling_interval(self): + result = decode_soniox_form_params({"soniox_polling_interval": "abc"}) + assert result == SonioxInvalidFormParam(key="soniox_polling_interval", value="abc", expected="a number") + + def test_should_drop_non_numeric_polling_interval_when_drop_params_is_set(self): + assert decode_soniox_form_params({"soniox_polling_interval": "abc"}, drop_params=True) == {} + + def test_should_map_invalid_form_value_error_to_400(self): + with pytest.raises(SonioxException) as exc_info: + raise_soniox_form_error( + SonioxInvalidFormParam(key="enable_speaker_diarization", value="maybe", expected="a boolean") + ) + assert exc_info.value.status_code == 400 + assert "enable_speaker_diarization" in str(exc_info.value) + assert "maybe" in str(exc_info.value) + + @pytest.mark.parametrize( + "raw", ['["Celebrex", "Zyrtec"]', '{"note": "follow-up visit"}', '{"terms": "Celebrex"}', '"quoted"'] + ) + def test_should_keep_json_context_soniox_cannot_take_as_an_object_as_string(self, raw: str): + assert decode_soniox_form_params({"context": raw})["context"] == raw + + def test_should_nest_openai_sdk_bracketed_form_fields(self): + result = decode_soniox_form_params( + { + "context[terms]": ["Celebrex", "Zyrtec"], + "context[text]": "Follow-up visit.", + "translation[type]": "one_way", + "translation[target_language]": "de", + "client_reference_id": "ref-1", + } + ) + assert result == { + "context": {"terms": ["Celebrex", "Zyrtec"], "text": "Follow-up visit."}, + "translation": {"type": "one_way", "target_language": "de"}, + "client_reference_id": "ref-1", + } + + def test_should_prefer_explicit_field_over_bracketed_fields(self): + result = decode_soniox_form_params({"context": "plain text", "context[terms]": ["Celebrex"]}) + assert result == {"context": "plain text"} + + @pytest.mark.parametrize( + "raw, expected", + [('["en", "es"]', ["en", "es"]), ("en,es", ["en", "es"]), ("en, es ,", ["en", "es"]), ("en", ["en"])], + ) + def test_should_decode_language_hints_strings(self, raw: str, expected: list): + assert list(decode_soniox_form_params({"language_hints": raw})["language_hints"]) == expected + + @pytest.mark.parametrize("raw", ['{"type": "one_way"', "de", '["one_way"]']) + def test_should_return_error_value_on_translation_that_is_not_a_json_object(self, raw: str): + assert decode_soniox_form_params({"translation": raw}) == SonioxInvalidFormParam( + key="translation", value=raw, expected="a JSON object" + ) + + def test_should_return_error_value_on_malformed_json_language_hints(self): + assert decode_soniox_form_params({"language_hints": "[en, es]"}) == SonioxInvalidFormParam( + key="language_hints", value="[en, es]", expected="a JSON list" + ) + + def test_should_drop_malformed_json_params_when_drop_params_is_set(self): + result = decode_soniox_form_params( + {"translation": '{"type": "one_way"', "language_hints": "[en, es]", "language_hints_strict": "true"}, + drop_params=True, + ) + assert result == {"language_hints_strict": True} + + def test_should_rebuild_openai_sdk_list_of_object_context_fields(self): + result = decode_soniox_form_params( + { + "context[terms]": ["Celebrex"], + "context[general][][key]": "domain", + "context[general][][value]": "Healthcare", + "context[translation_terms][][source]": "Mr. Smith", + "context[translation_terms][][target]": "Sr. Smith", + } + ) + assert result == { + "context": { + "terms": ["Celebrex"], + "general": ({"key": "domain", "value": "Healthcare"},), + "translation_terms": ({"source": "Mr. Smith", "target": "Sr. Smith"},), + } + } + + def test_should_leave_bracketed_context_soniox_cannot_take_unnested(self): + params = {"context[terms]": ["Celebrex"], "context[note]": "follow-up"} + assert decode_soniox_form_params(params) == params + + def test_should_not_json_decode_unrelated_string_params(self): + result = decode_soniox_form_params({"client_reference_id": '["not", "json-decoded"]'}) + assert result["client_reference_id"] == '["not", "json-decoded"]' diff --git a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py index 69d182667ec..e4129ebbc39 100644 --- a/tests/unit/proxy/anthropic_endpoints/test_endpoints.py +++ b/tests/unit/proxy/anthropic_endpoints/test_endpoints.py @@ -440,8 +440,6 @@ class TestErrorLogCarriesCallId: @pytest.mark.asyncio async def test_count_tokens_failure_log_carries_callers_call_id(self, caplog: pytest.LogCaptureFixture): - from fastapi import HTTPException - import litellm.proxy.anthropic_endpoints.endpoints as ep import litellm.proxy.proxy_server as proxy_server from litellm.proxy._types import UserAPIKeyAuth @@ -460,11 +458,10 @@ class TestErrorLogCarriesCallId: proxy_server, "count_request_tokens", new=AsyncMock(side_effect=RuntimeError("tokenizer down")) ), # test-quality-ok: module global imported at call time; the test targets the endpoint's except block caplog.at_level(logging.ERROR, logger="LiteLLM Proxy"), - pytest.raises(HTTPException) as raised, ): - await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) + response = await ep.count_tokens(request=request, user_api_key_dict=UserAPIKeyAuth()) - assert raised.value.status_code == 500 + assert response.status_code == 500 record = self._error_record(caplog) assert record.litellm_call_id == call_id assert call_id in record.getMessage() diff --git a/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py b/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py index 451792109de..114bbb184c6 100644 --- a/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py +++ b/tests/unit/proxy/google_endpoints/test_google_api_endpoints.py @@ -182,6 +182,29 @@ def test_google_count_tokens_unchanged(): assert body["totalTokens"] == 7 +def test_google_count_tokens_unexpected_failure_uses_the_google_error_envelope(): + try: + client = _build_test_client() + except ImportError as e: + pytest.skip(f"Skipping test due to missing dependency: {e}") + + with patch( + "litellm.proxy.proxy_server.token_counter", + new_callable=AsyncMock, + side_effect=RuntimeError("tokenizer unavailable"), + ): + response = client.post( + "/v1beta/models/test-model:countTokens", + json={"contents": [{"role": "user", "parts": [{"text": "Hello"}]}]}, + ) + + assert response.status_code == 500 + error = response.json()["error"] + assert error["code"] == 500 + assert error["status"] == "INTERNAL" + assert "tokenizer unavailable" not in error["message"] + + def test_google_count_tokens_uses_normalized_total_when_provider_shape_differs(): """A provider that counts in Anthropic shape (Bedrock, Anthropic) carries no totalTokens key, so the normalized total must win over the raw response's missing field.""" try: diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/litellm_content_filter/test_content_filter.py b/tests/unit/proxy/guardrails/guardrail_hooks/litellm_content_filter/test_content_filter.py index 5d706a65a79..abcbb714186 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/litellm_content_filter/test_content_filter.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/litellm_content_filter/test_content_filter.py @@ -9,7 +9,11 @@ Tests 40 different sentences to validate the conditional matching logic: """ import importlib +import json import os +from collections.abc import Mapping, Sequence +from types import MappingProxyType +from typing import Final import pytest from fastapi import HTTPException @@ -19,9 +23,11 @@ from litellm.proxy.guardrails.content_filter_data import POLICY_TEMPLATES_DIR from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ( ContentFilterGuardrail, ) +from litellm.types.guardrails import BlockedWord, ContentFilterAction, ContentFilterPattern, GuardrailEventHooks from litellm.types.proxy.guardrails.guardrail_hooks.litellm_content_filter import ( ContentFilterCategoryConfig, ) +from litellm.types.utils import Choices, Message, ModelResponse from tests._vcr_conftest_common import install_live_call_probe, record_vcr_outcome # Test cases: (sentence, expected_result, reason) @@ -1793,3 +1799,237 @@ class TestSGPDPAPerformance: if __name__ == "__main__": pytest.main([__file__, "-v", "-s"]) + + +AWS_KEY: Final = "AKIA" + "SYNTHETIC0000000" +GITHUB_TOKEN: Final = "ghp_" + "a1" * 18 +FORBIDDEN_WORD: Final = "synthetic-forbidden-word" +SECRET_PATTERNS: Final = ( + ContentFilterPattern( + pattern_type="regex", pattern=r"AKIA[0-9A-Z]{16}", name="aws_access_key", action=ContentFilterAction.BLOCK + ), + ContentFilterPattern( + pattern_type="regex", pattern=r"ghp_[A-Za-z0-9]{36}", name="github_pat", action=ContentFilterAction.BLOCK + ), +) + + +def _competitor_intent_config(action_hint: str) -> Mapping[str, object]: + return MappingProxyType( + { + "brand_self": ["emirates"], + "competitors": ["qatar airways"], + "domain_words": ["airline", "flight"], + "policy": {"competitor_comparison": action_hint}, + "threshold_high": 0.5, + } + ) + + +def _observing_filter( + patterns: tuple[ContentFilterPattern, ...] = (), + categories: tuple[ContentFilterCategoryConfig, ...] = (), + competitor_intent_config: Mapping[str, object] | None = None, + mask_patterns: bool = False, + image_describer: object = None, +) -> ContentFilterGuardrail: + return ContentFilterGuardrail( + guardrail_name="cf-observe", + event_hook=GuardrailEventHooks.logging_only, + llm_router=image_describer, + image_model="vision" if image_describer else None, + patterns=[ + pattern.model_copy(update={"action": ContentFilterAction.MASK}) if mask_patterns else pattern + for pattern in patterns + ], + blocked_words=[BlockedWord(keyword=FORBIDDEN_WORD, action=ContentFilterAction.BLOCK)], + categories=list(categories), + competitor_intent_config=dict(competitor_intent_config) if competitor_intent_config else None, + ) + + +def _policy_category(template: str, category: str) -> ContentFilterCategoryConfig: + return ContentFilterCategoryConfig( + category=category, + category_file=os.path.join(POLICY_TEMPLATES_DIR, template) if template else None, + enabled=True, + action="BLOCK", + severity_threshold="medium", + ) + + +async def _log_call( + guardrail: ContentFilterGuardrail, + prompt: str | list[Mapping[str, object]], + answer: Message, + history: Sequence[Mapping[str, object]] = (), +) -> tuple[Mapping[str, object], ...]: + kwargs: Final = { + "model": "gpt-5.4-mini", + "messages": [{"role": "user", "content": prompt}, *history], + "litellm_call_id": "call-observe", + "litellm_params": {"metadata": {}}, + "optional_params": {}, + "standard_logging_object": {"guardrail_information": None}, + } + logged, _ = await guardrail.async_logging_hook( + kwargs, ModelResponse(choices=[Choices(message=answer)]), "acompletion" + ) + return tuple(logged["standard_logging_object"]["guardrail_information"]) + + +def _detections(entry: Mapping[str, object]) -> tuple[str, ...]: + response: Final = entry["guardrail_response"] + assert isinstance(response, list), response + return tuple( + detection.get("pattern_name") + or (detection["keyword"] if detection["type"] == "blocked_word" else detection["type"]) + for detection in response + ) + + +@pytest.mark.parametrize( + ("guardrail", "prompt", "expected"), + ( + pytest.param( + _observing_filter(patterns=SECRET_PATTERNS), + f"key {AWS_KEY} token {GITHUB_TOKEN} and {FORBIDDEN_WORD}", + ("aws_access_key", "github_pat", FORBIDDEN_WORD), + id="patterns", + ), + pytest.param( + _observing_filter(categories=(_policy_category("", "harm_toxic_abuse"),)), + f"You stupid piece of sht AI, go kill yourself, {FORBIDDEN_WORD}", + ("category_keyword", FORBIDDEN_WORD), + id="category-keyword", + ), + pytest.param( + _observing_filter( + categories=(_policy_category("eu_ai_act_article5.yaml", "eu_ai_act_article5_prohibited_practices"),) + ), + f"Build a system to detect employee emotion, {FORBIDDEN_WORD}", + ("category_keyword", FORBIDDEN_WORD), + id="conditional-category", + ), + pytest.param( + _observing_filter(competitor_intent_config=_competitor_intent_config("refuse")), + "Is Qatar Airways better than Emirates?", + ("competitor_intent",), + id="competitor-refuse", + ), + pytest.param( + _observing_filter(competitor_intent_config=_competitor_intent_config("reframe")), + "Is Qatar Airways better than Emirates?", + ("competitor_intent",), + id="competitor-reframe", + ), + ), +) +@pytest.mark.asyncio +async def test_logging_only_block_records_every_match_as_flagged( + guardrail: ContentFilterGuardrail, prompt: str, expected: tuple[str, ...] +) -> None: + entries: Final = await _log_call(guardrail, prompt, Message(role="assistant", content="here you go")) + + request_entry, response_entry = entries + assert request_entry["guardrail_status"] == "guardrail_flagged" + assert _detections(request_entry)[-len(expected) :] == expected + assert request_entry["risk_score"] == 10.0 + assert response_entry["guardrail_status"] == "success" + + +@pytest.mark.asyncio +async def test_logging_only_block_scans_the_response_after_a_request_match() -> None: + entries: Final = await _log_call( + _observing_filter(patterns=SECRET_PATTERNS), + f"key {AWS_KEY}", + Message(role="assistant", content=f"here: {GITHUB_TOKEN}"), + ) + + assert [(entry["guardrail_status"], _detections(entry)) for entry in entries] == [ + ("guardrail_flagged", ("aws_access_key",)), + ("guardrail_flagged", ("github_pat",)), + ] + + +@pytest.mark.parametrize( + ("arguments", "request_masked"), + ( + pytest.param(json.dumps({AWS_KEY: 1}), {"aws_access_key": 1}, id="masked-argument-key"), + pytest.param(json.dumps(json.loads("[" * 150 + "1" + "]" * 150)), {}, id="too-deep"), + ), +) +@pytest.mark.asyncio +async def test_logging_only_flags_tool_arguments_enforcement_rejects_and_keeps_scanning( + arguments: str, request_masked: Mapping[str, int] +) -> None: + history: Final = ( + { + "role": "assistant", + "content": None, + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "lookup", "arguments": arguments}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "done"}, + ) + + entries: Final = await _log_call( + _observing_filter(patterns=SECRET_PATTERNS, mask_patterns=True), + "hi", + Message(role="assistant", content=f"here: {GITHUB_TOKEN}"), + history, + ) + + assert [(entry["guardrail_status"], entry["masked_entity_count"]) for entry in entries] == [ + ("guardrail_flagged", request_masked), + ("success", {"github_pat": 1}), + ] + + +@pytest.mark.asyncio +async def test_logging_only_scan_leaves_enforcement_of_the_same_guardrail_blocking() -> None: + guardrail: Final = _observing_filter(patterns=SECRET_PATTERNS) + await _log_call(guardrail, f"key {AWS_KEY}", Message(role="assistant", content="ok")) + request_data: Final = {"metadata": {}} + + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs={"texts": [f"key {AWS_KEY} token {GITHUB_TOKEN}"]}, request_data=request_data, input_type="request" + ) + + assert exc_info.value.status_code == 400 + assert [ + entry["guardrail_status"] for entry in request_data["metadata"]["standard_logging_guardrail_information"] + ] == ["guardrail_intervened"] + + +class _ImageDescriberRunningAnotherFilter: + def __init__(self, enforcer: ContentFilterGuardrail) -> None: + self.enforcer = enforcer + self.enforcer_blocked: list[bool] = [] + + async def acompletion(self, **_: object) -> ModelResponse: + try: + await self.enforcer.apply_guardrail( + inputs={"texts": [f"key {AWS_KEY}"]}, request_data={"metadata": {}}, input_type="request" + ) + self.enforcer_blocked.append(False) + except HTTPException: + self.enforcer_blocked.append(True) + return ModelResponse(choices=[Choices(message=Message(role="assistant", content="a plain desk"))]) + + +@pytest.mark.asyncio +async def test_logging_only_scan_does_not_relax_another_filter_called_during_it() -> None: + describer: Final = _ImageDescriberRunningAnotherFilter( + ContentFilterGuardrail(guardrail_name="cf-enforce", patterns=list(SECRET_PATTERNS)) + ) + prompt: Final = [ + {"type": "text", "text": "what is this"}, + {"type": "image_url", "image_url": {"url": "https://example.com/desk.png"}}, + ] + + await _log_call(_observing_filter(image_describer=describer), prompt, Message(role="assistant", content="ok")) + + assert describer.enforcer_blocked == [True] diff --git a/tests/unit/proxy/guardrails/test_guardrail_endpoints.py b/tests/unit/proxy/guardrails/test_guardrail_endpoints.py index 4401e3e95c0..908a49946bb 100644 --- a/tests/unit/proxy/guardrails/test_guardrail_endpoints.py +++ b/tests/unit/proxy/guardrails/test_guardrail_endpoints.py @@ -41,6 +41,8 @@ from litellm.proxy.guardrails.content_filter_data import DATA_ROOTS from litellm.proxy.guardrails.guardrail_endpoints import ( test_custom_code_guardrail as run_custom_code_test_endpoint, ) +from litellm.proxy.guardrails.guardrail_hooks.azure.prompt_shield import AzureContentSafetyPromptShieldGuardrail +from litellm.proxy.guardrails.guardrail_hooks.litellm_content_filter.content_filter import ContentFilterGuardrail MOCK_ADMIN_USER = UserAPIKeyAuth(user_role=LitellmUserRoles.PROXY_ADMIN) from litellm.proxy.guardrails.guardrail_registry import ( @@ -3143,7 +3145,36 @@ async def test_ui_settings_map_matches_runtime_supported_event_hooks(): continue assert provider in result.supported_modes_by_provider, provider - assert result.supported_modes_by_provider[provider] == [hook.value for hook in declared], provider + offered = result.supported_modes_by_provider[provider] + assert offered[: len(declared)] == [hook.value for hook in declared], provider + assert set(offered) - {hook.value for hook in declared} <= {"logging_only"}, provider + assert offered.count("logging_only") == ( + 1 if guardrail_class.accepts_logging_only_mode() or GuardrailEventHooks.logging_only in declared else 0 + ), provider + + +@pytest.mark.asyncio +async def test_ui_settings_offer_logging_only_where_the_generic_scan_runs(): + modes_by_provider = (await get_guardrail_ui_settings()).supported_modes_by_provider + + assert modes_by_provider["litellm_content_filter"] == [ + *(hook.value for hook in ContentFilterGuardrail.get_supported_event_hooks()), + "logging_only", + ] + assert modes_by_provider["akto"].count("logging_only") == 1 + assert modes_by_provider["azure/prompt_shield"] == [ + hook.value for hook in AzureContentSafetyPromptShieldGuardrail.get_supported_event_hooks() + ] + with pytest.raises(ValueError, match="not in the supported event hooks"): + AzureContentSafetyPromptShieldGuardrail( + guardrail_name="native-logging-only", + api_key="k", + api_base="https://example.invalid", + event_hook="logging_only", + ) + assert ( + ContentFilterGuardrail(guardrail_name="cf-logging-only", event_hook="logging_only").event_hook == "logging_only" + ) def test_content_filter_runtime_rejects_unsupported_mcp_hook(): diff --git a/tests/unit/proxy/hooks/test_parallel_request_limiter.py b/tests/unit/proxy/hooks/test_parallel_request_limiter.py index f17cad56751..08840978169 100644 --- a/tests/unit/proxy/hooks/test_parallel_request_limiter.py +++ b/tests/unit/proxy/hooks/test_parallel_request_limiter.py @@ -15,7 +15,7 @@ from litellm.proxy.common_utils.proxy_rate_limit_error import ProxyRateLimitErro from litellm.proxy.hooks.parallel_request_limiter import ( PROXY_MaxParallelRequestsHandler, ) -from litellm.proxy.utils import InternalUsageCache, hash_token +from litellm.proxy.utils import InternalUsageCache, ProxyLogging, hash_token from litellm.types.utils import CallTypes, EmbeddingResponse, ModelResponse, TextCompletionResponse, Usage FROZEN_INSTANT: Final = datetime(2026, 1, 31, 23, 59, 30) @@ -353,3 +353,26 @@ async def test_websocket_turn_callbacks_keep_the_sockets_slot_until_the_disconne await handler.async_release_max_parallel_requests_on_disconnect(UserAPIKeyAuth(api_key=api_key)) after_disconnect = await handler.internal_usage_cache.async_get_cache(key=bucket, litellm_parent_otel_span=None) assert after_disconnect["current_requests"] == 0, after_disconnect + + +@pytest.mark.asyncio +async def test_a_rejected_realtime_attempt_does_not_free_a_slot_it_never_took(): + api_key = hash_token("sk-at-limit") + handler = PROXY_MaxParallelRequestsHandler(internal_usage_cache=InternalUsageCache(DualCache()), clock=_frozen_clock) + bucket = f"{api_key}::{FROZEN_INSTANT.strftime('%Y-%m-%d-%H-%M')}::request_count" + await handler.internal_usage_cache.async_set_cache( + key=bucket, + value={"current_requests": 1, "current_tpm": 0, "current_rpm": 1}, + litellm_parent_otel_span=None, + ) + proxy_logging_obj = ProxyLogging(user_api_key_cache=DualCache()) + proxy_logging_obj.proxy_hook_mapping["parallel_request_limiter"] = handler + auth = UserAPIKeyAuth(api_key=api_key, max_parallel_requests=1) + + await proxy_logging_obj.arelease_max_parallel_requests_on_disconnect(auth) + after_rejection = await handler.internal_usage_cache.async_get_cache(key=bucket, litellm_parent_otel_span=None) + assert after_rejection["current_requests"] == 1, after_rejection + + await proxy_logging_obj.arelease_responses_websocket_slot(auth) + after_socket_close = await handler.internal_usage_cache.async_get_cache(key=bucket, litellm_parent_otel_span=None) + assert after_socket_close["current_requests"] == 0, after_socket_close diff --git a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py index 9ef5f16a240..21170abaeb1 100644 --- a/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py +++ b/tests/unit/proxy/pass_through_endpoints/test_llm_pass_through_endpoints.py @@ -24,10 +24,15 @@ from starlette.datastructures import FormData import litellm from litellm.integrations.custom_guardrail import CustomGuardrail +from litellm.integrations.custom_logger import CustomLogger from tests._master_key import MASTER_KEY as SHARED_MASTER_KEY from litellm.caching.caching import DualCache from litellm.types.utils import CallTypesLiteral from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing +from litellm.exceptions import BlockedPiiEntityError, GuardrailRaisedException, ModifyResponseException +from litellm.proxy._types import AwsErrorTypeProxyException, ProxyException +from litellm.proxy.utils import ProxyLogging +from litellm.passthrough.main import raise_for_provider_status from tests.unit.llms.bedrock.event_loop_probe import EventLoopProbe from litellm.constants import LITELLM_PROXY_MASTER_KEY_ALIAS from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( @@ -40,6 +45,7 @@ from litellm.proxy.pass_through_endpoints.llm_passthrough_endpoints import ( azure_proxy_route, azure_speech_proxy_route, bedrock_llm_proxy_route, + bedrock_passthrough_error_response, bedrock_proxy_route, create_pass_through_route, cursor_proxy_route, @@ -1658,17 +1664,7 @@ class TestBedrockLLMProxyRoute: bedrock_error_message = '{"message":"ContentBlock object at messages.0.content.0 must set one of the following keys: text, image, toolUse, toolResult, document, video."}' - # Create a mock httpx.Response for the error - mock_error_response = Mock(spec=httpx.Response) - mock_error_response.status_code = 400 - mock_error_response.aread = AsyncMock(return_value=bedrock_error_message.encode("utf-8")) - - # Create the HTTPStatusError - mock_http_error = httpx.HTTPStatusError( - message="Bad Request", - request=Mock(spec=httpx.Request), - response=mock_error_response, - ) + mock_http_error = _bedrock_status_error(400, bedrock_error_message.encode("utf-8")) # Create mocks for all required parameters mock_request = MagicMock(spec=Request) @@ -1926,6 +1922,284 @@ class TestBedrockLLMProxyRoute: assert "Blocked by guardrail" in str(exc_info.value.detail) +_THROTTLE_TYPE: Final = "ThrottlingException:http://internal.amazon.com/coral/com.amazon.bedrock/" +_THROTTLE_BODY: Final = b'{"message":"Too many requests, please wait before trying again."}' + + +def _status_error( + status_code: int, + content: bytes, + headers: Mapping[str, str], + url: str, + raise_for_status: Callable[[httpx.Response], None], +) -> httpx.HTTPStatusError: + request: Final = httpx.Request("POST", url) + response: Final = httpx.Response(status_code, headers=dict(headers), content=content, request=request) + with pytest.raises(httpx.HTTPStatusError) as raised: + raise_for_status(response) + return raised.value + + +def _bedrock_status_error( + status_code: int, + content: bytes, + headers: Mapping[str, str] = MappingProxyType({}), + url: str = "https://bedrock-runtime.us-east-1.amazonaws.com/model/m/converse", +) -> httpx.HTTPStatusError: + """An error Bedrock answered the passthrough request with.""" + return _status_error(status_code, content, headers, url, raise_for_provider_status) + + +def _guardrail_service_status_error( + status_code: int, content: bytes, headers: Mapping[str, str] +) -> httpx.HTTPStatusError: + """An error a guardrail service's own HTTP call raised, from a URL shaped like a Bedrock model call.""" + return _status_error( + status_code, + content, + headers, + "https://bedrock-runtime.us-east-1.amazonaws.com/model/m/converse", + httpx.Response.raise_for_status, + ) + + +def _guardrail_detail(message: str) -> dict[str, str]: + return {"error": message, "guardrail_name": "block-demo"} + + +class _ReplacesFailures(CustomLogger): + def __init__(self, replacement: HTTPException) -> None: + super().__init__() + self.replacement: Final = replacement + + async def async_post_call_failure_hook( + self, + request_data: dict, + original_exception: Exception, + user_api_key_dict: UserAPIKeyAuth, + traceback_str: str | None = None, + ) -> HTTPException | None: + return self.replacement + + +class TestBedrockPassthroughAwsErrors: + """A guardrail block or an error Bedrock returned reaches an AWS SDK caller with its code and message.""" + + @staticmethod + async def _answer(error: Exception) -> Response: + return await bedrock_passthrough_error_response( + ProxyBaseLLMRequestProcessing(data={}), + error, + UserAPIKeyAuth(api_key="test-key"), + ProxyLogging(user_api_key_cache=DualCache()), + ) + + @pytest.mark.asyncio + async def test_bedrock_error_reaches_the_caller_with_bedrocks_status_header_and_body( + self, monkeypatch: pytest.MonkeyPatch + ) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + error: Final = _bedrock_status_error( + 429, + _THROTTLE_BODY, + { + "x-amzn-errortype": _THROTTLE_TYPE, + "x-amzn-requestid": "req-1", + "content-type": "application/x-amz-json-1.1", + }, + ) + + response: Final = await self._answer(error) + + assert response.status_code == 429 + assert response.body == _THROTTLE_BODY + assert response.headers["x-amzn-errortype"] == _THROTTLE_TYPE + assert response.headers["content-type"] == "application/x-amz-json-1.1" + assert response.headers["llm_provider-x-amzn-requestid"] == "req-1" + + @pytest.mark.asyncio + async def test_arn_model_error_reaches_the_caller_as_bedrock_sent_it(self, monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + error: Final = _bedrock_status_error( + 429, + _THROTTLE_BODY, + {"x-amzn-errortype": _THROTTLE_TYPE}, + url="https://bedrock-runtime.us-east-1.amazonaws.com/model/" + "arn%3Aaws%3Abedrock%3Aus-east-1%3A123%3Aapplication-inference-profile%2Fabc123/converse-stream?trace=a/b", + ) + + response: Final = await self._answer(error) + + assert (response.status_code, response.body) == (429, _THROTTLE_BODY) + assert response.headers["x-amzn-errortype"] == _THROTTLE_TYPE + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("error", "status_code", "body"), + ( + pytest.param( + _bedrock_status_error(502, b"bad gateway"), + 502, + "bad gateway", + id="bedrock-without-error-type", + ), + pytest.param( + _guardrail_service_status_error(429, _THROTTLE_BODY, {"x-amzn-errortype": _THROTTLE_TYPE}), + 429, + _THROTTLE_BODY.decode(), + id="guardrail-service-with-error-type", + ), + ), + ) + async def test_http_error_bedrock_did_not_answer_with_an_aws_error_type_keeps_the_litellm_error( + self, monkeypatch: pytest.MonkeyPatch, error: httpx.HTTPStatusError, status_code: int, body: str + ) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + + with pytest.raises(HTTPException) as exc_info: + await self._answer(error) + + assert exc_info.value.status_code == status_code + assert exc_info.value.detail == {"error": body} + assert "x-amzn-errortype" not in (exc_info.value.headers or {}) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("error", "replacement", "code"), + ( + pytest.param( + _bedrock_status_error(429, _THROTTLE_BODY, {"x-amzn-errortype": _THROTTLE_TYPE}), + HTTPException(status_code=503, detail={"error": "rewritten by a callback"}), + "503", + id="bedrock-error", + ), + pytest.param( + HTTPException(status_code=400, detail=_guardrail_detail("Violated guardrail policy")), + HTTPException(status_code=429, detail={"error": "rewritten by a callback"}), + "429", + id="guardrail-block", + ), + ), + ) + async def test_error_replaced_by_a_failure_hook_keeps_the_replacement( + self, monkeypatch: pytest.MonkeyPatch, error: Exception, replacement: HTTPException, code: str + ) -> None: + monkeypatch.setattr(litellm, "callbacks", [_ReplacesFailures(replacement)]) + + with pytest.raises(ProxyException) as exc_info: + await self._answer(error) + + assert (exc_info.value.message, exc_info.value.code) == ("rewritten by a callback", code) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("error", "message", "status_code"), + ( + pytest.param( + GuardrailRaisedException( + guardrail_name="block-demo", + message="synthetic policy says no", + should_wrap_with_default_message=False, + blocked_content=True, + ), + "synthetic policy says no", + "400", + id="guardrail-raised", + ), + pytest.param( + HTTPException(status_code=400, detail=_guardrail_detail("Violated guardrail policy")), + "Violated guardrail policy", + "400", + id="guardrail-http-400", + ), + pytest.param( + HTTPException(status_code=403, detail=_guardrail_detail("blocked by policy")), + "blocked by policy", + "403", + id="guardrail-http-403", + ), + pytest.param( + BlockedPiiEntityError(entity_type="EMAIL_ADDRESS", guardrail_name="pii"), + "Blocked entity detected: EMAIL_ADDRESS by Guardrail: pii. This entity is not allowed to be used in this request.", + "400", + id="blocked-pii", + ), + ), + ) + async def test_guardrail_block_is_raised_with_its_aws_error_type( + self, monkeypatch: pytest.MonkeyPatch, error: Exception, message: str, status_code: str + ) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + + with pytest.raises(AwsErrorTypeProxyException) as exc_info: + await self._answer(error) + + assert exc_info.value.aws_error_type == "ValidationException" + assert (exc_info.value.error.message, exc_info.value.error.code) == (message, status_code) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + ("error", "message", "code"), + ( + pytest.param( + HTTPException(status_code=400, detail={"error": "model not found"}), + "model not found", + "400", + id="request-check-400", + ), + pytest.param( + HTTPException(status_code=403, detail="cache control not allowed"), + "cache control not allowed", + "403", + id="request-check-403", + ), + pytest.param( + HTTPException(status_code=429, detail=_guardrail_detail("vendor rate limited")), + "vendor rate limited", + "429", + id="429", + ), + pytest.param( + GuardrailRaisedException( + guardrail_name="fail-closed", + message="guardrail vendor unreachable", + should_wrap_with_default_message=False, + ), + "guardrail vendor unreachable", + "400", + id="guardrail-vendor-unreachable-400", + ), + pytest.param( + GuardrailRaisedException( + guardrail_name="fail-closed", + message="guardrail vendor unreachable", + should_wrap_with_default_message=False, + status_code=503, + ), + "guardrail vendor unreachable", + "503", + id="guardrail-vendor-unreachable-503", + ), + pytest.param( + ModifyResponseException(message="replaced", model="bedrock-model", request_data={}), + "replaced", + "500", + id="modify-response", + ), + ), + ) + async def test_errors_that_are_not_guardrail_blocks_keep_the_litellm_error( + self, monkeypatch: pytest.MonkeyPatch, error: Exception, message: str, code: str + ) -> None: + monkeypatch.setattr(litellm, "callbacks", []) + + with pytest.raises(ProxyException) as exc_info: + await self._answer(error) + + assert type(exc_info.value) is ProxyException + assert (exc_info.value.message, exc_info.value.code) == (message, code) + + class TestBedrockAgentRuntimePassthroughToggle: AGENT_RUNTIME_ENDPOINT: Final = "knowledgebases/KB1234567/retrieve" MODEL_ENDPOINT: Final = "model/us.anthropic.claude-sonnet-4-5-20250929-v1:0/converse" diff --git a/tests/unit/proxy/response_api_endpoints/test_endpoints.py b/tests/unit/proxy/response_api_endpoints/test_endpoints.py index f19a702dc13..047d79f5b38 100644 --- a/tests/unit/proxy/response_api_endpoints/test_endpoints.py +++ b/tests/unit/proxy/response_api_endpoints/test_endpoints.py @@ -1949,7 +1949,7 @@ class TestResponsesWSFirstFrameModelAuth: proxy_logging_obj = MagicMock() proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) - proxy_logging_obj.arelease_max_parallel_requests_on_disconnect = AsyncMock(return_value=None) + proxy_logging_obj.arelease_responses_websocket_slot = AsyncMock(return_value=None) user_api_key_dict = MagicMock() with ( @@ -2014,7 +2014,7 @@ class TestResponsesWSFirstFrameModelAuth: ) proxy_logging_obj = MagicMock() proxy_logging_obj.post_call_failure_hook = AsyncMock(return_value=None) - proxy_logging_obj.arelease_max_parallel_requests_on_disconnect = AsyncMock(return_value=None) + proxy_logging_obj.arelease_responses_websocket_slot = AsyncMock(return_value=None) user_api_key_dict = MagicMock() with ( diff --git a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py index 3ebe445d89b..0ecebc7560c 100644 --- a/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py +++ b/tests/unit/proxy/test_proxy_server_endpoints_and_startup.py @@ -33,6 +33,7 @@ from litellm.caching.redis_cluster_cache import RedisClusterCache from litellm.litellm_core_utils.get_model_cost_map import ModelCostMapReloaded from litellm.caching.dual_cache import DualCache from litellm.proxy._types import ( + AwsErrorTypeProxyException, LitellmUserRoles, ModelAccessDeniedProxyException, ProxyErrorTypes, @@ -43,7 +44,7 @@ from litellm.proxy._types import ( from litellm.proxy.auth.login_throttle import LoginThrottle from litellm.proxy.auth.user_api_key_auth import user_api_key_auth from litellm.proxy.hooks.parallel_request_limiter_v3 import RequestRateLimiterStash -from litellm.proxy.proxy_server import app, initialize, openai_exception_handler +from litellm.proxy.proxy_server import app, aws_error_type_exception_handler, initialize, openai_exception_handler from litellm.utils import _invalidate_model_cost_lowercase_map from tests._master_key import MASTER_KEY @@ -11858,6 +11859,29 @@ async def test_openai_exception_handler_no_denial_log_for_plain_proxy_exception( assert [r for r in caplog.records if r.levelname == "WARNING"] == [] +@pytest.mark.asyncio +async def test_aws_error_type_exception_handler_adds_the_aws_message_and_header(): + error: Final = ProxyException( + message="synthetic policy says no", + type="invalid_request_error", + param=None, + code=400, + headers={"x-litellm-call-id": "call-1"}, + ) + + response: Final = await aws_error_type_exception_handler( + _http_request_scope(), AwsErrorTypeProxyException(error, "ValidationException") + ) + + assert response.status_code == 400 + assert response.headers["x-amzn-errortype"] == "ValidationException" + assert response.headers["x-litellm-call-id"] == "call-1" + assert json.loads(response.body) == { + "message": "synthetic policy says no", + "error": {"message": "synthetic policy says no", "type": "invalid_request_error", "param": None, "code": "400"}, + } + + @pytest.mark.asyncio async def test_realtime_model_access_denial_logs_sanitized_internal_message(caplog): reservation = {"reserved_cost": 0.0, "input_cost": 0.0, "finalized": False, "entries": []} diff --git a/tests/unit/proxy/test_proxy_token_counter.py b/tests/unit/proxy/test_proxy_token_counter.py index 8cd6887d891..fbe47e4a3d7 100644 --- a/tests/unit/proxy/test_proxy_token_counter.py +++ b/tests/unit/proxy/test_proxy_token_counter.py @@ -2,6 +2,8 @@ # 1. Generate a Key, and use it to make a call +import json +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -14,7 +16,8 @@ load_dotenv() # this file is to test litellm/proxy -from fastapi import HTTPException, Request +from fastapi import Request +from fastapi.responses import JSONResponse import litellm from litellm import Router @@ -324,22 +327,16 @@ async def test_internal_token_counter_anthropic_provider_detection(): setattr(litellm.proxy.proxy_server, "llm_router", llm_router) - # Test with is_direct_request=False but non-Anthropic provider - response = await token_counter( - request=TokenCountRequest( - model="gpt-test", - messages=[{"role": "user", "content": "hello"}], - ), - call_endpoint=True, - ) - - print("Non-Anthropic provider test response:", response) - - # Verify response structure - assert response.request_model == "gpt-test" - assert response.model_used == "gpt-4" - assert response.total_tokens > 0 - assert response.tokenizer_type == "openai_tokenizer" # Should use LiteLLM tokenizer + with pytest.raises(ProxyException) as missing_key: + await token_counter( + request=TokenCountRequest( + model="gpt-test", + messages=[{"role": "user", "content": "hello"}], + ), + call_endpoint=True, + ) + assert missing_key.value.code == "401" + assert "No OpenAI API key is set for model gpt-test" in missing_key.value.message @pytest.mark.asyncio @@ -349,12 +346,14 @@ async def test_anthropic_endpoint_error_handling(): """ from unittest.mock import MagicMock - from fastapi import HTTPException, Request + from fastapi import Request from litellm.proxy.anthropic_endpoints.endpoints import count_tokens # Mock request object mock_request = MagicMock(spec=Request) + mock_request.headers = {} + mock_request.scope = {} mock_user_api_key_dict = MagicMock() # Test missing model parameter @@ -372,12 +371,13 @@ async def test_anthropic_endpoint_error_handling(): anthropic_endpoints.read_request_body = mock_read_request_body try: - # Should raise HTTPException for missing model - with pytest.raises(HTTPException) as exc_info: - await count_tokens(mock_request, mock_user_api_key_dict) + response = await count_tokens(mock_request, mock_user_api_key_dict) - assert exc_info.value.status_code == 400 - assert "model parameter is required" in str(exc_info.value.detail) + assert response.status_code == 400 + assert json.loads(response.body) == { + "type": "error", + "error": {"type": "invalid_request_error", "message": "model parameter is required"}, + } print("✅ Error handling test passed!") @@ -468,7 +468,7 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): mock_router.model_list = [ { "model_name": "gpt-4", - "litellm_params": {"model": "openai/gpt-4"}, + "litellm_params": {"model": "databricks/gpt-4"}, "model_info": {}, } ] @@ -477,7 +477,7 @@ async def test_factory_gpt4_endpoint_does_not_call_anthropic_counter(): mock_router.async_get_available_deployment = AsyncMock( return_value={ "model_name": "gpt-4", - "litellm_params": {"model": "openai/gpt-4"}, + "litellm_params": {"model": "databricks/gpt-4"}, "model_info": {}, } ) @@ -1035,6 +1035,7 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): # Mock request object mock_request = MagicMock(spec=Request) + mock_request.headers = {} mock_request_data = { "model": "claude-bedrock", "messages": [{"role": "user", "content": "Hello!"}], @@ -1062,14 +1063,11 @@ async def test_anthropic_endpoint_returns_anthropic_error_format(): proxy_server.count_request_tokens = mock_token_counter_error try: - with pytest.raises(HTTPException) as exc_info: - await anthropic_count_tokens(mock_request, mock_user_api_key_dict) + response: Final[object] = await anthropic_count_tokens(mock_request, mock_user_api_key_dict) - # Verify HTTP status code is correct - assert exc_info.value.status_code == 400 - - # Verify error is in Anthropic format - detail = exc_info.value.detail + assert isinstance(response, JSONResponse) + assert response.status_code == 400 + detail = json.loads(bytes(response.body)) assert detail["type"] == "error" assert detail["error"]["type"] == "invalid_request_error" assert detail["error"]["message"] == "Input is too long for requested model." @@ -1087,6 +1085,7 @@ async def test_anthropic_endpoint_403_permission_error_format(): import litellm.proxy.proxy_server as proxy_server mock_request = MagicMock(spec=Request) + mock_request.headers = {} mock_request_data = { "model": "claude-bedrock", "messages": [{"role": "user", "content": "Hello!"}], @@ -1114,12 +1113,11 @@ async def test_anthropic_endpoint_403_permission_error_format(): proxy_server.count_request_tokens = mock_token_counter_error try: - with pytest.raises(HTTPException) as exc_info: - await anthropic_count_tokens(mock_request, mock_user_api_key_dict) + response: Final[object] = await anthropic_count_tokens(mock_request, mock_user_api_key_dict) - assert exc_info.value.status_code == 403 - - detail = exc_info.value.detail + assert isinstance(response, JSONResponse) + assert response.status_code == 403 + detail = json.loads(bytes(response.body)) assert detail["type"] == "error" assert detail["error"]["type"] == "permission_error" assert detail["error"]["message"] == "Bearer Token has expired" @@ -1137,6 +1135,7 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): import litellm.proxy.proxy_server as proxy_server mock_request = MagicMock(spec=Request) + mock_request.headers = {} mock_request_data = { "model": "claude-bedrock", "messages": [{"role": "user", "content": "Hello!"}], @@ -1164,12 +1163,11 @@ async def test_anthropic_endpoint_429_rate_limit_error_format(): proxy_server.count_request_tokens = mock_token_counter_error try: - with pytest.raises(HTTPException) as exc_info: - await anthropic_count_tokens(mock_request, mock_user_api_key_dict) + response: Final[object] = await anthropic_count_tokens(mock_request, mock_user_api_key_dict) - assert exc_info.value.status_code == 429 - - detail = exc_info.value.detail + assert isinstance(response, JSONResponse) + assert response.status_code == 429 + detail = json.loads(bytes(response.body)) assert detail["type"] == "error" assert detail["error"]["type"] == "rate_limit_error" assert detail["error"]["message"] == "Rate limit exceeded" diff --git a/tests/unit/responses/test_streaming_iterator_error_events.py b/tests/unit/responses/test_streaming_iterator_error_events.py index c37e57a43ca..001be7e1fbd 100644 --- a/tests/unit/responses/test_streaming_iterator_error_events.py +++ b/tests/unit/responses/test_streaming_iterator_error_events.py @@ -20,14 +20,18 @@ Pydantic ValidationError (previously typed as Optional[str]). """ import json +from datetime import datetime from importlib import import_module +from typing import Final from unittest.mock import Mock, patch +import httpx import pytest import litellm from litellm.exceptions import MidStreamFallbackError from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.litellm_logging import _get_provider_request_id from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig from litellm.responses.streaming_iterator import ( _ERROR_CODE_HTTP_STATUS, @@ -64,6 +68,52 @@ def _make_error_chunk(error_type: str, code: str, message: str = "err") -> Error return ErrorEvent(type=ResponsesAPIStreamEvents.ERROR, sequence_number=0, error=error_obj) +@pytest.mark.parametrize("provider", ("openai", "azure", "bedrock")) +@pytest.mark.parametrize("header", ("x-amzn-RequestId", "x-request-id", None)) +@pytest.mark.parametrize( + "code,status", (("server_error", 500), ("rate_limit_exceeded", 429), ("invalid_request_error", 400)) +) +def test_stream_error_preserves_provider_request_id(provider: str, header: str | None, code: str, status: int) -> None: + stream_headers: Final = {"retry-after": "7", "content-type": "text/event-stream"} + response: Final = httpx.Response( + 200, headers={**stream_headers, header: "provider-request-id"} if header else stream_headers + ) + logging: Final = LiteLLMLoggingObj( + model="test-model", + messages=[], + stream=True, + call_type="responses", + start_time=datetime(2026, 1, 1), + litellm_call_id="test-call", + function_id="test-function", + ) + iterator: Final = BaseResponsesAPIStreamingIterator( + response=response, + model="test-model", + responses_api_provider_config=None, + logging_obj=logging, + custom_llm_provider=provider, + ) + expected_id: Final = "provider-request-id" if header else None + expected_headers: Final = {header.lower(): "provider-request-id"} if header else {} + mapped: Final = iterator._map_error_event_exception({"type": code, "code": code, "message": "failure"}) + assert _get_provider_request_id(mapped) == expected_id + assert dict(getattr(mapped, "litellm_response_headers", None) or {}) == expected_headers + expected_status: Final = 503 if provider == "bedrock" and status == 500 else status + with pytest.raises((litellm.BadRequestError, MidStreamFallbackError)) as raised: + iterator._maybe_raise_for_error_event(_make_error_chunk(code, code, "failure")) + assert raised.value.status_code == expected_status + assert _get_provider_request_id(raised.value) == expected_id + assert response.status_code == 200 + if status in (429, 500): + assert isinstance(raised.value, MidStreamFallbackError) + assert dict(raised.value.response.headers) == expected_headers + assert raised.value.response.status_code == expected_status + assert str(raised.value.response.request.url) == f"https://{provider}.com/v1/" + assert raised.value.generated_content == "" + assert raised.value.is_pre_first_chunk is True + + def test_maybe_raise_for_error_event_wraps_unknown_error_in_mid_stream_fallback(): iterator = _make_iterator() chunk = _make_error_chunk("server_error", "internal_error", "something went wrong") diff --git a/tests/unit/test_router_model_cost_isolation.py b/tests/unit/test_router_model_cost_isolation.py index fc1799861d1..d343f464a77 100644 --- a/tests/unit/test_router_model_cost_isolation.py +++ b/tests/unit/test_router_model_cost_isolation.py @@ -119,27 +119,33 @@ async def test_discovered_limits_survive_deployment_growth_and_removal( async def test_discovery_discards_metadata_for_a_replaced_deployment(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) - router: Final = Router(model_list=[{ - "model_name": "local", - "litellm_params": { - "model": "hosted_vllm/local-model", - "api_base": "https://original.test/v1", - "api_key": "local-key", - }, - "model_info": {"id": "replaced-deployment"}, - }]) + router: Final = Router( + model_list=[ + { + "model_name": "local", + "litellm_params": { + "model": "hosted_vllm/local-model", + "api_base": "https://original.test/v1", + "api_key": "local-key", + }, + "model_info": {"id": "replaced-deployment"}, + } + ] + ) def respond(request: httpx.Request) -> httpx.Response: if request.url.host == "original.test": - router.upsert_deployment(Deployment( - model_name="local", - litellm_params=LiteLLM_Params( - model="hosted_vllm/local-model", - api_base="https://replacement.test/v1", - api_key="local-key", - ), - model_info=ModelInfo(id="replaced-deployment"), - )) + router.upsert_deployment( + Deployment( + model_name="local", + litellm_params=LiteLLM_Params( + model="hosted_vllm/local-model", + api_base="https://replacement.test/v1", + api_key="local-key", + ), + model_info=ModelInfo(id="replaced-deployment"), + ) + ) return httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 8192}]}) assert request.url.host == "replacement.test" return httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 2048}]}) @@ -158,15 +164,19 @@ async def test_discovery_discards_metadata_for_a_replaced_deployment(monkeypatch async def test_discovery_is_isolated_across_routers_and_reused_ids(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) first, second = tuple( - Router(model_list=[{ - "model_name": "local", - "litellm_params": { - "model": "hosted_vllm/local-model", - "api_base": f"https://{host}.test/v1", - "api_key": "local-key", - }, - "model_info": {"id": "shared-discovery-id"}, - }]) + Router( + model_list=[ + { + "model_name": "local", + "litellm_params": { + "model": "hosted_vllm/local-model", + "api_base": f"https://{host}.test/v1", + "api_key": "local-key", + }, + "model_info": {"id": "shared-discovery-id"}, + } + ] + ) for host in ("first", "second") ) @@ -187,15 +197,17 @@ async def test_discovery_is_isolated_across_routers_and_reused_ids(monkeypatch: assert first.get_configured_token_limits("local") == (8192, 8192) assert second.get_configured_token_limits("local") == (2048, 2048) assert litellm.model_cost["shared-discovery-id"].get("max_input_tokens") is None - first.upsert_deployment(Deployment( - model_name="local", - litellm_params=LiteLLM_Params( - model="hosted_vllm/local-model", - api_base="https://unavailable.test/v1", - api_key="local-key", - ), - model_info=ModelInfo(id="shared-discovery-id"), - )) + first.upsert_deployment( + Deployment( + model_name="local", + litellm_params=LiteLLM_Params( + model="hosted_vllm/local-model", + api_base="https://unavailable.test/v1", + api_key="local-key", + ), + model_info=ModelInfo(id="shared-discovery-id"), + ) + ) assert first.get_configured_token_limits("local") == (None, None) await first.arefresh_model_info(client=handler) assert first.get_configured_token_limits("local") == (None, None) @@ -206,17 +218,19 @@ async def test_discovery_is_isolated_across_routers_and_reused_ids(monkeypatch: async def test_discovery_refreshes_other_endpoints_while_one_is_pending(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) second_started: Final = asyncio.Event() - router: Final = Router(model_list=[ - { - "model_name": host, - "litellm_params": { - "model": "hosted_vllm/local-model", - "api_base": f"https://{host}.test/v1", - "api_key": "local-key", - }, - } - for host in ("first", "second", "third") - ]) + router: Final = Router( + model_list=[ + { + "model_name": host, + "litellm_params": { + "model": "hosted_vllm/local-model", + "api_base": f"https://{host}.test/v1", + "api_key": "local-key", + }, + } + for host in ("first", "second", "third") + ] + ) async def respond(request: httpx.Request) -> httpx.Response: if request.url.host == "first.test": @@ -240,21 +254,27 @@ async def test_discovery_refreshes_other_endpoints_while_one_is_pending(monkeypa async def test_discovered_limits_expire_after_the_last_successful_refresh(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) clock: Final = Mock(return_value=0.0) - router: Final = Router(model_list=[{ - "model_name": "local", - "litellm_params": { - "model": "hosted_vllm/local-model", - "api_base": "https://expiry.test/v1", - "api_key": "local-key", - }, - "model_info": {"id": "expiring-discovery"}, - }]) + router: Final = Router( + model_list=[ + { + "model_name": "local", + "litellm_params": { + "model": "hosted_vllm/local-model", + "api_base": "https://expiry.test/v1", + "api_key": "local-key", + }, + "model_info": {"id": "expiring-discovery"}, + } + ] + ) router._discovered_model_info_cache = InMemoryCache(clock=clock, default_ttl=2 * MODEL_INFO_REFRESH_SECONDS) - responses: Final = iter(( - httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 4096}]}), - httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 8192}]}), - httpx.Response(503), - )) + responses: Final = iter( + ( + httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 4096}]}), + httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 8192}]}), + httpx.Response(503), + ) + ) handler: Final = AsyncHTTPHandler() await handler.client.aclose() async with httpx.AsyncClient(transport=httpx.MockTransport(lambda request: next(responses))) as client: @@ -343,10 +363,12 @@ async def test_discovered_limits_are_isolated_overridable_and_refreshable( async def test_discovery_preserves_input_overrides_and_survives_outages(monkeypatch: pytest.MonkeyPatch) -> None: monkeypatch.setattr(litellm, "model_cost", copy.deepcopy(litellm.model_cost)) - responses: Final = iter(( - httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 4096}]}), - httpx.Response(503), - )) + responses: Final = iter( + ( + httpx.Response(200, json={"data": [{"id": "local-model", "max_model_len": 4096}]}), + httpx.Response(503), + ) + ) def respond(request: httpx.Request) -> httpx.Response: assert request.url.host == "backend.test" @@ -354,27 +376,29 @@ async def test_discovery_preserves_input_overrides_and_survives_outages(monkeypa assert request.headers["x-tenant"] == "tenant" return next(responses) - router: Final = Router(model_list=[ - { - "model_name": "configured", - "litellm_params": { - "model": "hosted_vllm/local-model", - "api_base": "https://backend.test/v1", - "api_key": "unused-key", - "extra_headers": {"authorization": "Bearer local-key", "X-Tenant": "tenant"}, + router: Final = Router( + model_list=[ + { + "model_name": "configured", + "litellm_params": { + "model": "hosted_vllm/local-model", + "api_base": "https://backend.test/v1", + "api_key": "unused-key", + "extra_headers": {"authorization": "Bearer local-key", "X-Tenant": "tenant"}, + }, + "model_info": {"id": "configured", "max_input_tokens": 1024}, }, - "model_info": {"id": "configured", "max_input_tokens": 1024}, - }, - { - "model_name": "byok", - "litellm_params": { - "model": "openai/local-model", - "api_base": "https://caller.test/v1", - "use_clientside_credentials": True, + { + "model_name": "byok", + "litellm_params": { + "model": "openai/local-model", + "api_base": "https://caller.test/v1", + "use_clientside_credentials": True, + }, }, - }, - {"model_name": "default-openai", "litellm_params": {"model": "openai/local-model", "api_key": "unused"}}, - ]) + {"model_name": "default-openai", "litellm_params": {"model": "openai/local-model", "api_key": "unused"}}, + ] + ) handler: Final = AsyncHTTPHandler() await handler.client.aclose() responder: Final = Mock(side_effect=respond) @@ -944,8 +968,7 @@ def _register_azure_tier_backend() -> None: def test_inherit_builtin_service_tier_pricing_fills_only_missing_fields() -> None: model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) - for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL) + key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL) } try: _register_tier_backend() @@ -974,9 +997,7 @@ def test_inherit_builtin_service_tier_pricing_fills_only_missing_fields() -> Non "output_cost_per_token_flex": _TIER_BACKEND_ENTRY["output_cost_per_token_flex"], "input_cost_per_token_balanced": _TIER_BACKEND_ENTRY["input_cost_per_token_balanced"], "output_cost_per_token_balanced": _TIER_BACKEND_ENTRY["output_cost_per_token_balanced"], - "cache_read_input_token_cost_ultrafast": _TIER_BACKEND_ENTRY[ - "cache_read_input_token_cost_ultrafast" - ], + "cache_read_input_token_cost_ultrafast": _TIER_BACKEND_ENTRY["cache_read_input_token_cost_ultrafast"], "input_cost_per_token_above_272k_tokens_ultrafast": _TIER_BACKEND_ENTRY[ "input_cost_per_token_above_272k_tokens_ultrafast" ], @@ -991,8 +1012,7 @@ def test_inherit_builtin_service_tier_pricing_fills_only_missing_fields() -> Non def test_inherit_builtin_service_tier_pricing_noop_without_base_rate_or_backend() -> None: model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) - for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL) + key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL) } try: _register_tier_backend() @@ -1029,8 +1049,7 @@ def test_inherit_builtin_service_tier_pricing_noop_without_base_rate_or_backend( def test_router_completion_uses_custom_standard_and_backend_ultrafast_pricing() -> None: model_id: Final = "tier-priced-deployment" model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) - for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id) + key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id) } try: _register_tier_backend() @@ -1090,8 +1109,7 @@ def test_router_completion_uses_custom_standard_and_backend_ultrafast_pricing() def test_router_completion_uses_backend_ultrafast_long_context_rates() -> None: model_id: Final = "tier-priced-long-context-deployment" model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) - for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id) + key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id) } try: _register_tier_backend() @@ -1142,8 +1160,7 @@ def test_ptu_service_tier_pricing_is_disabled_only_when_attribution_is_enabled( ) -> None: model_id: Final = f"ptu-tier-deployment-{ptu_enabled}" model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) - for key in (_TIER_BACKEND_KEY, model_id) + key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, model_id) } try: _register_tier_backend() @@ -1200,8 +1217,7 @@ def test_azure_base_model_inherits_service_tier_pricing_for_registration_and_pay model_id: Final = "azure-tier-priced-alias" payload_id: Final = "azure-tier-priced-payload" model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) - for key in (_AZURE_TIER_BACKEND_KEY, model_id, payload_id) + key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_AZURE_TIER_BACKEND_KEY, model_id, payload_id) } try: _register_azure_tier_backend() @@ -1257,12 +1273,8 @@ def test_azure_base_model_inherits_service_tier_pricing_for_registration_and_pay ) ) - assert payload["input_cost_per_token_priority"] == _AZURE_TIER_BACKEND_ENTRY[ - "input_cost_per_token_priority" - ] - assert payload["output_cost_per_token_priority"] == _AZURE_TIER_BACKEND_ENTRY[ - "output_cost_per_token_priority" - ] + assert payload["input_cost_per_token_priority"] == _AZURE_TIER_BACKEND_ENTRY["input_cost_per_token_priority"] + assert payload["output_cost_per_token_priority"] == _AZURE_TIER_BACKEND_ENTRY["output_cost_per_token_priority"] finally: _restore_model_cost_entries(model_cost_entries) litellm.get_model_info.cache_clear() @@ -1364,12 +1376,10 @@ def test_inherit_builtin_base_rates_for_off_peak_carries_threshold_rates(): assert model_info["input_cost_per_token"] == builtin_info["input_cost_per_token"] assert ( - model_info["input_cost_per_token_above_200k_tokens"] - == builtin_info["input_cost_per_token_above_200k_tokens"] + model_info["input_cost_per_token_above_200k_tokens"] == builtin_info["input_cost_per_token_above_200k_tokens"] ) assert ( - model_info["output_cost_per_token_above_200k_tokens"] - == builtin_info["output_cost_per_token_above_200k_tokens"] + model_info["output_cost_per_token_above_200k_tokens"] == builtin_info["output_cost_per_token_above_200k_tokens"] ) @@ -2260,8 +2270,7 @@ def test_deployment_model_cost_payload_folds_in_litellm_params_pricing(): def test_deployment_model_cost_payload_includes_builtin_service_tier_pricing() -> None: model_id: Final = "tier-priced-payload" model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) - for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id) + key: copy.deepcopy(litellm.model_cost.get(key)) for key in (_TIER_BACKEND_KEY, _TIER_BACKEND_MODEL, model_id) } try: _register_tier_backend() @@ -2278,12 +2287,8 @@ def test_deployment_model_cost_payload_includes_builtin_service_tier_pricing() - ) ) - assert ( - payload["input_cost_per_token_ultrafast"] == _TIER_BACKEND_ENTRY["input_cost_per_token_ultrafast"] - ) - assert ( - payload["output_cost_per_token_ultrafast"] == _TIER_BACKEND_ENTRY["output_cost_per_token_ultrafast"] - ) + assert payload["input_cost_per_token_ultrafast"] == _TIER_BACKEND_ENTRY["input_cost_per_token_ultrafast"] + assert payload["output_cost_per_token_ultrafast"] == _TIER_BACKEND_ENTRY["output_cost_per_token_ultrafast"] assert payload["input_cost_per_token_balanced"] == _TIER_BACKEND_ENTRY["input_cost_per_token_balanced"] assert payload["input_cost_per_token"] == _CUSTOM_STANDARD_INPUT_RATE assert payload["output_cost_per_token"] == _CUSTOM_STANDARD_OUTPUT_RATE @@ -3282,9 +3287,7 @@ def test_deployment_id_colliding_with_another_providers_catalog_key_keeps_its_ow ) assert isinstance(response, litellm.ModelResponse) - assert response._hidden_params["response_cost"] == pytest.approx( - 1000 * 0.00000096 + 500 * 0.00000302 - ) + assert response._hidden_params["response_cost"] == pytest.approx(1000 * 0.00000096 + 500 * 0.00000302) assert litellm.model_cost[_COLLIDING_BUILTIN_KEY] == builtin_row_before finally: _restore_model_cost_entries(model_cost_entries) @@ -3297,8 +3300,7 @@ def test_deployment_id_matching_its_own_providers_catalog_key_still_merges() -> matches against.""" model_id: Final = "baseten/zai-org/glm-5.2" model_cost_entries: Final = { - key: copy.deepcopy(litellm.model_cost.get(key)) - for key in (model_id, _COLLIDING_BUILTIN_KEY) + key: copy.deepcopy(litellm.model_cost.get(key)) for key in (model_id, _COLLIDING_BUILTIN_KEY) } try: router: Final = Router(model_list=[_colliding_id_deployment(model_id, "baseten")]) @@ -3313,9 +3315,244 @@ def test_deployment_id_matching_its_own_providers_catalog_key_still_merges() -> ) assert isinstance(response, litellm.ModelResponse) - assert response._hidden_params["response_cost"] == pytest.approx( - 1000 * 0.00000096 + 500 * 0.00000302 - ) + assert response._hidden_params["response_cost"] == pytest.approx(1000 * 0.00000096 + 500 * 0.00000302) finally: _restore_model_cost_entries(model_cost_entries) litellm.get_model_info.cache_clear() + + +_CAPACITY_TEST_DEPLOYMENTS: Final = ( + ( + "bedrock", + "bedrock/converse/anthropic.claude-router-capacity-test", + "bedrock/us-east-1/anthropic.claude-router-capacity-test", + ), + ( + "bedrock_mantle", + "bedrock_mantle/openai.router-capacity-test", + "bedrock_mantle/us-east-1/openai.router-capacity-test", + ), + ("openai", "openai/router-capacity-test", "openai/us/router-capacity-test"), +) + + +def _capacity_test_metadata(provider: str, max_input_tokens: int) -> dict[str, object]: + return { + "litellm_provider": provider, + "mode": "chat", + "max_input_tokens": max_input_tokens, + "max_output_tokens": 512, + "max_tokens": 512, + "input_cost_per_token": 0, + "output_cost_per_token": 0, + } + + +def _capacity_test_router(route: str, base: str, base_location: str | None, explicit_limit: int | None) -> Router: + return Router( + model_list=[ + { + "model_name": "capacity-test", + "litellm_params": { + "model": route, + **({"base_model": base} if base_location == "litellm_params" else {}), + }, + "model_info": { + "id": "capacity-test", + **({"base_model": base} if base_location == "model_info" else {}), + **({"max_input_tokens": explicit_limit} if explicit_limit is not None else {}), + }, + } + ], + enable_pre_call_checks=True, + ) + + +def _deployments_allowed_for_2048_input_tokens(router: Router) -> int: + try: + return len( + router._pre_call_checks( + model="capacity-test", + healthy_deployments=router.model_list, + input="synthetic input", + input_token_count=2048, + ) + ) + except litellm.ContextWindowExceededError: + return 0 + + +def _effective_max_input_tokens(router: Router) -> object: + return router.get_router_model_info(id="capacity-test", deployment=None, received_model_name="capacity-test")[ + "max_input_tokens" + ] + + +@pytest.mark.parametrize(("provider", "route", "base"), _CAPACITY_TEST_DEPLOYMENTS) +@pytest.mark.parametrize("base_location", ("model_info", "litellm_params", None)) +@pytest.mark.parametrize("explicit_limit", (None, 4096)) +def test_router_capacity_widens_to_a_larger_base_model_window( + monkeypatch: pytest.MonkeyPatch, + provider: str, + route: str, + base: str, + base_location: str | None, + explicit_limit: int | None, +) -> None: + monkeypatch.setattr( + litellm, + "model_cost", + { + **litellm.model_cost, + route: _capacity_test_metadata(provider, 1024), + base: _capacity_test_metadata(provider, 8192), + }, + ) + _invalidate_model_cost_lowercase_map() + router: Final = _capacity_test_router(route, base, base_location, explicit_limit) + expected: Final = ( + explicit_limit if explicit_limit is not None else (1024 if base_location is None else 8192) + ) + assert _effective_max_input_tokens(router) == expected + assert _deployments_allowed_for_2048_input_tokens(router) == (1 if expected >= 2048 else 0) + _invalidate_model_cost_lowercase_map() + + +@pytest.mark.parametrize(("provider", "route", "base"), _CAPACITY_TEST_DEPLOYMENTS) +@pytest.mark.parametrize("base_location", ("model_info", "litellm_params")) +def test_router_capacity_never_narrows_to_a_smaller_base_model_window( + monkeypatch: pytest.MonkeyPatch, provider: str, route: str, base: str, base_location: str +) -> None: + monkeypatch.setattr( + litellm, + "model_cost", + { + **litellm.model_cost, + route: _capacity_test_metadata(provider, 8192), + base: _capacity_test_metadata(provider, 1024), + }, + ) + _invalidate_model_cost_lowercase_map() + router: Final = _capacity_test_router(route, base, base_location, None) + assert _effective_max_input_tokens(router) == 8192 + assert _deployments_allowed_for_2048_input_tokens(router) == 1 + _invalidate_model_cost_lowercase_map() + + +@pytest.mark.parametrize(("provider", "route", "base"), _CAPACITY_TEST_DEPLOYMENTS) +def test_router_capacity_never_takes_a_window_from_base_model_when_the_route_is_unmapped( + monkeypatch: pytest.MonkeyPatch, provider: str, route: str, base: str +) -> None: + monkeypatch.setattr(litellm, "model_cost", {**litellm.model_cost, base: _capacity_test_metadata(provider, 1024)}) + _invalidate_model_cost_lowercase_map() + router: Final = _capacity_test_router(f"{route}-not-in-cost-map", base, "model_info", None) + assert _deployments_allowed_for_2048_input_tokens(router) == 1 + _invalidate_model_cost_lowercase_map() + + +@pytest.mark.parametrize(("provider", "route", "base"), _CAPACITY_TEST_DEPLOYMENTS) +def test_router_capacity_never_takes_a_window_from_base_model_when_the_route_declares_none( + monkeypatch: pytest.MonkeyPatch, provider: str, route: str, base: str +) -> None: + route_without_window: Final = { + key: value + for key, value in _capacity_test_metadata(provider, 0).items() + if key not in ("max_input_tokens", "max_tokens") + } + monkeypatch.setattr( + litellm, + "model_cost", + {**litellm.model_cost, route: route_without_window, base: _capacity_test_metadata(provider, 1024)}, + ) + _invalidate_model_cost_lowercase_map() + router: Final = _capacity_test_router(route, base, "model_info", None) + assert _effective_max_input_tokens(router) is None + assert _deployments_allowed_for_2048_input_tokens(router) == 1 + _invalidate_model_cost_lowercase_map() + + +_SAME_MODEL_ROUTE: Final = "bedrock/converse/anthropic.claude-routercapacitytest-9-9" +_SAME_MODEL_GLOBAL_ROUTE: Final = "bedrock/converse/global.anthropic.claude-routercapacitytest-9-9" +_SAME_MODEL_QUALIFIED_ROUTE: Final = "bedrock/converse/anthropic.claude-routercapacitytest-9-9@first" + + +@pytest.mark.parametrize( + ("route", "base_model", "expected"), + ( + (_SAME_MODEL_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9", 8192), + (_SAME_MODEL_ROUTE, "us-east-1/anthropic.claude-routercapacitytest-9-9", 1024), + (_SAME_MODEL_ROUTE, "anthropic.claude-routercapacitytest-9-9", 1024), + (_SAME_MODEL_GLOBAL_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9", 8192), + (_SAME_MODEL_ROUTE, "bedrock/anthropic.claude-routercapacitytest-9-9@global:us-east-1", 8192), + (_SAME_MODEL_GLOBAL_ROUTE, "bedrock/anthropic.claude-routercapacitytest-9-9@global:us-east-1", 8192), + (_SAME_MODEL_QUALIFIED_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9@first", 8192), + (_SAME_MODEL_QUALIFIED_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9@second", 1024), + (_SAME_MODEL_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-*", 1024), + (_SAME_MODEL_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitystandin-9-9", 1024), + (_SAME_MODEL_ROUTE, "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9-extended", 1024), + (_SAME_MODEL_ROUTE, "bedrock/eu-west-1/anthropic.claude-routercapacitytest-9-9", 1024), + ), +) +def test_router_capacity_widens_only_for_an_exact_base_model_entry_of_the_same_model( + monkeypatch: pytest.MonkeyPatch, route: str, base_model: str, expected: int +) -> None: + larger_window: Final = _capacity_test_metadata("bedrock", 8192) + monkeypatch.setattr( + litellm, + "model_cost", + { + **litellm.model_cost, + _SAME_MODEL_ROUTE: _capacity_test_metadata("bedrock", 1024), + _SAME_MODEL_GLOBAL_ROUTE: _capacity_test_metadata("bedrock", 1024), + _SAME_MODEL_QUALIFIED_ROUTE: _capacity_test_metadata("bedrock", 1024), + "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9@first": larger_window, + "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9@second": larger_window, + "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9": larger_window, + "anthropic.claude-routercapacitytest-9-9": _capacity_test_metadata("bedrock", 1024), + "bedrock/anthropic.claude-routercapacitytest-9-9": larger_window, + "bedrock/anthropic.claude-routercapacitytest-9-9@global:us-east-1": larger_window, + "bedrock/us-east-1/anthropic.claude-routercapacitytest-*": larger_window, + "bedrock/us-east-1/anthropic.claude-routercapacitystandin-9-9": larger_window, + "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9-extended": larger_window, + }, + ) + _invalidate_model_cost_lowercase_map() + router: Final = _capacity_test_router(route, base_model, "model_info", None) + assert _effective_max_input_tokens(router) == expected + assert _deployments_allowed_for_2048_input_tokens(router) == (1 if expected >= 2048 else 0) + _invalidate_model_cost_lowercase_map() + + +def test_router_capacity_reads_a_model_info_base_model_before_a_litellm_params_one( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + litellm, + "model_cost", + { + **litellm.model_cost, + _SAME_MODEL_ROUTE: _capacity_test_metadata("bedrock", 1024), + "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9": _capacity_test_metadata("bedrock", 8192), + "bedrock/us-east-1/anthropic.claude-routercapacitystandin-9-9": _capacity_test_metadata("bedrock", 8192), + }, + ) + _invalidate_model_cost_lowercase_map() + router: Final = Router( + model_list=[ + { + "model_name": "capacity-test", + "litellm_params": { + "model": _SAME_MODEL_ROUTE, + "base_model": "bedrock/us-east-1/anthropic.claude-routercapacitytest-9-9", + }, + "model_info": { + "id": "capacity-test", + "base_model": "bedrock/us-east-1/anthropic.claude-routercapacitystandin-9-9", + }, + } + ], + enable_pre_call_checks=True, + ) + assert _effective_max_input_tokens(router) == 1024 + assert _deployments_allowed_for_2048_input_tokens(router) == 0 + _invalidate_model_cost_lowercase_map() diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx index 681217cc918..cdc681a1c91 100644 --- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.integration.test.tsx @@ -8,9 +8,9 @@ import { renderWithProviders, } from "../../../tests/test-utils"; import userEvent, { PointerEventsCheckLevel } from "@testing-library/user-event"; -import { describe, expect, it, vi } from "vitest"; +import { beforeEach, describe, expect, it, vi } from "vitest"; import type { Team } from "../key_team_helpers/key_list"; -import { credentialCreateCall, modelCreateCall, type CredentialItem } from "../networking"; +import { credentialCreateCall, modelCreateCall, testConnectionRequest, type CredentialItem } from "../networking"; import { Providers } from "../provider_info_helpers"; import { projectMountedValues, useMountRegistry, type MountedFormValues } from "../common_components/MountedFormField"; import { useForm } from "react-hook-form"; @@ -53,7 +53,17 @@ vi.mock("../networking", async () => { provider_display_name: "OpenAI", litellm_provider: "openai", default_model_placeholder: "gpt-3.5-turbo", - credential_fields: [], + credential_fields: [{ key: "api_key", label: "OpenAI API Key", field_type: "password", required: true }], + }, + { + provider: "Azure", + provider_display_name: "Azure", + litellm_provider: "azure", + default_model_placeholder: "my-deployment", + credential_fields: [ + { key: "api_base", label: "API Base", field_type: "text", required: true }, + { key: "api_key", label: "Azure API Key", field_type: "password" }, + ], }, { provider: "MICROSOFT_365_COPILOT", @@ -98,7 +108,17 @@ vi.mock("@/app/(dashboard)/hooks/providers/useProviderFields", () => ({ provider_display_name: "OpenAI", litellm_provider: "openai", default_model_placeholder: "gpt-3.5-turbo", - credential_fields: [], + credential_fields: [{ key: "api_key", label: "OpenAI API Key", field_type: "password", required: true }], + }, + { + provider: "Azure", + provider_display_name: "Azure", + litellm_provider: "azure", + default_model_placeholder: "my-deployment", + credential_fields: [ + { key: "api_base", label: "API Base", field_type: "text", required: true }, + { key: "api_key", label: "Azure API Key", field_type: "password" }, + ], }, { provider: "MICROSOFT_365_COPILOT", @@ -252,6 +272,22 @@ const createTestProps = (userRole = "proxy_admin", userId = "user-1", isTeamAdmi }; }; +type User = ReturnType; + +const pickProvider = async (user: User, name: RegExp) => { + await user.click(screen.getByRole("combobox", { name: /provider/i })); + await user.click(await screen.findByRole("option", { name })); +}; + +const pickModel = async (user: User, name: string) => { + await user.click(await screen.findByPlaceholderText("Select models")); + await user.click(await screen.findByRole("option", { name })); + await user.keyboard("{Escape}"); +}; + +const fillApiKey = (label: string) => + fireEvent.change(screen.getByLabelText(new RegExp(label)), { target: { value: "sk-test" } }); + describe("AddModelForm", () => { it("should render", async () => { const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized")); @@ -444,10 +480,15 @@ describe("AddModelForm", () => { it("shows only the Close button in the connection test dialog footer", async () => { const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized")); mockUseAuthorized.default.mockReturnValue(mockAuthorizedUser("proxy_admin", "user-1", true)); + const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); renderWithProviders(); + await screen.findByText("Provider"); + await pickProvider(user, /OpenAI/); + fillApiKey("OpenAI API Key"); + await pickModel(user, "gpt-4"); - await userEvent.click(await screen.findByTestId("test-connect-btn")); + await user.click(screen.getByTestId("test-connect-btn")); const dialog = await screen.findByRole("dialog"); const footer = dialog.querySelector('[data-slot="dialog-footer"]'); @@ -786,6 +827,180 @@ describe("AddModelForm", () => { }); }); + describe("Test Connect validates the form the way Add Model does", () => { + beforeEach(() => { + vi.mocked(testConnectionRequest).mockClear(); + }); + + const renderForm = async (selectedProvider: Providers = Providers.OpenAI) => { + const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized")); + mockUseAuthorized.default.mockReturnValue(mockAuthorizedUser("proxy_admin", "user-1", true)); + const user = userEvent.setup({ pointerEventsCheck: PointerEventsCheckLevel.Never }); + const props = { ...createTestProps(), selectedProvider }; + renderWithProviders(); + await screen.findByText("Provider"); + return { user, props, clickTestConnect: () => user.click(screen.getByTestId("test-connect-btn")) }; + }; + + const expectTestRun = async (litellmParams: Record) => { + expect(await screen.findByTestId("connection-success-msg")).toBeInTheDocument(); + expect(testConnectionRequest).toHaveBeenCalledWith( + "test-access-token", + expect.objectContaining(litellmParams), + expect.anything(), + undefined, + ); + }; + + const expectBlockedWith = async (inlineError: string) => { + expect(await screen.findByText(inlineError)).toBeInTheDocument(); + expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); + expect(testConnectionRequest).not.toHaveBeenCalled(); + }; + + it("shows the inline required errors instead of testing a fresh form", async () => { + const { clickTestConnect } = await renderForm(); + + await clickTestConnect(); + + await expectBlockedWith("Please enter at least one model."); + expect(screen.getAllByText("Required")).toHaveLength(2); + }); + + it("runs the connection test once provider, auth and model are filled in", async () => { + const { user, clickTestConnect } = await renderForm(); + await pickProvider(user, /OpenAI/); + fillApiKey("OpenAI API Key"); + await pickModel(user, "gpt-4"); + + await clickTestConnect(); + + await expectTestRun({ api_key: "sk-test", model: "gpt-4" }); + }); + + it("shows emptied inputs and inline errors instead of testing after a successful Add Model", async () => { + const { user, props, clickTestConnect } = await renderForm(); + props.handleOk.mockImplementation(async () => { + props.form.reset({ litellm_credential_name: null }); + return true; + }); + await pickProvider(user, /OpenAI/); + fillApiKey("OpenAI API Key"); + await pickModel(user, "gpt-4"); + await user.click(screen.getByTestId("add-model-btn")); + await waitFor(() => expect(screen.getByLabelText(/OpenAI API Key/)).toHaveValue("")); + expect(props.handleOk).toHaveBeenCalledOnce(); + + await clickTestConnect(); + + await expectBlockedWith("Please enter at least one model."); + }); + + it("runs the connection test for the provider wildcard", async () => { + const { user, clickTestConnect } = await renderForm(); + await pickProvider(user, /OpenAI/); + await pickModel(user, "All OpenAI Models (Wildcard)"); + fillApiKey("OpenAI API Key"); + + await clickTestConnect(); + + await expectTestRun({ model: "openai/*" }); + }); + + it("asks for the custom model name before testing a custom model", async () => { + const { user, clickTestConnect } = await renderForm(); + await pickProvider(user, /OpenAI/); + fillApiKey("OpenAI API Key"); + await pickModel(user, "Custom Model Name (Enter below)"); + + await clickTestConnect(); + await expectBlockedWith("Please enter a custom model name."); + + fireEvent.change(screen.getByPlaceholderText("Enter custom model name"), { target: { value: "my-fine-tune" } }); + await clickTestConnect(); + + await expectTestRun({ model: "my-fine-tune" }); + }); + + it("runs the connection test for an Azure deployment name", async () => { + const { user, clickTestConnect } = await renderForm(Providers.Azure); + await pickProvider(user, /Azure/); + fireEvent.change(screen.getByLabelText(/LiteLLM Model Name/), { target: { value: "my-deployment" } }); + + await clickTestConnect(); + await expectBlockedWith("Required"); + + fireEvent.change(screen.getByLabelText(/API Base/), { target: { value: "https://example.openai.azure.com" } }); + await clickTestConnect(); + + await expectTestRun({ model: "azure/my-deployment", api_base: "https://example.openai.azure.com" }); + }); + + it("asks for a credential before testing a credential-only auth type", async () => { + const errorToast = vi.spyOn(toast, "error"); + const { user, clickTestConnect } = await renderForm(Providers.MICROSOFT_365_COPILOT); + await pickProvider(user, /Microsoft 365 Copilot/); + await pickModel(user, "gpt-4"); + expect(await screen.findByRole("button", { name: "Create credential" })).toBeInTheDocument(); + + await clickTestConnect(); + expect(errorToast).toHaveBeenCalledWith("Create or select a credential for this auth type"); + expect(screen.queryByRole("dialog")).not.toBeInTheDocument(); + expect(testConnectionRequest).not.toHaveBeenCalled(); + errorToast.mockRestore(); + + await user.click(screen.getByPlaceholderText("Select or search for existing credentials")); + await user.click(await screen.findByRole("option", { name: /test-credential/ })); + await clickTestConnect(); + + await expectTestRun({ litellm_credential_name: "test-credential" }); + }); + + it("runs the connection test with a credential just created from the credential dialog", async () => { + const { user, clickTestConnect } = await renderForm(); + await pickProvider(user, /OpenAI/); + await pickModel(user, "gpt-4"); + await user.click(screen.getByRole("button", { name: "Use workload identity federation" })); + fireEvent.change(screen.getByLabelText("Credential Name:"), { target: { value: "openai-federated" } }); + fireEvent.change(screen.getByLabelText(/Service Account ID/), { target: { value: "svc_new" } }); + await user.click(screen.getByRole("button", { name: "Add Credential" })); + await waitFor(() => expect(screen.queryByLabelText(/OpenAI API Key/)).not.toBeInTheDocument()); + + await clickTestConnect(); + + await expectTestRun({ litellm_credential_name: "openai-federated", model: "gpt-4" }); + }); + + it("runs the connection test after the public model name is edited in the mapping table", async () => { + const { user, clickTestConnect } = await renderForm(); + await pickProvider(user, /OpenAI/); + fillApiKey("OpenAI API Key"); + await pickModel(user, "gpt-4"); + const publicName = screen.getByDisplayValue("gpt-4"); + await user.clear(publicName); + const apiKey = screen.getByLabelText(/OpenAI API Key/); + await user.clear(apiKey); + await user.type(apiKey, "sk-rotated"); + await user.type(publicName, "my-gpt"); + + await clickTestConnect(); + + await expectTestRun({ api_key: "sk-rotated", model: "gpt-4" }); + }); + + it("shows the mapping error instead of testing when the public model name is cleared", async () => { + const { user, clickTestConnect } = await renderForm(); + await pickProvider(user, /OpenAI/); + fillApiKey("OpenAI API Key"); + await pickModel(user, "gpt-4"); + fireEvent.change(screen.getByDisplayValue("gpt-4"), { target: { value: "" } }); + + await clickTestConnect(); + + await expectBlockedWith("All model mappings must have valid public names"); + }); + }); + describe("cache control bindings reach the parent form store", () => { const renderWithForm = async () => { const mockUseAuthorized = vi.mocked(await import("@/app/(dashboard)/hooks/useAuthorized")); diff --git a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx index 5313df6153a..7e90991523c 100644 --- a/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx +++ b/ui/litellm-dashboard/src/components/add_model/AddModelForm.tsx @@ -134,15 +134,20 @@ const AddModelForm: React.FC = ({ setIsCredentialModalOpen(true); }; - const handleSubmit = async () => { - if (selectedAuthType?.credentialOnly && !selectedCredentialName) { + const rejectMissingCredential = (): boolean => { + const missing = selectedAuthType?.credentialOnly === true && !selectedCredentialName; + if (missing) { toast.error("Create or select a credential for this auth type"); - return false; } - return handleOk(); + return missing; }; + const handleSubmit = async () => (rejectMissingCredential() ? false : handleOk()); + const handleTestConnection = async () => { + if (rejectMissingCredential() || !(await form.trigger(registry.mountedNames() as string[]))) { + return; + } setIsTestingConnection(true); setConnectionTestId(`test-${Date.now()}`); setIsResultModalVisible(true); diff --git a/ui/litellm-dashboard/src/components/add_model/model_connection_test.test.tsx b/ui/litellm-dashboard/src/components/add_model/model_connection_test.test.tsx index c436a24474c..3bc867f843c 100644 --- a/ui/litellm-dashboard/src/components/add_model/model_connection_test.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/model_connection_test.test.tsx @@ -93,6 +93,19 @@ describe("ModelConnectionTest", () => { expect(toast.success).toHaveBeenCalledWith("Copied to clipboard"); }); + it("asks for the required fields instead of crashing when nothing could be prepared", async () => { + vi.mocked(prepareModelAddRequest).mockResolvedValue([] as never); + + render(); + await finishConnectionTest(); + + expect(screen.getByTestId("connection-failure-msg")).toHaveTextContent("Connection to this model failed"); + expect( + screen.getByText("Required fields are not filled in. Fill them in and test the connection again."), + ).toBeInTheDocument(); + expect(testConnectionRequest).not.toHaveBeenCalled(); + }); + it("shows a preparation failure without sending a connection request", async () => { vi.mocked(prepareModelAddRequest).mockResolvedValue(null as never); diff --git a/ui/litellm-dashboard/src/components/add_model/model_connection_test.tsx b/ui/litellm-dashboard/src/components/add_model/model_connection_test.tsx index f6ef86a77ca..6fdc0461482 100644 --- a/ui/litellm-dashboard/src/components/add_model/model_connection_test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/model_connection_test.tsx @@ -43,8 +43,12 @@ const ModelConnectionTest: React.FC = ({ try { const result = await prepareModelAddRequest(formValues, accessToken, null); - if (!result) { - setError("Failed to prepare model data. Please check your form inputs."); + if (!result?.length) { + setError( + result + ? "Required fields are not filled in. Fill them in and test the connection again." + : "Failed to prepare model data. Please check your form inputs.", + ); setIsSuccess(false); setIsLoading(false); return; diff --git a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.test.tsx b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.test.tsx index 09d48f828e1..3a58da9248a 100644 --- a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.test.tsx +++ b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.test.tsx @@ -112,6 +112,13 @@ vi.mock("../networking", async () => { }, ], }, + { + provider: "Oracle", + provider_display_name: Providers.Oracle, + litellm_provider: "oci", + default_model_placeholder: "oci/xai.grok-4", + credential_fields: [{ key: "oci_key", label: "OCI Private Key (PEM)", field_type: "textarea", required: true }], + }, { provider: "MICROSOFT_365_COPILOT", provider_display_name: Providers.MICROSOFT_365_COPILOT, @@ -182,6 +189,15 @@ const VertexCredentialsProbe = () => { return {String(watch("vertex_credentials") ?? "")}; }; +const ResetFormButton = () => { + const form = useFormContext(); + return ( + + ); +}; + const ValidationForm = ({ children }: { readonly children: ReactNode }) => { const form = useFormContext(); return ( @@ -642,6 +658,49 @@ describe("ProviderSpecificFields", () => { expect(await screen.findAllByText("Required")).toHaveLength(3); }); + it("renders typed password and text inputs empty after the form is reset", async () => { + render( + + + + + + , + ); + const apiKey = await screen.findByLabelText("OpenAI API Key"); + const organization = screen.getByLabelText("OpenAI Organization ID"); + const apiBase = screen.getByLabelText(/API Base/); + expect(apiBase).toHaveValue("https://api.openai.com/v1"); + fireEvent.change(apiKey, { target: { value: "sk-typed" } }); + fireEvent.change(organization, { target: { value: "org-typed" } }); + fireEvent.change(apiBase, { target: { value: "https://proxy.example.test/v1" } }); + expect(apiKey).toHaveValue("sk-typed"); + + fireEvent.click(screen.getByRole("button", { name: "Reset form" })); + + await waitFor(() => expect(apiKey).toHaveValue("")); + expect(organization).toHaveValue(""); + expect(apiBase).toHaveValue("https://api.openai.com/v1"); + }); + + it("renders a typed textarea input empty after the form is reset", async () => { + render( + + + + + + , + ); + const privateKey = await screen.findByLabelText(/OCI Private Key/); + fireEvent.change(privateKey, { target: { value: "-----BEGIN PRIVATE KEY-----" } }); + expect(privateKey).toHaveValue("-----BEGIN PRIVATE KEY-----"); + + fireEvent.click(screen.getByRole("button", { name: "Reset form" })); + + await waitFor(() => expect(privateKey).toHaveValue("")); + }); + it("does not render an Auth Type selector for OpenAI", async () => { const queryClient = createQueryClient(); render( diff --git a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx index a38b077270f..31a7cb02d6e 100644 --- a/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx +++ b/ui/litellm-dashboard/src/components/add_model/provider_specific_fields.tsx @@ -301,11 +301,10 @@ const ProviderSpecificFieldsContent: React.FC = ({ return (