diff --git a/litellm/_internal_context.py b/litellm/_internal_context.py index df549701cf4..c4527e1d099 100644 --- a/litellm/_internal_context.py +++ b/litellm/_internal_context.py @@ -36,6 +36,8 @@ REDIS_FAMILIES_METADATA_KEY: Final = "families" _service_caller: Final[ContextVar[str | None]] = ContextVar("service_caller", default=None) +_resent_on: Final[ContextVar[Callable[[Exception], bool] | None]] = ContextVar("resent_on", default=None) + @contextmanager def post_response_phase() -> Generator[None]: @@ -113,6 +115,23 @@ def current_service_caller() -> str | None: return _service_caller.get() +@contextmanager +def resent_on(rejection: Callable[[Exception], bool]) -> Generator[None]: + """An attempt its caller sends again when it fails with ``rejection``. That failure is not the request's + outcome, so it is neither logged as the request's failure nor retried, and its budget reservation and rate-limit + slot stay with the attempt sent again.""" + token: Final = _resent_on.set(rejection) + try: + yield + finally: + _restore(_resent_on, token) + + +def is_resent_failure(error: Exception) -> bool: + rejection: Final = _resent_on.get() + return rejection is not None and rejection(error) + + @contextmanager def pinned_billing_time(moment: datetime) -> Generator[None]: """Price every rate lookup inside this block at ``moment`` rather than at each one's own clock read.""" diff --git a/litellm/completion_extras/litellm_responses_transformation/handler.py b/litellm/completion_extras/litellm_responses_transformation/handler.py index 3f7017eccc2..c47bb046021 100644 --- a/litellm/completion_extras/litellm_responses_transformation/handler.py +++ b/litellm/completion_extras/litellm_responses_transformation/handler.py @@ -182,6 +182,9 @@ class ResponsesToCompletionBridgeHandler: model_response: Final = validated_kwargs["model_response"] logging_obj: Final = validated_kwargs["logging_obj"] custom_llm_provider: Final = validated_kwargs["custom_llm_provider"] + from .transformation import serves_anthropic_messages + + serves_messages_client: Final = serves_anthropic_messages(logging_obj) if kwargs.get("stream") is True and "stream" not in optional_params: optional_params = {**optional_params, "stream": True} @@ -222,6 +225,7 @@ class ResponsesToCompletionBridgeHandler: encoding=validated_kwargs["encoding"], api_key=kwargs.get("api_key"), json_mode=kwargs.get("json_mode"), + keeps_every_reasoning_item=serves_messages_client, ) elif isinstance(result, ModelResponse): if not stream: @@ -247,6 +251,7 @@ class ResponsesToCompletionBridgeHandler: encoding=validated_kwargs["encoding"], api_key=kwargs.get("api_key"), json_mode=kwargs.get("json_mode"), + keeps_every_reasoning_item=serves_messages_client, ) else: if self._is_preformatted_cached_chat_stream(result): @@ -260,6 +265,7 @@ class ResponsesToCompletionBridgeHandler: structured_output=is_structured_output_request( cast("Mapping[str, object]", optional_params) # cast-ok: optional_params is an untyped dict ), + streams_reasoning_signatures=serves_messages_client, ) streamwrapper: Final = CustomStreamWrapper( completion_stream=completion_stream, @@ -282,6 +288,9 @@ class ResponsesToCompletionBridgeHandler: model_response: Final = validated_kwargs["model_response"] logging_obj: Final = validated_kwargs["logging_obj"] custom_llm_provider: Final = validated_kwargs["custom_llm_provider"] + from .transformation import serves_anthropic_messages + + serves_messages_client: Final = serves_anthropic_messages(logging_obj) if kwargs.get("stream") is True and "stream" not in optional_params: optional_params = {**optional_params, "stream": True} @@ -325,6 +334,7 @@ class ResponsesToCompletionBridgeHandler: encoding=validated_kwargs["encoding"], api_key=kwargs.get("api_key"), json_mode=kwargs.get("json_mode"), + keeps_every_reasoning_item=serves_messages_client, ) elif isinstance(result, ModelResponse): if not stream: @@ -350,6 +360,7 @@ class ResponsesToCompletionBridgeHandler: encoding=validated_kwargs["encoding"], api_key=kwargs.get("api_key"), json_mode=kwargs.get("json_mode"), + keeps_every_reasoning_item=serves_messages_client, ) else: if self._is_preformatted_cached_chat_stream(result): @@ -363,6 +374,7 @@ class ResponsesToCompletionBridgeHandler: structured_output=is_structured_output_request( cast("Mapping[str, object]", optional_params) # cast-ok: optional_params is an untyped dict ), + streams_reasoning_signatures=serves_messages_client, ) streamwrapper: Final = CustomStreamWrapper( completion_stream=completion_stream, diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index 51f84860c3e..28c5f03da5e 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -31,9 +31,11 @@ from litellm.litellm_core_utils.allowed_tools_tool_choice import ( parse_allowed_tools_tool_choice, responses_allowed_tools_tool_choice, ) +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_or_none from litellm.litellm_core_utils.hidden_params import get_hidden_params, get_or_create_hidden_params from litellm.litellm_core_utils.prompt_templates.common_utils import ( responses_reasoning_items_from_thinking_blocks, + thinking_blocks_from_reasoning_items, with_prompt_cache_breakpoint, ) from litellm.llms.base_llm.base_model_iterator import BaseModelResponseIterator @@ -50,6 +52,8 @@ from litellm.responses.utils import ResponsesAPIRequestUtils, normalize_response from litellm.types.llms.openai import ( ChatCompletionAnnotation, ChatCompletionReasoningItem, + ChatCompletionRedactedThinkingBlock, + ChatCompletionThinkingBlock, ChatCompletionToolCallChunk, ChatCompletionToolCallFunctionChunk, ChatCompletionToolParamFunctionChunk, @@ -59,7 +63,7 @@ from litellm.types.llms.openai import ( ResponsesAPIResponse, ResponsesAPIStreamEvents, ) -from litellm.types.utils import GenericStreamingChunk, ModelResponseStream +from litellm.types.utils import CallTypes, Delta, GenericStreamingChunk, ModelResponseStream, StreamingChoices if TYPE_CHECKING: from openai.types.responses import ResponseInputImageParam, ResponseOutputItem @@ -75,8 +79,6 @@ if TYPE_CHECKING: AllMessageValues, ChatCompletionFileObject, ChatCompletionImageObject, - ChatCompletionRedactedThinkingBlock, - ChatCompletionThinkingBlock, ChatCompletionToolReferenceObject, OpenAIMessageContentListBlock, ) @@ -230,6 +232,10 @@ class _BuiltReasoningItem(TypedDict): summary: Sequence[_ReasoningSummaryText] +def serves_anthropic_messages(logging_obj: object) -> bool: + return attribute_or_none(logging_obj, "call_type") == CallTypes.anthropic_messages.value + + def _get_reasoning_items( msg: "AllMessageValues", ) -> list[ChatCompletionReasoningItem]: @@ -448,6 +454,77 @@ def _as_chat_reasoning_items( return cast(list[ChatCompletionReasoningItem], list(reasoning_items)) +def _with_stream_block_index( + block: Mapping[str, object], + index: int, + summary_streamed: bool, +) -> ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock: + if block.get("type") == "redacted_thinking": + data: Final = block.get("data") + return ChatCompletionRedactedThinkingBlock( + type="redacted_thinking", + data=data if isinstance(data, str) else "", + index=index, + ) + thinking: Final = block.get("thinking") + signature: Final = block.get("signature") + return ChatCompletionThinkingBlock( + type="thinking", + thinking=thinking if isinstance(thinking, str) and not summary_streamed else "", + signature=signature if isinstance(signature, str) else None, + index=index, + ) + + +def _stream_output_index(chunk: Mapping[str, object]) -> int | None: + output_index: Final = chunk.get("output_index") + return output_index if isinstance(output_index, int) and not isinstance(output_index, bool) else None + + +def _reasoning_done_thinking_blocks( + item: Mapping[str, object], + output_index: int | None, + summary_streamed: bool, +) -> tuple[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, ...]: + index: Final = output_index if output_index is not None else 0 + return tuple( + _with_stream_block_index(block, index, summary_streamed) + for block in thinking_blocks_from_reasoning_items((item,)) + ) + + +def _reasoning_done_stream_chunk( + chunk: Mapping[str, object], + summary_streamed_outputs: frozenset[int | None], +) -> ModelResponseStream | None: + event_type: Final = chunk.get("type") + event_name: Final = ( + event_type.value + if isinstance(event_type, ResponsesAPIStreamEvents) + else event_type + if isinstance(event_type, str) + else None + ) + if event_name != ResponsesAPIStreamEvents.OUTPUT_ITEM_DONE: + return None + item: Final = as_str_mapping(chunk.get("item")) + if item is None or item.get("type") != "reasoning": + return None + output_index: Final = _stream_output_index(chunk) + signed_blocks: Final = _reasoning_done_thinking_blocks(item, output_index, output_index in summary_streamed_outputs) + if not signed_blocks: + return None + return ModelResponseStream( + choices=[ + StreamingChoices( + index=0, + delta=Delta(thinking_blocks=list(signed_blocks)), + finish_reason=None, + ) + ] + ) + + def _shift_annotation(annotation: ChatCompletionAnnotation, offset: int) -> ChatCompletionAnnotation: shifted: Final = { key: value + offset if key in ("start_index", "end_index") and isinstance(value, int) else value @@ -481,7 +558,7 @@ def _fold_choices_into_one( choices: Sequence["Choices"], tool_calls: Sequence[Mapping[str, object]], trailing_reasoning_content: str | None, - trailing_reasoning_item: _BuiltReasoningItem | None, + trailing_reasoning_items: Sequence[_BuiltReasoningItem], tool_search_items: Sequence[Mapping[str, object]], structured_output: bool = False, ) -> "Choices": @@ -500,7 +577,7 @@ def _fold_choices_into_one( reasoning_contents: Final = (*(_message_reasoning_content(m) for m in messages), trailing_reasoning_content) reasoning_items: Final = ( *chain.from_iterable(_message_reasoning_items(m) for m in messages), - *(() if trailing_reasoning_item is None else (trailing_reasoning_item,)), + *trailing_reasoning_items, ) return Choices( message=Message( @@ -1169,6 +1246,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): output_items: Sequence[object], handle_raw_dict_callback: Callable[..., tuple["Choices | None", int]] | None = None, structured_output: bool = False, + keeps_every_reasoning_item: bool = False, ) -> list["Choices"]: """ Convert Responses API output items to chat completion choices. @@ -1178,6 +1256,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): handle_raw_dict_callback: Optional callback for handling raw dict items structured_output: The chat request carried a json response_format, so only the first message item becomes content + keeps_every_reasoning_item: Keep each reasoning item before an answer, not only the last Returns: List of Choices objects @@ -1200,7 +1279,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): choices: Final[list[Choices]] = [] index = 0 reasoning_content: str | None = None - pending_reasoning_item: _BuiltReasoningItem | None = None + pending_reasoning_items: tuple[_BuiltReasoningItem, ...] = () # Collect all tool calls to put them in a single choice # (Chat Completions API expects all tool calls in one message) @@ -1209,12 +1288,19 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): for item in output_items: if isinstance(item, ResponseReasoningItem): - pending_reasoning_item = _build_reasoning_item( - item_id=item.id, - encrypted_content=getattr(item, "encrypted_content", None), - summary_raw=item.summary, + pending_reasoning_items = ( + *(pending_reasoning_items if keeps_every_reasoning_item else ()), + _build_reasoning_item( + item_id=item.id, + encrypted_content=getattr(item, "encrypted_content", None), + summary_raw=item.summary, + ), + ) + reasoning_content = " ".join( + s["text"] + for s in chain.from_iterable(built["summary"] for built in pending_reasoning_items) + if s.get("text") ) - reasoning_content = " ".join(s["text"] for s in pending_reasoning_item["summary"] if s.get("text")) elif isinstance(item, ResponseOutputMessage): for content in item.content: @@ -1231,7 +1317,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): annotations=annotations, reasoning_items=cast( list[ChatCompletionReasoningItem] | None, - ([pending_reasoning_item] if pending_reasoning_item is not None else None), + (list(pending_reasoning_items) or None), ), ) @@ -1244,7 +1330,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ) reasoning_content = None # flush - pending_reasoning_item = None # flush + pending_reasoning_items = () # flush index += 1 elif isinstance(item, ResponseFunctionToolCall): @@ -1309,7 +1395,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): choices=choices, tool_calls=accumulated_tool_calls, trailing_reasoning_content=reasoning_content, - trailing_reasoning_item=pending_reasoning_item, + trailing_reasoning_items=pending_reasoning_items, tool_search_items=tool_search_items, structured_output=structured_output, ) @@ -1416,6 +1502,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): encoding: object, api_key: str | None = None, json_mode: bool | None = None, + keeps_every_reasoning_item: bool = False, ) -> "ModelResponse": """Transform Responses API response to chat completion response""" from litellm.responses.utils import ResponseAPILoggingUtils @@ -1447,6 +1534,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): structured_output=is_structured_output_request( cast("Mapping[str, object]", optional_params) # cast-ok: optional_params is an untyped dict ), + keeps_every_reasoning_item=keeps_every_reasoning_item, ) response_is_incomplete: Final = raw_response.status == "incomplete" or ( @@ -1510,9 +1598,14 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): sync_stream: bool, json_mode: bool | None = False, structured_output: bool = False, + streams_reasoning_signatures: bool = False, ) -> BaseModelResponseIterator: return OpenAiResponsesToChatCompletionStreamIterator( - streaming_response, sync_stream, json_mode, structured_output=structured_output + streaming_response, + sync_stream, + json_mode, + structured_output=structured_output, + streams_reasoning_signatures=streams_reasoning_signatures, ) def _convert_content_str_to_input_text(self, content: str, role: str) -> dict[str, object]: @@ -1891,13 +1984,16 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): sync_stream: bool, json_mode: bool | None = False, structured_output: bool = False, + streams_reasoning_signatures: bool = False, ): super().__init__(streaming_response, sync_stream, json_mode) + self._streams_reasoning_signatures: Final = streams_reasoning_signatures self._chat_completion_id: str | None = None self._served_service_tier: str | None = None self._tool_call_index_map: dict[int, int] = {} # mutable-ok: per-stream accumulator state self._structured_output: Final = structured_output self._first_message_item_done: bool = False + self._summary_streamed_outputs: frozenset[int | None] = frozenset() def _handle_string_chunk( self, str_line: Union[str, "BaseModel"] @@ -2109,7 +2205,6 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): ) ] ) - elif event_type == "response.output_text.delta": # Content part added to output content_part = parsed_chunk.get("delta", None) @@ -2213,15 +2308,24 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): """ verbose_logger.debug("Chat provider: transform_streaming_response called with chunk: %s", chunk) self._remember_served_service_tier(chunk) + reasoning_chunk: Final = ( + _reasoning_done_stream_chunk(chunk, self._summary_streamed_outputs) + if self._streams_reasoning_signatures + else None + ) translated: Final = ( self._empty_content_chunk() if self._is_text_delta_after_structured_answer(chunk) + else reasoning_chunk + if reasoning_chunk is not None else OpenAiResponsesToChatCompletionStreamIterator.translate_responses_chunk_to_openai_stream( chunk, tool_call_index_map=self._tool_call_index_map ) ) if self._structured_output and _is_message_item_done(chunk): self._first_message_item_done = True + if _responses_event_type(chunk) == "response.reasoning_summary_text.delta" and chunk.get("delta"): + self._summary_streamed_outputs = self._summary_streamed_outputs | {_stream_output_index(chunk)} return self._with_served_service_tier(self._with_stream_scoped_id(translated)) def _is_text_delta_after_structured_answer(self, chunk: Mapping[str, object]) -> bool: diff --git a/litellm/litellm_core_utils/get_llm_provider_logic.py b/litellm/litellm_core_utils/get_llm_provider_logic.py index a67c1bb7df8..a9be736360b 100644 --- a/litellm/litellm_core_utils/get_llm_provider_logic.py +++ b/litellm/litellm_core_utils/get_llm_provider_logic.py @@ -158,6 +158,14 @@ def declared_authenticating_provider(model: str | None, custom_llm_provider: str return declared if declared in PROVIDERS_THAT_AUTHENTICATE_ON_PROVIDER_INFO else None +def resolve_model_and_provider_for_metadata(model: str, custom_llm_provider: str | None = None) -> tuple[str, str]: + declared: Final = declared_authenticating_provider(model, custom_llm_provider) + if declared is not None: + return model.removeprefix(f"{declared}/"), declared + resolved_model, resolved_provider, _, _ = get_llm_provider(model=model, custom_llm_provider=custom_llm_provider) + return resolved_model, resolved_provider + + def inferred_provider(model: str | None) -> str | None: if not model: return None diff --git a/litellm/litellm_core_utils/litellm_logging.py b/litellm/litellm_core_utils/litellm_logging.py index 95323c2ef5e..8cd8e55c292 100644 --- a/litellm/litellm_core_utils/litellm_logging.py +++ b/litellm/litellm_core_utils/litellm_logging.py @@ -4593,6 +4593,7 @@ class Logging(LiteLLMLoggingBaseClass): from litellm.completion_extras.litellm_responses_transformation.transformation import ( LiteLLMResponsesTransformationHandler, responses_provider_metadata, + serves_anthropic_messages, ) served_id: Final = _provider_response_id(result) @@ -4607,6 +4608,7 @@ class Logging(LiteLLMLoggingBaseClass): optional_params={}, litellm_params={}, encoding=litellm.encoding, + keeps_every_reasoning_item=serves_anthropic_messages(self), ) translated.id = served_id or translated.id return translated diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index ac5791957b9..f9c4ebca350 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -20,6 +20,7 @@ from pydantic import TypeAdapter import litellm from litellm import verbose_logger +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.router_utils.batch_utils import InMemoryFile from litellm.types.llms.openai import ( AllMessageValues, @@ -2289,6 +2290,89 @@ def encrypted_content_of_block(block: Mapping[str, object]) -> str | None: return encrypted_content_from_signature(_encrypted_reasoning_field(block)) +_REASONING_SUMMARY_SEPARATOR: Final = "\n\n" + +_OBJECT_TUPLE: Final = TypeAdapter(tuple[object, ...]) + + +def _object_sequence(value: object) -> tuple[object, ...]: + if not isinstance(value, Sequence) or isinstance(value, (str, bytes)): + return () + return _OBJECT_TUPLE.validate_python(value) + + +def _iter_mappings(value: object) -> tuple[Mapping[str, object], ...]: + return tuple(mapping for item in _object_sequence(value) if (mapping := as_str_mapping(item)) is not None) + + +def _reasoning_summary_part_text(part: object) -> str: + mapping: Final = as_str_mapping(part) + if mapping is not None: + text: Final = mapping.get("text") + return text if isinstance(text, str) else "" + text_attr: Final = getattr(part, "text", None) + return text_attr if isinstance(text_attr, str) else "" + + +def _reasoning_item_summary_text(item: Mapping[str, object]) -> str: + parts: Final = tuple(_reasoning_summary_part_text(part) for part in _object_sequence(item.get("summary"))) + return _REASONING_SUMMARY_SEPARATOR.join(part for part in parts if part) + + +def thinking_block_from_reasoning_item( + item: object, +) -> ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock | None: + """The chat thinking block that carries one Responses `reasoning` item back to an Anthropic client. + + `encrypted_content` rides the same `litellm_encrypted_reasoning:` signature the Responses + adapter uses. A summary becomes thinking text. Bytes with no summary become `redacted_thinking`, + because an empty thinking string is not something the client can replay. + """ + mapping: Final = as_str_mapping(item) + if mapping is None or mapping.get("type") not in (None, "reasoning"): + return None + encrypted_content: Final = mapping.get("encrypted_content") + if not isinstance(encrypted_content, str) or not encrypted_content: + return None + signature: Final = encrypted_reasoning_signature(encrypted_content) + text: Final = _reasoning_item_summary_text(mapping) + if not text: + return ChatCompletionRedactedThinkingBlock(type="redacted_thinking", data=signature) + return ChatCompletionThinkingBlock(type="thinking", thinking=text, signature=signature) + + +def thinking_blocks_from_reasoning_items( + items: object, +) -> tuple[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, ...]: + return tuple( + block for item in _object_sequence(items) if (block := thinking_block_from_reasoning_item(item)) is not None + ) + + +def _block_encrypted_content(block: object) -> str | None: + mapping: Final = as_str_mapping(block) + if mapping is None: + return None + return encrypted_content_of_block(mapping) + + +def thinking_blocks_with_encrypted_reasoning( + thinking_blocks: object, + reasoning_items: object, +) -> tuple[Mapping[str, object], ...]: + """Thinking blocks already on the message, plus any encrypted reasoning item they do not already carry.""" + direct: Final = _iter_mappings(thinking_blocks) + represented: Final = frozenset( + content for block in direct if (content := _block_encrypted_content(block)) is not None + ) + extra: Final = tuple( + block + for block in thinking_blocks_from_reasoning_items(_iter_mappings(reasoning_items)) + if _block_encrypted_content(block) not in represented + ) + return (*extra, *direct) + + def is_encrypted_reasoning_block(block: object) -> bool: """A thinking or redacted_thinking block carrying Responses API encrypted reasoning. @@ -2380,6 +2464,14 @@ def anthropic_content_lists(messages: Sequence[object]) -> Iterator[object]: ) +def _content_carries_encrypted_reasoning(content: object) -> bool: + return any(is_encrypted_reasoning_block(block) for block in _object_sequence(content)) + + +def messages_carry_encrypted_reasoning(messages: Sequence[object]) -> bool: + return any(_content_carries_encrypted_reasoning(content) for content in anthropic_content_lists(messages)) + + def _strip_encrypted_reasoning_from_blocks( content: object, *, diff --git a/litellm/litellm_core_utils/prompt_templates/factory.py b/litellm/litellm_core_utils/prompt_templates/factory.py index 7b3885eac6e..767f3b4c30d 100644 --- a/litellm/litellm_core_utils/prompt_templates/factory.py +++ b/litellm/litellm_core_utils/prompt_templates/factory.py @@ -4599,6 +4599,16 @@ def get_assistant_message_block_or_continue_message( raise ValueError(f"Unsupported content type: {type(content_block)}") +def _as_thinking_block(item: object) -> ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock | None: + match item: + case {"type": "thinking" | "redacted_thinking"}: + return cast( # cast-ok: the mapping pattern matched a thinking block's type tag + ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, item + ) + case _: + return None + + class BedrockConverseMessagesProcessor: @staticmethod def _initial_message_setup( @@ -4647,11 +4657,15 @@ class BedrockConverseMessagesProcessor: deployment_model: str | None = None, base_model: str | None = None, ) -> list[BedrockMessageBlock]: - from litellm.llms.bedrock.common_utils import bedrock_converse_accepts_redacted_thinking + from litellm.llms.bedrock.common_utils import ( + bedrock_converse_accepts_redacted_thinking, + bedrock_converse_rejects_signed_thinking, + ) contents: list[BedrockMessageBlock] = [] msg_i = 0 to_claude: Final = bedrock_converse_accepts_redacted_thinking(model, base_model) + drop_signed_thinking: Final = bedrock_converse_rejects_signed_thinking(model, base_model) messages = BedrockConverseMessagesProcessor._initial_message_setup( messages, model, llm_provider, user_continue_message @@ -4803,7 +4817,7 @@ class BedrockConverseMessagesProcessor: if thinking_blocks is not None: converted_thinking_blocks = ( BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks( - thinking_blocks, to_claude=to_claude + thinking_blocks, to_claude=to_claude, drop_signed=drop_signed_thinking ) ) assistant_content = BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content( @@ -4817,7 +4831,7 @@ class BedrockConverseMessagesProcessor: if isinstance(element, dict): if element["type"] == "thinking" or element["type"] == "redacted_thinking": thinking_block = BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks( - thinking_blocks=[element], to_claude=to_claude + thinking_blocks=[element], to_claude=to_claude, drop_signed=drop_signed_thinking ) assistants_parts = ( BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content( @@ -4844,7 +4858,9 @@ class BedrockConverseMessagesProcessor: model=model, deployment_model=deployment_model, ) - if _cache_point_block is not None: + if _cache_point_block is not None and BedrockConverseMessagesProcessor.keeps_cache_point( + element, drop_signed=drop_signed_thinking + ): assistants_parts.append(_cache_point_block) assistant_content.extend(assistants_parts) elif _assistant_content is not None and isinstance(_assistant_content, str): @@ -4892,6 +4908,7 @@ class BedrockConverseMessagesProcessor: thinking_blocks: Sequence[ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock], *, to_claude: bool = False, + drop_signed: bool = False, ) -> list[BedrockContentBlock]: return [ BedrockContentBlock( @@ -4901,8 +4918,25 @@ class BedrockConverseMessagesProcessor: ) for thinking_block in thinking_blocks if not (to_claude and is_encrypted_reasoning_block(thinking_block)) + and not (drop_signed and BedrockConverseMessagesProcessor._is_signed(thinking_block)) ] + @staticmethod + def keeps_cache_point(element: object, *, drop_signed: bool) -> bool: + match _as_thinking_block(element): + case None: + return True + case thinking_block: + return not (drop_signed and BedrockConverseMessagesProcessor._is_signed(thinking_block)) + + @staticmethod + def _is_signed(thinking_block: ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock) -> bool: + match thinking_block: + case {"signature": signature} if signature: + return True + case _: + return False + @staticmethod def _to_reasoning_content_block( thinking_block: ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock, @@ -5059,11 +5093,15 @@ def _bedrock_converse_messages_pt( - Please ensure that function response turn comes immediately after a function call turn - Conversation blocks and tool result blocks cannot be provided in the same turn. Issue: https://github.com/BerriAI/litellm/issues/6053 """ - from litellm.llms.bedrock.common_utils import bedrock_converse_accepts_redacted_thinking + from litellm.llms.bedrock.common_utils import ( + bedrock_converse_accepts_redacted_thinking, + bedrock_converse_rejects_signed_thinking, + ) contents: list[BedrockMessageBlock] = [] msg_i = 0 to_claude: Final = bedrock_converse_accepts_redacted_thinking(model, base_model) + drop_signed_thinking: Final = bedrock_converse_rejects_signed_thinking(model, base_model) messages = BedrockConverseMessagesProcessor._initial_message_setup( messages, model, llm_provider, user_continue_message @@ -5216,7 +5254,7 @@ def _bedrock_converse_messages_pt( if thinking_blocks is not None: converted_thinking_blocks = ( BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks( - thinking_blocks, to_claude=to_claude + thinking_blocks, to_claude=to_claude, drop_signed=drop_signed_thinking ) ) assistant_content = BedrockConverseMessagesProcessor.add_thinking_blocks_to_assistant_content( @@ -5231,7 +5269,7 @@ def _bedrock_converse_messages_pt( if element["type"] == "thinking" or element["type"] == "redacted_thinking": thinking_block = ( BedrockConverseMessagesProcessor.translate_thinking_blocks_to_reasoning_content_blocks( - thinking_blocks=[element], to_claude=to_claude + thinking_blocks=[element], to_claude=to_claude, drop_signed=drop_signed_thinking ) ) assistants_parts = ( @@ -5260,7 +5298,9 @@ def _bedrock_converse_messages_pt( model=model, deployment_model=deployment_model, ) - if _cache_point_block is not None: + if _cache_point_block is not None and BedrockConverseMessagesProcessor.keeps_cache_point( + element, drop_signed=drop_signed_thinking + ): assistants_parts.append(_cache_point_block) assistant_content.extend(assistants_parts) elif _assistant_content is not None and isinstance(_assistant_content, str): diff --git a/litellm/litellm_core_utils/streaming_handler.py b/litellm/litellm_core_utils/streaming_handler.py index 7135b9c1071..d91a67a2d13 100644 --- a/litellm/litellm_core_utils/streaming_handler.py +++ b/litellm/litellm_core_utils/streaming_handler.py @@ -912,6 +912,7 @@ class CustomStreamWrapper: and model_response.choices[0].delta.role is not None ) or (getattr(model_response.choices[0].delta, "reasoning_items", None) is not None) + or getattr(model_response.choices[0].delta, "thinking_blocks", None) ): return True else: diff --git a/litellm/llms/anthropic/pass_through/adapters/handler.py b/litellm/llms/anthropic/pass_through/adapters/handler.py index 8c0c065c457..a818d6f3565 100644 --- a/litellm/llms/anthropic/pass_through/adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/adapters/handler.py @@ -1,5 +1,8 @@ -from collections.abc import AsyncIterator, Coroutine, Iterator, Mapping, Sequence +import re +from collections.abc import AsyncIterator, Callable, Coroutine, Iterator, Mapping, Sequence +from contextlib import AbstractContextManager, nullcontext from dataclasses import dataclass +from functools import partial from typing import ( TYPE_CHECKING, Final, @@ -11,16 +14,25 @@ from pydantic import TypeAdapter from typing_extensions import TypedDict import litellm +from litellm._internal_context import resent_on from litellm._logging import verbose_logger from litellm.litellm_core_utils.asyncify import run_async_function +from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObject +from litellm.litellm_core_utils.model_param_helper import peek_reasoning_summary_aliases +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + is_encrypted_reasoning_block, + thinking_blocks_with_encrypted_reasoning, +) from litellm.litellm_core_utils.tool_fields import ( ROUTES_WITHOUT_MCP_SERVERS, drops_mcp_servers, tools_without_mcp_servers, ) +from litellm.llms.anthropic.common_utils import strip_encrypted_reasoning_blocks_from_anthropic_messages from litellm.llms.anthropic.pass_through.adapters.transformation import ( AnthropicAdapter, + openai_thinking_routes_to_responses, ) from litellm.llms.anthropic.pass_through.context_management import ( AnthropicContextManagementError, @@ -48,9 +60,9 @@ from litellm.types.llms.anthropic_messages.anthropic_response import ( ) from litellm.types.llms.openai import OpenAIWebSearchOptions from litellm.types.utils import ModelResponse -from litellm.utils import get_model_info if TYPE_CHECKING: + from litellm.chat_completions.dispatch import ChatResult from litellm.router import Router from litellm.types.proxy.auth.user_api_key_auth import UserAPIKeyAuth @@ -68,22 +80,156 @@ class _CompletionKwargs(TypedDict, total=False, extra_items=object): custom_llm_provider: str +_ANTHROPIC_MESSAGES: Final = TypeAdapter(list[dict[str, object]]) +_UNVERIFIED_ENCRYPTED_CONTENT: Final = re.compile(r"encrypted content \S+ could not be verified", re.IGNORECASE) +_PrepareReplay: TypeAlias = "Callable[[_AnthropicMessages], _CompletionKwargs]" + + @dataclass(frozen=True, slots=True) class _ChatTargetReasoning: effort: str summary: str | None = None +def _provider_error_code(body: object) -> str | None: + mapping: Final = as_str_mapping(body) + if mapping is None: + return None + nested: Final = as_str_mapping(mapping.get("error")) + source: Final = nested if nested is not None else mapping + code: Final = source.get("code") + return code if isinstance(code, str) else None + + +def _logged_rejection_code(error: litellm.BadRequestError) -> str: + if error.code == "invalid_encrypted_content" or _provider_error_code(error.body) == "invalid_encrypted_content": + return "invalid_encrypted_content" + if "invalid_encrypted_content" in str(error): + return "invalid_encrypted_content" + return "could_not_be_verified" + + +def _completion_text(completion_kwargs: _CompletionKwargs, key: str) -> str: + value: Final = completion_kwargs.get(key) + return value if isinstance(value, str) else "" + + +def _message_carries_encrypted_reasoning(message: object) -> bool: + mapping: Final = as_str_mapping(message) + if mapping is None: + return False + blocks: Final = thinking_blocks_with_encrypted_reasoning( + mapping.get("thinking_blocks"), mapping.get("reasoning_items") + ) + return any(is_encrypted_reasoning_block(block) for block in blocks) + + +def _sends_encrypted_reasoning(completion_kwargs: _CompletionKwargs) -> bool: + messages: Final = completion_kwargs.get("messages") + sent: Final = iter_items(messages) if messages is not None else () + return any(_message_carries_encrypted_reasoning(message) for message in sent) + + +def _rejects_encrypted_reasoning(error: litellm.BadRequestError) -> bool: + error_message: Final = str(error) + return ( + "invalid_encrypted_content" in error_message or _UNVERIFIED_ENCRYPTED_CONTENT.search(error_message) is not None + ) + + +def _is_rejected_reasoning(error: Exception) -> bool: + return isinstance(error, litellm.BadRequestError) and _rejects_encrypted_reasoning(error) + + +def resent_without_rejected_reasoning(sent_encrypted_reasoning: bool) -> AbstractContextManager[None]: + return resent_on(_is_rejected_reasoning) if sent_encrypted_reasoning else nullcontext() + + +def messages_without_rejected_reasoning( + messages: Sequence[Mapping[str, object]], + error: litellm.BadRequestError, + *, + sent_encrypted_reasoning: bool, + provider: str, + model: str, +) -> "_AnthropicMessages | None": + if not sent_encrypted_reasoning or not _rejects_encrypted_reasoning(error): + return None + verbose_logger.warning( + "Backend rejected replayed encrypted reasoning, retrying once without it. " + "Enable the encrypted_content_affinity pre-call check to drop such blocks before the call. " + "provider=%s model=%s status=%s code=%s", + provider, + model, + error.status_code, + _logged_rejection_code(error), + ) + return _ANTHROPIC_MESSAGES.validate_python( + strip_encrypted_reasoning_blocks_from_anthropic_messages([dict(message) for message in messages]) + ) + + +def _messages_without_rejected_reasoning( + messages: _AnthropicMessages, + error: litellm.BadRequestError, + completion_kwargs: _CompletionKwargs, + sent_encrypted_reasoning: bool, +) -> "_AnthropicMessages | None": + return messages_without_rejected_reasoning( + messages, + error, + sent_encrypted_reasoning=sent_encrypted_reasoning, + provider=_completion_text(completion_kwargs, "custom_llm_provider"), + model=_completion_text(completion_kwargs, "model"), + ) + + +async def _acompletion_without_rejected_reasoning( + completion_kwargs: _CompletionKwargs, + messages: _AnthropicMessages, + prepare_replay: _PrepareReplay, +) -> "ChatResult": + sent_encrypted_reasoning: Final = _sends_encrypted_reasoning(completion_kwargs) + try: + with resent_without_rejected_reasoning(sent_encrypted_reasoning): + return await litellm.acompletion(**completion_kwargs) + except litellm.BadRequestError as error: + replayed_messages: Final = _messages_without_rejected_reasoning( + messages, error, completion_kwargs, sent_encrypted_reasoning + ) + if replayed_messages is None: + raise + return await litellm.acompletion(**prepare_replay(replayed_messages)) + + +def _completion_without_rejected_reasoning( + completion_kwargs: _CompletionKwargs, + messages: _AnthropicMessages, + prepare_replay: _PrepareReplay, +) -> "ChatResult | Coroutine[object, object, ChatResult]": + sent_encrypted_reasoning: Final = _sends_encrypted_reasoning(completion_kwargs) + try: + with resent_without_rejected_reasoning(sent_encrypted_reasoning): + return litellm.completion(**completion_kwargs) + except litellm.BadRequestError as error: + replayed_messages: Final = _messages_without_rejected_reasoning( + messages, error, completion_kwargs, sent_encrypted_reasoning + ) + if replayed_messages is None: + raise + return litellm.completion(**prepare_replay(replayed_messages)) + + +def _is_compaction_block(block: object) -> bool: + mapping: Final = as_str_mapping(block) + return mapping is not None and mapping.get("type") == "compaction" + + def _messages_have_compaction_block(messages: _AnthropicMessages) -> bool: """Return True when any message carries a ``compaction`` content block.""" - for msg in messages: - content = msg.get("content") - if not isinstance(content, list): - continue - for block in content: - if isinstance(block, dict) and block.get("type") == "compaction": - return True - return False + return any( + any(map(_is_compaction_block, content)) for msg in messages if isinstance(content := msg.get("content"), list) + ) def _proxy_router_fallback() -> "Router | None": @@ -413,34 +559,14 @@ class LiteLLMMessagesToCompletionTransformationHandler: If the user provides a `summary` field in the thinking dict, it is passed through to the OpenAI reasoning params (opt-in per OpenAI spec). """ - custom_llm_provider = completion_kwargs.get("custom_llm_provider") - if custom_llm_provider is None: - try: - _, inferred_provider, _, _ = litellm.utils.get_llm_provider( - model=cast(str, completion_kwargs.get("model")) - ) - custom_llm_provider = inferred_provider - except Exception: - custom_llm_provider = None - - if custom_llm_provider != "openai": - return - - if not isinstance(thinking, dict) or thinking.get("type") != "enabled": - return - model: Final = completion_kwargs.get("model") - try: - model_info: Final = get_model_info(model=cast(str, model), custom_llm_provider=custom_llm_provider) - if model_info and model_info.get("supports_reasoning") is False: - # Model doesn't support reasoning/responses API, don't route - return - except Exception: - pass + if not isinstance(model, str) or thinking is None: + return + if not openai_thinking_routes_to_responses(model, completion_kwargs.get("custom_llm_provider"), thinking): + return - if isinstance(model, str) and model and "responses/" not in model: - local_model: Final = model.removeprefix(f"{custom_llm_provider}/") - completion_kwargs["model"] = f"{custom_llm_provider}/responses/{local_model}" + if "responses/" not in model: + completion_kwargs["model"] = f"openai/responses/{model.removeprefix('openai/')}" auto_summary: Final = is_reasoning_auto_summary_enabled() @@ -646,12 +772,15 @@ class LiteLLMMessagesToCompletionTransformationHandler: if "output_config" in extra_kwargs: request_data["output_config"] = extra_kwargs["output_config"] + api_base: Final = provider_kwargs.get("api_base") ( openai_request, tool_name_mapping, ) = ANTHROPIC_ADAPTER.translate_completion_input_params_with_tool_mapping( request_data, custom_llm_provider=target_provider, + api_base=api_base if isinstance(api_base, str) else None, + reasoning_summary=peek_reasoning_summary_aliases(provider_kwargs), ) if openai_request is None: @@ -779,12 +908,9 @@ class LiteLLMMessagesToCompletionTransformationHandler: parent_trace_id=sub_call_trace_id, ) - ( - completion_kwargs, - tool_name_mapping, - ) = LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs( + prepare: Final = partial( + LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs, max_tokens=max_tokens, - messages=effective_messages, model=model, metadata=metadata, stop_sequences=stop_sequences, @@ -799,8 +925,10 @@ class LiteLLMMessagesToCompletionTransformationHandler: output_format=output_format, extra_kwargs=kwargs, ) - - completion_response: Final = await litellm.acompletion(**completion_kwargs) + initial_kwargs, tool_name_mapping = prepare(messages=effective_messages) + completion_response: Final = await _acompletion_without_rejected_reasoning( + initial_kwargs, effective_messages, lambda replayed_messages: prepare(messages=replayed_messages)[0] + ) logging_obj: Final = litellm_logging_obj_from_kwargs(kwargs) if stream: @@ -924,12 +1052,9 @@ class LiteLLMMessagesToCompletionTransformationHandler: effective_messages: Final = polyfill_result.messages if polyfill_result is not None else messages effective_system: Final = polyfill_result.system if polyfill_result is not None else system - ( - completion_kwargs, - tool_name_mapping, - ) = LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs( + prepare: Final = partial( + LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs, max_tokens=max_tokens, - messages=effective_messages, model=model, metadata=metadata, stop_sequences=stop_sequences, @@ -944,8 +1069,10 @@ class LiteLLMMessagesToCompletionTransformationHandler: output_format=output_format, extra_kwargs=kwargs, ) - - completion_response: Final = litellm.completion(**completion_kwargs) + initial_kwargs, tool_name_mapping = prepare(messages=effective_messages) + completion_response: Final = _completion_without_rejected_reasoning( + initial_kwargs, effective_messages, lambda replayed_messages: prepare(messages=replayed_messages)[0] + ) logging_obj: Final = litellm_logging_obj_from_kwargs(kwargs) if stream: diff --git a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py index 18db0f28338..6f094256197 100644 --- a/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py +++ b/litellm/llms/anthropic/pass_through/adapters/streaming_iterator.py @@ -24,6 +24,7 @@ from typing_extensions import assert_never from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.exceptions import ContextWindowExceededError, MidStreamFallbackError +from litellm.litellm_core_utils.core_helpers import as_str_mapping from litellm.litellm_core_utils.hidden_params import set_hidden_params from litellm.llms.anthropic.pass_through.safeguards import StreamedSafeguardResults from litellm.llms.anthropic.pass_through.utils import anthropic_sse_frame @@ -71,6 +72,21 @@ def _optional_attr_sequence(obj: object, name: str) -> Sequence[object]: return value if value else () +def _first_thinking_block(chunk: "ModelResponseStream") -> Mapping[str, object] | None: + delta: Final = _optional_attr(chunk.choices[0], "delta") if chunk.choices else None + blocks: Final = _optional_attr_sequence(delta, "thinking_blocks") + return as_str_mapping(blocks[0]) if blocks else None + + +def _thinking_block_index(block: Mapping[str, object] | None) -> int | None: + index: Final = block.get("index") if block is not None else None + return index if isinstance(index, int) else None + + +def _closes_thinking_block(block: Mapping[str, object] | None) -> bool: + return block is not None and _thinking_block_index(block) is not None and bool(block.get("signature")) + + def _error_status_and_message(exc: Exception) -> tuple[int, str]: if isinstance(exc, (BaseLLMException, MidStreamFallbackError, ContextWindowExceededError)): return exc.status_code, exc.message @@ -458,12 +474,14 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): sent_first_chunk: bool = False sent_content_block_start: bool = False sent_content_block_finish: bool = False - current_content_block_type: Literal["text", "tool_use", "thinking"] = "text" + current_content_block_type: Literal["text", "tool_use", "thinking", "redacted_thinking"] = "text" sent_last_message: bool = False holding_chunk: ContentBlockDelta | None = None holding_stop_reason_chunk: MessageBlockDelta | None = None queued_usage_chunk: bool = False current_content_block_index: int = 0 + current_thinking_block_index: int | None = None + current_thinking_block_closed: bool = False def __init__( self, @@ -1352,11 +1370,31 @@ class AnthropicStreamWrapper(AdapterCompletionStreamWrapper): original_name: Final = self.tool_name_mapping.get(truncated_name, truncated_name) tool_block["name"] = original_name + thinking_block: Final = _first_thinking_block(chunk) + thinking_block_index: Final = _thinking_block_index(thinking_block) + previous_thinking_block_index: Final = self.current_thinking_block_index + follows_closed_thinking_block: Final = self.current_thinking_block_closed and block_type == "thinking" + self.current_thinking_block_closed = _closes_thinking_block(thinking_block) + if ( + thinking_block_index is not None + or block_type != self.current_content_block_type + or follows_closed_thinking_block + ): + self.current_thinking_block_index = thinking_block_index + if block_type != self.current_content_block_type: self.current_content_block_type = block_type self.current_content_block_start = content_block_start return True + if follows_closed_thinking_block or ( + thinking_block_index is not None + and previous_thinking_block_index is not None + and thinking_block_index != previous_thinking_block_index + ): + self.current_content_block_start = content_block_start + return True + # For parallel tool calls, we'll necessarily have a new content block # if we get a function name since it signals a new tool call if block_type == "tool_use": diff --git a/litellm/llms/anthropic/pass_through/adapters/transformation.py b/litellm/llms/anthropic/pass_through/adapters/transformation.py index d1560c5b2ed..17e50498d46 100644 --- a/litellm/llms/anthropic/pass_through/adapters/transformation.py +++ b/litellm/llms/anthropic/pass_through/adapters/transformation.py @@ -2,12 +2,14 @@ import copy import hashlib import json from collections.abc import AsyncIterator, Iterator, Mapping, Sequence +from dataclasses import dataclass from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, TypeAlias, TypeVar, cast from pydantic import JsonValue, TypeAdapter import litellm +from litellm.litellm_core_utils.core_helpers import as_str_mapping, attribute_items from litellm.litellm_core_utils.hidden_params import get_hidden_params from litellm.llms.anthropic.pass_through.utils import ( is_reasoning_auto_summary_enabled, @@ -122,6 +124,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( anthropic_image_source_to_openai_url, parse_tool_call_arguments, reasoning_content_from_thinking_blocks, + thinking_blocks_with_encrypted_reasoning, with_prompt_cache_breakpoint, ) from litellm.litellm_core_utils.prompt_templates.factory import ( @@ -219,6 +222,72 @@ def target_supports_mid_conversation_system(model: str | None, custom_llm_provid return supports_mid_conversation_system(model=model, custom_llm_provider=custom_llm_provider) +@dataclass(frozen=True, slots=True) +class ChatRouteInputs: + tools: Sequence[Mapping[str, object]] | None = None + reasoning_effort: str | Mapping[str, object] | None = None + api_base: str | None = None + thinking: Mapping[str, object] | None = None + reasoning_summary: object | None = None + + +_NO_ROUTE_INPUTS: Final = ChatRouteInputs() + + +def _inferred_provider(model: str) -> str | None: + try: + return litellm.utils.get_llm_provider(model=model)[1] + except Exception: # noqa: BLE001 # an unknown model name has no inferred provider + return None + + +def openai_thinking_routes_to_responses( + model: str | None, + custom_llm_provider: str | None, + thinking: Mapping[str, object] | None, +) -> bool: + if not model or thinking is None or thinking.get("type") != "enabled": + return False + provider: Final = custom_llm_provider or _inferred_provider(model) + if provider != "openai": + return False + try: + model_info: Final = litellm.get_model_info(model=model, custom_llm_provider=provider) + except Exception: # noqa: BLE001 # a model missing from the cost map still routes, as before + return True + return model_info.get("supports_reasoning") is not False + + +def _keeps_encrypted_reasoning_for_responses( + model: str | None, + custom_llm_provider: str | None, + route: ChatRouteInputs, +) -> bool: + from litellm.responses.bridge_check import chat_completion_ends_in_responses_api + + return openai_thinking_routes_to_responses( + model, custom_llm_provider, route.thinking + ) or chat_completion_ends_in_responses_api( + model, + custom_llm_provider, + tools=route.tools, + reasoning_effort=route.reasoning_effort, + reasoning_summary=route.reasoning_summary, + api_base=route.api_base, + ) + + +def _anthropic_thinking_blocks(source: object) -> tuple[Mapping[str, object], ...]: + return thinking_blocks_with_encrypted_reasoning( + getattr(source, "thinking_blocks", None), + getattr(source, "reasoning_items", None), + ) + + +def _streamed_thinking_blocks(delta: object) -> tuple[Mapping[str, object], ...]: + return tuple(block for item in attribute_items(delta, "thinking_blocks") if (block := as_str_mapping(item))) + + def _chat_tool_param(function_chunk: ChatCompletionToolParamFunctionChunk, tool: object) -> ChatCompletionToolParam: eager_input_streaming: Final = eager_input_streaming_flag(tool) if eager_input_streaming is None: @@ -246,7 +315,12 @@ class AnthropicAdapter: return result def translate_completion_input_params_with_tool_mapping( - self, kwargs, *, custom_llm_provider: str | None = None + self, + kwargs, + *, + custom_llm_provider: str | None = None, + api_base: str | None = None, + reasoning_summary: object | None = None, ) -> tuple[ChatCompletionRequest | None, dict[str, str]]: """ Translate Anthropic request params to OpenAI format, returning tool name mapping. @@ -280,6 +354,8 @@ class AnthropicAdapter: ) = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( anthropic_message_request=request_body, custom_llm_provider=custom_llm_provider, + api_base=api_base, + reasoning_summary=reasoning_summary, ) return translated_body, tool_name_mapping @@ -465,9 +541,14 @@ class LiteLLMAnthropicMessagesAdapter: *, custom_llm_provider: str | None = None, preserve_midturn_system: bool = False, + route: ChatRouteInputs = _NO_ROUTE_INPUTS, ) -> list[AllMessageValues]: new_messages: Final[list[AllMessageValues]] = [] - replayable_messages: Final = strip_encrypted_reasoning_blocks_from_anthropic_messages(messages) + replayable_messages: Final = ( + messages + if _keeps_encrypted_reasoning_for_responses(model, custom_llm_provider, route) + else strip_encrypted_reasoning_blocks_from_anthropic_messages(messages) + ) leading_count: Final = next( (i for i, m in enumerate(replayable_messages) if not is_system_role_message(m)), len(replayable_messages), @@ -1221,6 +1302,8 @@ class LiteLLMAnthropicMessagesAdapter: anthropic_message_request: AnthropicMessagesRequest, *, custom_llm_provider: str | None = None, + api_base: str | None = None, + reasoning_summary: object | None = None, preserve_midturn_system: bool = False, ) -> tuple[ChatCompletionRequest, dict[str, str]]: """ @@ -1232,26 +1315,15 @@ class LiteLLMAnthropicMessagesAdapter: for tools that exceeded OpenAI's 64-char limit """ # Debug: Processing Anthropic message request - new_messages: list[AllMessageValues] = [] tool_name_mapping: dict[str, str] = {} - ## CONVERT ANTHROPIC MESSAGES TO OPENAI messages_list: Final[list[AllAnthropicPassThroughMessageValues]] = cast( list[AllAnthropicPassThroughMessageValues], anthropic_message_request["messages"], ) - new_messages = self.translate_anthropic_messages_to_openai( - messages=messages_list, - model=anthropic_message_request.get("model"), - custom_llm_provider=custom_llm_provider, - preserve_midturn_system=preserve_midturn_system, - ) - ## ADD SYSTEM MESSAGE TO MESSAGES - self._add_system_message_to_messages(new_messages, anthropic_message_request) - new_kwargs: Final[ChatCompletionRequest] = { "model": anthropic_message_request["model"], - "messages": new_messages, + "messages": [], } ## CONVERT METADATA (user_id + litellm metadata) self._translate_metadata_to_openai( @@ -1275,6 +1347,23 @@ class LiteLLMAnthropicMessagesAdapter: new_kwargs=new_kwargs, custom_llm_provider=custom_llm_provider, ) + ## CONVERT ANTHROPIC MESSAGES TO OPENAI, after tools and thinking pick the route + new_messages: Final = self.translate_anthropic_messages_to_openai( + messages=messages_list, + model=anthropic_message_request.get("model"), + custom_llm_provider=custom_llm_provider, + preserve_midturn_system=preserve_midturn_system, + route=ChatRouteInputs( + tools=new_kwargs.get("tools"), + reasoning_effort=new_kwargs.get("reasoning_effort"), + api_base=api_base, + thinking=anthropic_message_request.get("thinking"), + reasoning_summary=reasoning_summary, + ), + ) + ## ADD SYSTEM MESSAGE TO MESSAGES + self._add_system_message_to_messages(new_messages, anthropic_message_request) + new_kwargs["messages"] = new_messages ## CONVERT STOP_SEQUENCES self._translate_stop_sequences_to_openai( anthropic_message_request=anthropic_message_request, @@ -1358,9 +1447,9 @@ class LiteLLMAnthropicMessagesAdapter: for choice in choices ): new_content.extend(dict(_COMPACTION_BLOCK.validate_python(block)) for block in compaction_blocks) - # Handle thinking blocks first - if hasattr(choice.message, "thinking_blocks") and choice.message.thinking_blocks: - for thinking_block in choice.message.thinking_blocks: + signed_thinking_blocks = _anthropic_thinking_blocks(choice.message) + if signed_thinking_blocks: + for thinking_block in signed_thinking_blocks: if thinking_block.get("type") == "thinking": if is_empty_unsigned_thinking_block(thinking_block): continue @@ -1627,7 +1716,7 @@ class LiteLLMAnthropicMessagesAdapter: def translate_streaming_openai_chunk_to_anthropic_content_block( self, choices: list[OpenAIStreamingChoice | StreamingChoices] ) -> tuple[ - Literal["text", "tool_use", "thinking"], + Literal["text", "tool_use", "thinking", "redacted_thinking"], "ContentBlockContentBlockDict", ]: from litellm._uuid import uuid @@ -1660,20 +1749,24 @@ class LiteLLMAnthropicMessagesAdapter: choice.delta ) is not None: return "text", TextBlock(type="text", text="") - elif isinstance(choice, StreamingChoices) and hasattr(choice.delta, "thinking_blocks"): - thinking_blocks = choice.delta.thinking_blocks or [] - if len(thinking_blocks) > 0: - thinking_block = thinking_blocks[0] - if thinking_block["type"] == "thinking": - thinking = thinking_block.get("thinking") or "" - signature = thinking_block.get("signature") or "" + elif isinstance(choice, StreamingChoices) and (thinking_blocks := _streamed_thinking_blocks(choice.delta)): + thinking_block = thinking_blocks[0] + if thinking_block["type"] == "redacted_thinking": + redacted_data = thinking_block.get("data") + return "redacted_thinking", ChatCompletionRedactedThinkingBlock( + type="redacted_thinking", + data=redacted_data if isinstance(redacted_data, str) else "", + ) + if thinking_block["type"] == "thinking": + thinking = thinking_block.get("thinking") or "" + signature = thinking_block.get("signature") or "" - assert isinstance(thinking, str) - assert isinstance(signature, str) + assert isinstance(thinking, str) + assert isinstance(signature, str) - return "thinking", ChatCompletionThinkingBlock( - type="thinking", thinking=thinking, signature=signature - ) + return "thinking", ChatCompletionThinkingBlock( + type="thinking", thinking=thinking, signature=signature + ) # OpenAI-compatible reasoning backends (e.g. vLLM/SGLang reasoning # parsers) populate ``reasoning_content`` without ``thinking_blocks``. # ``Delta`` deletes the ``thinking_blocks`` attribute when unset, so the @@ -1706,19 +1799,17 @@ class LiteLLMAnthropicMessagesAdapter: for tool in choice.delta.tool_calls: if tool.function is not None and tool.function.arguments is not None: partial_json = (partial_json or "") + tool.function.arguments - elif isinstance(choice, StreamingChoices) and hasattr(choice.delta, "thinking_blocks"): - thinking_blocks = choice.delta.thinking_blocks or [] - if len(thinking_blocks) > 0: - for thinking_block in thinking_blocks: - if thinking_block["type"] == "thinking": - thinking = thinking_block.get("thinking") or "" - signature = thinking_block.get("signature") or "" + elif isinstance(choice, StreamingChoices) and (thinking_blocks := _streamed_thinking_blocks(choice.delta)): + for thinking_block in thinking_blocks: + if thinking_block["type"] == "thinking": + thinking = thinking_block.get("thinking") or "" + signature = thinking_block.get("signature") or "" - assert isinstance(thinking, str) - assert isinstance(signature, str) + assert isinstance(thinking, str) + assert isinstance(signature, str) - reasoning_content += thinking - reasoning_signature += signature + reasoning_content += thinking + reasoning_signature += signature # Handle reasoning_content when thinking_blocks is not present # This handles providers like OpenRouter that return reasoning_content elif isinstance(choice, StreamingChoices) and hasattr(choice.delta, "reasoning_content"): diff --git a/litellm/llms/anthropic/pass_through/messages/handler.py b/litellm/llms/anthropic/pass_through/messages/handler.py index eeff6302ee8..949d81f3c1b 100644 --- a/litellm/llms/anthropic/pass_through/messages/handler.py +++ b/litellm/llms/anthropic/pass_through/messages/handler.py @@ -27,6 +27,7 @@ from litellm.llms.base_llm.anthropic_messages.transformation import ( from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler from litellm.llms.custom_httpx.llm_http_handler import BaseLLMHTTPHandler +from litellm.responses.mcp.tool_precheck import may_name_litellm_mcp_gateway from litellm.types.llms.anthropic import AllAnthropicToolsValues from litellm.types.llms.anthropic_messages.anthropic_request import AnthropicMetadata from litellm.types.llms.anthropic_messages.anthropic_response import ( @@ -569,38 +570,34 @@ def anthropic_messages_handler( # downstream path (native passthrough and both bridges) gets real tools rather than a # reference the provider cannot resolve. Popped from kwargs so it never reaches the provider. skip_mcp_handler: Final = kwargs.pop("_skip_mcp_handler", False) - if not skip_mcp_handler and tools: + if not skip_mcp_handler and tools and may_name_litellm_mcp_gateway(tools): # pyright: ignore[reportUnknownArgumentType] # untyped tools list from litellm.llms.anthropic.pass_through.messages.mcp_handler import ( anthropic_messages_with_mcp, ) - from litellm.responses.mcp.litellm_proxy_mcp_handler import ( - LiteLLM_Proxy_MCP_Handler, - ) - if LiteLLM_Proxy_MCP_Handler.should_use_litellm_mcp_gateway(tools=tools): - return anthropic_messages_with_mcp( - max_tokens=max_tokens, - messages=messages, - model=model, - metadata=metadata, - stop_sequences=stop_sequences, - stream=stream, - system=cast( # cast-ok: anthropic accepts list system blocks; handler signature is narrower - "str | None", system - ), - temperature=temperature, - thinking=thinking, - tool_choice=tool_choice, - tools=tools, - top_k=top_k, - top_p=top_p, - container=container, - api_key=api_key, - api_base=api_base, - client=client, - custom_llm_provider=custom_llm_provider, - **kwargs, - ) + return anthropic_messages_with_mcp( + max_tokens=max_tokens, + messages=messages, + model=model, + metadata=metadata, + stop_sequences=stop_sequences, + stream=stream, + system=cast( # cast-ok: anthropic accepts list system blocks; handler signature is narrower + "str | None", system + ), + temperature=temperature, + thinking=thinking, + tool_choice=tool_choice, + tools=tools, + top_k=top_k, + top_p=top_p, + container=container, + api_key=api_key, + api_base=api_base, + client=client, + custom_llm_provider=custom_llm_provider, + **kwargs, + ) anthropic_messages_provider_config: Final = _native_messages_config( model, custom_llm_provider, kwargs.get("model_info") diff --git a/litellm/llms/anthropic/pass_through/responses_adapters/handler.py b/litellm/llms/anthropic/pass_through/responses_adapters/handler.py index 3eeaee8ad43..abe0f2f419d 100644 --- a/litellm/llms/anthropic/pass_through/responses_adapters/handler.py +++ b/litellm/llms/anthropic/pass_through/responses_adapters/handler.py @@ -14,9 +14,10 @@ from litellm.integrations.anthropic_cache_control_hook import ( AnthropicCacheControlHook, supports_openai_prompt_cache_breakpoint, ) -from litellm.litellm_core_utils.core_helpers import normalize_drop_params +from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items, normalize_drop_params from litellm.litellm_core_utils.prompt_templates.common_utils import cache_control_as_prompt_cache_breakpoint from litellm.litellm_core_utils.tool_fields import drops_mcp_servers, tools_without_mcp_servers +from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator from litellm.types.integrations.custom_logger import CONVERTED_STREAM_KEYS from litellm.types.llms.anthropic import ( AllAnthropicMessageValues, @@ -31,7 +32,11 @@ from litellm.types.llms.anthropic_messages.anthropic_response import ( from litellm.types.llms.openai import ResponsesAPIResponse from litellm.utils import ProviderConfigManager -from ..adapters.handler import request_safeguards_evaluator +from ..adapters.handler import ( + messages_without_rejected_reasoning, + request_safeguards_evaluator, + resent_without_rejected_reasoning, +) from ..safeguards import with_safeguard_results from ..utils import litellm_logging_obj_from_kwargs, local_model_name from .streaming_iterator import AnthropicResponsesStreamWrapper @@ -254,6 +259,78 @@ def _build_responses_kwargs( return responses_kwargs +_ResponsesKwargsFor: TypeAlias = Callable[[AnthropicRequestMessages], Mapping[str, object]] +_ResponsesResult: TypeAlias = ResponsesAPIResponse | BaseResponsesAPIStreamingIterator +_SyncResponsesResult: TypeAlias = _ResponsesResult | Coroutine[object, object, _ResponsesResult] + + +def _is_encrypted_reasoning_item(item: object) -> bool: + mapping: Final = as_str_mapping(item) + return mapping is not None and mapping.get("type") == "reasoning" and bool(mapping.get("encrypted_content")) + + +def _sends_encrypted_reasoning(responses_kwargs: Mapping[str, object]) -> bool: + input_items: Final = responses_kwargs.get("input") + sent: Final = iter_items(input_items) if input_items is not None else () + return any(map(_is_encrypted_reasoning_item, sent)) + + +def _kwarg_text(responses_kwargs: Mapping[str, object], key: str) -> str: + value: Final = responses_kwargs.get(key) + return value if isinstance(value, str) else "" + + +def _messages_without_rejected_reasoning( + messages: AnthropicRequestMessages, + error: litellm.BadRequestError, + responses_kwargs: Mapping[str, object], + sent_encrypted_reasoning: bool, +) -> AnthropicRequestMessages | None: + return messages_without_rejected_reasoning( + messages, + error, + sent_encrypted_reasoning=sent_encrypted_reasoning, + provider=_kwarg_text(responses_kwargs, "custom_llm_provider"), + model=_kwarg_text(responses_kwargs, "model"), + ) + + +async def _aresponses_without_rejected_reasoning( + messages: AnthropicRequestMessages, responses_kwargs_for: _ResponsesKwargsFor +) -> tuple[Mapping[str, object], _ResponsesResult]: + responses_kwargs: Final = responses_kwargs_for(messages) + sent_encrypted_reasoning: Final = _sends_encrypted_reasoning(responses_kwargs) + try: + with resent_without_rejected_reasoning(sent_encrypted_reasoning): + return responses_kwargs, await litellm.aresponses(**responses_kwargs) + except litellm.BadRequestError as error: + replayed_messages: Final = _messages_without_rejected_reasoning( + messages, error, responses_kwargs, sent_encrypted_reasoning + ) + if replayed_messages is None: + raise + replay_kwargs: Final = responses_kwargs_for(replayed_messages) + return replay_kwargs, await litellm.aresponses(**replay_kwargs) + + +def _responses_without_rejected_reasoning( + messages: AnthropicRequestMessages, responses_kwargs_for: _ResponsesKwargsFor +) -> tuple[Mapping[str, object], _SyncResponsesResult]: + responses_kwargs: Final = responses_kwargs_for(messages) + sent_encrypted_reasoning: Final = _sends_encrypted_reasoning(responses_kwargs) + try: + with resent_without_rejected_reasoning(sent_encrypted_reasoning): + return responses_kwargs, litellm.responses(**responses_kwargs) + except litellm.BadRequestError as error: + replayed_messages: Final = _messages_without_rejected_reasoning( + messages, error, responses_kwargs, sent_encrypted_reasoning + ) + if replayed_messages is None: + raise + replay_kwargs: Final = responses_kwargs_for(replayed_messages) + return replay_kwargs, litellm.responses(**replay_kwargs) + + class LiteLLMMessagesToResponsesAPIHandler: """ Handles Anthropic /v1/messages requests for OpenAI / Azure models by @@ -281,28 +358,30 @@ class LiteLLMMessagesToResponsesAPIHandler: output_format: AnthropicOutputSchema | None = None, **kwargs: object, ) -> AnthropicMessagesResponse | AsyncIterator[bytes]: - responses_kwargs: Final = _build_responses_kwargs( - max_tokens=max_tokens, - messages=messages, - model=model, - context_management=context_management, - metadata=metadata, - output_config=output_config, - stop_sequences=stop_sequences, - stream=stream, - system=system, - temperature=temperature, - thinking=thinking, - tool_choice=tool_choice, - tools=tools, - top_k=top_k, - top_p=top_p, - output_format=output_format, - extra_kwargs=kwargs, - ) + + def responses_kwargs_for(sent_messages: AnthropicRequestMessages) -> Mapping[str, object]: + return _build_responses_kwargs( + max_tokens=max_tokens, + messages=sent_messages, + model=model, + context_management=context_management, + metadata=metadata, + output_config=output_config, + stop_sequences=stop_sequences, + stream=stream, + system=system, + temperature=temperature, + thinking=thinking, + tool_choice=tool_choice, + tools=tools, + top_k=top_k, + top_p=top_p, + output_format=output_format, + extra_kwargs=kwargs, + ) safeguards_evaluator: Final = request_safeguards_evaluator(kwargs, messages) - result: Final = await litellm.aresponses(**responses_kwargs) + responses_kwargs, result = await _aresponses_without_rejected_reasoning(messages, responses_kwargs_for) if stream: wrapper: Final = AnthropicResponsesStreamWrapper( @@ -366,28 +445,28 @@ class LiteLLMMessagesToResponsesAPIHandler: **kwargs, ) - # Sync path - responses_kwargs: Final = _build_responses_kwargs( - max_tokens=max_tokens, - messages=messages, - model=model, - context_management=context_management, - metadata=metadata, - output_config=output_config, - stop_sequences=stop_sequences, - stream=stream, - system=system, - temperature=temperature, - thinking=thinking, - tool_choice=tool_choice, - tools=tools, - top_k=top_k, - top_p=top_p, - output_format=output_format, - extra_kwargs=kwargs, - ) + def responses_kwargs_for(sent_messages: AnthropicRequestMessages) -> Mapping[str, object]: + return _build_responses_kwargs( + max_tokens=max_tokens, + messages=sent_messages, + model=model, + context_management=context_management, + metadata=metadata, + output_config=output_config, + stop_sequences=stop_sequences, + stream=stream, + system=system, + temperature=temperature, + thinking=thinking, + tool_choice=tool_choice, + tools=tools, + top_k=top_k, + top_p=top_p, + output_format=output_format, + extra_kwargs=kwargs, + ) - result: Final = litellm.responses(**responses_kwargs) + responses_kwargs, result = _responses_without_rejected_reasoning(messages, responses_kwargs_for) if stream: wrapper: Final = AnthropicResponsesStreamWrapper( diff --git a/litellm/llms/bedrock/common_utils.py b/litellm/llms/bedrock/common_utils.py index 30a7ddb2224..d872b4f11c7 100644 --- a/litellm/llms/bedrock/common_utils.py +++ b/litellm/llms/bedrock/common_utils.py @@ -1336,15 +1336,29 @@ def bedrock_converse_accepts_redacted_thinking(model: str, base_model: str | Non its own. An ARN that hides its model, like an application inference profile, falls back to the deployment's ``base_model`` and counts as non-Claude without one. """ + return _is_claude(_converse_target_model(model, base_model)) + + +def bedrock_converse_rejects_signed_thinking(model: str, base_model: str | None) -> bool: + if base_model is None and "." not in get_bedrock_base_model(model): + return False + return not _is_claude_behind_any_region_prefix(_converse_target_model(model, base_model)) + + +def _converse_target_model(model: str, base_model: str | None) -> str: if base_model is None or "." in get_bedrock_base_model(model): - return _is_claude(model) - return _is_claude(base_model) + return model + return base_model def _is_claude(model: str) -> bool: return get_bedrock_base_model(model).startswith(("anthropic", "claude")) +def _is_claude_behind_any_region_prefix(model: str) -> bool: + return _is_claude(model) or get_bedrock_base_model(model).partition(".")[2].startswith("anthropic.") + + def bedrock_converse_supports_parallel_tool_use_config(model: str) -> bool: return any( (litellm.model_cost.get(candidate) or {}).get("supports_parallel_tool_use_config") is True diff --git a/litellm/llms/bedrock/passthrough/transformation.py b/litellm/llms/bedrock/passthrough/transformation.py index d7d6341fe7a..453e730e7fc 100644 --- a/litellm/llms/bedrock/passthrough/transformation.py +++ b/litellm/llms/bedrock/passthrough/transformation.py @@ -13,6 +13,10 @@ from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.passthrough.transformation import BasePassthroughConfig, PassthroughStreamCollector from litellm.llms.bedrock.chat.converse_to_openai import as_json_array, as_json_object +from litellm.llms.bedrock.chat.invoke_transformations.base_invoke_transformation import ( + bedrock_invoke_guardrail_headers, +) +from litellm.llms.bedrock.request_metadata import merge_bedrock_invoke_headers from litellm.types.llms.openai import ChatCompletionSystemMessage, ChatCompletionUserMessage from litellm.types.utils import ModelResponseStream @@ -38,8 +42,35 @@ def is_bedrock_streaming_endpoint(endpoint: str) -> bool: _CONVERSE_ACTIONS: Final = frozenset({"converse", "converse-stream"}) -def _is_converse_url(url: str) -> bool: - return httpx.URL(url).path.rstrip("/").rsplit("/", 1)[-1] in _CONVERSE_ACTIONS +_INVOKE_ACTIONS: Final = frozenset({"invoke", "invoke-with-response-stream"}) +_BEDROCK_INVOKE_GUARDRAIL_HEADERS: Final = frozenset( + {"x-amzn-bedrock-guardrailidentifier", "x-amzn-bedrock-guardrailversion", "x-amzn-bedrock-trace"} +) + + +def _converse_guardrail_config(action: str, guardrail_config: object) -> object: + fields: Final = as_json_object(guardrail_config) + if action == "converse-stream" or fields is None: + return guardrail_config + return {key: value for key, value in fields.items() if key != "streamProcessingMode"} + + +def _with_deployment_guardrail( + action: str, + headers: Mapping[str, str], + body: Mapping[str, object], + guardrail_config: object, +) -> tuple[Mapping[str, str], Mapping[str, object]]: + if guardrail_config is None: + return headers, body + if action in _CONVERSE_ACTIONS: + return headers, {**body, "guardrailConfig": _converse_guardrail_config(action, guardrail_config)} + if action in _INVOKE_ACTIONS: + guardrail_headers: Final = tuple(bedrock_invoke_guardrail_headers(guardrail_config).items()) + return merge_bedrock_invoke_headers( + dict(headers), (), guardrail_headers, _BEDROCK_INVOKE_GUARDRAIL_HEADERS + ), body + return headers, body def _is_cache_point(entry: object) -> bool: @@ -363,13 +394,20 @@ class BedrockPassthroughConfig(BaseAWSLLM, BedrockModelInfo, BedrockEventStreamD api_base: str, model: str | None = None, ) -> tuple[dict, bytes | None]: + action: Final = httpx.URL(api_base).path.rstrip("/").rsplit("/", 1)[-1] body: Final[Mapping[str, object]] = request_data or {} optional_params: Final = litellm_params.copy() + guarded_headers, guarded_body = _with_deployment_guardrail( + action, + headers, + _converse_body_to_send(body) if action in _CONVERSE_ACTIONS else body, + litellm_params.get("guardrailConfig"), + ) return self._sign_request( service_name="bedrock", - headers=headers, + headers=dict(guarded_headers), optional_params=optional_params, - request_data=dict(_converse_body_to_send(body) if _is_converse_url(api_base) else body), + request_data=dict(guarded_body), api_base=api_base, model=model, api_key=self.get_bedrock_bearer_token(optional_params), diff --git a/litellm/llms/bedrock/request_metadata.py b/litellm/llms/bedrock/request_metadata.py index 1f4e5886508..685b52e3b1c 100644 --- a/litellm/llms/bedrock/request_metadata.py +++ b/litellm/llms/bedrock/request_metadata.py @@ -179,12 +179,11 @@ def merge_bedrock_invoke_headers( """ Merge the ``X-Amzn-*`` headers the Invoke paths derive from params. - ``caller_owned`` (the guardrail headers) defers to a header the caller already set, which is - the long-standing behaviour for those. ``proxy_owned_names`` are dropped from the caller's - headers unconditionally and re-supplied only from ``proxy_owned``, because those names carry - proxy-authenticated identity into an AWS billing record that the caller must not be able to - write. Names are compared case-insensitively so a caller cannot leave a second spelling in - the dict and let the transport pick the winner. + ``caller_owned`` defers to a header the caller already set. ``proxy_owned_names`` are dropped + from the caller's headers unconditionally and re-supplied only from ``proxy_owned``, for names + the caller must not be able to write, such as proxy-authenticated identity headed for an AWS + billing record or a deployment's guardrail. Names are compared case-insensitively so a caller + cannot leave a second spelling in the dict and let the transport pick the winner. """ if not caller_owned and not proxy_owned and not proxy_owned_names: return headers diff --git a/litellm/llms/dashscope/videos/transformation.py b/litellm/llms/dashscope/videos/transformation.py index 17b685a0e1a..625ce3bab7e 100644 --- a/litellm/llms/dashscope/videos/transformation.py +++ b/litellm/llms/dashscope/videos/transformation.py @@ -35,6 +35,7 @@ from litellm.types.videos.utils import ( decode_video_id_with_provider, encode_video_id_with_provider, extract_original_video_id, + id_safe_name, ) from litellm.utils import get_model_info @@ -599,17 +600,13 @@ def _stamped_poll_models(metadata: object) -> tuple[str, str | None] | None: fields: Final = _PARAMETERS_ADAPTER.validate_python(metadata) except ValidationError: return None - requested: Final = _id_safe_name(fields.get("model_group")) - served: Final = _id_safe_name(fields.get("deployment_model_name")) + requested: Final = id_safe_name(fields.get("model_group")) + served: Final = id_safe_name(fields.get("deployment_model_name")) if requested is None: return None return requested, None if served is None or served == requested or "*" in served else served -def _id_safe_name(name: object) -> str | None: - return name if isinstance(name, str) and name and ";" not in name else None - - class DashScopeVideoConfig(BaseVideoConfig): def get_supported_openai_params(self, model: str) -> list[str]: # mutable-ok: BaseVideoConfig contract returns list return [ diff --git a/litellm/llms/openai/responses/transformation.py b/litellm/llms/openai/responses/transformation.py index 464dc799c16..811b5df9933 100644 --- a/litellm/llms/openai/responses/transformation.py +++ b/litellm/llms/openai/responses/transformation.py @@ -63,7 +63,9 @@ _NO_TOOL_UPDATE: Final[Mapping[str, object]] = MappingProxyType({}) _MODEL_FAMILIES_REJECTING_TOP_LEVEL_SCHEMA_COMBINATORS: Final = ("gpt-4", "gpt-3.5", "chatgpt-4o", "o1", "o3", "o4") _PROVIDERS_WITH_OPENAI_SCHEMA_VALIDATOR: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI}) _PROVIDERS_VALIDATING_TOOL_CALL_ITEM_IDS: Final = frozenset({LlmProviders.AZURE, LlmProviders.OPENAI}) -_PROVIDERS_REPLAYING_ONLY_THEIR_OWN_REASONING: Final = _PROVIDERS_VALIDATING_TOOL_CALL_ITEM_IDS +_PROVIDERS_REPLAYING_ONLY_THEIR_OWN_REASONING: Final = frozenset( + {LlmProviders.AZURE, LlmProviders.BEDROCK, LlmProviders.BEDROCK_MANTLE, LlmProviders.OPENAI} +) _ROLES_READING_PROMPT_CACHE_BREAKPOINTS: Final = frozenset({"system", "developer", "user"}) _PART_TYPES_READING_PROMPT_CACHE_BREAKPOINTS: Final = frozenset({"input_text", "input_image", "input_file"}) _INPUT_ITEM: Final[TypeAdapter[Mapping[str, object]]] = TypeAdapter(Mapping[str, object]) diff --git a/litellm/llms/vertex_ai/batches/handler.py b/litellm/llms/vertex_ai/batches/handler.py index de5fe39cfc9..a039347e824 100644 --- a/litellm/llms/vertex_ai/batches/handler.py +++ b/litellm/llms/vertex_ai/batches/handler.py @@ -14,8 +14,10 @@ from litellm.litellm_core_utils.url_utils import ( ) from litellm.llms.custom_httpx.http_handler import ( HTTPHandler, + MaskedHTTPStatusError, get_async_httpx_client, get_httpx_client, + mask_sensitive_info, ) from litellm.llms.vertex_ai.common_utils import VertexAIError, get_vertex_base_url from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import VertexLLM @@ -80,6 +82,15 @@ def _vertex_batch_list_payload(response: _VertexBatchListJsonSource) -> dict[str return response.json() +def _vertex_batch_create_error(error: httpx.HTTPStatusError) -> MaskedHTTPStatusError: + body: Final = mask_sensitive_info(error.response.text) + return MaskedHTTPStatusError( + httpx.HTTPStatusError(f"{error}\n{body}", request=error.request, response=error.response), + message=body, + text=body, + ) + + class VertexAIBatchPrediction(VertexLLM): def __init__(self, gcs_bucket_name: str, *args, **kwargs): super().__init__(*args, **kwargs) @@ -165,11 +176,19 @@ class VertexAIBatchPrediction(VertexLLM): headers=headers, ) - response: Final = sync_handler.post( - url=api_base, - headers=headers, - data=json.dumps(vertex_batch_request), - ) + try: + response: Final = sync_handler.post( + url=api_base, + headers=headers, + data=json.dumps(vertex_batch_request), + ) + except httpx.HTTPStatusError as e: + litellm.verbose_logger.error( + "Vertex AI batch create failed: status=%s, body=%s", + e.response.status_code, + e.response.text[:1000], + ) + raise _vertex_batch_create_error(e) from e payload_view: Final[_VertexBatchPayloadView] = {"payload": response.json()} _json_response: Final = payload_view["payload"] @@ -265,13 +284,12 @@ class VertexAIBatchPrediction(VertexLLM): data=json.dumps(vertex_batch_request), ) except httpx.HTTPStatusError as e: - error_body: Final = e.response.text litellm.verbose_logger.error( "Vertex AI batch create failed: status=%s, body=%s", e.response.status_code, - error_body[:1000], + e.response.text[:1000], ) - raise + raise _vertex_batch_create_error(e) from e payload_view: Final[_VertexBatchPayloadView] = {"payload": response.json()} _json_response: Final = payload_view["payload"] diff --git a/litellm/llms/vertex_ai/batches/transformation.py b/litellm/llms/vertex_ai/batches/transformation.py index a8d02045740..2460278294c 100644 --- a/litellm/llms/vertex_ai/batches/transformation.py +++ b/litellm/llms/vertex_ai/batches/transformation.py @@ -9,14 +9,41 @@ from litellm._logging import verbose_logger from litellm._uuid import uuid from litellm.llms.vertex_ai.common_utils import ( VertexAIError, + VertexAIModelRoute, convert_vertex_datetime_to_openai_datetime, + get_vertex_ai_model_route, ) +from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import VertexAIAnthropicConfig from litellm.types.llms.openai import BatchJobStatus, CreateBatchRequest from litellm.types.llms.vertex_ai import * -from litellm.types.llms.vertex_ai import GenerateContentResponseBody +from litellm.types.llms.vertex_ai import GenerateContentResponseBody, VertexPartnerProvider from litellm.types.utils import LiteLLMBatch, ModelInfo, Usage _NATIVE_VERTEX_RESPONSE: Final = TypeAdapter(GenerateContentResponseBody) +_JSON_OBJECT: Final = TypeAdapter(Mapping[str, object]) +VERTEX_ANTHROPIC_VERSION: Final = "vertex-2023-10-16" +_ANTHROPIC_PUBLISHER_MODEL_PATH: Final = "publishers/anthropic/models/" + + +def is_vertex_anthropic_batch_model(model: str) -> bool: + """ + Whether a Vertex batch model (bare, `vertex_ai/`-prefixed or a publisher path) is a Claude model, decided the + way Vertex chat routes Claude. Vertex batches Claude only as `publishers/anthropic/models/` with rows in + Anthropic's Messages shape + """ + bare_model: Final = model.rsplit("/", 1)[-1] + return _ANTHROPIC_PUBLISHER_MODEL_PATH in model or ( + get_vertex_ai_model_route(bare_model) is VertexAIModelRoute.PARTNER_MODELS + and VertexPartnerProvider.claude.value in bare_model + ) + + +def as_json_object(value: object) -> Mapping[str, object] | None: + return _JSON_OBJECT.validate_python(value) if isinstance(value, dict) else None + + +def is_vertex_anthropic_batch_output_response(response_body: Mapping[str, object]) -> bool: + return response_body.get("type") == "message" and isinstance(response_body.get("usage"), dict) def _int_field(mapping: Mapping[str, object], key: str) -> int: @@ -43,8 +70,22 @@ def is_vertex_embedding_batch_output_response(response_body: Mapping[str, object return isinstance(response_body.get("embedding"), dict) +def is_vertex_anthropic_batch_output_row(vertex_output_row: Mapping[str, object]) -> bool: + """ + Whether a Vertex batch output row came from a Claude (Anthropic Messages) request. Vertex echoes only the + request's top-level `custom_id` on these rows, with an Anthropic `message` response on success and an empty + response plus a `status` on failure, so a Claude input file read back is never mistaken for output + """ + if not ("custom_id" in vertex_output_row and "processed_time" in vertex_output_row): + return False + response: Final = as_json_object(vertex_output_row.get("response")) + return (response is not None and is_vertex_anthropic_batch_output_response(response)) or bool( + vertex_output_row.get("status") + ) + + def is_native_vertex_batch_output_row(row: Mapping[str, object]) -> bool: - return isinstance(row.get("request"), dict) + return isinstance(row.get("request"), dict) or is_vertex_anthropic_batch_output_row(row) class NativeVertexBatchCostCalculator(Protocol): @@ -72,6 +113,9 @@ def _native_vertex_row_usage( response_body: Mapping[str, object], calculate_usage: Callable[[GenerateContentResponseBody], Usage], ) -> Usage | None: + anthropic_usage: Final = as_json_object(response_body.get("usage")) + if is_vertex_anthropic_batch_output_response(response_body) and anthropic_usage is not None: + return VertexAIAnthropicConfig().calculate_usage(usage_object=anthropic_usage, reasoning_content=None) if "usageMetadata" not in response_body: if not is_vertex_embedding_batch_output_response(response_body): return None @@ -97,21 +141,24 @@ def native_vertex_batch_row_stats( """ Usage and cost of one native Vertex predictions.jsonl row, a `{"request": ..., "response": {"candidates": [...], "usageMetadata": {...}, "modelVersion": ...}}` - generateContent object or a `{"request": ..., "response": {"embedding": {...}, "usageMetadata": {...}}}` - embedding object (an embedding row without `usageMetadata` is billed from its documented `tokenCount`). - `model_name` (the deployment model) prices the row unless it is a wildcard, else its own `modelVersion` - does, else the wildcard name so explicit deployment prices still apply; a row without a response, a - generateContent row without `response.usageMetadata`, and a row whose response fails validation are + generateContent object, a `{"request": ..., "response": {"embedding": {...}, "usageMetadata": {...}}}` + embedding object (an embedding row without `usageMetadata` is billed from its documented `tokenCount`) or a + `{"custom_id": ..., "request": ..., "response": {"type": "message", "model": ..., "usage": {...}}}` Claude object. + `model_name` (the deployment model) prices the row unless it is a wildcard, else its own `modelVersion` (Claude + rows: `model`) does, else the wildcard name so explicit deployment prices still apply; a row without a response, + a generateContent row without `response.usageMetadata`, and a row whose response fails validation are None (failed). """ - response_body: Final = row.get("response") - if not isinstance(response_body, dict): + response_body: Final = as_json_object(row.get("response")) + if response_body is None: return None usage: Final = _native_vertex_row_usage(response_body, calculate_usage) if usage is None: return None total_tokens: Final = usage.total_tokens or (usage.prompt_tokens + usage.completion_tokens) - model_version: Final = response_body.get("modelVersion") + model_version: Final = response_body.get( + "model" if is_vertex_anthropic_batch_output_response(response_body) else "modelVersion" + ) deployment_model: Final = model_name if model_name and "*" not in model_name else None model: Final = deployment_model or (model_version if isinstance(model_version, str) else model_name) if model is None: diff --git a/litellm/llms/vertex_ai/files/transformation.py b/litellm/llms/vertex_ai/files/transformation.py index 396acd270fa..b64d81e4630 100644 --- a/litellm/llms/vertex_ai/files/transformation.py +++ b/litellm/llms/vertex_ai/files/transformation.py @@ -19,6 +19,7 @@ from typing_extensions import ReadOnly, Required import litellm from litellm._uuid import uuid +from litellm.constants import DEFAULT_CHAT_COMPLETION_PARAM_VALUES, RESPONSE_FORMAT_TOOL_NAME from litellm.files.types import FileContentStreamingResult from litellm.files.utils import FilesAPIUtils from litellm.litellm_core_utils.cloud_storage_security import ( @@ -30,11 +31,13 @@ from litellm.litellm_core_utils.cloud_storage_security import ( split_configured_cloud_bucket_name, validate_managed_cloud_file_id, ) +from litellm.litellm_core_utils.core_helpers import as_str_mapping, iter_items, normalize_drop_params from litellm.litellm_core_utils.litellm_logging import Logging from litellm.litellm_core_utils.prompt_templates.common_utils import ( extract_file_data, extract_file_metadata, ) +from litellm.llms.anthropic.chat.transformation import ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY from litellm.llms.base_llm.base_utils import map_developer_role_to_system_role from litellm.llms.base_llm.chat.transformation import BaseLLMException from litellm.llms.base_llm.files.batch_records import responses_batch_body_to_chat_body @@ -43,7 +46,14 @@ from litellm.llms.base_llm.files.transformation import ( BaseFileUploadStream, LiteLLMLoggingObj, ) -from litellm.llms.vertex_ai.batches.transformation import vertex_embedding_prompt_token_count +from litellm.llms.vertex_ai.batches.transformation import ( + VERTEX_ANTHROPIC_VERSION, + as_json_object, + is_vertex_anthropic_batch_model, + is_vertex_anthropic_batch_output_response, + is_vertex_anthropic_batch_output_row, + vertex_embedding_prompt_token_count, +) from litellm.llms.vertex_ai.common_utils import ( convert_vertex_datetime_to_openai_datetime, get_vertex_ai_fine_tuned_endpoint_id, @@ -55,6 +65,7 @@ from litellm.llms.vertex_ai.gemini.vertex_and_google_ai_studio_gemini import ( from litellm.llms.vertex_ai.gemini_embeddings.batch_embed_content_transformation import ( transform_openai_input_gemini_embed_content, ) +from litellm.llms.vertex_ai.vertex_ai_partner_models.anthropic.transformation import VertexAIAnthropicConfig from litellm.types.files import StreamingMediaUploadConfig from litellm.types.llms.openai import ( AllMessageValues, @@ -75,6 +86,7 @@ from litellm.types.utils import ( ModelResponse, Usage, ) +from litellm.utils import get_non_default_completion_params, get_optional_params from ..common_utils import VertexAIError from ..vertex_llm_base import VertexBase @@ -83,6 +95,7 @@ _GCP_LABEL_VALUE_MAX_LEN: Final = 63 _CUSTOM_ID_RAW_LABEL_PREFIX: Final = "b32_" _VERTEX_BATCH_KEY_FIELD: Final = "key" _MANAGED_GCS_MODEL_PATH_PATTERN: Final = re.compile(r"publishers/[^/]+/models/([^/?]+)") +_PUBLISHER_MODEL_PATH_PATTERN: Final = re.compile(r"publishers/[^/]+/models/") _EMBED_REQUEST_FIELD_BY_GEMINI_PARAM: Final = ( ("outputDimensionality", "output_dimensionality"), ("taskType", "task_type"), @@ -220,10 +233,14 @@ def _get_litellm_batch_custom_id(vertex_output_row: Mapping[str, object]) -> str """ Resolve the OpenAI `custom_id` for a Vertex batch output row. - Embedding rows carry it in the top-level `key` field that Vertex echoes back; - `generateContent` rows have no such field, so it is smuggled through request - labels instead (see `_set_litellm_batch_custom_id_labels`). + Claude rows carry it in their own top-level `custom_id` field and embedding rows in + the top-level `key` field that Vertex echoes back; `generateContent` rows have no + such field, so it is smuggled through request labels instead (see + `_set_litellm_batch_custom_id_labels`). """ + anthropic_custom_id: Final = vertex_output_row.get("custom_id") + if anthropic_custom_id is not None: + return str(anthropic_custom_id) key = vertex_output_row.get(_VERTEX_BATCH_KEY_FIELD) if key is not None: return unquote(str(key)) @@ -284,6 +301,12 @@ def _is_vertex_generate_content_batch_output_row(vertex_output_row: Mapping[str, ) +def _is_vertex_chat_batch_output_row(vertex_output_row: Mapping[str, object]) -> bool: + return _is_vertex_generate_content_batch_output_row(vertex_output_row) or is_vertex_anthropic_batch_output_row( + vertex_output_row + ) + + def _try_parse_vertex_batch_output_row(line: bytes) -> _VertexBatchRow | None: try: row: Final = _parse_vertex_batch_output_row(line.decode("utf-8")) @@ -405,6 +428,37 @@ def _openai_batch_output_row( } +def _vertex_anthropic_batch_output_row_to_openai( + custom_id: str, + response: object, + mock_httpx_response: httpx.Response, +) -> _OpenAIBatchOutputRow: + """ + One OpenAI batch output row from a Vertex Claude row's Anthropic `message` response + """ + anthropic_response: Final = as_json_object(response) + if anthropic_response is None or not is_vertex_anthropic_batch_output_response(anthropic_response): + return _openai_batch_output_row( + custom_id=custom_id, + error_code="vertex_ai_error", + error_message="Claude batch row has no message response", + ) + try: + model_response: Final = VertexAIAnthropicConfig().transform_parsed_response( + completion_response=dict(anthropic_response), + raw_response=mock_httpx_response, + model_response=ModelResponse(), + json_mode=True, + ) + except Exception as e: # noqa: BLE001 # one malformed row must fail alone, not the whole output file + return _openai_batch_output_row( + custom_id=custom_id, + error_code="transformation_error", + error_message=f"Failed to transform response: {e}", + ) + return _openai_batch_output_row(custom_id=custom_id, body=model_response.model_dump()) + + def _split_vertex_batch_key(vertex_output_row: Mapping[str, object]) -> tuple[str, int, int]: """ Resolve `(custom_id, index within that custom_id, group size)` for a Vertex batch @@ -679,21 +733,123 @@ def _openai_batch_jsonl_entry_to_vertex_embeddings_rows( ) +_CLAUDE_BATCH_ROW_FIXED_PARAMS: Final = frozenset( + ("model", "messages", "custom_llm_provider", "drop_params", "stream", "stream_options") +) +_CLAUDE_BATCH_OPENAI_PARAMS: Final = frozenset(DEFAULT_CHAT_COMPLETION_PARAM_VALUES) - {"metadata"} +_CLAUDE_BATCH_DEPLOYMENT_PARAMS: Final = ( + _CLAUDE_BATCH_OPENAI_PARAMS - _CLAUDE_BATCH_ROW_FIXED_PARAMS - {"extra_headers"} +) | {"drop_params"} + + +def _claude_batch_deployment_params(litellm_params: Mapping[str, object]) -> Mapping[str, object]: + return MappingProxyType( + {key: value for key, value in litellm_params.items() if key in _CLAUDE_BATCH_DEPLOYMENT_PARAMS} + ) + + +def _names_the_response_format_tool(tool: object) -> bool: + match tool: + case {"function": {"name": str(name)}}: + return name == RESPONSE_FORMAT_TOOL_NAME + case _: + return False + + +def _openai_batch_jsonl_entry_to_vertex_anthropic_row( + custom_id: object, + messages: list[AllMessageValues], + request_params: Mapping[str, object], + model: str, + deployment_params: Mapping[str, object], +) -> Mapping[str, object]: + """ + Transforms an OpenAI chat batch entry into the Anthropic Messages row Vertex batches Claude with, + `{"custom_id": ..., "request": {"messages": [...], "anthropic_version": "vertex-2023-10-16", "max_tokens": ...}}` + """ + if custom_id is None: + raise ValueError("custom_id is required on Claude lines, since Claude output rows are matched back only by it") + config: Final = VertexAIAnthropicConfig() + bare_model: Final = model.rsplit("/", 1)[-1] + params: Final = {**deployment_params, **request_params} + drop_params: Final = litellm.drop_params is True or normalize_drop_params(params.get("drop_params")) is True + chat_params: Final = { + **get_non_default_completion_params(params), + **{key: value for key, value in params.items() if key in _CLAUDE_BATCH_OPENAI_PARAMS}, + } + if any(map(_names_the_response_format_tool, iter_items(chat_params.get("tools") or ()))): + raise ValueError( + f"tool name {RESPONSE_FORMAT_TOOL_NAME!r} is reserved for response_format on Claude, " + "and batch output cannot tell a user tool by that name apart, so rename it" + ) + optional_params: Final[Mapping[str, object]] = get_optional_params( + model=bare_model, + custom_llm_provider="vertex_ai", + drop_params=drop_params, + messages=messages, + **{ # pyright: ignore[reportArgumentType] # batch JSONL bodies are untyped; get_optional_params checks each param as chat does + key: value for key, value in chat_params.items() if key not in _CLAUDE_BATCH_ROW_FIXED_PARAMS + }, + ) + request_state: Final[dict[str, object]] = {"drop_params": drop_params} + anthropic_request: Final[Mapping[str, object]] = config.transform_request( + model=bare_model, + messages=messages, + optional_params={ + **{key: value for key, value in optional_params.items() if key not in ("json_mode", "extra_body")}, + "anthropic_version": VERTEX_ANTHROPIC_VERSION, + }, + litellm_params=request_state, + headers={}, + ) + renamed_tools: Final = as_str_mapping(request_state.get(ANTHROPIC_TOOL_NAME_REVERSE_MAP_KEY)) + if renamed_tools: + raise ValueError( + f"tool names {sorted(map(str, renamed_tools.values()))} must match ^[a-zA-Z0-9_-]{{1,128}}$ for Claude on Vertex, " + "and batch output cannot map a renamed tool back, so rename these tools" + ) + extra_body: Final = as_json_object(optional_params.get("extra_body")) + request: Final[Mapping[str, object]] = ( + anthropic_request + if extra_body is None + else { + **anthropic_request, + **config.transform_extra_body( + extra_body=extra_body, request=anthropic_request, model=bare_model, litellm_params={} + ), + } + ) + return {"custom_id": str(custom_id), "request": request} + + def _openai_batch_jsonl_entry_to_vertex_rows( openai_entry: dict[str, Any], map_openai_to_vertex_params: Callable[[dict[str, object]], dict[str, object]], + batch_model: str | None = None, + deployment_params: Mapping[str, object] = MappingProxyType({}), ) -> tuple[Mapping[str, object], ...]: """ Transforms a single OpenAI JSONL batch entry into the Vertex rows it maps to. + `batch_model` is the deployment model the batch runs against and wins over the entry's own `body.model`, + which picks the row shape: Anthropic Messages for Claude, `GenerateContentRequest` otherwise. + jsonl body for vertex is {"request": } Example Vertex jsonl {"request":{"contents": [{"role": "user", "parts": [{"text": "What is the relation between the following video and image samples?"}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/video/animals.mp4", "mimeType": "video/mp4"}}, {"fileData": {"fileUri": "gs://cloud-samples-data/generative-ai/image/cricket.jpeg", "mimeType": "image/jpeg"}}]}]}} """ + entry_body: Final = openai_entry.get("body") + model: Final = batch_model or (str(entry_body.get("model") or "") if isinstance(entry_body, dict) else "") + is_anthropic_model: Final = is_vertex_anthropic_batch_model(model) if _is_embeddings_batch_entry(openai_entry): + if is_anthropic_model: + raise VertexAIError( + status_code=400, + message=f"Vertex AI batch embeddings are not supported for Claude models ('{model}')", + ) return _openai_batch_jsonl_entry_to_vertex_embeddings_rows(openai_entry) - if not isinstance(openai_entry.get("body"), dict): + if not isinstance(entry_body, dict): raise TypeError("`body` on batch requests must be a JSON object, but was missing or not an object") openai_request_body: Final = openai_entry["body"] chat_request_body: Final = ( @@ -701,8 +857,21 @@ def _openai_batch_jsonl_entry_to_vertex_rows( if _is_responses_batch_entry(openai_entry) else openai_request_body ) + chat_messages: Final = map_developer_role_to_system_role(chat_request_body.get("messages", [])) + if is_anthropic_model: + return ( + _openai_batch_jsonl_entry_to_vertex_anthropic_row( + custom_id=openai_entry.get("custom_id"), + messages=chat_messages, + request_params={ + key: value for key, value in chat_request_body.items() if key not in ("model", "messages") + }, + model=model, + deployment_params=deployment_params, + ), + ) vertex_request_body: Final = transform_request_body( - messages=map_developer_role_to_system_role(chat_request_body.get("messages", [])), + messages=chat_messages, model=chat_request_body.get("model", ""), optional_params=map_openai_to_vertex_params(chat_request_body), custom_llm_provider="vertex_ai", @@ -854,15 +1023,24 @@ class _OpenAIToVertexBatchUploadStream(BaseFileUploadStream): self, openai_file_content: FileTypes, map_openai_to_vertex_params: Callable[[dict[str, object]], dict[str, object]], + batch_model: str | None = None, + deployment_params: Mapping[str, object] = MappingProxyType({}), ) -> None: self._openai_file_content = openai_file_content self._map_openai_to_vertex_params = map_openai_to_vertex_params + self._batch_model = batch_model + self._deployment_params = deployment_params def _iter_vertex_jsonl_chunks(self) -> Iterator[bytes]: first = True for lineno, entry in _iter_openai_jsonl_entries(self._openai_file_content): try: - vertex_rows = _openai_batch_jsonl_entry_to_vertex_rows(entry, self._map_openai_to_vertex_params) + vertex_rows = _openai_batch_jsonl_entry_to_vertex_rows( + entry, + self._map_openai_to_vertex_params, + batch_model=self._batch_model, + deployment_params=self._deployment_params, + ) except Exception as e: raise VertexAIError( status_code=400, @@ -918,16 +1096,31 @@ class _RawFileUploadStream(BaseFileUploadStream): return _iter_raw_file_chunks(self._file_content) +def _deployment_batch_model(litellm_params: Mapping[str, object]) -> str | None: + model: Final = litellm_params.get("model") + return model.removeprefix("vertex_ai/") if isinstance(model, str) and model else None + + +def _publisher_model_path(raw_model: str) -> str: + if _PUBLISHER_MODEL_PATH_PATTERN.search(raw_model): + return raw_model + publisher: Final = "anthropic" if is_vertex_anthropic_batch_model(raw_model) else "google" + return f"publishers/{publisher}/models/{raw_model}" + + def _managed_batch_object_name(raw_model: str, *, passthrough: bool) -> str: endpoint_id: Final = get_vertex_ai_fine_tuned_endpoint_id(raw_model) - model_path: Final = ( - f"endpoints/{endpoint_id}" - if endpoint_id is not None - else (raw_model if "publishers/google/models" in raw_model else f"publishers/google/models/{raw_model}") - ) - safe_model_path: Final = sanitize_cloud_object_path(model_path, fallback="model") + model_path: Final = f"endpoints/{endpoint_id}" if endpoint_id is not None else _publisher_model_path(raw_model) prefix: Final = _PASSTHROUGH_MANAGED_GCS_PREFIX if passthrough else VERTEX_AI_MANAGED_GCS_PREFIX - return f"{prefix}{safe_model_path}/{uuid.uuid4()}" + return f"{prefix}{_safe_managed_model_path(model_path)}/{uuid.uuid4()}" + + +def _safe_managed_model_path(model_path: str) -> str: + if not is_vertex_anthropic_batch_model(model_path): + return sanitize_cloud_object_path(model_path, fallback="model") + return "@".join( + sanitize_cloud_object_path(version_part, fallback="model") for version_part in model_path.split("@") + ) class VertexAIFilesConfig(VertexBase, BaseFilesConfig): @@ -1030,10 +1223,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): fallback_filename="file", ) - def _get_configured_bucket_name(self, litellm_params: dict) -> str: + def _get_configured_bucket_name(self, litellm_params: Mapping[str, object]) -> str: + configured_bucket_name: Final = litellm_params.get("gcs_bucket_name") or litellm_params.get("bucket_name") bucket_name: Final = ( - litellm_params.get("gcs_bucket_name") - or litellm_params.get("bucket_name") + (configured_bucket_name if isinstance(configured_bucket_name, str) else None) or os.getenv("GCS_BATCH_BUCKET_NAME") or os.getenv("GCS_BUCKET_NAME") ) @@ -1047,7 +1240,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): api_key: str | None, model: str, optional_params: dict, - litellm_params: dict, + litellm_params: Mapping[str, object], data: CreateFileRequest, ) -> str: """ @@ -1071,11 +1264,10 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): raise ValueError("file is required") if purpose is None: raise ValueError("purpose is required") - configured_model: Final = litellm_params.get("model") object_name = self.get_object_name( file_data, purpose, - deployment_model=configured_model if isinstance(configured_model, str) else None, + deployment_model=_deployment_batch_model(litellm_params), passthrough=is_passthrough_batch_upload(data, litellm_params), ) if object_prefix: @@ -1126,7 +1318,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): model: str, create_file_data: CreateFileRequest, optional_params: dict, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> bytes | str | dict: """ 2 Cases: @@ -1157,6 +1349,8 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): body_stream=_OpenAIToVertexBatchUploadStream( file_data, self._map_openai_to_vertex_params, + batch_model=_deployment_batch_model(litellm_params), + deployment_params=_claude_batch_deployment_params(litellm_params), ), content_type="application/json", ) @@ -1207,7 +1401,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): def get_error_class(self, error_message: str, status_code: int, headers: dict | Headers) -> BaseLLMException: return VertexAIError(status_code=status_code, message=error_message, headers=headers) - def _parse_gcs_uri(self, file_id: str, litellm_params: dict | None = None) -> tuple[str, str]: + def _parse_gcs_uri(self, file_id: str, litellm_params: Mapping[str, object] | None = None) -> tuple[str, str]: """ Validate a managed GCS file_id and return (bucket, url-encoded-object-path). """ @@ -1225,7 +1419,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): self, file_id: str, optional_params: dict, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> tuple[str, dict]: bucket, encoded_object = self._parse_gcs_uri(file_id, litellm_params) url: Final = f"https://storage.googleapis.com/storage/v1/b/{bucket}/o/{encoded_object}" @@ -1255,7 +1449,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): self, file_id: str, optional_params: dict, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> tuple[str, dict]: bucket, encoded_object = self._parse_gcs_uri(file_id, litellm_params) url: Final = f"https://storage.googleapis.com/storage/v1/b/{bucket}/o/{encoded_object}" @@ -1298,7 +1492,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): self, file_content_request, optional_params: dict, - litellm_params: dict, + litellm_params: Mapping[str, object], ) -> tuple[str, dict]: file_id: Final = file_content_request.get("file_id", "") bucket, encoded_object = self._parse_gcs_uri(file_id, litellm_params) @@ -1395,7 +1589,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): headers=MappingProxyType({**headers, "content-length": str(len(transformed_content))}), ) - if not _is_vertex_generate_content_batch_output_row(first_row): + if not _is_vertex_chat_batch_output_row(first_row): return FileContentStreamingResult(stream_iterator=replayed_stream, headers=headers) return FileContentStreamingResult( @@ -1479,10 +1673,7 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): # first line is not valid UTF-8/JSON) raises and falls through to the # passthrough below, leaving the content untouched. first_row: Final = _parse_vertex_batch_output_row(first_line) - if not ( - _is_vertex_embeddings_batch_output_row(first_row) - or _is_vertex_generate_content_batch_output_row(first_row) - ): + if not (_is_vertex_embeddings_batch_output_row(first_row) or _is_vertex_chat_batch_output_row(first_row)): return content context: Final = _new_vertex_batch_output_row_transform_context() @@ -1547,6 +1738,13 @@ class VertexAIFilesConfig(VertexBase, BaseFilesConfig): error_message=status, ) + if is_vertex_anthropic_batch_output_row(vertex_output): + return _vertex_anthropic_batch_output_row_to_openai( + custom_id=custom_id, + response=vertex_output.get("response"), + mock_httpx_response=mock_httpx_response, + ) + # Transform successful response using existing transformation vertex_response: Final = vertex_output.get("response", {}) diff --git a/litellm/main.py b/litellm/main.py index b03df777ef4..181b33cf342 100644 --- a/litellm/main.py +++ b/litellm/main.py @@ -27,6 +27,7 @@ from functools import partial from types import MappingProxyType from typing import TYPE_CHECKING, Any, Final, Literal, Optional, Protocol, Union, cast, get_args +from litellm._internal_context import is_resent_failure from litellm._logging import redact_string from litellm._uuid import uuid @@ -140,6 +141,7 @@ from litellm.llms.vertex_ai.common_utils import ( get_vertex_ai_model_route, ) from litellm.realtime_api.main import realtime_health_check +from litellm.responses.mcp.tool_precheck import may_name_litellm_mcp_gateway from litellm.secret_managers.main import get_secret_bool, get_secret_str from litellm.types.completion import ( CompletionDispatchContext, @@ -5381,63 +5383,55 @@ def completion( control_options: Final = _resolve_control_options(kwargs, model) skip_mcp_handler: Final = kwargs.pop("_skip_mcp_handler", False) - if not skip_mcp_handler and tools: + if not skip_mcp_handler and tools and may_name_litellm_mcp_gateway(tools): # pyright: ignore[reportUnknownArgumentType] # untyped tools list from litellm.responses.mcp.chat_completions_handler import acompletion_with_mcp - from litellm.responses.mcp.litellm_proxy_mcp_handler import ( - LiteLLM_Proxy_MCP_Handler, - ) - from litellm.types.llms.openai import ToolParam - # Check if MCP tools are present (following responses pattern) - # Cast tools to Optional[Iterable[ToolParam]] for type checking - tools_for_mcp: Final = cast(Iterable[ToolParam] | None, tools) - if LiteLLM_Proxy_MCP_Handler.should_use_litellm_mcp_gateway(tools=tools_for_mcp): - return acompletion_with_mcp( # pyright: ignore[reportReturnType] # MCP path returns a coroutine that acompletion() awaits; completion()'s sync return type omits it - model=model, - messages=messages, - functions=functions, - function_call=function_call, - timeout=timeout, - temperature=temperature, - top_p=top_p, - n=n, - stream=stream, - stream_options=stream_options, - stop=stop, - max_tokens=max_tokens, - max_completion_tokens=max_completion_tokens, - modalities=modalities, - prediction=prediction, - audio=audio, - presence_penalty=presence_penalty, - frequency_penalty=frequency_penalty, - logit_bias=logit_bias, - user=user, - response_format=response_format, - seed=seed, - tools=tools, - tool_choice=tool_choice, - parallel_tool_calls=parallel_tool_calls, - logprobs=logprobs, - top_logprobs=top_logprobs, - deployment_id=deployment_id, - reasoning_effort=reasoning_effort, - verbosity=verbosity, - safety_identifier=safety_identifier, - service_tier=service_tier, - store=store, - prompt_cache_key=prompt_cache_key, - base_url=base_url, - api_version=api_version, - api_key=api_key, - model_list=model_list, - extra_headers=extra_headers, - thinking=thinking, - web_search_options=web_search_options, - shared_session=shared_session, - enable_json_schema_validation=enable_json_schema_validation, - **kwargs, - ) + return acompletion_with_mcp( # pyright: ignore[reportReturnType] # MCP path returns a coroutine that acompletion() awaits; completion()'s sync return type omits it + model=model, + messages=messages, + functions=functions, + function_call=function_call, + timeout=timeout, + temperature=temperature, + top_p=top_p, + n=n, + stream=stream, + stream_options=stream_options, + stop=stop, + max_tokens=max_tokens, + max_completion_tokens=max_completion_tokens, + modalities=modalities, + prediction=prediction, + audio=audio, + presence_penalty=presence_penalty, + frequency_penalty=frequency_penalty, + logit_bias=logit_bias, + user=user, + response_format=response_format, + seed=seed, + tools=tools, + tool_choice=tool_choice, + parallel_tool_calls=parallel_tool_calls, + logprobs=logprobs, + top_logprobs=top_logprobs, + deployment_id=deployment_id, + reasoning_effort=reasoning_effort, + verbosity=verbosity, + safety_identifier=safety_identifier, + service_tier=service_tier, + store=store, + prompt_cache_key=prompt_cache_key, + base_url=base_url, + api_version=api_version, + api_key=api_key, + model_list=model_list, + extra_headers=extra_headers, + thinking=thinking, + web_search_options=web_search_options, + shared_session=shared_session, + enable_json_schema_validation=enable_json_schema_validation, + **kwargs, + ) api_base = kwargs.get("api_base", None) mock_response: Final[MOCK_RESPONSE_TYPE | None] = kwargs.get("mock_response", None) mock_tool_calls: Final = kwargs.get("mock_tool_calls", None) @@ -6277,6 +6271,10 @@ def completion( ) +def _is_not_resent_failure(error: BaseException) -> bool: + return isinstance(error, Exception) and not is_resent_failure(error) + + def completion_with_retries(*args, **kwargs): """ Executes a litellm.completion() with 3 retries @@ -6296,10 +6294,15 @@ def completion_with_retries(*args, **kwargs): retryer = tenacity.Retrying( wait=tenacity.wait_exponential(multiplier=1, max=10), stop=tenacity.stop_after_attempt(num_retries), + retry=tenacity.retry_if_exception(_is_not_resent_failure), reraise=True, ) else: - retryer = tenacity.Retrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True) + retryer = tenacity.Retrying( + stop=tenacity.stop_after_attempt(num_retries), + retry=tenacity.retry_if_exception(_is_not_resent_failure), + reraise=True, + ) return retryer(original_function, *args, **kwargs) @@ -6322,10 +6325,15 @@ async def acompletion_with_retries(*args, **kwargs): retryer = tenacity.AsyncRetrying( wait=tenacity.wait_exponential(multiplier=1, max=10), stop=tenacity.stop_after_attempt(num_retries), + retry=tenacity.retry_if_exception(_is_not_resent_failure), reraise=True, ) else: - retryer = tenacity.AsyncRetrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True) + retryer = tenacity.AsyncRetrying( + stop=tenacity.stop_after_attempt(num_retries), + retry=tenacity.retry_if_exception(_is_not_resent_failure), + reraise=True, + ) return await retryer(original_function, *args, **kwargs) @@ -6350,10 +6358,15 @@ def responses_with_retries(*args, **kwargs): retryer = tenacity.Retrying( wait=tenacity.wait_exponential(multiplier=1, max=10), stop=tenacity.stop_after_attempt(num_retries), + retry=tenacity.retry_if_exception(_is_not_resent_failure), reraise=True, ) else: - retryer = tenacity.Retrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True) + retryer = tenacity.Retrying( + stop=tenacity.stop_after_attempt(num_retries), + retry=tenacity.retry_if_exception(_is_not_resent_failure), + reraise=True, + ) return retryer(original_function, *args, **kwargs) @@ -6377,10 +6390,15 @@ async def aresponses_with_retries(*args, **kwargs): retryer = tenacity.AsyncRetrying( wait=tenacity.wait_exponential(multiplier=1, max=10), stop=tenacity.stop_after_attempt(num_retries), + retry=tenacity.retry_if_exception(_is_not_resent_failure), reraise=True, ) else: - retryer = tenacity.AsyncRetrying(stop=tenacity.stop_after_attempt(num_retries), reraise=True) + retryer = tenacity.AsyncRetrying( + stop=tenacity.stop_after_attempt(num_retries), + retry=tenacity.retry_if_exception(_is_not_resent_failure), + reraise=True, + ) return await retryer(original_function, *args, **kwargs) diff --git a/litellm/passthrough/main.py b/litellm/passthrough/main.py index 744f9e1aa4c..4e94b1bb910 100644 --- a/litellm/passthrough/main.py +++ b/litellm/passthrough/main.py @@ -437,6 +437,7 @@ def llm_passthrough_route( params: QueryParamTypes | None = None, cookies: CookieTypes | None = None, client: HTTPHandler | AsyncHTTPHandler | None = None, + guardrailConfig: object = None, **kwargs, ) -> ( httpx.Response @@ -547,7 +548,7 @@ def llm_passthrough_route( _request_data: dict | None = data if isinstance(data, dict) else (json if isinstance(json, dict) else None) headers, signed_json_body = provider_config.sign_request( headers=headers, - litellm_params=litellm_params_dict, + litellm_params={**litellm_params_dict, "guardrailConfig": guardrailConfig}, request_data=_request_data, api_base=str(updated_url), model=model, diff --git a/litellm/proxy/auth/auth_utils.py b/litellm/proxy/auth/auth_utils.py index 98bb0985f4a..d012bbe51f3 100644 --- a/litellm/proxy/auth/auth_utils.py +++ b/litellm/proxy/auth/auth_utils.py @@ -4,6 +4,7 @@ import re import sys from collections.abc import Collection, Iterator, Mapping from functools import lru_cache +from itertools import accumulate, chain, repeat from logging import Logger from types import MappingProxyType from typing import Any, Final, Protocol @@ -39,6 +40,7 @@ from litellm.types.passthrough_endpoints.pass_through_endpoints import ( from litellm.types.router import CONFIGURABLE_CLIENTSIDE_AUTH_PARAMS, Deployment, server_owned_wif_fields_present from litellm.types.router import reject_server_owned_wif_params as _reject_server_owned_wif_params from litellm.types.utils import CustomPricingLiteLLMParams, oauth_token_exchange_litellm_params +from litellm.types.videos.main import DecodedVideoId def is_invalid_virtual_key_error(exception: BaseException | None) -> bool: @@ -566,6 +568,31 @@ def fallback_target_model_name(target: object) -> str | None: return None +def _direct_fallback_model_names(model: str, llm_router: Router) -> frozenset[str]: + from litellm.proxy.auth.model_checks import get_all_fallbacks + + return frozenset( + name for target in get_all_fallbacks(model, llm_router) if (name := fallback_target_model_name(target)) + ) + + +def _fallback_hop( + frontier_and_reached: tuple[frozenset[str], frozenset[str]], llm_router: Router +) -> tuple[frozenset[str], frozenset[str]]: + frontier, reached = frontier_and_reached + next_hop: Final = ( + frozenset(chain.from_iterable(_direct_fallback_model_names(model, llm_router) for model in frontier)) - reached + ) + return next_hop, reached | next_hop + + +def router_fallback_model_names(model: str, llm_router: Router) -> frozenset[str]: + """The groups a failure on `model` can reach through the router's fallbacks, chained and generic (`*`) ones.""" + hops: Final = accumulate(repeat(llm_router), _fallback_hop, initial=(frozenset({model}), frozenset({model}))) + _, reached = next(hop for hop in hops if not hop[0]) + return reached - {model} + + def request_fallback_model_names(request_body: Mapping[str, object]) -> tuple[str, ...]: return tuple( dict.fromkeys( @@ -1975,10 +2002,10 @@ def _extract_models_from_managed_resource_id( ) if resource_id_field == "video_id": - model_id = decode_video_id_with_provider(resource_id).get("model_id") + decoded_video: Final = decode_video_id_with_provider(resource_id) _append_model_candidates( candidates=candidates, - value=_resolve_model_id_with_router(model_id, llm_router), + value=_video_access_group(decoded_video, llm_router) or decoded_video.get("model_id"), ) else: model_id = decode_character_id_with_provider(resource_id).get("model_id") @@ -1992,6 +2019,25 @@ def _extract_models_from_managed_resource_id( return _dedupe_model_candidates(candidates) +def _video_access_group(decoded_video: DecodedVideoId, llm_router: Router | None) -> str | None: + """ + The group a key must be allowed to call to use a video: the creating deployment's group, or, while router + fallbacks skip model access checks, the group whose router fallback created it. + """ + from litellm.proxy.auth.fallback_model_access import router_fallback_access_check + from litellm.proxy.video_endpoints.utils import video_model_group + + creating_group: Final = video_model_group(decoded_video, llm_router) + requested_model: Final = decoded_video.get("requested_model") + if llm_router is None or creating_group is None or requested_model is None: + return creating_group + if router_fallback_access_check.is_enforced(): + return creating_group + if creating_group in router_fallback_model_names(requested_model, llm_router): + return requested_model + return creating_group + + def _resolve_model_id_with_router(model_id: str | None, llm_router: Router | None) -> str | None: if model_id is None or llm_router is None: return model_id diff --git a/litellm/proxy/video_endpoints/endpoints.py b/litellm/proxy/video_endpoints/endpoints.py index c4f33eae114..d4643994043 100644 --- a/litellm/proxy/video_endpoints/endpoints.py +++ b/litellm/proxy/video_endpoints/endpoints.py @@ -20,9 +20,14 @@ from litellm.proxy.common_utils.openai_endpoint_utils import ( ) from litellm.proxy.image_endpoints.endpoints import batch_to_bytesio from litellm.proxy.video_endpoints.utils import ( + VideoRequestRoute, encode_character_id_in_response, + encode_video_id_in_response, extract_model_from_target_model_names, + fallback_requested_model, get_custom_provider_from_data, + routable_model_name, + video_model_group, video_reference_to_id, ) from litellm.types.videos.utils import ( @@ -84,6 +89,7 @@ async def video_generation( # Read request body data: Final = await read_request_body(request=request) + requested: Final = VideoRequestRoute.model_validate(data) if input_reference is not None: input_reference_file: Final = await batch_to_bytesio([input_reference]) if input_reference_file: @@ -118,7 +124,13 @@ async def video_generation( version=version, ) else: - return generated + requested_model: Final = routable_model_name(requested.model, llm_router, requested.custom_llm_provider) + return encode_video_id_in_response( + response=generated, + fallback_provider=requested.custom_llm_provider, + fallback_model_id=requested_model, + requested_model=fallback_requested_model(requested_model, generated, llm_router), + ) @router.get( @@ -269,12 +281,9 @@ async def video_status( if custom_llm_provider: data["custom_llm_provider"] = custom_llm_provider - # Resolve model_name from model_id if available - # This allows the router to automatically inject litellm_params from the model config - if model_id_from_decoded and llm_router: - resolved_model: Final = llm_router.resolve_model_name_from_model_id(model_id_from_decoded, provider_from_id) - if resolved_model: - data["model"] = resolved_model + resolved_model: Final = video_model_group(decoded, llm_router) + if resolved_model: + data["model"] = resolved_model # Process request using ProxyBaseLLMRequestProcessing processor: Final = ProxyBaseLLMRequestProcessing(data=data, served_model=decoded.get("served_model")) @@ -305,7 +314,12 @@ async def video_status( version=version, ) else: - return status + return encode_video_id_in_response( + response=status, + fallback_provider=custom_llm_provider, + fallback_model_id=model_id_from_decoded, + polled_video=decoded, + ) @router.get( @@ -358,7 +372,6 @@ async def video_content( decoded: Final = decode_video_id_with_provider(video_id) provider_from_id: Final = decoded.get("custom_llm_provider") - model_id_from_decoded: Final = decoded.get("model_id") custom_llm_provider: Final = ( get_custom_llm_provider_from_request_headers(request=request) @@ -369,12 +382,9 @@ async def video_content( if custom_llm_provider: data["custom_llm_provider"] = custom_llm_provider - # Resolve model_name from model_id if available - # This allows the router to automatically inject litellm_params from the model config - if model_id_from_decoded and llm_router: - resolved_model: Final = llm_router.resolve_model_name_from_model_id(model_id_from_decoded, provider_from_id) - if resolved_model: - data["model"] = resolved_model + resolved_model: Final = video_model_group(decoded, llm_router) + if resolved_model: + data["model"] = resolved_model # Process request using ProxyBaseLLMRequestProcessing processor: Final = ProxyBaseLLMRequestProcessing(data=data, served_model=decoded.get("served_model")) try: @@ -462,6 +472,7 @@ async def video_remix( ) data: Final = await read_request_body(request=request) + requested: Final = VideoRequestRoute.model_validate(data) data["video_id"] = video_id decoded: Final = decode_video_id_with_provider(video_id) @@ -471,18 +482,15 @@ async def video_remix( custom_llm_provider: Final = ( get_custom_llm_provider_from_request_headers(request=request) or get_custom_llm_provider_from_request_query(request=request) - or data.get("custom_llm_provider") + or requested.custom_llm_provider or provider_from_id ) if custom_llm_provider: data["custom_llm_provider"] = custom_llm_provider - # Resolve model_name from model_id if available - # This allows the router to automatically inject litellm_params from the model config - if model_id_from_decoded and llm_router: - resolved_model: Final = llm_router.resolve_model_name_from_model_id(model_id_from_decoded, provider_from_id) - if resolved_model: - data["model"] = resolved_model + resolved_model: Final = video_model_group(decoded, llm_router) + if resolved_model: + data["model"] = resolved_model # Process request using ProxyBaseLLMRequestProcessing processor: Final = ProxyBaseLLMRequestProcessing(data=data, served_model=decoded.get("served_model")) @@ -513,7 +521,13 @@ async def video_remix( version=version, ) else: - return remixed + return encode_video_id_in_response( + response=remixed, + fallback_provider=custom_llm_provider, + fallback_model_id=model_id_from_decoded + or routable_model_name(requested.model, llm_router, custom_llm_provider), + requested_model=decoded.get("requested_model"), + ) @router.post( @@ -771,6 +785,7 @@ async def video_edit( ) data: Final = await read_request_body(request=request) + requested: Final = VideoRequestRoute.model_validate(data) uploaded_video: Final = data.pop("video", None) if isinstance(uploaded_video, StarletteUploadFile): video_files: Final = await batch_to_bytesio((uploaded_video,)) @@ -793,10 +808,9 @@ async def video_edit( ) data["custom_llm_provider"] = custom_llm_provider - if model_id_from_decoded and llm_router: - resolved_model: Final = llm_router.resolve_model_name_from_model_id(model_id_from_decoded, provider_from_id) - if resolved_model: - data["model"] = resolved_model + resolved_model: Final = video_model_group(decoded, llm_router) + if resolved_model: + data["model"] = resolved_model processor: Final = ProxyBaseLLMRequestProcessing(data=data, served_model=decoded.get("served_model")) try: @@ -826,7 +840,13 @@ async def video_edit( version=version, ) else: - return edited + return encode_video_id_in_response( + response=edited, + fallback_provider=custom_llm_provider, + fallback_model_id=model_id_from_decoded + or routable_model_name(requested.model, llm_router, custom_llm_provider), + requested_model=decoded.get("requested_model"), + ) @router.post( @@ -875,6 +895,7 @@ async def video_extension( ) data: Final = await read_request_body(request=request) + requested: Final = VideoRequestRoute.model_validate(data) data["video_id"] = video_reference_to_id(data.pop("video", None)) decoded: Final = decode_video_id_with_provider(data["video_id"]) @@ -890,10 +911,9 @@ async def video_extension( ) data["custom_llm_provider"] = custom_llm_provider - if model_id_from_decoded and llm_router: - resolved_model: Final = llm_router.resolve_model_name_from_model_id(model_id_from_decoded, provider_from_id) - if resolved_model: - data["model"] = resolved_model + resolved_model: Final = video_model_group(decoded, llm_router) + if resolved_model: + data["model"] = resolved_model processor: Final = ProxyBaseLLMRequestProcessing(data=data, served_model=decoded.get("served_model")) try: @@ -923,4 +943,10 @@ async def video_extension( version=version, ) else: - return extended + return encode_video_id_in_response( + response=extended, + fallback_provider=custom_llm_provider, + fallback_model_id=model_id_from_decoded + or routable_model_name(requested.model, llm_router, custom_llm_provider), + requested_model=decoded.get("requested_model"), + ) diff --git a/litellm/proxy/video_endpoints/utils.py b/litellm/proxy/video_endpoints/utils.py index 554f0f6ef53..a2be42f1b76 100644 --- a/litellm/proxy/video_endpoints/utils.py +++ b/litellm/proxy/video_endpoints/utils.py @@ -1,8 +1,18 @@ -from typing import Any, Final +from typing import TYPE_CHECKING, Any, Final import orjson +from pydantic import BaseModel, ConfigDict -from litellm.types.videos.utils import encode_character_id_with_provider +from litellm.types.videos.main import DecodedVideoId, VideoObject +from litellm.types.videos.utils import ( + decode_video_id_with_provider, + encode_character_id_with_provider, + encode_video_id_with_provider, + id_safe_name, +) + +if TYPE_CHECKING: + from litellm.router import Router def extract_model_from_target_model_names(target_model_names: Any) -> str | None: @@ -65,3 +75,104 @@ def encode_character_id_in_response(response: Any, custom_llm_provider: str, mod model_id=model_id, ) return response + + +class VideoRequestRoute(BaseModel): + model_config = ConfigDict(frozen=True) + + model: object = None + custom_llm_provider: object = None + + +def routable_model_name(model: object, llm_router: "Router | None", custom_llm_provider: object) -> str | None: + name: Final = id_safe_name(model) + if name is None or llm_router is None: + return name + provider: Final = custom_llm_provider if isinstance(custom_llm_provider, str) else None + return name if llm_router.resolve_model_name_from_model_id(name, provider) else None + + +def video_model_group(decoded: DecodedVideoId, llm_router: "Router | None") -> str | None: + """ + The group a video id names. A provider model that several groups share resolves to the group of the + deployment that created the video, under its team's public name for a team model, while a name the router + routes by itself (the id's model, bare or provider-prefixed), or one that routes to that deployment's group + (an alias of it), is kept as written. + """ + if llm_router is None: + return None + model_id: Final = decoded.get("model_id") + provider: Final = decoded.get("custom_llm_provider") + resolved: Final = llm_router.resolve_model_name_from_model_id(model_id, provider) if model_id else None + deployment_id: Final = decoded.get("deployment_id") + if not deployment_id or resolved in (model_id, f"{provider}/{model_id}"): + return resolved + deployment: Final = llm_router.get_deployment(model_id=deployment_id) + if deployment is None: + return resolved + creating_group: Final = deployment.model_info.team_public_model_name or deployment.model_name + if resolved is not None and llm_router.resolve_served_model_name(resolved, creating_group) == creating_group: + return resolved + return creating_group + + +def fallback_requested_model(requested_model: str | None, response: object, llm_router: "Router | None") -> str | None: + """The group a create asked for when one of that group's router fallbacks served it.""" + from litellm.proxy.auth.auth_utils import router_fallback_model_names + + if requested_model is None or llm_router is None or not isinstance(response, VideoObject): + return None + deployment_id: Final = id_safe_name(response.hidden_params.get("model_id")) + deployment: Final = llm_router.get_deployment(model_id=deployment_id) if deployment_id else None + if deployment is None: + return None + creating_group: Final = deployment.model_info.team_public_model_name or deployment.model_name + return requested_model if creating_group in router_fallback_model_names(requested_model, llm_router) else None + + +def encode_video_id_in_response( + response: object, + fallback_provider: object, + fallback_model_id: object, + polled_video: DecodedVideoId | None = None, + requested_model: str | None = None, +) -> object: + """ + Stamps the provider, model and the deployment that answered into a returned video id. A poll of an existing + video (`polled_video`) answered without a routed deployment keeps the deployment its id already named, and + every returned id keeps the group a router fallback was serving (`requested_model`). + """ + if not isinstance(response, VideoObject) or not response.id: + return response + + decoded: Final = decode_video_id_with_provider(response.id) + hidden_params: Final = response.hidden_params + routed_deployment_id: Final = id_safe_name(hidden_params.get("model_id")) + deployment_id: Final = routed_deployment_id or (polled_video or decoded).get("deployment_id") + kept_requested_model: Final = ( + id_safe_name(requested_model) or decoded.get("requested_model") or (polled_video or {}).get("requested_model") + ) + model_id: Final = decoded.get("model_id") or id_safe_name(fallback_model_id) or routed_deployment_id + provider: Final = ( + id_safe_name(decoded.get("custom_llm_provider")) + or id_safe_name(hidden_params.get("custom_llm_provider")) + or id_safe_name(fallback_provider) + ) + if not model_id or not provider: + return response + if ( + model_id == decoded.get("model_id") + and deployment_id == decoded.get("deployment_id") + and kept_requested_model == decoded.get("requested_model") + ): + return response + + restamped_id: Final = encode_video_id_with_provider( + video_id=decoded.get("video_id") or response.id, + provider=provider, + model_id=model_id, + served_model=decoded.get("served_model"), + deployment_id=deployment_id, + requested_model=kept_requested_model, + ) + return response.model_copy(update={"id": restamped_id}) diff --git a/litellm/responses/bridge_check.py b/litellm/responses/bridge_check.py index d5e3dfc0c02..6b80132a077 100644 --- a/litellm/responses/bridge_check.py +++ b/litellm/responses/bridge_check.py @@ -4,7 +4,10 @@ from urllib.parse import urlsplit import litellm from litellm._logging import verbose_logger -from litellm.litellm_core_utils.get_llm_provider_logic import get_llm_provider +from litellm.litellm_core_utils.get_llm_provider_logic import ( + get_llm_provider, + resolve_model_and_provider_for_metadata, +) from litellm.llms.azure_ai.chat.transformation import AzureAIGPT5Config from litellm.llms.azure_ai.common_utils import ( azure_ai_supports_native_responses, @@ -69,7 +72,7 @@ def responses_api_bridge_check( reasoning_effort: str | Mapping[str, object] | None = None, reasoning_summary: object | None = None, api_base: str | None = None, -) -> tuple[dict, str]: +) -> tuple[dict[str, object], str]: model_info: dict[str, object] = {} # Global flag: route ALL OpenAI chat completions through Responses API. @@ -186,3 +189,38 @@ def responses_api_bridge_check( model = model.replace("responses/", "") return model_info, model + + +def chat_completion_ends_in_responses_api( + model: str | None, + custom_llm_provider: str | None, + *, + tools: Sequence[Mapping[str, object]] | None = None, + reasoning_effort: str | Mapping[str, object] | None = None, + reasoning_summary: object | None = None, + api_base: str | None = None, +) -> bool: + """Whether chat completion for this model is served by a Responses API. + + Encrypted reasoning can be replayed only there. A normal chat provider cannot + verify the block, so the messages adapter strips it. ``chatgpt``, a ``responses/`` + model, and a request whose tools or reasoning effort the bridge sends to Responses keep it. + """ + if model is None and custom_llm_provider is None: + return False + try: + resolved_model, provider = resolve_model_and_provider_for_metadata(model or "", custom_llm_provider) + except Exception as error: # noqa: BLE001 # unknown model names must not fail message translation + verbose_logger.debug("responses replay check failed for %s: %s", model, error) + return custom_llm_provider == LlmProviders.CHATGPT.value or ( + isinstance(model, str) and (model == "chatgpt" or model.startswith("chatgpt/")) + ) + model_info, _resolved = responses_api_bridge_check( + model=resolved_model, + custom_llm_provider=provider, + tools=tools, + reasoning_effort=reasoning_effort, + reasoning_summary=reasoning_summary, + api_base=api_base, + ) + return model_info.get("mode") == "responses" diff --git a/litellm/responses/main.py b/litellm/responses/main.py index f030105151d..809d933acfd 100644 --- a/litellm/responses/main.py +++ b/litellm/responses/main.py @@ -47,6 +47,7 @@ from litellm.responses.litellm_completion_transformation.handler import ( LiteLLMCompletionTransformationHandler, ) from litellm.responses.mcp.request_context import MCPRequestContext +from litellm.responses.mcp.tool_precheck import may_name_litellm_mcp_gateway from litellm.responses.utils import ResponsesAPIRequestUtils from litellm.types.llms.openai import ( PromptObject, @@ -1134,11 +1135,7 @@ def _responses_try_dispatch_mcp_gateway( | None ): """Return a response when MCP gateway handles the call; otherwise None.""" - from litellm.responses.mcp.litellm_proxy_mcp_handler import ( - LiteLLM_Proxy_MCP_Handler, - ) - - if skip_mcp_handler or not LiteLLM_Proxy_MCP_Handler.should_use_litellm_mcp_gateway(tools=tools): + if skip_mcp_handler or not may_name_litellm_mcp_gateway(tools): return None mcp_call_kwargs: Final = { "model": model, diff --git a/litellm/responses/mcp/litellm_proxy_mcp_handler.py b/litellm/responses/mcp/litellm_proxy_mcp_handler.py index 163caf63f62..e6b2f9ead68 100644 --- a/litellm/responses/mcp/litellm_proxy_mcp_handler.py +++ b/litellm/responses/mcp/litellm_proxy_mcp_handler.py @@ -1,4 +1,3 @@ -import re import traceback from collections.abc import Awaitable, Callable, Collection, Iterable, Mapping, Sequence from dataclasses import dataclass, field @@ -31,6 +30,13 @@ from litellm.proxy._experimental.mcp_server.utils import ( strip_known_server_prefix, ) from litellm.responses.main import aresponses # noqa: TID251 # inner call must skip the MCP gateway that invoked it +from litellm.responses.mcp.tool_precheck import ( + LITELLM_PROXY_MCP_SERVER_URL_PREFIX, + may_name_litellm_mcp_gateway, + names_gateway_explicitly, + names_whole_gateway, + proxy_path_mcp_name, +) from litellm.responses.streaming_iterator import BaseResponsesAPIStreamingIterator from litellm.types.llms.openai import ( ResponseInputParam, @@ -76,44 +82,6 @@ class _AresponsesOptions(ExtraItemsTypedDict, total=False, extra_items=object): pass -LITELLM_PROXY_MCP_SERVER_URL: Final = "litellm_proxy" -LITELLM_PROXY_MCP_SERVER_URL_PREFIX: Final = f"{LITELLM_PROXY_MCP_SERVER_URL}/mcp/" - -_PROXY_MCP_PATH_RE: Final = re.compile(r"^https?://.+/mcp/([^/]+)$") - - -def _mcp_server_url(tool: ToolParam) -> str | None: - if not isinstance(tool, dict) or tool.get("type") != "mcp": - return None - server_url: Final = tool.get("server_url") - return server_url if isinstance(server_url, str) else None - - -def _names_gateway_explicitly(tool: ToolParam) -> bool: - return (_mcp_server_url(tool) or "").startswith(LITELLM_PROXY_MCP_SERVER_URL) - - -def _proxy_path_mcp_name(tool: ToolParam) -> str | None: - server_url: Final = _mcp_server_url(tool) - match: Final = None if server_url is None else _PROXY_MCP_PATH_RE.match(server_url) - return None if match is None else match.group(1) - - -def _server_url_of(tool: object) -> str | None: - match tool: - case {"type": "mcp", "server_url": str() as server_url} if isinstance(tool, dict): - return server_url - case _: - return None - - -def _names_whole_gateway(tool: object) -> bool: - server_url: Final = _server_url_of(tool) or "" - return server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL) and not server_url.startswith( - LITELLM_PROXY_MCP_SERVER_URL_PREFIX - ) - - def _registered_mcp_servers() -> Collection[MCPServer]: from litellm.proxy._experimental.mcp_server.mcp_server_manager import ( global_mcp_server_manager, @@ -159,7 +127,7 @@ async def _gateway_served_names( async def _served_mcp_path_names( tools: Collection[ToolParam], served_names: Callable[[Collection[str]], Awaitable[frozenset[str]]] ) -> frozenset[str]: - names: Final = frozenset(name for name in map(_proxy_path_mcp_name, tools) if name is not None) + names: Final = frozenset(name for name in map(proxy_path_mcp_name, tools) if name is not None) return await served_names(names) if names else frozenset[str]() @@ -229,17 +197,15 @@ class LiteLLM_Proxy_MCP_Handler: @staticmethod def should_use_litellm_mcp_gateway(tools: Iterable[ToolParam] | None) -> bool: - """True when a tool may name this gateway: server_url "litellm_proxy..." or an http(s) URL ending in - /mcp/. `_split_mcp_tools` then settles which of the latter the gateway actually serves.""" - return any(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) is not None for tool in tools or ()) + return may_name_litellm_mcp_gateway(tools) _should_use_litellm_mcp_gateway = should_use_litellm_mcp_gateway @staticmethod def _parse_mcp_tools(tools: Iterable[Mapping[str, object]] | None) -> SplitTools: items: Final = tuple(tools or ()) - gateway_tools: Final[list[ToolParam]] = [tool for tool in items if _names_gateway_explicitly(tool)] - other_tools: Final[list[ToolParam]] = [tool for tool in items if not _names_gateway_explicitly(tool)] + gateway_tools: Final[list[ToolParam]] = [tool for tool in items if names_gateway_explicitly(tool)] + other_tools: Final[list[ToolParam]] = [tool for tool in items if not names_gateway_explicitly(tool)] return gateway_tools, other_tools @staticmethod @@ -252,7 +218,7 @@ class LiteLLM_Proxy_MCP_Handler: return LiteLLM_Proxy_MCP_Handler._parse_mcp_tools( [ {**tool, "server_url": f"{LITELLM_PROXY_MCP_SERVER_URL_PREFIX}{name}"} - if (name := _proxy_path_mcp_name(tool)) in served + if (name := proxy_path_mcp_name(tool)) in served else tool for tool in items ] @@ -268,7 +234,7 @@ class LiteLLM_Proxy_MCP_Handler: ) -> tuple[bool, ...]: items: Final = tuple(tools or ()) served: Final = await _served_mcp_path_names(items, served_names) - return tuple(_names_gateway_explicitly(tool) or _proxy_path_mcp_name(tool) in served for tool in items) + return tuple(names_gateway_explicitly(tool) or proxy_path_mcp_name(tool) in served for tool in items) @staticmethod async def _apply_toolset_permissions( @@ -690,7 +656,7 @@ class LiteLLM_Proxy_MCP_Handler: server_url: Final = tool.get("server_url") if not isinstance(label, str) or not isinstance(server_url, str): return None - if _names_whole_gateway(tool): + if names_whole_gateway(tool): return _LabelScope(label=label, whole_gateway=True) kind, value = await LiteLLM_Proxy_MCP_Handler._resolve_proxy_mcp_name( server_url.split("/")[-1], user_api_key_auth, granted_toolsets diff --git a/litellm/responses/mcp/tool_precheck.py b/litellm/responses/mcp/tool_precheck.py new file mode 100644 index 00000000000..1fa95749ef6 --- /dev/null +++ b/litellm/responses/mcp/tool_precheck.py @@ -0,0 +1,50 @@ +import importlib +import re +from collections.abc import Iterable +from typing import Final + +LITELLM_PROXY_MCP_SERVER_URL: Final = "litellm_proxy" +LITELLM_PROXY_MCP_SERVER_URL_PREFIX: Final = f"{LITELLM_PROXY_MCP_SERVER_URL}/mcp/" + +_PROXY_MCP_PATH_RE: Final = re.compile(r"^https?://.+/mcp/([^/]+)$") +_GATEWAY_REGISTRY_MODULE: Final = "litellm.proxy._experimental.mcp_server.mcp_server_manager" + + +def _mcp_server_url(tool: object) -> str | None: + match tool: + case {"type": "mcp", "server_url": str() as server_url} if isinstance(tool, dict): + return server_url + case _: + return None + + +def names_gateway_explicitly(tool: object) -> bool: + return (_mcp_server_url(tool) or "").startswith(LITELLM_PROXY_MCP_SERVER_URL) + + +def names_whole_gateway(tool: object) -> bool: + server_url: Final = _mcp_server_url(tool) or "" + return server_url.startswith(LITELLM_PROXY_MCP_SERVER_URL) and not server_url.startswith( + LITELLM_PROXY_MCP_SERVER_URL_PREFIX + ) + + +def proxy_path_mcp_name(tool: object) -> str | None: + server_url: Final = _mcp_server_url(tool) + match: Final = None if server_url is None else _PROXY_MCP_PATH_RE.match(server_url) + return None if match is None else match.group(1) + + +def _gateway_registry_loads() -> bool: + try: + importlib.import_module(_GATEWAY_REGISTRY_MODULE) + except ImportError: + return False + return True + + +def may_name_litellm_mcp_gateway(tools: Iterable[object] | None) -> bool: + items: Final = tuple(tools or ()) + if any(map(names_gateway_explicitly, items)): + return True + return any(proxy_path_mcp_name(tool) is not None for tool in items) and _gateway_registry_loads() diff --git a/litellm/router.py b/litellm/router.py index 244fac94d42..9880ca6fdd0 100644 --- a/litellm/router.py +++ b/litellm/router.py @@ -34,6 +34,7 @@ from collections.abc import ( MutableMapping, Sequence, ) +from contextvars import ContextVar from dataclasses import dataclass from datetime import datetime, timezone from functools import lru_cache, partial @@ -223,6 +224,7 @@ from litellm.router_utils.cooldown_handlers import ( is_background_response_poll_not_found, is_caller_scoped_auth_failure, is_caller_timeout_408, + is_video_lookup_not_found, set_cooldown_deployments, ) from litellm.router_utils.deployment_pricing import deployment_pricing_models @@ -370,6 +372,7 @@ from litellm.types.utils import ( shared_backend_model_info, ) from litellm.types.utils import ModelInfo as ModelMapInfo +from litellm.types.videos.utils import VIDEO_ID_CALL_TYPES, decode_video_id_with_provider from litellm.utils import ( CustomStreamWrapper, EmbeddingResponse, @@ -488,6 +491,7 @@ def model_info_is_active_for_environment(model_info: Mapping[str, object] | None _PreRoutingStrategyT = TypeVar("_PreRoutingStrategyT") _CallbackT = TypeVar("_CallbackT") +_DeploymentT = TypeVar("_DeploymentT", bound=Mapping[str, object]) _ALIAS_PARAMS_NEVER_FORWARDED: Final = frozenset({"model", "api_base", "api_key", "api_version"}) _ALIAS_MARKER_FORWARDED_PARAMS_KWARG: Final = "_alias_marker_forwarded_params" @@ -496,6 +500,7 @@ _CLAUDE_CODE_SESSION_ROUTER_TTL_SECONDS: Final = 3600 CLAUDE_CODE_SESSION_ROUTER_BINDING_TARGET: Final = "claude_code_session_router_binding" _PRE_CALL_CHECK_FILTERED_PARAMS: Final = frozenset({"response_format"}) +_video_call_deployment_id: Final[ContextVar[str | None]] = ContextVar("video_call_deployment_id", default=None) _RUNTIME_TOGGLEABLE_PRE_CALL_CHECKS: Final[Mapping[str, type[CustomLogger]]] = MappingProxyType( { "prompt_caching": PromptCachingDeploymentCheck, @@ -5717,17 +5722,19 @@ class Router: passthrough_on_no_deployment: Final = kwargs.pop("passthrough_on_no_deployment", False) function_name: Final = "_ageneric_api_call_with_fallbacks" + handler_name: Final[str | None] = getattr(original_generic_function, "__name__", None) deployment = None # rebind-ok: pre-init so the except block can stamp a failure with no deployment picked try: parent_otel_span: Final = get_parent_otel_span_from_kwargs(kwargs) try: - deployment = await self.async_get_available_deployment( # rebind-ok: set on success, see pre-init above - model=model, - request_kwargs=kwargs, - messages=kwargs.get("messages", None), - input=kwargs.get("input", None), - specific_deployment=kwargs.pop("specific_deployment", None), - ) + with self._video_call_pin(handler_name, kwargs): + deployment = await self.async_get_available_deployment( # rebind-ok: set on success, see pre-init + model=model, + request_kwargs=kwargs, + messages=kwargs.get("messages", None), + input=kwargs.get("input", None), + specific_deployment=kwargs.pop("specific_deployment", None), + ) except Exception as e: if passthrough_on_no_deployment: return await original_generic_function(model=model, **kwargs) @@ -6409,6 +6416,37 @@ class Router: ) return await self._ageneric_api_call_with_fallbacks(original_function=original_function, **kwargs) + @staticmethod + @contextlib.contextmanager + def _video_call_pin(handler_name: str | None, request_kwargs: Mapping[str, object]) -> Generator[None]: + """ + Scopes the deployment a video id names to the routing of one call on that video (status, content, remix, + edit, extension). Any other call routes as if no id were sent, whatever its body carries. + """ + video_id: Final = request_kwargs.get("video_id") + deployment_id: Final = ( + decode_video_id_with_provider(video_id).get("deployment_id") + if handler_name in VIDEO_ID_CALL_TYPES and isinstance(video_id, str) + else None + ) + token: Final = _video_call_deployment_id.set(deployment_id) + try: + yield + finally: + _video_call_deployment_id.reset(token) + + @staticmethod + def _narrow_to_video_deployment(healthy_deployments: Sequence[_DeploymentT]) -> Sequence[_DeploymentT]: + """ + A video lives on the provider account of the deployment that created it, so a call on an id naming a + deployment that every routing filter kept goes to that deployment, else to whichever one routing picks. + """ + deployment_id: Final = _video_call_deployment_id.get() + creating: Final = tuple( + deployment for deployment in healthy_deployments if deployment_id in Router._deployment_ids((deployment,)) + ) + return creating or healthy_deployments + def _generic_api_call_with_fallbacks(self, model: str, original_function: Callable, **kwargs): """ Make a generic LLM API call through the router, this allows you to use retries/fallbacks with litellm router @@ -6430,13 +6468,14 @@ class Router: kwargs=kwargs, metadata_variable_name=metadata_variable_name, ) - deployment: Final = self.get_available_deployment( - model=model, - messages=kwargs.get("messages", None), - input=kwargs.get("input", None), - specific_deployment=kwargs.pop("specific_deployment", None), - request_kwargs=kwargs, - ) + with self._video_call_pin(handler_name, kwargs): + deployment: Final = self.get_available_deployment( + model=model, + messages=kwargs.get("messages", None), + input=kwargs.get("input", None), + specific_deployment=kwargs.pop("specific_deployment", None), + request_kwargs=kwargs, + ) self._update_kwargs_with_deployment(deployment=deployment, kwargs=kwargs, function_name="generic_api_call") data: Final = deployment["litellm_params"].copy() @@ -8914,6 +8953,13 @@ class Router: ) return False + if is_video_lookup_not_found(exception, kwargs): + verbose_router_logger.debug( + "Router: Exiting 'deployment_callback_on_failure' without cooldown. " + "Provider 404 says the video is not on this deployment's account, not that the deployment is down." + ) + return False + exception_status: Final = getattr(exception, "status_code", "") if is_caller_timeout_408(kwargs, exception_status): @@ -12337,7 +12383,8 @@ class Router: 1. First, check if model_id directly matches a model_name or deployment ID 2. If not, search through router's model_list to find a match by litellm_params.model 3. If not, keep model_id when the router routes that name itself (an alias, a routing group, or a wildcard - whose deployment serves `custom_llm_provider`) + whose deployment serves `custom_llm_provider`), else name it `custom_llm_provider/model_id` when only that + form routes to such a wildcard 4. Return the model_name if found, None otherwise Args: @@ -12385,8 +12432,15 @@ class Router: if self.is_recognized_model(model_id): return model_id - if custom_llm_provider is not None and self.pattern_router.routes_to_provider(model_id, custom_llm_provider): + if custom_llm_provider is None: + return None + if self.pattern_router.routes_to_provider(model_id, custom_llm_provider): return model_id + provider_prefixed: Final = f"{custom_llm_provider}/{model_id}" + if not model_id.startswith(f"{custom_llm_provider}/") and self.pattern_router.routes_to_provider( + provider_prefixed, custom_llm_provider + ): + return provider_prefixed return None def resolve_served_model_name(self, model_name: str, served_model: str) -> str: @@ -12915,13 +12969,13 @@ class Router: return [{**DeploymentTypedDict(**m), "model_name": model_name} for m in matches] def wildcard_route_deployments(self, model: str, team_id: str | None) -> Sequence[Mapping[str, object]]: - """The deployments routing serves ``model`` from through a wildcard: the most specific - proxy-wide pattern matching it, else ``team_id``'s own pattern, never both.""" - proxy_wide: Final[Sequence[Mapping[str, object]]] = self.pattern_router.get_deployments_by_pattern(model=model) + """The deployments routing serves ``model`` from through a wildcard: ``team_id``'s own pattern + matching it, else the most specific proxy-wide pattern, never both.""" team_router: Final = self.team_pattern_routers.get(team_id) if team_id is not None else None - if proxy_wide or team_router is None: - return proxy_wide - return team_router.get_deployments_by_pattern(model=model) + team_owned: Final[Sequence[Mapping[str, object]]] = ( + team_router.get_deployments_by_pattern(model=model) if team_router is not None else () + ) + return team_owned or self.pattern_router.get_deployments_by_pattern(model=model) def team_ptu_ceiling(self, team_id: str, model_group: str) -> PTUTeamCeiling | None: """The per-minute ceiling ``team_id``'s PTU shares set on the group serving ``model_group``, @@ -13897,10 +13951,10 @@ class Router: def _filter_deployments_by_model_access_groups( self, model: str, - healthy_deployments: list, + healthy_deployments: list[DeploymentTypedDict], request_kwargs: Mapping[str, object] | None, request_team_id: str | None, - ) -> list: + ) -> list[DeploymentTypedDict]: """ Restrict candidate deployments to caller-authorized model access groups. @@ -14084,7 +14138,8 @@ class Router: ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) _target_order: Final = (request_kwargs or {}).pop("_target_order", None) healthy_deployments = litellm.utils.get_order_filtered_deployments( - cast(list[dict], healthy_deployments), target_order=_target_order + list(self._narrow_to_video_deployment(cast(list[dict[str, object]], healthy_deployments))), + target_order=_target_order, ) ## WEIGHTED FAILOVER EXCLUSION ## -> drop deployments already tried in @@ -15271,7 +15326,7 @@ class Router: ## ORDER FILTERING ## -> if user set 'order' in deployments, return deployments with lowest order (e.g. order=1 > order=2) _target_order: Final = (request_kwargs or {}).pop("_target_order", None) healthy_deployments = litellm.utils.get_order_filtered_deployments( - healthy_deployments, target_order=_target_order + list(self._narrow_to_video_deployment(healthy_deployments)), target_order=_target_order ) ## WEIGHTED FAILOVER EXCLUSION ## -> drop deployments already tried in diff --git a/litellm/router_utils/cooldown_handlers.py b/litellm/router_utils/cooldown_handlers.py index cdca5009b0e..20ad8a6e163 100644 --- a/litellm/router_utils/cooldown_handlers.py +++ b/litellm/router_utils/cooldown_handlers.py @@ -35,6 +35,7 @@ from litellm.types.utils import ( BACKGROUND_RESPONSE_STATUS_POLL_CALL_ORIGIN, InternalCallOrigin, ) +from litellm.types.videos.utils import VIDEO_ID_CALL_TYPES from .router_callbacks.track_deployment_metrics import ( get_deployment_failures_for_current_minute, @@ -89,6 +90,11 @@ def is_background_response_poll_not_found(exception: Exception, litellm_params: ) +def is_video_lookup_not_found(exception: object, failure_kwargs: Mapping[str, object]) -> bool: + """Whether a call on an existing video id failed with a provider 404: the video is not on that account.""" + return getattr(exception, "status_code", None) == 404 and failure_kwargs.get("call_type") in VIDEO_ID_CALL_TYPES + + _EXCEPTION_POLICY_FIELDS: Final[tuple[tuple[type, str], ...]] = ( # ContentPolicyViolationError subclasses BadRequestError, so it must be checked first. (litellm.ContentPolicyViolationError, "ContentPolicyViolationErrorAllowedFails"), diff --git a/litellm/router_utils/pattern_match_deployments.py b/litellm/router_utils/pattern_match_deployments.py index 62ba5592f90..daa052a46e1 100644 --- a/litellm/router_utils/pattern_match_deployments.py +++ b/litellm/router_utils/pattern_match_deployments.py @@ -173,7 +173,7 @@ class PatternMatchRouter: return None # No matching pattern found def routes_to_provider(self, request: str, custom_llm_provider: str) -> bool: - deployments: Final = _ROUTED_DEPLOYMENTS.validate_python(self.route(request) or ()) + deployments: Final = _ROUTED_DEPLOYMENTS.validate_python(self.get_pattern(request) or ()) return any( (deployment.litellm_params.custom_llm_provider or inferred_provider(deployment.litellm_params.model)) == custom_llm_provider diff --git a/litellm/types/llms/anthropic.py b/litellm/types/llms/anthropic.py index 53abac75162..2f59c1765a5 100644 --- a/litellm/types/llms/anthropic.py +++ b/litellm/types/llms/anthropic.py @@ -545,7 +545,9 @@ class ContentBlockStartText(TypedDict): content_block: TextBlock -ContentBlockContentBlockDict = ToolUseBlock | TextBlock | ChatCompletionThinkingBlock +ContentBlockContentBlockDict = ( + ToolUseBlock | TextBlock | ChatCompletionThinkingBlock | ChatCompletionRedactedThinkingBlock +) ContentBlockStart = ContentBlockStartToolUse | ContentBlockStartText diff --git a/litellm/types/videos/main.py b/litellm/types/videos/main.py index 05444b1e7ec..4d6f8fc6a30 100644 --- a/litellm/types/videos/main.py +++ b/litellm/types/videos/main.py @@ -113,6 +113,8 @@ class DecodedVideoId(TypedDict, total=False): model_id: str | None video_id: str served_model: ReadOnly[str | None] + deployment_id: ReadOnly[str | None] + requested_model: ReadOnly[str | None] class CharacterObject(LiteLLMBaseModel): diff --git a/litellm/types/videos/utils.py b/litellm/types/videos/utils.py index 9a07c9ba802..0fc2b8ede1b 100644 --- a/litellm/types/videos/utils.py +++ b/litellm/types/videos/utils.py @@ -6,6 +6,7 @@ Format: vid_{base64_encoded_string} """ import base64 +from collections.abc import Sequence from typing import Final from litellm._logging import verbose_logger @@ -14,6 +15,22 @@ from litellm.types.videos.main import DecodedVideoId VIDEO_ID_PREFIX: Final = "video_" VIDEO_ID_SERVED_MODEL_PREFIX: Final = "served_model:" +VIDEO_ID_DEPLOYMENT_PREFIX: Final = "deployment_id:" +VIDEO_ID_REQUESTED_MODEL_PREFIX: Final = "requested_model:" +VIDEO_ID_CALL_TYPES: Final = frozenset( + { + "avideo_status", + "video_status", + "avideo_content", + "video_content", + "avideo_remix", + "video_remix", + "avideo_edit", + "video_edit", + "avideo_extension", + "video_extension", + } +) CHARACTER_ID_PREFIX: Final = "character_" CHARACTER_ID_TEMPLATE: Final = "litellm:custom_llm_provider:{};model_id:{};character_id:{}" @@ -37,7 +54,12 @@ def _add_base64_padding(value: str) -> str: def encode_video_id_with_provider( - video_id: str, provider: str, model_id: str | None = None, served_model: str | None = None + video_id: str, + provider: str, + model_id: str | None = None, + served_model: str | None = None, + deployment_id: str | None = None, + requested_model: str | None = None, ) -> str: """Encode provider and model_id into video_id using base64.""" if not provider or not video_id: @@ -54,12 +76,24 @@ def encode_video_id_with_provider( # ID is not encoded (even if it starts with video_), so encode it assembled_id = str(SpecialEnums.LITELLM_MANAGED_VIDEO_COMPLETE_STR.value).format(provider, model_id or "", video_id) served_part: Final = f";{VIDEO_ID_SERVED_MODEL_PREFIX}{served_model}" if served_model else "" + deployment_part: Final = f";{VIDEO_ID_DEPLOYMENT_PREFIX}{deployment_id}" if deployment_id else "" + requested_part: Final = f";{VIDEO_ID_REQUESTED_MODEL_PREFIX}{requested_model}" if requested_model else "" - base64_encoded_id: Final[str] = base64.b64encode((assembled_id + served_part).encode("utf-8")).decode("utf-8") + base64_encoded_id: Final[str] = base64.b64encode( + (assembled_id + served_part + deployment_part + requested_part).encode("utf-8") + ).decode("utf-8") return f"{VIDEO_ID_PREFIX}{base64_encoded_id}" +def id_safe_name(name: object) -> str | None: + return name if isinstance(name, str) and name and ";" not in name else None + + +def _trailing_part(parts: Sequence[str], prefix: str) -> str | None: + return next((part.removeprefix(prefix) or None for part in parts[3:] if part.startswith(prefix)), None) + + def decode_video_id_with_provider(encoded_video_id: str) -> DecodedVideoId: """Decode provider and model_id from encoded video_id.""" if not encoded_video_id: @@ -80,43 +114,6 @@ def decode_video_id_with_provider(encoded_video_id: str) -> DecodedVideoId: cleaned_id = encoded_video_id.replace(VIDEO_ID_PREFIX, "") cleaned_id = _add_base64_padding(cleaned_id) decoded_id: Final = base64.b64decode(cleaned_id.encode("utf-8")).decode("utf-8") - - if ";" not in decoded_id: - return DecodedVideoId( - custom_llm_provider=None, - model_id=None, - video_id=encoded_video_id, - ) - - parts: Final = decoded_id.split(";") - - custom_llm_provider = None - model_id = None - decoded_video_id = encoded_video_id - - if len(parts) >= 3: - custom_llm_provider_part: Final = parts[0] - model_id_part: Final = parts[1] - video_id_part: Final = parts[2] - - custom_llm_provider = custom_llm_provider_part.replace("litellm:custom_llm_provider:", "") - model_id = model_id_part.replace("model_id:", "") - decoded_video_id = video_id_part.replace("video_id:", "") - - served_model: Final = ( - parts[3].removeprefix(VIDEO_ID_SERVED_MODEL_PREFIX) or None - if len(parts) >= 4 and parts[3].startswith(VIDEO_ID_SERVED_MODEL_PREFIX) - else None - ) - - if served_model is None: - return DecodedVideoId(custom_llm_provider=custom_llm_provider, model_id=model_id, video_id=decoded_video_id) - return DecodedVideoId( - custom_llm_provider=custom_llm_provider, - model_id=model_id, - video_id=decoded_video_id, - served_model=served_model, - ) except Exception as e: verbose_logger.debug("Error decoding video_id '%s': %s", encoded_video_id, e) return DecodedVideoId( @@ -125,6 +122,42 @@ def decode_video_id_with_provider(encoded_video_id: str) -> DecodedVideoId: video_id=encoded_video_id, ) + if ";" not in decoded_id: + return DecodedVideoId( + custom_llm_provider=None, + model_id=None, + video_id=encoded_video_id, + ) + + parts: Final = decoded_id.split(";") + + custom_llm_provider = None + model_id = None + decoded_video_id = encoded_video_id + + if len(parts) >= 3: + custom_llm_provider_part: Final = parts[0] + model_id_part: Final = parts[1] + video_id_part: Final = parts[2] + + custom_llm_provider = custom_llm_provider_part.replace("litellm:custom_llm_provider:", "") + model_id = model_id_part.replace("model_id:", "") + decoded_video_id = video_id_part.replace("video_id:", "") + + served_model: Final = _trailing_part(parts, VIDEO_ID_SERVED_MODEL_PREFIX) + deployment_id: Final = _trailing_part(parts, VIDEO_ID_DEPLOYMENT_PREFIX) + requested_model: Final = _trailing_part(parts, VIDEO_ID_REQUESTED_MODEL_PREFIX) + decoded_parts: Final = DecodedVideoId( + custom_llm_provider=custom_llm_provider, model_id=model_id, video_id=decoded_video_id + ) + with_served: Final[DecodedVideoId] = ( + decoded_parts if served_model is None else {**decoded_parts, "served_model": served_model} + ) + with_deployment: Final[DecodedVideoId] = ( + with_served if deployment_id is None else {**with_served, "deployment_id": deployment_id} + ) + return with_deployment if requested_model is None else {**with_deployment, "requested_model": requested_model} + def extract_original_video_id(encoded_video_id: str) -> str: """Extract original video ID without encoding.""" diff --git a/litellm/utils.py b/litellm/utils.py index 1033ddc1d3d..35bf2518d35 100644 --- a/litellm/utils.py +++ b/litellm/utils.py @@ -51,7 +51,7 @@ import litellm.litellm_core_utils # audio_utils.utils is lazy-loaded - only imported when needed for transcription calls import litellm.litellm_core_utils.json_validation_rule -from litellm._internal_context import is_internal_call +from litellm._internal_context import is_internal_call, is_resent_failure from litellm._lazy_imports import ( _get_modified_max_tokens, get_default_encoding, @@ -1985,12 +1985,13 @@ def client(original_function): retry_policy=kwargs.get("retry_policy"), ) kwargs["retry_policy"] = reset_retry_policy() # prevent infinite loops - litellm.num_retries = None # set retries to None to prevent infinite loops + if not is_resent_failure(e): + litellm.num_retries = None # set retries to None to prevent infinite loops context_window_fallback_dict: Final = kwargs.get("context_window_fallback_dict", {}) is_completion_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=False) if ( - num_retries and not is_completion_litellm_router_call + num_retries and not is_completion_litellm_router_call and not is_resent_failure(e) ): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying if ( isinstance(e, openai.APIError) @@ -2020,11 +2021,12 @@ def client(original_function): retry_policy=kwargs.get("retry_policy"), ) kwargs["retry_policy"] = reset_retry_policy() # prevent infinite loops - litellm.num_retries = None # set retries to None to prevent infinite loops + if not is_resent_failure(e): + litellm.num_retries = None # set retries to None to prevent infinite loops is_responses_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=False) if ( - num_retries and not is_responses_litellm_router_call + num_retries and not is_responses_litellm_router_call and not is_resent_failure(e) ): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying if ( isinstance(e, openai.APIError) @@ -2037,7 +2039,7 @@ def client(original_function): end_time = datetime.datetime.now() # LOG FAILURE - handle streaming failure logging in the _next_ object, remove `handle_failure` once it's deprecated - if logging_obj: + if logging_obj and not is_resent_failure(e): logging_obj.failure_handler( e, traceback_exception, start_time, end_time ) # DO NOT MAKE THREADED - router retry fallback relies on this! @@ -2284,7 +2286,7 @@ def client(original_function): # Reuse the timestamp taken right when the deployment call itself failed, before # the failure hook ran, so a slow callback doesn't inflate the reported duration. end_time = _deployment_call_end_time if _deployment_call_end_time is not None else datetime.datetime.now() # noqa: DTZ005 # matches the naive datetimes this whole function already times start_time/end_time with - if logging_obj and not _is_litellm_internal_call: + if logging_obj and not _is_litellm_internal_call and not is_resent_failure(e): unbind_budget_reservation_from_callbacks(logging_obj.litellm_params) try: logging_obj.failure_handler( @@ -2305,7 +2307,10 @@ def client(original_function): is_acompletion_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=True) if ( - num_retries and not is_acompletion_litellm_router_call and not isinstance(e, ImportError) + num_retries + and not is_acompletion_litellm_router_call + and not isinstance(e, ImportError) + and not is_resent_failure(e) ): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying try: litellm.num_retries = None # set retries to None to prevent infinite loops @@ -2316,8 +2321,9 @@ def client(original_function): elif isinstance(e, openai.APIError): # generic api error kwargs["retry_strategy"] = "constant_retry" result = await litellm.acompletion_with_retries(*args, **kwargs) - except Exception: - pass + except Exception as retry_error: + if is_resent_failure(retry_error): + raise else: return result elif ( @@ -2336,7 +2342,10 @@ def client(original_function): is_aresponses_litellm_router_call: Final = _is_litellm_router_call(kwargs, is_async=True) if ( - num_retries and not is_aresponses_litellm_router_call and not isinstance(e, ImportError) + num_retries + and not is_aresponses_litellm_router_call + and not isinstance(e, ImportError) + and not is_resent_failure(e) ): # only enter this if call is not from litellm router/proxy. router has it's own logic for retrying try: litellm.num_retries = None # set retries to None to prevent infinite loops @@ -2347,8 +2356,9 @@ def client(original_function): elif isinstance(e, openai.APIError): # generic api error kwargs["retry_strategy"] = "constant_retry" result = await litellm.aresponses_with_retries(*args, **kwargs) - except Exception: - pass + except Exception as retry_error: + if is_resent_failure(retry_error): + raise else: return result diff --git a/pyproject.toml b/pyproject.toml index fe156695894..84200cb91a3 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "litellm" -version = "1.108.2" +version = "1.108.3" description = "Library to easily interface with LLM API providers" readme = "README.md" requires-python = ">=3.10, <3.15" @@ -396,7 +396,7 @@ litellm-enterprise = { workspace = true } members = ["enterprise", "litellm-proxy-extras"] [tool.commitizen] -version = "1.108.2" +version = "1.108.3" version_files = [ "pyproject.toml:^version", ] diff --git a/tests/base_sdk_tests/check_base_sdk_install.py b/tests/base_sdk_tests/check_base_sdk_install.py index 3abea2a9707..4abd71e39a6 100644 --- a/tests/base_sdk_tests/check_base_sdk_install.py +++ b/tests/base_sdk_tests/check_base_sdk_install.py @@ -102,6 +102,29 @@ def check_completion() -> str: return "mock completion round-trips" +def check_completion_with_function_tool() -> str: + import litellm + + response: Final = litellm.completion( + model="gpt-4o", + messages=[{"role": "user", "content": "ping"}], + tools=[ + { + "type": "function", + "function": { + "name": "noop", + "description": "No operation", + "parameters": {"type": "object", "properties": {}}, + }, + } + ], + mock_response="pong", + ) + content: Final = response.choices[0].message.content + _require(content == "pong", f"mock completion with a function tool returned {content!r}") + return "mock completion with a function tool round-trips" + + def check_mcp_install_guidance() -> str: try: import litellm.experimental_mcp_client @@ -209,6 +232,7 @@ CHECKS: tuple[tuple[str, Callable[[], str]], ...] = ( ("import litellm", check_import), ("optional MCP installation guidance", check_mcp_install_guidance), ("chat completion", check_completion), + ("chat completion with a function tool", check_completion_with_function_tool), ("embedding", check_embedding), ("bundled model metadata", check_bundled_model_metadata), ("token counter", check_token_counter), diff --git a/tests/base_sdk_tests/check_sdk_http.py b/tests/base_sdk_tests/check_sdk_http.py index c93306cdced..f7a1e30ac99 100644 --- a/tests/base_sdk_tests/check_sdk_http.py +++ b/tests/base_sdk_tests/check_sdk_http.py @@ -89,7 +89,6 @@ def payload(response: object) -> Mapping[str, JsonValue]: return JSON_OBJECT.validate_json(response.model_dump_json(exclude_none=True)) -@skipIf(importlib.util.find_spec("mcp") is None, "BUG: LIT-9333 completion(tools=...) imports optional mcp") def check_tools(base: str, asynchronous: bool = False) -> None: tools: Final = [ { @@ -122,7 +121,6 @@ def check_tools(base: str, asynchronous: bool = False) -> None: assert json.loads(body) == {"model": "gpt-6", "messages": MESSAGES, "tools": tools} -@skipIf(importlib.util.find_spec("mcp") is None, "BUG: LIT-9333 responses() imports optional mcp") def check_responses(base: str, asynchronous: bool = False) -> None: output: Final[Sequence[JsonValue]] = [ { @@ -144,14 +142,62 @@ def check_responses(base: str, asynchronous: bool = False) -> None: "usage": {"input_tokens": 3, "output_tokens": 2, "total_tokens": 5}, } ) - arguments: Final = dict(model="openai/gpt-6", input="ping", api_key="test-key", api_base=base + "/v1") + tools: Final = [ + {"type": "mcp", "server_label": "docs", "server_url": "https://mcp.example.com/sse"}, + {"type": "mcp", "server_label": "repos", "server_url": "https://mcp.example.com/mcp/repos"}, + ] + arguments: Final = dict(model="openai/gpt-6", input="ping", tools=tools, api_key="test-key", api_base=base + "/v1") response: Final = asyncio.run(litellm.aresponses(**arguments)) if asynchronous else litellm.responses(**arguments) assert payload(response)["output"] == output assert payload(response)["usage"] == {"input_tokens": 3, "output_tokens": 2, "total_tokens": 5} path, headers, body = REQUESTS.get(timeout=10) assert path == "/v1/responses" assert headers["authorization"] == "Bearer test-key" - assert json.loads(body) == {"model": "gpt-6", "input": "ping"} + assert json.loads(body) == {"model": "gpt-6", "input": "ping", "tools": tools} + + +def check_messages_tools(base: str) -> None: + tools: Final = [ + { + "name": "lookup", + "description": "Look up the forecast for a city", + "input_schema": {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]}, + } + ] + content: Final = [{"type": "tool_use", "id": "toolu-http-check", "name": "lookup", "input": {"city": "Paris"}}] + enqueue( + { + "id": "msg-tools-http-check", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": content, + "stop_reason": "tool_use", + "stop_sequence": None, + "usage": {"input_tokens": 3, "output_tokens": 2}, + } + ) + + async def send() -> object: + create: Final = cast( # cast-ok: SDK facade exposes untyped kwargs + Callable[..., Awaitable[object]], litellm.anthropic.messages.acreate + ) + return await create( + model="anthropic/claude-sonnet-5-5", + messages=MESSAGES, + tools=tools, + max_tokens=16, + api_key="test-key", + api_base=base, + ) + + response: Final = JSON_OBJECT.validate_python(asyncio.run(send())) + assert response["content"] == content, response + path, headers, body = REQUESTS.get(timeout=10) + assert path == "/v1/messages" + assert headers["x-api-key"] == "test-key" + sent: Final = JSON_OBJECT.validate_json(body) + assert (sent["model"], sent["messages"], sent["tools"]) == ("claude-sonnet-5-5", MESSAGES, tools), sent @skipIf(importlib.util.find_spec("fastapi") is None, "BUG: LIT-9333 cache hits import proxy-only fastapi") @@ -392,6 +438,7 @@ def main() -> None: partial(check_tools, asynchronous=True), check_responses, partial(check_responses, asynchronous=True), + check_messages_tools, check_cache_hit, check_messages_stream_and_callback, ): diff --git a/tests/code_coverage_tests/router_code_coverage.py b/tests/code_coverage_tests/router_code_coverage.py index 9bd077ef3a5..8e84cf2e3e9 100644 --- a/tests/code_coverage_tests/router_code_coverage.py +++ b/tests/code_coverage_tests/router_code_coverage.py @@ -99,6 +99,8 @@ ignored_function_names = [ "_embedding", "_aembedding", "_stream_pre_content_error", # Tested through the non-retriable retry error tests in test_router.py + "_narrow_to_video_deployment", # Tested through video_status and avideo_status polls in test_router_provider_endpoints.py + "_video_call_pin", # Tested through the video_status pin tests and the chat video_id test in test_router_provider_endpoints.py "_deployment_num_retries", # Tested through the deployment num_retries mid-stream budget test in test_router.py "_request_fallback_list", # Tested through every mid-stream retry test in test_router.py "_request_model_group", # Tested through test_anthropic_messages_retry_budget_precedence_direct_call diff --git a/tests/e2e/llm_translation/test_together_ai_e2e.py b/tests/e2e/llm_translation/test_together_ai_e2e.py index 2d5d5f2d93a..7775c5c6ed2 100644 --- a/tests/e2e/llm_translation/test_together_ai_e2e.py +++ b/tests/e2e/llm_translation/test_together_ai_e2e.py @@ -7,9 +7,9 @@ backends are pinned because the registry has no flag for what they prove: ``enab and the ``{"reasoning": {"enabled": false}}`` toggle that ``reasoning_effort="none"`` maps to are Qwen hybrid-model contracts, and the structured-output case runs on that hybrid model with reasoning off, since a reasoning-only model can spend the whole token budget -thinking and return no content. MiniMax-M3 is the serverless model whose +thinking and return no content. MiniMax-M3.1-Flash is the serverless model whose template renders a replayed ``reasoning_content`` back into the prompt (Qwen and -DeepSeek silently drop it). MiniMax-M3 honors that replayed field on nearly every call, not +DeepSeek silently drop it). MiniMax-M3.1-Flash honors that replayed field on nearly every call, not every call (one miss in dozens of otherwise identical calls), so the replay case asks up to ``REPLAY_ATTEMPTS`` times and fails only when no answer carries the secret, which a proxy that strips the field guarantees. Requires TOGETHER_API_KEY on the proxy; no @@ -56,7 +56,7 @@ from pydantic import BaseModel pytestmark = pytest.mark.e2e HYBRID_REASONING_BACKEND = "together_ai/Qwen/Qwen3.5-9B" -REASONING_REPLAY_BACKEND = "together_ai/MiniMaxAI/MiniMax-M3" +REASONING_REPLAY_BACKEND = "together_ai/MiniMaxAI/MiniMax-M3.1-Flash" SECRET_PROMPT = "Remember this for later and reply with just OK." SECRET_REASONING = "The user told me their favorite color is chartreuse. I must remember it." diff --git a/tests/integration/_support/client.py b/tests/integration/_support/client.py index 053318c2451..fe1db22987d 100644 --- a/tests/integration/_support/client.py +++ b/tests/integration/_support/client.py @@ -95,6 +95,10 @@ class Gateway: client: httpx.Client key: str upstream_url: str + database_url: str | None = None + + def rows(self, query: str, parameters: tuple[str, ...]) -> list[dict[str, JsonValue]]: + return read_rows(query, parameters, database_url=self.database_url) def request( self, @@ -267,12 +271,14 @@ class Scenario: ) assert read_rows('SELECT model_id FROM "LiteLLM_ProxyModelTable" WHERE model_id = %s', (identity,)) == [] - def model(self, *, model_info: Mapping[str, JsonValue] | None = None, **parameters: JsonValue) -> str: - name: Final = f"integration-{uuid.uuid4().hex}" + def model( + self, *, name: str | None = None, model_info: Mapping[str, JsonValue] | None = None, **parameters: JsonValue + ) -> str: + model_name: Final = name if name is not None else f"integration-{uuid.uuid4().hex}" created: Final = self.gateway.post( "/model/new", { - "model_name": name, + "model_name": model_name, "litellm_params": { "model": "openai/gpt-4o-mini", "api_key": "integration-provider-key", @@ -284,7 +290,7 @@ class Scenario: ) identity: Final = string_value(object_value(created["model_info"])["id"]) self.cleanups.callback(self.delete_model, identity) - return name + return model_name @contextmanager diff --git a/tests/integration/mcp/test_mcp_llm_endpoints.py b/tests/integration/mcp/test_mcp_llm_endpoints.py index c3bccec1f8e..c1517a48c1c 100644 --- a/tests/integration/mcp/test_mcp_llm_endpoints.py +++ b/tests/integration/mcp/test_mcp_llm_endpoints.py @@ -468,6 +468,61 @@ def test_server_scoped_gateway_url_exposes_only_that_servers_tools(gateway: Gate assert [call["body"]["params"]["name"] for call in _peer_add_calls(rig.peer)] == ["add"] +@pytest.mark.parametrize("surface", ("chat", "responses", "messages")) +def test_proxy_path_url_of_a_served_server_is_listed_and_executed_by_the_gateway( + gateway: Gateway, surface: Surface +) -> None: + with _rig(gateway, surface) as rig: + server_url: Final = f"{str(gateway.client.base_url).rstrip('/')}/mcp/{rig.alias}" + response: Final = rig.send(_granted_key(rig), [{**AUTO, "server_url": server_url}]) + assert response.status_code == 200, response.text + requests: Final = rig.upstream_tools() + assert len(requests) == 2 and all(rig.tool in names for names in requests), requests + peer_calls: Final = tuple(JSON_OBJECT.validate_python(call["body"]) for call in _peer_add_calls(rig.peer)) + assert [object_value(call["params"])["name"] for call in peer_calls] == ["add"], peer_calls + assert rig.final_text(JSON_OBJECT.validate_json(response.content)) == ANSWER, response.text + + +LOOKUP_PARAMETERS: Final = {"type": "object", "properties": {"city": {"type": "string"}}, "required": ["city"]} +CALLER_FUNCTION_TOOLS: Final[Mapping[Surface, Mapping[str, object]]] = { + "chat": {"type": "function", "function": {"name": "lookup", "parameters": LOOKUP_PARAMETERS}}, + "responses": {"type": "function", "name": "lookup", "parameters": LOOKUP_PARAMETERS}, + "messages": {"name": "lookup", "input_schema": LOOKUP_PARAMETERS}, +} + + +@pytest.mark.parametrize("surface", ("chat", "responses", "messages")) +def test_tools_that_do_not_name_the_gateway_reach_the_provider_unchanged_and_never_the_peer( + gateway: Gateway, surface: Surface +) -> None: + with _rig(gateway, surface) as rig: + remote_mcp_tools: Final = [ + { + "type": "mcp", + "server_label": "docs", + "server_url": "https://mcp.example.com/sse", + "require_approval": "never", + }, + { + "type": "mcp", + "server_label": "repos", + "server_url": f"https://mcp.example.com/mcp/unserved{uuid.uuid4().hex[:8]}", + "require_approval": "never", + }, + ] + response: Final = rig.send(_granted_key(rig), [CALLER_FUNCTION_TOOLS[surface], *remote_mcp_tools]) + assert response.status_code == 200, response.text + bodies: Final = tuple( + JSON_OBJECT.validate_json(request.body) for request in rig.wire.drain() if request.method == "POST" + ) + assert len(bodies) == 1, bodies + forwarded: Final = bodies[0].get("tools") + assert isinstance(forwarded, list) and len(forwarded) == 3, forwarded + assert [tool for tool in forwarded if isinstance(tool, dict) and tool.get("type") == "mcp"] == remote_mcp_tools + assert "lookup" in json.dumps([tool for tool in forwarded if tool not in remote_mcp_tools]), forwarded + assert tool_calls(rig.peer.drain()) == (), "a tool that does not name the gateway reached the peer" + + def test_streaming_chat_executes_the_tool_once_and_streams_the_follow_up(gateway: Gateway) -> None: with _rig(gateway, "chat") as rig: key: Final = _granted_key(rig) diff --git a/tests/integration/mcp/test_mcp_transports.py b/tests/integration/mcp/test_mcp_transports.py index ef240cba570..be9e016e947 100644 --- a/tests/integration/mcp/test_mcp_transports.py +++ b/tests/integration/mcp/test_mcp_transports.py @@ -678,7 +678,7 @@ def test_sdk_progress_and_cancellation_reach_upstream_without_losing_session( from starlette.routing import Mount from integration._support.asgi import asgi_server - from integration._support.mcp import _drain_sse_streams + from integration._support.mcp import _acknowledge_notifications_after_dispatch, _drain_sse_streams from integration._support.process import owned_proxy upstream: Final = MCPServer("progress-cancellation-peer") @@ -770,7 +770,9 @@ def test_sdk_progress_and_cancellation_reach_upstream_without_losing_session( watchers = tuple(task for task in asyncio.all_tasks() if "_shutdown_watcher" in repr(task.get_coro())) await asyncio.wait_for(asyncio.gather(*watchers), 2) - hosted: Final = Starlette(routes=[Mount("/", app=app)], lifespan=lifespan) + hosted: Final = Starlette( + routes=[Mount("/", app=_acknowledge_notifications_after_dispatch(app))], lifespan=lifespan + ) with asgi_server(hosted) as upstream_url, httpx.Client() as client: config: Final = tmp_path / "proxy.yaml" config.write_text( diff --git a/tests/integration/providers/test_bedrock_converse_redacted_thinking_wire.py b/tests/integration/providers/test_bedrock_converse_redacted_thinking_wire.py index 79ec0545c44..1da52b0b6dc 100644 --- a/tests/integration/providers/test_bedrock_converse_redacted_thinking_wire.py +++ b/tests/integration/providers/test_bedrock_converse_redacted_thinking_wire.py @@ -414,3 +414,74 @@ def test_converse_redacted_thinking_replays_to_the_anthropic_api_as_claude_nativ {"type": "redacted_thinking", "data": NATIVE_REDACTED}, {"type": "text", "text": FIRST_ANSWER}, ) + + +_SIGNED_THINKING: Final[dict[str, JsonValue]] = { + "type": "thinking", + "thinking": "Checking the forecast.", + "signature": "scripted-signature-" + "s" * 32, +} +_CONVERSE_SIGNED_PART: Final[dict[str, JsonValue]] = { + "reasoningContent": { + "reasoningText": {"text": _SIGNED_THINKING["thinking"], "signature": _SIGNED_THINKING["signature"]} + } +} + + +@pytest.mark.parametrize( + ("deployment", "replayed_thinking"), + [ + pytest.param({"model": f"bedrock/converse/{NOVA_MODEL_ID}"}, (), id="nova"), + pytest.param( + { + "model": f"bedrock/converse/{APPLICATION_PROFILE_ARN}", + "model_info": {"base_model": f"bedrock/{NOVA_MODEL_ID}"}, + }, + (), + id="application-profile-arn-with-nova-base-model", + ), + pytest.param( + {"model": f"bedrock/converse/{APPLICATION_PROFILE_ARN}"}, + (_CONVERSE_SIGNED_PART,), + id="application-profile-arn-without-base-model", + ), + pytest.param( + { + "model": f"bedrock/converse/{APPLICATION_PROFILE_ARN}", + "model_info": {"base_model": f"bedrock/{CLAUDE_MODEL_ID}"}, + }, + (_CONVERSE_SIGNED_PART,), + id="application-profile-arn-with-claude-base-model", + ), + pytest.param(_CLAUDE, (_CONVERSE_SIGNED_PART,), id="claude"), + ], +) +def test_signed_history_keeps_its_thinking_unless_the_converse_target_is_known_not_to_be_claude( + gateway: Gateway, deployment: dict[str, JsonValue], replayed_thinking: tuple[dict[str, JsonValue], ...] +) -> None: + model_info: Final = deployment.get("model_info") + parameters: Final = {key: value for key, value in deployment.items() if key != "model_info"} + with wire_server(_bedrock_peer()) as wire, gateway.scenario() as scenario: + model: Final = scenario.model( + model_info=model_info if isinstance(model_info, dict) else None, + api_key=TOKEN, + aws_region_name="us-east-1", + aws_bedrock_runtime_endpoint=wire.url, + **parameters, + ) + + reply: Final = _chat( + gateway, + { + "model": model, + "max_tokens": 256, + "messages": [ + {"role": "user", "content": QUESTION}, + {"role": "assistant", "content": FIRST_ANSWER, "thinking_blocks": [_SIGNED_THINKING]}, + {"role": "user", "content": FOLLOW_UP}, + ], + }, + ) + + assert reply.content == FINAL_ANSWER + assert _replayed_assistant_turn(wire).content == (*replayed_thinking, {"text": FIRST_ANSWER}) diff --git a/tests/integration/providers/test_bedrock_passthrough_guardrail_config_wire.py b/tests/integration/providers/test_bedrock_passthrough_guardrail_config_wire.py new file mode 100644 index 00000000000..f995a7ac83f --- /dev/null +++ b/tests/integration/providers/test_bedrock_passthrough_guardrail_config_wire.py @@ -0,0 +1,531 @@ +import base64 +import json +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 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 +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.upstream import _aws_event_frame +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 180) + +_MODEL_ID: Final = "anthropic.claude-sonnet-5-v1:0" +_EVENT_STREAM: Final = "application/vnd.amazon.eventstream" +_BLOCKED_WORD: Final = "zzzforbidden" +_BLOCKED_TEXT: Final = "Blocked by guardrail" +_VERSIONLESS_GUARDRAIL_MESSAGE: Final = "guardrailIdentifier and guardrailVersion must be provided together." +_ACTIONS: Final = ("converse", "converse-stream", "invoke", "invoke-with-response-stream") +_DEPLOYMENT_GUARDRAIL: Final[Mapping[str, str]] = { + "guardrailIdentifier": "deployment-guardrail", + "guardrailVersion": "3", + "trace": "enabled", + "streamProcessingMode": "sync", +} +_CALLER_GUARDRAIL: Final[Mapping[str, str]] = { + "guardrailIdentifier": "caller-guardrail", + "guardrailVersion": "9", + "trace": "disabled", +} +_MANTLE_KEY: Final = "scripted-mantle-key" +_CALLER_INVOKE_HEADERS: Final[Mapping[str, str]] = { + "x-pass-x-amzn-bedrock-guardrailidentifier": "caller-guardrail", + "x-pass-x-amzn-bedrock-guardrailversion": "9", + "x-pass-x-amzn-bedrock-trace": "DISABLED", +} + + +@dataclass(frozen=True, slots=True) +class Applied: + """The guardrail a request asked Bedrock for, read from where Bedrock reads it for that action.""" + + identifier: str + version: str | None + trace: str | None + stream_processing_mode: str | None + + +@dataclass(frozen=True, slots=True) +class Outcome: + """What the AWS SDK caller reads back.""" + + text: str + intervened: bool + traced_guardrail: str | None + + +def _action(request: Request) -> str: + return request.target.partition("?")[0].rsplit("/", 1)[-1] + + +def _applied(request: Request) -> Applied | None: + if _action(request).startswith("converse"): + config: Final = json.loads(request.body).get("guardrailConfig") + if config is None: + return None + return Applied( + config["guardrailIdentifier"], + config.get("guardrailVersion"), + config.get("trace"), + config.get("streamProcessingMode"), + ) + identifier: Final = request.headers.get("x-amzn-bedrock-guardrailidentifier") + if identifier is None: + return None + return Applied( + identifier, + request.headers.get("x-amzn-bedrock-guardrailversion"), + request.headers.get("x-amzn-bedrock-trace"), + None, + ) + + +def _identifier(request: Request) -> str | None: + applied: Final = _applied(request) + return None if applied is None else applied.identifier + + +def _assessment(applied: Applied) -> Mapping[str, JsonValue]: + return { + applied.identifier: { + "wordPolicy": {"customWords": [{"match": _BLOCKED_WORD, "action": "BLOCKED", "detected": True}]}, + "appliedGuardrailDetails": {"guardrailId": applied.identifier, "guardrailVersion": applied.version}, + } + } + + +def _bedrock_with_guardrails(request: Request) -> Reply: + """Bedrock's runtime as the caller sees it: the guardrail the request names blocks a prompt holding the blocked + word, answering with its message and, when the request enabled tracing, a trace naming that guardrail. A guardrail + named without a version is refused with a 400 ValidationException, as real Bedrock in us-east-1 refused it on + every one of these actions on 2026-10-19.""" + applied: Final = _applied(request) + if applied is not None and applied.version is None: + return Reply( + status=400, + body=json.dumps({"message": _VERSIONLESS_GUARDRAIL_MESSAGE}).encode(), + headers={"x-amzn-errortype": "ValidationException"}, + ) + blocked: Final = applied is not None and _BLOCKED_WORD.encode() in request.body + traced: Final = blocked and applied is not None and (applied.trace or "").lower() == "enabled" + text: Final = _BLOCKED_TEXT if blocked else "ok" + stop_reason: Final = "guardrail_intervened" if blocked else "end_turn" + usage: Final = {"inputTokens": 5, "outputTokens": 1, "totalTokens": 6} + converse_trace: Final = {"trace": {"guardrail": {"inputAssessment": _assessment(applied)}}} if traced else {} + action: Final = _action(request) + if action == "converse": + return Reply( + body=json.dumps( + { + "output": {"message": {"role": "assistant", "content": [{"text": text}]}}, + "stopReason": stop_reason, + "usage": usage, + "metrics": {"latencyMs": 3}, + **converse_trace, + } + ).encode() + ) + if action == "converse-stream": + return Reply( + chunks=( + _aws_event_frame("messageStart", {"role": "assistant"}, "sc", "u"), + _aws_event_frame("contentBlockDelta", {"delta": {"text": text}, "contentBlockIndex": 0}, "sc", "u"), + _aws_event_frame("contentBlockStop", {"contentBlockIndex": 0}, "sc", "u"), + _aws_event_frame("messageStop", {"stopReason": stop_reason}, "sc", "u"), + _aws_event_frame( + "metadata", {"usage": usage, "metrics": {"latencyMs": 3}, **converse_trace}, "sc", "u" + ), + ), + content_type=_EVENT_STREAM, + ) + message: Final = { + "id": "msg_bedrock", + "type": "message", + "role": "assistant", + "model": "claude", + "content": [{"type": "text", "text": text}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 5, "output_tokens": 1}, + **({"amazon-bedrock-guardrailAction": "INTERVENED"} if blocked else {}), + **({"amazon-bedrock-trace": {"guardrail": {"input": _assessment(applied)}}} if traced else {}), + } + if action == "invoke": + return Reply(body=json.dumps(message).encode()) + chunk: Final = {"bytes": base64.b64encode(json.dumps(message).encode()).decode()} + return Reply(chunks=(_aws_event_frame("chunk", chunk, "sc", "u"),), content_type=_EVENT_STREAM) + + +def _bedrock_outage(request: Request) -> Reply: + return Reply( + status=500, + body=json.dumps({"message": "scripted outage"}).encode(), + headers={"x-amzn-errortype": "InternalServerException"}, + ) + + +def _traced_guardrail(assessment: object) -> str | None: + return None if not isinstance(assessment, dict) else next(iter(assessment)) + + +def _send(runtime: object, action: str, model: str, prompt: str, caller_guardrail: bool) -> Outcome: + messages: Final = [{"role": "user", "content": [{"text": prompt}]}] + invoke_body: Final = json.dumps( + {"anthropic_version": "bedrock-2023-05-31", "max_tokens": 16, "messages": [{"role": "user", "content": prompt}]} + ) + invoke_guardrail: Final = ( + {"guardrailIdentifier": "caller-guardrail", "guardrailVersion": "9", "trace": "DISABLED"} + if caller_guardrail + else {} + ) + if action == "converse": + response: Final = getattr(runtime, "converse")( + modelId=model, + messages=messages, + **({"guardrailConfig": dict(_CALLER_GUARDRAIL)} if caller_guardrail else {}), + ) + return Outcome( + text=response["output"]["message"]["content"][0]["text"], + intervened=response["stopReason"] == "guardrail_intervened", + traced_guardrail=_traced_guardrail(response.get("trace", {}).get("guardrail", {}).get("inputAssessment")), + ) + if action == "converse-stream": + events: Final = list( + getattr(runtime, "converse_stream")( + modelId=model, + messages=messages, + **( + {"guardrailConfig": {**_CALLER_GUARDRAIL, "streamProcessingMode": "async"}} + if caller_guardrail + else {} + ), + )["stream"] + ) + (stop,) = tuple(event["messageStop"]["stopReason"] for event in events if "messageStop" in event) + (metadata,) = tuple(event["metadata"] for event in events if "metadata" in event) + return Outcome( + text="".join( + event["contentBlockDelta"]["delta"]["text"] for event in events if "contentBlockDelta" in event + ), + intervened=stop == "guardrail_intervened", + traced_guardrail=_traced_guardrail(metadata.get("trace", {}).get("guardrail", {}).get("inputAssessment")), + ) + if action == "invoke": + message: Final = json.loads( + getattr(runtime, "invoke_model")(modelId=model, body=invoke_body, **invoke_guardrail)["body"].read() + ) + else: + (message,) = tuple( + json.loads(event["chunk"]["bytes"]) + for event in getattr(runtime, "invoke_model_with_response_stream")( + modelId=model, body=invoke_body, **invoke_guardrail + )["body"] + ) + return Outcome( + text=message["content"][0]["text"], + intervened=message.get("amazon-bedrock-guardrailAction") == "INTERVENED", + traced_guardrail=_traced_guardrail(message.get("amazon-bedrock-trace", {}).get("guardrail", {}).get("input")), + ) + + +def _signed_header_names(request: Request) -> frozenset[str]: + return frozenset(request.headers["authorization"].split("SignedHeaders=")[1].split(",")[0].split(";")) + + +def _deployment(wire: Wire, guardrail_config: Mapping[str, str] | None = None) -> Mapping[str, JsonValue]: + return { + "model": f"bedrock/{_MODEL_ID}", + "api_base": wire.url, + "aws_access_key_id": "AKIASCRIPTEDPROVIDER", + "aws_secret_access_key": "scripted-secret", + "aws_region_name": "us-east-1", + "use_in_pass_through": True, + **({} if guardrail_config is None else {"guardrailConfig": dict(guardrail_config)}), + } + + +@dataclass(frozen=True, slots=True) +class Rig: + proxy: Gateway + wire: Wire + guarded_model: str + unguarded_model: str + balanced_model: str + balanced_wires: Mapping[str, Wire] + failing_model: str + failing_wire: Wire + fallback_wire: Wire + mantle_model: str + versionless_model: str + key: str + + def runtime( # noqa: ANN201 # boto3 clients are untyped + self, monkeypatch: pytest.MonkeyPatch, extra_headers: Mapping[str, str] = MappingProxyType({}) + ): + monkeypatch.setenv("AWS_BEARER_TOKEN_BEDROCK", self.key) + client: Final = boto3.client( + "bedrock-runtime", + endpoint_url=f"{str(self.proxy.client.base_url).rstrip('/')}/bedrock", + region_name="us-east-1", + config=Config(read_timeout=case_timeout_seconds(), retries={"total_max_attempts": 1}), + ) + client.meta.events.register( + "before-send.bedrock-runtime.*", + lambda request, **_: request.headers.update(extra_headers), + ) + return client + + def received(self, wire: Wire, marker: str) -> tuple[Request, ...]: + return tuple(request for request in wire.drain() if marker.encode() in request.body) + + +@pytest.fixture(scope="module") +def rig(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[Rig]: + suffix: Final = uuid.uuid4().hex[:12] + guarded_model: Final = f"bedrock-guarded-{suffix}" + unguarded_model: Final = f"bedrock-unguarded-{suffix}" + balanced_model: Final = f"bedrock-balanced-{suffix}" + failing_model: Final = f"bedrock-failing-{suffix}" + fallback_model: Final = f"bedrock-fallback-{suffix}" + mantle_model: Final = f"bedrock-mantle-{suffix}" + versionless_model: Final = f"bedrock-versionless-{suffix}" + directory: Final = tmp_path_factory.mktemp("bedrock-guardrail-config") + base_config: Final = yaml.safe_load(Path("tests/integration/proxy_config.yaml").read_text()) + with ( + wire_server(_bedrock_with_guardrails) as wire, + wire_server(_bedrock_with_guardrails) as wire_a, + wire_server(_bedrock_with_guardrails) as wire_b, + wire_server(_bedrock_outage) as failing_wire, + wire_server(_bedrock_with_guardrails) as fallback_wire, + ): + config: Final = { + **base_config, + "model_list": [ + *base_config.get("model_list", []), + { + "model_name": guarded_model, + "litellm_params": _deployment(wire, _DEPLOYMENT_GUARDRAIL), + }, + {"model_name": unguarded_model, "litellm_params": _deployment(wire)}, + { + "model_name": balanced_model, + "litellm_params": _deployment( + wire_a, {"guardrailIdentifier": "guardrail-a", "guardrailVersion": "1"} + ), + }, + { + "model_name": balanced_model, + "litellm_params": _deployment( + wire_b, {"guardrailIdentifier": "guardrail-b", "guardrailVersion": "2"} + ), + }, + { + "model_name": failing_model, + "litellm_params": _deployment( + failing_wire, + {"guardrailIdentifier": "failing-guardrail", "guardrailVersion": "5", "trace": "enabled"}, + ), + }, + { + "model_name": fallback_model, + "litellm_params": _deployment( + fallback_wire, + {"guardrailIdentifier": "fallback-guardrail", "guardrailVersion": "6", "trace": "enabled"}, + ), + }, + { + "model_name": versionless_model, + "litellm_params": _deployment( + wire, {"guardrailIdentifier": "versionless-guardrail", "trace": "enabled"} + ), + }, + { + "model_name": mantle_model, + "litellm_params": { + "model": "bedrock_mantle/anthropic.claude-sonnet-5", + "api_base": wire.url, + "api_key": _MANTLE_KEY, + "aws_region_name": "us-east-1", + "use_in_pass_through": True, + "guardrailConfig": { + "guardrailIdentifier": "mantle-guardrail", + "guardrailVersion": "4", + "trace": "enabled", + }, + }, + }, + ], + "router_settings": { + **base_config.get("router_settings", {}), + "fallbacks": [{failing_model: [fallback_model]}], + }, + } + config_path: Final = directory / "bedrock-guardrail-config.yaml" + config_path.write_text(yaml.safe_dump(config)) + with ( + owned_proxy(session_gateway, directory, {}, config=config_path, workers=2) as proxy, + proxy.scenario() as scenario, + ): + yield Rig( + proxy=proxy, + wire=wire, + guarded_model=guarded_model, + unguarded_model=unguarded_model, + balanced_model=balanced_model, + balanced_wires={"guardrail-a": wire_a, "guardrail-b": wire_b}, + failing_model=failing_model, + failing_wire=failing_wire, + fallback_wire=fallback_wire, + mantle_model=mantle_model, + versionless_model=versionless_model, + key=scenario.key(key_alias=f"guardrail-key-{suffix}"), + ) + + +def _expected_on_the_wire(action: str, request: Request) -> Applied: + if action == "converse": + return Applied("deployment-guardrail", "3", "enabled", None) + if action == "converse-stream": + return Applied("deployment-guardrail", "3", "enabled", "sync") + assert {"x-amzn-bedrock-guardrailidentifier", "x-amzn-bedrock-guardrailversion", "x-amzn-bedrock-trace"} <= ( + _signed_header_names(request) + ), request.headers["authorization"] + return Applied("deployment-guardrail", "3", "ENABLED", None) + + +@pytest.mark.parametrize("action", _ACTIONS) +def test_passthrough_on_a_guardrailed_deployment_blocks_with_the_deployment_guardrail( + rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str +) -> None: + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch) + + blocked: Final = _send(runtime, action, rig.guarded_model, f"{marker} say {_BLOCKED_WORD}", caller_guardrail=False) + allowed: Final = _send(runtime, action, rig.guarded_model, f"{marker} say hello", caller_guardrail=False) + + assert blocked == Outcome(_BLOCKED_TEXT, intervened=True, traced_guardrail="deployment-guardrail"), blocked + assert allowed == Outcome("ok", intervened=False, traced_guardrail=None), allowed + received: Final = rig.received(rig.wire, marker) + assert [_action(request) for request in received] == [action, action], received + assert [_applied(request) for request in received] == [ + _expected_on_the_wire(action, request) for request in received + ], [(request.headers, request.body) for request in received] + + +@pytest.mark.parametrize("action", _ACTIONS) +def test_caller_guardrail_is_replaced_by_the_deployment_guardrail( + rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str +) -> None: + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch, _CALLER_INVOKE_HEADERS) + + outcome: Final = _send(runtime, action, rig.guarded_model, f"{marker} say {_BLOCKED_WORD}", caller_guardrail=True) + + assert outcome == Outcome(_BLOCKED_TEXT, intervened=True, traced_guardrail="deployment-guardrail"), outcome + (request,) = rig.received(rig.wire, marker) + assert _applied(request) == _expected_on_the_wire(action, request), (request.headers, request.body) + assert "caller-guardrail" not in json.dumps({**request.headers, "body": request.body.decode()}), request + + +@pytest.mark.parametrize("action", _ACTIONS) +def test_mantle_deployment_passthrough_sends_its_guardrail_with_bearer_auth( + rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str +) -> None: + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch) + + outcome: Final = _send(runtime, action, rig.mantle_model, f"{marker} say {_BLOCKED_WORD}", caller_guardrail=False) + + assert outcome == Outcome(_BLOCKED_TEXT, intervened=True, traced_guardrail="mantle-guardrail"), outcome + (request,) = rig.received(rig.wire, marker) + assert request.headers["authorization"] == f"Bearer {_MANTLE_KEY}", request.headers + assert _applied(request) == Applied( + "mantle-guardrail", "4", "ENABLED" if action.startswith("invoke") else "enabled", None + ), (request.headers, request.body) + + +@pytest.mark.parametrize("action", _ACTIONS) +def test_deployment_guardrail_without_a_version_returns_the_bedrock_rejection( + rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str +) -> None: + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch) + + with pytest.raises(ClientError) as raised: + _send(runtime, action, rig.versionless_model, f"{marker} say hello", caller_guardrail=False) + + error: Final = raised.value.response.get("Error", {}) + assert (error.get("Code"), error.get("Message")) == ("ValidationException", _VERSIONLESS_GUARDRAIL_MESSAGE), error + received: Final = rig.received(rig.wire, marker) + assert {_applied(request) for request in received} == { + Applied("versionless-guardrail", None, "ENABLED" if action.startswith("invoke") else "enabled", None) + }, [(request.headers, request.body) for request in received] + + +@pytest.mark.parametrize("action", _ACTIONS) +def test_deployment_without_a_guardrail_leaves_the_caller_guardrail_as_sent( + rig: Rig, monkeypatch: pytest.MonkeyPatch, action: str +) -> None: + if action.startswith("invoke"): + pytest.skip("BUG: the router-model invoke pass-through drops the caller's X-Amzn-Bedrock-Guardrail* headers") + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch) + + outcome: Final = _send(runtime, action, rig.unguarded_model, f"{marker} say {_BLOCKED_WORD}", caller_guardrail=True) + + assert outcome == Outcome(_BLOCKED_TEXT, intervened=True, traced_guardrail=None), outcome + (request,) = rig.received(rig.wire, marker) + assert _applied(request) == Applied( + "caller-guardrail", "9", "disabled", "async" if action == "converse-stream" else None + ), request.body + + +def test_each_load_balanced_request_carries_the_guardrail_of_the_deployment_it_reached( + rig: Rig, monkeypatch: pytest.MonkeyPatch +) -> None: + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch) + + outcomes: Final = tuple( + _send(runtime, action, rig.balanced_model, f"{marker}-{attempt} say {_BLOCKED_WORD}", caller_guardrail=False) + for attempt, action in enumerate(_ACTIONS * 6) + ) + + assert all(outcome.intervened for outcome in outcomes), outcomes + reached: Final = {name: rig.received(wire, marker) for name, wire in rig.balanced_wires.items()} + assert sum(len(requests) for requests in reached.values()) == len(outcomes), reached + assert all(reached.values()), {name: len(requests) for name, requests in reached.items()} + assert {name: {_identifier(request) for request in requests} for name, requests in reached.items()} == { + "guardrail-a": {"guardrail-a"}, + "guardrail-b": {"guardrail-b"}, + }, reached + + +def test_a_fallback_sends_the_guardrail_of_the_deployment_it_fell_back_to( + rig: Rig, monkeypatch: pytest.MonkeyPatch +) -> None: + marker: Final = uuid.uuid4().hex + runtime: Final = rig.runtime(monkeypatch) + + outcomes: Final = tuple( + _send(runtime, action, rig.failing_model, f"{marker} say {_BLOCKED_WORD}", caller_guardrail=False) + for action in _ACTIONS + ) + + assert outcomes == (Outcome(_BLOCKED_TEXT, intervened=True, traced_guardrail="fallback-guardrail"),) * 4, outcomes + failed: Final = rig.received(rig.failing_wire, marker) + served: Final = rig.received(rig.fallback_wire, marker) + assert [_action(request) for request in served] == list(_ACTIONS), served + assert {_identifier(request) for request in failed} == {"failing-guardrail"}, failed + assert {_identifier(request) for request in served} == {"fallback-guardrail"}, served diff --git a/tests/integration/providers/test_responses_minted_reasoning_replay_wire.py b/tests/integration/providers/test_responses_minted_reasoning_replay_wire.py index 2253443dfc0..4ff5871ff97 100644 --- a/tests/integration/providers/test_responses_minted_reasoning_replay_wire.py +++ b/tests/integration/providers/test_responses_minted_reasoning_replay_wire.py @@ -77,7 +77,18 @@ _AZURE_AI_OPENAI_HOST: Final = _Deployment( "/openai/v1/responses?api-version=preview", strips_message_status=True, ) -_DROPPING: Final = (_OPENAI, _AZURE, _AZURE_AI_OPENAI_HOST) +_BEDROCK_MANTLE: Final = _Deployment( + "bedrock_mantle", "bedrock_mantle/openai.gpt-oss-120b", "synthetic-mantle-key", "/v1/responses" +) +_BEDROCK: Final = _Deployment( + "bedrock", + "bedrock/openai.gpt-oss-120b-1:0", + "synthetic-bedrock-key", + "/openai/v1/responses", + MappingProxyType({"aws_region_name": "us-east-1"}), + model_info=MappingProxyType({"supported_endpoints": ["/v1/responses"]}), +) +_DROPPING: Final = (_OPENAI, _AZURE, _AZURE_AI_OPENAI_HOST, _BEDROCK_MANTLE, _BEDROCK) _KEEPING: Final = ( _Deployment("litellm_proxy", f"litellm_proxy/{_GPT}", "synthetic-proxy-key", "/responses"), _Deployment("databricks", "databricks/gpt-5.6", "synthetic-databricks-key", "/responses"), @@ -95,15 +106,6 @@ _KEEPING: Final = ( "/v1/responses", types_untyped_items_as_messages=True, ), - _Deployment("bedrock_mantle", "bedrock_mantle/openai.gpt-oss-120b", "synthetic-mantle-key", "/v1/responses"), - _Deployment( - "bedrock", - "bedrock/openai.gpt-oss-120b-1:0", - "synthetic-bedrock-key", - "/openai/v1/responses", - MappingProxyType({"aws_region_name": "us-east-1"}), - model_info=MappingProxyType({"supported_endpoints": ["/v1/responses"]}), - ), *( _Deployment(slug, f"{slug}/{model}", f"synthetic-{slug}-key", "/responses") for slug, model in ( diff --git a/tests/integration/routing/test_video_poll_routing_wire.py b/tests/integration/routing/test_video_poll_routing_wire.py index e0d88a6b0ff..a3d767275a0 100644 --- a/tests/integration/routing/test_video_poll_routing_wire.py +++ b/tests/integration/routing/test_video_poll_routing_wire.py @@ -338,9 +338,9 @@ def test_an_id_without_a_served_group_still_polls_its_own_deployment(rig: _Rig) video_id: Final = created_id(create(rig.gateway, {"model": rig.member_a, "seconds": "5", "size": "1280x720"})) assert decoded(video_id).get("served_model") is None legacy: Final = forged_id("dashscope", rig.member_a, task_of(video_id), None) - assert legacy == video_id + assert {**decoded(legacy), "deployment_id": decoded(video_id).get("deployment_id")} == decoded(video_id) rig.drain() - _assert_served(_poll(rig, legacy), legacy) + _assert_served(_poll(rig, legacy), video_id) download: Final = _download(rig, legacy) assert download.status_code == 200 and download.content == MP4, download.text received: Final = rig.a.drain() diff --git a/tests/unit/batches/test_batch_utils.py b/tests/unit/batches/test_batch_utils.py index ed9bceba927..1ee85522ae3 100644 --- a/tests/unit/batches/test_batch_utils.py +++ b/tests/unit/batches/test_batch_utils.py @@ -2296,6 +2296,36 @@ async def test_native_vertex_rows_route_to_vertex_cost_path_without_flag(monkeyp assert {(call["model"], call["custom_llm_provider"]) for call in calls} == {("gemini-2.5-flash", "vertex_ai")} +@pytest.mark.asyncio +async def test_raw_vertex_claude_rows_are_priced_like_with_the_native_output_flag(monkeypatch): + claude_message: Final = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "Blue."}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 12, "output_tokens": 5}, + } + rows: Final = [ + {"custom_id": "ok", "processed_time": "2026-10-29T03:00:00Z", "response": claude_message, "status": ""}, + {"custom_id": "bad", "processed_time": "2026-10-29T03:00:00Z", "response": {}, "status": "Bad Request"}, + ] + + async def cost_with_flag(flag: bool) -> bu.BatchCostUsageResult: + monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", flag, raising=False) + return await bu.calculate_batch_cost_and_usage( + file_content_dictionary=rows, custom_llm_provider="vertex_ai", model_name="vertex_ai/claude-sonnet-4-6" + ) + + unflagged: Final = await cost_with_flag(False) + flagged: Final = await cost_with_flag(True) + + assert unflagged.cost > 0 + assert (unflagged.successful_requests, unflagged.failed_requests) == (1, 1) + assert (unflagged.cost, unflagged.usage.total_tokens) == (flagged.cost, flagged.usage.total_tokens) + + @pytest.mark.asyncio async def test_openai_shaped_vertex_rows_keep_the_generic_path_without_flag(monkeypatch): monkeypatch.setattr(litellm, "disable_vertex_batch_output_transformation", False, raising=False) diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py index c5d7ca96a21..ab31cee5e07 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_handler.py @@ -1,7 +1,12 @@ +import json from datetime import datetime +from pathlib import Path +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch +import httpx import pytest +import respx import litellm @@ -11,6 +16,7 @@ from litellm.completion_extras.litellm_responses_transformation.handler import ( from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper from litellm.types.router import GenericLiteLLMParams +from litellm.litellm_core_utils.prompt_templates.common_utils import encrypted_reasoning_signature from litellm.types.utils import ModelResponse @@ -335,3 +341,110 @@ async def test_acompletion_keeps_provider_native_model_id_through_responses( handed_model = responses_call.call_args.kwargs["model"] assert _upstream_model_for(handed_model, custom_llm_provider) == expected_upstream_model + + +_CHATGPT_RESPONSES_URL: Final = "https://chatgpt.com/backend-api/codex/responses" + + +def _reasoning_item(item_id: str) -> dict[str, object]: + return {"type": "reasoning", "id": item_id, "summary": [], "encrypted_content": f"gAAAA_{item_id}"} + + +def _weather_call(index: int, city: str) -> dict[str, object]: + return { + "type": "function_call", + "id": f"fc_{index}", + "status": "completed", + "arguments": json.dumps({"city": city}), + "call_id": f"call_{index}", + "name": "get_weather", + } + + +def _two_reasoning_tool_turn_stream() -> httpx.Response: + turn: Final = { + "id": "resp_two_reasoning", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.6-sol", + "output": [ + _reasoning_item("rs_first"), + _weather_call(1, "Oslo"), + _reasoning_item("rs_second"), + _weather_call(2, "Lyon"), + ], + "usage": {"input_tokens": 10, "output_tokens": 20, "total_tokens": 30}, + } + body: Final = f"data: {json.dumps({'type': 'response.completed', 'response': turn})}\n\ndata: [DONE]\n\n" + return httpx.Response(200, headers={"content-type": "text/event-stream"}, text=body) + + +_FAR_FUTURE_EXPIRY: Final = 4_102_444_800 + + +@pytest.fixture +def _chatgpt_backend(tmp_path: Path, monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter) -> None: + auth: Final = {"access_token": "chatgpt-test-token", "expires_at": _FAR_FUTURE_EXPIRY, "account_id": "acct_test"} + (tmp_path / "auth.json").write_text(json.dumps(auth)) + monkeypatch.setenv("CHATGPT_TOKEN_DIR", str(tmp_path)) + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post(_CHATGPT_RESPONSES_URL).mock(side_effect=lambda _request: _two_reasoning_tool_turn_stream()) + + +_WEATHER_QUESTION: Final = "Weather in Oslo and Lyon?" +_WEATHER_TOOL_INPUT: Final = {"type": "object", "properties": {"city": {"type": "string"}}} +_EVERY_REASONING_SIGNATURE: Final = [ + encrypted_reasoning_signature("gAAAA_rs_first"), + encrypted_reasoning_signature("gAAAA_rs_second"), +] + + +def _thinking_signatures(reply: object) -> list[object]: + assert isinstance(reply, dict) + return [block.get("data") for block in reply["content"] if block["type"] == "redacted_thinking"] + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_chatgpt_backend") +async def test_async_messages_client_gets_every_reasoning_item_of_a_non_streamed_turn() -> None: + reply: Final = await litellm.anthropic.messages.acreate( + model="chatgpt/gpt-5.6-sol", + max_tokens=64, + messages=[{"role": "user", "content": _WEATHER_QUESTION}], + tools=[{"name": "get_weather", "input_schema": _WEATHER_TOOL_INPUT}], + ) + + assert _thinking_signatures(reply) == _EVERY_REASONING_SIGNATURE + + +_CHAT_WEATHER_TOOL: Final = {"type": "function", "function": {"name": "get_weather", "parameters": _WEATHER_TOOL_INPUT}} + + +def _reasoning_item_ids(response: object) -> list[object]: + assert isinstance(response, ModelResponse) + return [item["id"] for item in response.choices[0].message.reasoning_items] + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_chatgpt_backend") +async def test_async_chat_client_keeps_only_the_last_reasoning_item_of_a_turn() -> None: + response: Final = await litellm.acompletion( + model="chatgpt/gpt-5.6-sol", + messages=[{"role": "user", "content": _WEATHER_QUESTION}], + tools=[_CHAT_WEATHER_TOOL], + ) + + assert _reasoning_item_ids(response) == ["rs_second"] + + +@pytest.mark.usefixtures("_chatgpt_backend") +def test_sync_chat_client_keeps_only_the_last_reasoning_item_of_a_turn() -> None: + response: Final = litellm.completion( + model="chatgpt/gpt-5.6-sol", + messages=[{"role": "user", "content": _WEATHER_QUESTION}], + tools=[_CHAT_WEATHER_TOOL], + ) + + assert _reasoning_item_ids(response) == ["rs_second"] diff --git a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py index 95037761b11..d56adc07368 100644 --- a/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py +++ b/tests/unit/completion_extras/litellm_responses_transformation/test_completion_extras_litellm_responses_transformation_transformation.py @@ -3272,6 +3272,402 @@ def test_reasoning_items_streaming_emitted_on_response_completed(): assert ri["summary"][0]["text"] == summary_text +def test_reasoning_output_item_done_emits_signed_thinking_blocks_the_stream_keeps(): + """A chatgpt /v1/messages stream drops the reasoning item unless this delta carries thinking_blocks. + + response.output_item.done is the only event that has encrypted_content. The chat delta has to + expose it as a signed thinking block, and that block alone has to count as stream content. + """ + from unittest.mock import MagicMock + + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + OpenAiResponsesToChatCompletionStreamIterator, + ) + from litellm.litellm_core_utils.prompt_templates.common_utils import encrypted_reasoning_signature + from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper + from litellm.llms.anthropic.pass_through.adapters.transformation import LiteLLMAnthropicMessagesAdapter + + encrypted: Final = "gAAAA_round_trip" + done: Final = { + "type": "response.output_item.done", + "output_index": 0, + "item": { + "type": "reasoning", + "id": "rs_done", + "encrypted_content": encrypted, + "summary": [], + }, + } + result: Final = OpenAiResponsesToChatCompletionStreamIterator( + streaming_response=None, sync_stream=True, streams_reasoning_signatures=True + ).chunk_parser(done) + blocks: Final = result.choices[0].delta.thinking_blocks + assert blocks is not None + assert blocks[0]["type"] == "redacted_thinking" + assert blocks[0]["data"] == encrypted_reasoning_signature(encrypted) + assert blocks[0]["index"] == 0 + assert getattr(result.choices[0].delta, "reasoning_content", None) in (None, "") + + wrapper: Final = CustomStreamWrapper( + completion_stream=None, + model="chatgpt/gpt-6.1-sol", + logging_obj=MagicMock(), + custom_llm_provider="chatgpt", + ) + assert wrapper.is_chunk_non_empty({}, result, {}) is True + + block_type, start = LiteLLMAnthropicMessagesAdapter().translate_streaming_openai_chunk_to_anthropic_content_block( + choices=result.choices + ) + assert block_type == "redacted_thinking" + assert start["data"] == encrypted_reasoning_signature(encrypted) + + +_SIGNED_REASONING_STREAM_BODY: Final[Mapping[str, object]] = { + "id": "resp_signed_reasoning", + "object": "response", + "created_at": 1789788253, + "status": "completed", + "model": "gpt-6-sol", + "output": [ + {"type": "reasoning", "id": "rs_signed", "summary": [], "encrypted_content": "gAAAA_consumer"}, + { + "type": "message", + "id": "msg_logged_call_type", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "pong", "annotations": []}], + }, + ], + "usage": {"input_tokens": 10, "output_tokens": 5, "total_tokens": 15}, +} + +_GET_WEATHER_TOOL: Final = { + "type": "function", + "function": {"name": "get_weather", "parameters": {"type": "object", "properties": {}}}, +} + + +def _signed_reasoning_sse() -> str: + events: Final = ( + { + "type": "response.created", + "response": {**_SIGNED_REASONING_STREAM_BODY, "status": "in_progress", "output": [], "usage": None}, + }, + {"type": "response.output_item.done", "output_index": 0, "item": _SIGNED_REASONING_STREAM_BODY["output"][0]}, + { + "type": "response.output_text.delta", + "item_id": "msg_logged_call_type", + "output_index": 1, + "content_index": 0, + "delta": "pong", + }, + {"type": "response.completed", "response": _SIGNED_REASONING_STREAM_BODY}, + ) + return "".join(f"event: {event['type']}\ndata: {json.dumps(event)}\n\n" for event in events) + + +_AZURE_BRIDGE_KWARGS: Final[Mapping[str, object]] = { + "api_base": "https://example.invalid", + "api_version": "2025-04-01-preview", + "api_key": "test-key", +} + + +@pytest.mark.asyncio +async def test_streamed_reasoning_signature_reaches_messages_clients_and_never_chat_completions_deltas( + monkeypatch: pytest.MonkeyPatch, +) -> None: + from litellm.litellm_core_utils.prompt_templates.common_utils import encrypted_reasoning_signature + + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + with respx.mock(assert_all_called=True) as router: + router.post(url__startswith="https://example.invalid/openai/responses").mock( + side_effect=lambda _request: httpx.Response( + 200, text=_signed_reasoning_sse(), headers={"content-type": "text/event-stream"} + ) + ) + chat_stream: Final = await litellm.acompletion( + model="azure/gpt-6-sol", + messages=[{"role": "user", "content": "ping"}], + tools=[_GET_WEATHER_TOOL], + stream=True, + **_AZURE_BRIDGE_KWARGS, + ) + chat_deltas: Final = [chunk.choices[0].delta async for chunk in chat_stream if chunk.choices] + messages_stream: Final = await litellm.anthropic.messages.acreate( + model="azure/gpt-6-sol", + messages=[{"role": "user", "content": "ping"}], + tools=[{"name": "get_weather", "input_schema": {"type": "object", "properties": {}}}], + max_tokens=64, + stream=True, + **_AZURE_BRIDGE_KWARGS, + ) + sse: Final = b"".join([chunk async for chunk in messages_stream]).decode() + + assert [getattr(delta, "thinking_blocks", None) for delta in chat_deltas if getattr(delta, "thinking_blocks", None)] == [] + assert "".join(delta.content or "" for delta in chat_deltas) == "pong" + messages_events: Final = [json.loads(line[len("data: ") :]) for line in sse.splitlines() if line.startswith("data: ")] + started: Final = [event["content_block"] for event in messages_events if event.get("type") == "content_block_start"] + assert [block["type"] for block in started] == ["redacted_thinking", "text"] + assert started[0]["data"] == encrypted_reasoning_signature("gAAAA_consumer") + + +def test_reasoning_output_item_done_with_summary_signs_the_thinking_text(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + OpenAiResponsesToChatCompletionStreamIterator, + ) + from litellm.litellm_core_utils.prompt_templates.common_utils import encrypted_reasoning_signature + + encrypted: Final = "gAAAA_with_summary" + done: Final = { + "type": "response.output_item.done", + "output_index": 2, + "item": { + "type": "reasoning", + "id": "rs_sum", + "encrypted_content": encrypted, + "summary": [{"type": "summary_text", "text": "Oslo is colder."}], + }, + } + result: Final = OpenAiResponsesToChatCompletionStreamIterator( + streaming_response=None, sync_stream=True, streams_reasoning_signatures=True + ).chunk_parser(done) + block: Final = result.choices[0].delta.thinking_blocks[0] + assert block["type"] == "thinking" + assert block["thinking"] == "Oslo is colder." + assert block["signature"] == encrypted_reasoning_signature(encrypted) + assert block["index"] == 2 + + +def test_streamed_reasoning_summary_reaches_messages_clients_once_with_its_signature(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + OpenAiResponsesToChatCompletionStreamIterator, + ) + from litellm.litellm_core_utils.prompt_templates.common_utils import encrypted_reasoning_signature + from litellm.llms.anthropic.pass_through.adapters.streaming_iterator import AnthropicStreamWrapper + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + encrypted: Final = "gAAAA_streamed_summary" + events: Final = ( + {"type": "response.reasoning_summary_part.added", "output_index": 0, "summary_index": 0}, + {"type": "response.reasoning_summary_text.delta", "output_index": 0, "delta": "Oslo is "}, + {"type": "response.reasoning_summary_text.delta", "output_index": 0, "delta": "colder."}, + { + "type": "response.output_item.done", + "output_index": 0, + "item": { + "type": "reasoning", + "id": "rs_streamed", + "encrypted_content": encrypted, + "summary": [{"type": "summary_text", "text": "Oslo is colder."}], + }, + }, + {"type": "response.output_text.delta", "output_index": 1, "delta": "Oslo."}, + ) + parser: Final = OpenAiResponsesToChatCompletionStreamIterator( + streaming_response=None, sync_stream=True, streams_reasoning_signatures=True + ) + chunks: Final = ( + *(parser.chunk_parser(event) for event in events), + ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")]), + ) + + sse_events: Final = list(AnthropicStreamWrapper(completion_stream=iter(chunks), model="chatgpt/gpt-5.5")) + deltas: Final = [event["delta"] for event in sse_events if event.get("type") == "content_block_delta"] + assert "".join(d.get("thinking", "") for d in deltas if d["type"] == "thinking_delta") == "Oslo is colder." + assert [d["signature"] for d in deltas if d["type"] == "signature_delta"] == [ + encrypted_reasoning_signature(encrypted) + ] + + +def test_reasoning_item_repeated_at_stream_end_reaches_messages_clients_once(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + OpenAiResponsesToChatCompletionStreamIterator, + ) + from litellm.litellm_core_utils.prompt_templates.common_utils import encrypted_reasoning_signature + from litellm.llms.anthropic.pass_through.adapters.streaming_iterator import AnthropicStreamWrapper + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + encrypted: Final = "gAAAA_completed_repeat" + reasoning_item: Final = {"type": "reasoning", "id": "rs_repeat", "summary": [], "encrypted_content": encrypted} + events: Final = ( + {"type": "response.output_item.done", "output_index": 0, "item": reasoning_item}, + {"type": "response.output_text.delta", "output_index": 1, "delta": "Oslo."}, + ) + parser: Final = OpenAiResponsesToChatCompletionStreamIterator( + streaming_response=None, sync_stream=True, streams_reasoning_signatures=True + ) + chunks: Final = ( + *(parser.chunk_parser(event) for event in events), + ModelResponseStream( + choices=[StreamingChoices(index=0, delta=Delta(reasoning_items=[reasoning_item]), finish_reason=None)] + ), + ModelResponseStream(choices=[StreamingChoices(index=0, delta=Delta(), finish_reason="stop")]), + ) + + sse_events: Final = list(AnthropicStreamWrapper(completion_stream=iter(chunks), model="azure/gpt-5.4-nano")) + started: Final = [event["content_block"] for event in sse_events if event.get("type") == "content_block_start"] + assert [block["type"] for block in started] == ["redacted_thinking", "text"] + assert started[0]["data"] == encrypted_reasoning_signature(encrypted) + + +def test_reasoning_output_item_done_without_encrypted_content_stays_empty(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + OpenAiResponsesToChatCompletionStreamIterator, + ) + + done: Final = { + "type": "response.output_item.done", + "output_index": 0, + "item": { + "type": "reasoning", + "id": "rs_plain", + "summary": [{"type": "summary_text", "text": "already streamed"}], + }, + } + result: Final = OpenAiResponsesToChatCompletionStreamIterator( + streaming_response=None, sync_stream=True, streams_reasoning_signatures=True + ).chunk_parser(done) + assert getattr(result.choices[0].delta, "thinking_blocks", None) in (None, []) + + +def test_transform_response_hands_encrypted_reasoning_to_messages_clients_only(): + from unittest.mock import Mock + + from litellm.litellm_core_utils.prompt_templates.common_utils import encrypted_reasoning_signature + from litellm.llms.anthropic.pass_through.adapters.transformation import ( + LiteLLMAnthropicMessagesAdapter, + ) + from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.types.utils import ModelResponse, Usage + + encrypted: Final = "gAAAA_non_stream" + reasoning: Final = ResponseReasoningItem( + id="rs_tool", + summary=[], + type="reasoning", + content=None, + encrypted_content=encrypted, + status=None, + ) + function_call: Final = ResponseFunctionToolCall( + id="fc_1", + type="function_call", + status="completed", + arguments='{"city":"Oslo"}', + call_id="call_1", + name="get_weather", + ) + raw: Final = ResponsesAPIResponse.model_construct( + id="resp_tool", + created_at=1, + error=None, + incomplete_details=None, + model="gpt-6.1-sol", + output=[reasoning, function_call], + status="completed", + usage=None, + ) + model_response: Final = ModelResponse( + id="chatcmpl-tool", + created=1, + model=None, + object="chat.completion", + system_fingerprint=None, + choices=[], + usage=Usage(completion_tokens=0, prompt_tokens=0, total_tokens=0), + ) + result: Final = LiteLLMResponsesTransformationHandler().transform_response( + model="gpt-6.1-sol", + raw_response=raw, + model_response=model_response, + logging_obj=Mock(), + request_data={"model": "gpt-6.1-sol"}, + messages=[{"role": "user", "content": "Which city?"}], + optional_params={}, + litellm_params={}, + encoding=Mock(), + ) + message: Final = result.choices[0].message + assert message.reasoning_items[0]["encrypted_content"] == encrypted + assert "thinking_blocks" not in result.model_dump()["choices"][0]["message"] + assert message.tool_calls[0].function.name == "get_weather" + + anthropic_response: Final = LiteLLMAnthropicMessagesAdapter().translate_openai_response_to_anthropic( + response=result, tool_name_mapping=None + ) + content: Final = anthropic_response["content"] + assert [block["type"] for block in content] == ["redacted_thinking", "tool_use"] + assert content[0]["data"] == encrypted_reasoning_signature(encrypted) + + +@pytest.mark.parametrize( + ("keeps_every_reasoning_item", "kept_ids"), + [ + pytest.param(True, ["rs_first", "rs_second"], id="messages_client_keeps_every_item"), + pytest.param(False, ["rs_second"], id="chat_client_keeps_the_last_item"), + ], +) +def test_non_streamed_tool_turn_hands_messages_clients_every_reasoning_item( + keeps_every_reasoning_item: bool, kept_ids: list[str] +) -> None: + from unittest.mock import Mock + + from litellm.types.llms.openai import ResponsesAPIResponse + from litellm.types.utils import ModelResponse, Usage + + def reasoning(item_id: str) -> ResponseReasoningItem: + return ResponseReasoningItem( + id=item_id, summary=[], type="reasoning", content=None, encrypted_content=f"gAAAA_{item_id}", status=None + ) + + def call(index: int, city: str) -> ResponseFunctionToolCall: + return ResponseFunctionToolCall( + id=f"fc_{index}", + type="function_call", + status="completed", + arguments=json.dumps({"city": city}), + call_id=f"call_{index}", + name="get_weather", + ) + + raw: Final = ResponsesAPIResponse.model_construct( + id="resp_two_reasoning", + created_at=1, + error=None, + incomplete_details=None, + model="gpt-5.4-nano", + output=[reasoning("rs_first"), call(1, "Oslo"), reasoning("rs_second"), call(2, "Lyon")], + status="completed", + usage=None, + ) + result: Final = LiteLLMResponsesTransformationHandler().transform_response( + model="gpt-5.4-nano", + raw_response=raw, + model_response=ModelResponse( + id="chatcmpl-two-reasoning", + created=1, + model=None, + object="chat.completion", + system_fingerprint=None, + choices=[], + usage=Usage(completion_tokens=0, prompt_tokens=0, total_tokens=0), + ), + logging_obj=Mock(), + request_data={"model": "gpt-5.4-nano"}, + messages=[{"role": "user", "content": "Oslo or Lyon?"}], + optional_params={}, + litellm_params={"custom_llm_provider": "azure"}, + encoding=Mock(), + keeps_every_reasoning_item=keeps_every_reasoning_item, + ) + + message: Final = result.choices[0].message + assert [item["id"] for item in message.reasoning_items] == kept_ids + assert [tool_call.function.name for tool_call in message.tool_calls] == ["get_weather", "get_weather"] + + def test_streaming_function_call_tool_id_for_degenerate_call_id(): """In streaming, Bedrock Mantle's function_call event carries a unique ``id`` (``fc_...``) and a non-unique, index-based ``call_id`` (``call_0``). For that diff --git a/tests/unit/conftest.py b/tests/unit/conftest.py index 8331fb1937b..9028b8461b0 100644 --- a/tests/unit/conftest.py +++ b/tests/unit/conftest.py @@ -93,7 +93,15 @@ CALLBACK_LISTS: Final = ( "_async_failure_callback", "_async_input_callback", ) -RESET_TO_NONE_GLOBALS: Final = ("model_fallbacks", "cache") +RESET_TO_NONE_GLOBALS: Final = ( + "model_fallbacks", + "cache", + "max_fallbacks", + "default_fallbacks", + "fallbacks", + "context_window_fallbacks", + "content_policy_fallbacks", +) RESTORED_GLOBALS: Final = ( "disable_aiohttp_transport", "force_ipv4", 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 813ebdffff8..b31f45e506e 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 @@ -68,7 +68,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.llms.vertex_ai.gemini.transformation import gemini_convert_messages_with_history from litellm.types.llms.bedrock import BedrockConverseReasoningContentBlock -from litellm.types.llms.openai import AllMessageValues +from litellm.types.llms.openai import AllMessageValues, ChatCompletionUserMessage READ_FILE: Final[list[object]] = [{"name": "read_file", "description": "Read the contents of a file."}] @@ -422,6 +422,46 @@ def test_anthropic_messages_pt_drops_the_stream_index_from_replayed_thinking_blo ] +@pytest.mark.parametrize( + ("model", "expected_content"), + ( + ( + "global.anthropic.claude-haiku-5-5", + [ + {"reasoningContent": {"reasoningText": {"text": "pondering", "signature": "signed-reasoning"}}}, + {"text": "hello"}, + ], + ), + ("global.zai.glm-5.3", [{"text": "hello"}]), + ("amazon.nova-pro-v1:0", [{"text": "hello"}]), + ("converse/us.openai.gpt-6-luna", [{"text": "hello"}]), + ), +) +def test_bedrock_converse_replays_signed_thinking_only_to_claude( + model: str, expected_content: Sequence[Mapping[str, object]] +) -> None: + body: Final = litellm.AmazonConverseConfig().transform_request( + model=model, + messages=[ + ChatCompletionUserMessage(role="user", content="hi"), + ChatCompletionAssistantMessage( + role="assistant", + content="hello", + thinking_blocks=[{"type": "thinking", "thinking": "pondering", "signature": "signed-reasoning"}], + ), + ChatCompletionUserMessage(role="user", content="continue"), + ], + optional_params={}, + litellm_params={}, + headers={}, + ) + assert body["messages"] == [ + {"role": "user", "content": [{"text": "hi"}]}, + {"role": "assistant", "content": expected_content}, + {"role": "user", "content": [{"text": "continue"}]}, + ] + + async def test_anthropic_bedrock_thinking_blocks_with_none_content(): """ Test the specific function that processes thinking_blocks when content is None @@ -695,6 +735,192 @@ async def test_bedrock_converse_keeps_dropping_redacted_thinking_for_non_claude_ assert assistant_content == [{"text": "hello"}] +_HIDDEN_MODEL_ARN: Final = "arn:aws:bedrock:us-west-2:123456789012:application-inference-profile/abc123" + + +async def _converse_messages( + messages: list[dict[str, object]], + model: str, + base_model: str | None, + use_async: bool, + assistant_continue_message: str | None = None, +) -> list[dict[str, object]]: + return ( + await BedrockConverseMessagesProcessor._bedrock_converse_messages_pt_async( + messages=messages, + model=model, + llm_provider="bedrock", + base_model=base_model, + assistant_continue_message=assistant_continue_message, + ) + if use_async + else _bedrock_converse_messages_pt( + messages=messages, + model=model, + llm_provider="bedrock", + base_model=base_model, + assistant_continue_message=assistant_continue_message, + ) + ) + + +def _assistant_with_thinking(placement: str, block: dict[str, object]) -> dict[str, object]: + return ( + {"role": "assistant", "content": "hello", "thinking_blocks": [block]} + if placement == "thinking_blocks" + else {"role": "assistant", "content": [block, {"type": "text", "text": "hello"}]} + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("placement", ["thinking_blocks", "content"]) +@pytest.mark.parametrize( + "signature", ["claude-signature", encrypted_reasoning_signature("gAAAA_1")], ids=["claude", "encrypted"] +) +@pytest.mark.parametrize( + ("model", "base_model"), + [ + ("converse/us.openai.gpt-6-luna", None), + ("us.amazon.nova-pro-v1:0", None), + ("us.amazon.nova-pro-v1:0", "anthropic.claude-sonnet-4-6"), + (_HIDDEN_MODEL_ARN, "openai.gpt-oss-120b-1:0"), + ], +) +async def test_bedrock_converse_drops_signed_thinking_for_a_known_non_claude_target( + use_async: bool, placement: str, signature: str, model: str, base_model: str | None +): + block: Final = {"type": "thinking", "thinking": "private plan", "signature": signature} + messages: Final = [ + {"role": "user", "content": "hi"}, + _assistant_with_thinking(placement, block), + {"role": "user", "content": "again"}, + ] + + assert await _converse_messages(messages, model, base_model, use_async) == [ + {"role": "user", "content": [{"text": "hi"}]}, + {"role": "assistant", "content": [{"text": "hello"}]}, + {"role": "user", "content": [{"text": "again"}]}, + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize("placement", ["thinking_blocks", "content"]) +@pytest.mark.parametrize( + ("model", "base_model"), + [ + ("us.anthropic.claude-sonnet-4-6", None), + ("ca.anthropic.claude-haiku-5-5", None), + ("converse/in.anthropic.claude-sonnet-5", None), + ("claude-sonnet-4-6", None), + (_HIDDEN_MODEL_ARN, None), + (_HIDDEN_MODEL_ARN, "anthropic.claude-sonnet-4-6"), + ], +) +async def test_bedrock_converse_keeps_signed_thinking_for_a_target_that_may_be_claude( + use_async: bool, placement: str, model: str, base_model: str | None +): + block: Final = {"type": "thinking", "thinking": "private plan", "signature": "claude-signature"} + messages: Final = [ + {"role": "user", "content": "hi"}, + _assistant_with_thinking(placement, block), + {"role": "user", "content": "again"}, + ] + + result: Final = await _converse_messages(messages, model, base_model, use_async) + + assistant_content: Final = next(m["content"] for m in result if m["role"] == "assistant") + assert assistant_content == [ + {"reasoningContent": {"reasoningText": {"text": "private plan", "signature": "claude-signature"}}}, + {"text": "hello"}, + ] + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +async def test_bedrock_converse_keeps_the_tool_call_when_dropping_signed_thinking(use_async: bool): + messages: Final = [ + {"role": "user", "content": "weather?"}, + { + "role": "assistant", + "content": None, + "thinking_blocks": [{"type": "thinking", "thinking": "private plan", "signature": "claude-signature"}], + "tool_calls": [ + {"id": "call_1", "type": "function", "function": {"name": "get_weather", "arguments": "{}"}} + ], + }, + {"role": "tool", "tool_call_id": "call_1", "content": "sunny"}, + ] + + result: Final = await _converse_messages(messages, "converse/us.openai.gpt-6-luna", None, use_async) + + assert [m["role"] for m in result] == ["user", "assistant", "user"] + assert [list(part) for part in result[1]["content"]] == [["toolUse"]] + assert result[1]["content"][0]["toolUse"]["toolUseId"] == "call_1" + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +@pytest.mark.parametrize( + ("assistant_continue_message", "expected"), + [ + (None, [{"role": "user", "content": [{"text": "hi"}, {"text": "again"}]}]), + ( + "Please continue.", + [ + {"role": "user", "content": [{"text": "hi"}]}, + {"role": "assistant", "content": [{"text": "Please continue."}]}, + {"role": "user", "content": [{"text": "again"}]}, + ], + ), + ], + ids=["merged_into_one_user_turn", "split_by_the_continue_message"], +) +async def test_bedrock_converse_drops_a_turn_that_held_only_signed_thinking( + use_async: bool, assistant_continue_message: str | None, expected: list[dict[str, object]] +): + messages: Final = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": None, + "thinking_blocks": [{"type": "thinking", "thinking": "private plan", "signature": "claude-signature"}], + }, + {"role": "user", "content": "again"}, + ] + + result: Final = await _converse_messages( + messages, "converse/us.openai.gpt-6-luna", None, use_async, assistant_continue_message + ) + + assert result == expected + + +@pytest.mark.asyncio +@pytest.mark.parametrize("use_async", [False, True]) +async def test_bedrock_converse_drops_the_cache_point_of_a_dropped_signed_thinking_block(use_async: bool): + messages: Final = [ + {"role": "user", "content": "hi"}, + { + "role": "assistant", + "content": [ + { + "type": "thinking", + "thinking": "private plan", + "signature": "claude-signature", + "cache_control": {"type": "ephemeral"}, + } + ], + }, + {"role": "user", "content": "again"}, + ] + + result: Final = await _converse_messages(messages, "converse/us.amazon.nova-2-lite-v1:0", None, use_async) + + assert result == [{"role": "user", "content": [{"text": "hi"}, {"text": "again"}]}] + + @pytest.mark.parametrize( "thinking_block", [ diff --git a/tests/unit/litellm_core_utils/test_streaming_handler.py b/tests/unit/litellm_core_utils/test_streaming_handler.py index 637811f6387..8dfce043fe1 100644 --- a/tests/unit/litellm_core_utils/test_streaming_handler.py +++ b/tests/unit/litellm_core_utils/test_streaming_handler.py @@ -246,6 +246,37 @@ def test_is_chunk_non_empty(initialized_custom_stream_wrapper: CustomStreamWrapp ) +def test_is_chunk_non_empty_keeps_thinking_blocks_without_other_content( + initialized_custom_stream_wrapper: CustomStreamWrapper, +): + """Encrypted reasoning arrives as thinking_blocks and nothing else, so that delta has to be yielded.""" + chunk: Final = ModelResponseStream( + id="chunk-thinking", + object="chat.completion.chunk", + created=1, + model="chatgpt/gpt-6.1-sol", + choices=[ + StreamingChoices( + index=0, + finish_reason=None, + delta=Delta( + thinking_blocks=[ + {"type": "redacted_thinking", "data": "litellm_encrypted_reasoning:gAAAA"} + ] + ), + ) + ], + ) + assert ( + initialized_custom_stream_wrapper.is_chunk_non_empty( + completion_obj={}, + model_response=chunk, + response_obj={}, + ) + is True + ) + + def test_is_chunk_non_empty_with_annotations( initialized_custom_stream_wrapper: CustomStreamWrapper, ): diff --git a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py index 2f0a7ad070b..f02e7762deb 100644 --- a/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py +++ b/tests/unit/llms/anthropic/pass_through/adapters/test_anthropic_experimental_pass_through_adapters_transformation.py @@ -1,15 +1,18 @@ import base64 import functools import json +import logging +from collections.abc import Mapping, Sequence +from contextlib import nullcontext from types import MappingProxyType from typing import Any, Final, cast +import httpx import pytest +import respx import litellm - - - +from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.prompt_templates.common_utils import ( TOOL_RESULT_IMAGE_PLACEHOLDER, encrypted_reasoning_signature, @@ -31,6 +34,7 @@ from litellm.llms.anthropic.pass_through.messages.mid_conversation_system import from litellm.llms.openai.chat.gpt_transformation import OpenAIGPTConfig from litellm.types.llms.anthropic import ( AnthopicMessagesAssistantMessageParam, + AnthropicMessagesRequest, AnthropicMessagesUserMessageParam, ) from litellm.types.llms.openai import ChatCompletionAssistantToolCall @@ -509,6 +513,218 @@ def test_translate_anthropic_messages_to_openai_drops_bridge_encrypted_reasoning assert [b["signature"] for b in result[2]["thinking_blocks"]] == ["EqQBCkYIAxgCIkA_signed"] +def test_chatgpt_messages_keep_encrypted_reasoning_and_replay_it(): + """chatgpt /v1/messages ends in the Responses API, so the signed block has to survive into the next request.""" + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + + summary_encrypted: Final = "gAAAA_summary" + bytes_encrypted: Final = "gAAAA_bytes" + messages: Final = [ + AnthropicMessagesUserMessageParam(role="user", content="Which city is colder?"), + AnthopicMessagesAssistantMessageParam( + role="assistant", + content=[ + { + "type": "thinking", + "thinking": "Oslo is colder.", + "signature": encrypted_reasoning_signature(summary_encrypted), + }, + {"type": "redacted_thinking", "data": encrypted_reasoning_signature(bytes_encrypted)}, + {"type": "tool_use", "id": "call_1", "name": "get_weather", "input": {"city": "Oslo"}}, + ], + ), + AnthropicMessagesUserMessageParam( + role="user", + content=[{"type": "tool_result", "tool_use_id": "call_1", "content": "-7 C"}], + ), + ] + + kept: Final = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( + messages=messages, + model="chatgpt/gpt-6.1-sol", + custom_llm_provider="chatgpt", + ) + assistant: Final = kept[1] + assert [block["type"] for block in assistant["thinking_blocks"]] == ["thinking", "redacted_thinking"] + assert assistant["thinking_blocks"][0]["signature"] == encrypted_reasoning_signature(summary_encrypted) + assert assistant["thinking_blocks"][1]["data"] == encrypted_reasoning_signature(bytes_encrypted) + + responses_input, _instructions = LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api( + kept + ) + reasoning: Final = tuple(item for item in responses_input if item.get("type") == "reasoning") + assert [item["encrypted_content"] for item in reasoning] == [summary_encrypted, bytes_encrypted] + types: Final = [item.get("type") for item in responses_input] + assert types.index("reasoning") < types.index("function_call") + + stripped: Final = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( + messages=messages, + model="gpt-4o", + custom_llm_provider="openai", + ) + assert not stripped[1].get("thinking_blocks") + + prefixed: Final = LiteLLMAnthropicMessagesAdapter().translate_anthropic_messages_to_openai( + messages=messages, + model="responses/gpt-5-mini", + custom_llm_provider="openai", + ) + assert prefixed[1]["thinking_blocks"][1]["data"] == encrypted_reasoning_signature(bytes_encrypted) + + +def test_translate_openai_reasoning_items_become_signed_thinking_blocks(): + encrypted: Final = "gAAAA_outbound" + openai_choices: Final = [ + Choices( + message=Message( + role="assistant", + content="Oslo.", + reasoning_items=[ + { + "type": "reasoning", + "id": "rs_1", + "encrypted_content": encrypted, + "summary": [{"type": "summary_text", "text": "Oslo is colder."}], + } + ], + tool_calls=[ + ChatCompletionAssistantToolCall( + id="call_1", + type="function", + function=Function(name="get_weather", arguments='{"city": "Oslo"}'), + ) + ], + ) + ) + ] + + result: Final = LiteLLMAnthropicMessagesAdapter()._translate_openai_content_to_anthropic(choices=openai_choices) + + assert result[0]["type"] == "thinking" + assert result[0]["thinking"] == "Oslo is colder." + assert result[0]["signature"] == encrypted_reasoning_signature(encrypted) + assert result[1]["type"] == "text" + assert result[2]["type"] == "tool_use" + + +def test_translate_openai_reasoning_items_without_summary_become_redacted_thinking(): + encrypted: Final = "gAAAA_bytes_only" + openai_choices: Final = [ + Choices( + message=Message( + role="assistant", + content=None, + reasoning_items=[ + {"type": "reasoning", "id": "rs_1", "encrypted_content": encrypted, "summary": []} + ], + ) + ) + ] + + result: Final = LiteLLMAnthropicMessagesAdapter()._translate_openai_content_to_anthropic(choices=openai_choices) + + assert result == [ + {"type": "redacted_thinking", "data": encrypted_reasoning_signature(encrypted)}, + ] + + +def test_translate_openai_reasoning_items_are_not_duplicated_when_thinking_blocks_already_carry_them(): + encrypted: Final = "gAAAA_once" + signature: Final = encrypted_reasoning_signature(encrypted) + openai_choices: Final = [ + Choices( + message=Message( + role="assistant", + content="Oslo.", + thinking_blocks=[{"type": "redacted_thinking", "data": signature}], + reasoning_items=[ + {"type": "reasoning", "id": "rs_1", "encrypted_content": encrypted, "summary": []} + ], + ) + ) + ] + + result: Final = LiteLLMAnthropicMessagesAdapter()._translate_openai_content_to_anthropic(choices=openai_choices) + + assert [block["type"] for block in result] == ["redacted_thinking", "text"] + + +def _replayed_reasoning_request(tools: list[dict[str, object]]) -> AnthropicMessagesRequest: + return AnthropicMessagesRequest( + model="azure/gpt-5.4-nano", + max_tokens=64, + messages=[ + AnthropicMessagesUserMessageParam(role="user", content="Which city is colder?"), + AnthopicMessagesAssistantMessageParam( + role="assistant", + content=[ + {"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAAA_azure_tools")}, + {"type": "tool_use", "id": "call_1", "name": "get_weather", "input": {"city": "Oslo"}}, + ], + ), + AnthropicMessagesUserMessageParam( + role="user", + content=[{"type": "tool_result", "tool_use_id": "call_1", "content": "-7 C"}], + ), + ], + tools=tools, + ) + + +@pytest.mark.parametrize( + ("tools", "kept"), + [ + pytest.param( + [{"name": "get_weather", "input_schema": {"type": "object", "properties": {}}}], + True, + id="function_tools_bridge_to_responses", + ), + pytest.param([], False, id="no_tools_stay_on_chat"), + ], +) +def test_encrypted_reasoning_is_kept_when_the_request_tools_route_it_to_responses( + tools: list[dict[str, object]], kept: bool +) -> None: + openai_request, _mapping = LiteLLMAnthropicMessagesAdapter().translate_anthropic_to_openai( + _replayed_reasoning_request(tools), + custom_llm_provider="azure", + api_base="https://keep-check.openai.azure.com", + ) + + assistant: Final = openai_request["messages"][1] + assert bool(assistant.get("thinking_blocks")) is kept + + +@pytest.mark.parametrize( + ("deployment_params", "kept"), + [ + pytest.param({"reasoning_summary": "auto"}, True, id="deployment_summary_bridges_to_responses"), + pytest.param({"extra_body": {"reasoningSummary": "auto"}}, True, id="extra_body_summary_bridges_to_responses"), + pytest.param({}, False, id="no_summary_stays_on_chat"), + ], +) +def test_encrypted_reasoning_is_kept_when_the_deployment_reasoning_summary_routes_it_to_responses( + deployment_params: dict[str, object], kept: bool +) -> None: + from litellm.llms.anthropic.pass_through.adapters.handler import ( + LiteLLMMessagesToCompletionTransformationHandler, + ) + + request: Final = _replayed_reasoning_request([]) + completion_kwargs, _mapping = LiteLLMMessagesToCompletionTransformationHandler._prepare_completion_kwargs( + max_tokens=64, + messages=cast(list[dict[str, object]], request["messages"]), + model="azure/gpt-5.1", + thinking={"type": "enabled", "budget_tokens": 2048}, + extra_kwargs={"api_base": "https://keep-check.openai.azure.com", **deployment_params}, + ) + + assistant: Final = cast(list[Mapping[str, object]], completion_kwargs["messages"])[1] + assert bool(assistant.get("thinking_blocks")) is kept + + def test_translate_anthropic_messages_to_openai_sets_reasoning_content(): """Reasoning-aware chat providers read reasoning_content, so thinking text must land there. @@ -5235,3 +5451,575 @@ def test_eager_input_streaming_tool_reaches_bedrock_converse_as_beta(): assert data["additionalModelRequestFields"]["anthropic_beta"] == ["fine-grained-tool-streaming-2025-05-14"] assert data["toolConfig"]["tools"][0]["toolSpec"]["inputSchema"]["json"] == EAGER_INPUT_SCHEMA + + +@pytest.fixture +def _chat_bridge_httpx(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + litellm.in_memory_llm_clients_cache.flush_cache() + + +_RETRY_API_BASE: Final = "https://encrypted-retry.openai.azure.com" +_RETRY_URL: Final = r".*encrypted-retry\.openai\.azure\.com/openai/.*responses.*" +_RETRY_TOOLS: Final = [{"name": "get_weather", "input_schema": {"type": "object", "properties": {}}}] +_RETRY_REJECTION: Final = httpx.Response( + 400, + json={ + "error": { + "message": "The encrypted content gAAA_foreign could not be verified.", + "type": "invalid_request_error", + "param": None, + "code": "invalid_encrypted_content", + } + }, +) +_UNRELATED_REJECTION: Final = httpx.Response( + 400, json={"error": {"message": "max_output_tokens is too small", "type": "invalid_request_error"}} +) +_UNRELATED_UNVERIFIED_REJECTION: Final = httpx.Response( + 400, json={"error": {"message": "The organization could not be verified.", "type": "invalid_request_error"}} +) +_RETRY_REPLY: Final = { + "id": "resp_retry", + "object": "response", + "created_at": 1700000000, + "status": "completed", + "model": "gpt-5.4-nano", + "output": [ + { + "type": "message", + "id": "msg_retry", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Oslo", "annotations": []}], + } + ], + "usage": {"input_tokens": 10, "output_tokens": 2, "total_tokens": 12}, + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], +} + + +def _history_with_foreign_reasoning() -> list[dict[str, object]]: + return [ + {"role": "user", "content": "Which city is colder?"}, + { + "role": "assistant", + "content": [ + {"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAA_foreign")}, + {"type": "tool_use", "id": "call_1", "name": "get_weather", "input": {"city": "Oslo"}}, + ], + }, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_1", "content": "-7 C"}]}, + ] + + +def _sent_reasoning_items(request: httpx.Request) -> list[object]: + sent: Final = json.loads(request.read()) + return [item for item in sent["input"] if item.get("type") == "reasoning"] + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_chat_bridge_httpx") +async def test_async_messages_retry_once_without_reasoning_the_backend_cannot_decrypt( + respx_mock: respx.MockRouter, + caplog: pytest.LogCaptureFixture, +) -> None: + route: Final = respx_mock.post(url__regex=_RETRY_URL).mock( + side_effect=[_RETRY_REJECTION, httpx.Response(200, json=_RETRY_REPLY)] + ) + + with caplog.at_level(logging.WARNING, logger="LiteLLM"): + response: Final = await litellm.anthropic.messages.acreate( + model="azure/gpt-5.4-nano", + api_base=_RETRY_API_BASE, + api_key="fake-key", + api_version="2025-04-01-preview", + max_tokens=64, + tools=_RETRY_TOOLS, + messages=_history_with_foreign_reasoning(), + ) + + warnings: Final = tuple( + record.getMessage() + for record in caplog.records + if record.levelno >= logging.WARNING and record.name == "LiteLLM" + ) + retry_warnings: Final = tuple(message for message in warnings if "retrying once without it" in message) + assert route.call_count == 2 + assert len(_sent_reasoning_items(route.calls[0].request)) == 1 + assert _sent_reasoning_items(route.calls[1].request) == [] + assert [block["text"] for block in response["content"] if block["type"] == "text"] == ["Oslo"] + assert retry_warnings == ( + "Backend rejected replayed encrypted reasoning, retrying once without it. " + "Enable the encrypted_content_affinity pre-call check to drop such blocks before the call. " + "provider=azure model=azure/gpt-5.4-nano status=400 code=invalid_encrypted_content", + ) + assert all( + "gAAA_foreign" not in message and "litellm_encrypted_reasoning:" not in message for message in warnings + ) + + +@pytest.mark.usefixtures("_chat_bridge_httpx") +def test_sync_messages_retry_once_without_reasoning_the_backend_cannot_decrypt( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(url__regex=_RETRY_URL).mock( + side_effect=[_RETRY_REJECTION, httpx.Response(200, json=_RETRY_REPLY)] + ) + + response: Final = litellm.anthropic.messages.create( + model="azure/gpt-5.4-nano", + api_base=_RETRY_API_BASE, + api_key="fake-key", + api_version="2025-04-01-preview", + max_tokens=64, + tools=_RETRY_TOOLS, + messages=_history_with_foreign_reasoning(), + ) + + assert route.call_count == 2 + assert _sent_reasoning_items(route.calls[1].request) == [] + assert [block["text"] for block in response["content"] if block["type"] == "text"] == ["Oslo"] + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_chat_bridge_httpx") +async def test_async_messages_with_num_retries_send_the_rejected_reasoning_only_once( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(url__regex=_RETRY_URL).mock( + side_effect=[_RETRY_REJECTION, httpx.Response(200, json=_RETRY_REPLY)] + ) + + await litellm.anthropic.messages.acreate( + model="azure/gpt-5.4-nano", + api_base=_RETRY_API_BASE, + api_key="fake-key", + api_version="2025-04-01-preview", + max_tokens=64, + tools=_RETRY_TOOLS, + messages=_history_with_foreign_reasoning(), + num_retries=2, + ) + + assert [len(_sent_reasoning_items(call.request)) for call in route.calls] == [1, 0] + + +@pytest.mark.usefixtures("_chat_bridge_httpx") +def test_sync_messages_with_global_num_retries_send_the_rejected_reasoning_once_and_keep_the_setting( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "num_retries", 2) + route: Final = respx_mock.post(url__regex=_RETRY_URL).mock( + side_effect=[_RETRY_REJECTION, httpx.Response(200, json=_RETRY_REPLY)] + ) + + litellm.anthropic.messages.create( + model="azure/gpt-5.4-nano", + api_base=_RETRY_API_BASE, + api_key="fake-key", + api_version="2025-04-01-preview", + max_tokens=64, + tools=_RETRY_TOOLS, + messages=_history_with_foreign_reasoning(), + ) + + assert [len(_sent_reasoning_items(call.request)) for call in route.calls] == [1, 0] + assert litellm.num_retries == 2 + + +_SERVER_ERROR: Final = httpx.Response( + 500, json={"error": {"message": "The server had an error while processing your request.", "type": "server_error"}} +) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_chat_bridge_httpx") +async def test_async_messages_retry_a_server_error_with_the_reasoning_and_drop_it_only_once_rejected( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(url__regex=_RETRY_URL).mock( + side_effect=[_SERVER_ERROR, _RETRY_REJECTION, httpx.Response(200, json=_RETRY_REPLY)] + ) + + response: Final = await litellm.anthropic.messages.acreate( + model="azure/gpt-5.4-nano", + api_base=_RETRY_API_BASE, + api_key="fake-key", + api_version="2025-04-01-preview", + max_tokens=64, + tools=_RETRY_TOOLS, + messages=_history_with_foreign_reasoning(), + num_retries=2, + ) + + assert [len(_sent_reasoning_items(call.request)) for call in route.calls] == [1, 1, 0] + assert [block["text"] for block in response["content"] if block["type"] == "text"] == ["Oslo"] + + +@pytest.mark.usefixtures("_chat_bridge_httpx") +def test_sync_messages_retry_a_server_error_with_the_reasoning_and_drop_it_only_once_rejected( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(url__regex=_RETRY_URL).mock( + side_effect=[_SERVER_ERROR, _RETRY_REJECTION, httpx.Response(200, json=_RETRY_REPLY)] + ) + + response: Final = litellm.anthropic.messages.create( + model="azure/gpt-5.4-nano", + api_base=_RETRY_API_BASE, + api_key="fake-key", + api_version="2025-04-01-preview", + max_tokens=64, + tools=_RETRY_TOOLS, + messages=_history_with_foreign_reasoning(), + num_retries=2, + ) + + assert [len(_sent_reasoning_items(call.request)) for call in route.calls] == [1, 1, 0] + assert [block["text"] for block in response["content"] if block["type"] == "text"] == ["Oslo"] + + +class _FailureLog(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.rejections: tuple[str, ...] = () + + def _record(self, kwargs: Mapping[str, object]) -> None: + error: Final = str(kwargs.get("exception")) + self.rejections = ( + *self.rejections, + next((phrase for phrase in _UPSTREAM_REJECTION_PHRASES if phrase in error), error), + ) + + def log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._record(kwargs) + + async def async_log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._record(kwargs) + + +_UPSTREAM_REJECTION_PHRASES: Final = ("could not be verified", "max_output_tokens is too small") +_OUTCOMES_OF_THE_TURN_SENT_AGAIN: Final = ( + pytest.param(httpx.Response(200, json=_RETRY_REPLY), (), id="served"), + pytest.param(_UNRELATED_REJECTION, ("max_output_tokens is too small",), id="rejected"), +) + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_chat_bridge_httpx") +@pytest.mark.parametrize(("sent_again", "logged_rejections"), _OUTCOMES_OF_THE_TURN_SENT_AGAIN) +async def test_async_messages_log_only_the_outcome_of_the_turn_sent_again( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + sent_again: httpx.Response, + logged_rejections: tuple[str, ...], +) -> None: + failure_log: Final = _FailureLog() + monkeypatch.setattr(litellm, "callbacks", [failure_log]) + respx_mock.post(url__regex=_RETRY_URL).mock(side_effect=[_RETRY_REJECTION, sent_again]) + + with pytest.raises(litellm.BadRequestError) if logged_rejections else nullcontext(): + await litellm.anthropic.messages.acreate( + model="azure/gpt-5.4-nano", + api_base=_RETRY_API_BASE, + api_key="fake-key", + api_version="2025-04-01-preview", + max_tokens=64, + tools=_RETRY_TOOLS, + messages=_history_with_foreign_reasoning(), + ) + + assert failure_log.rejections == logged_rejections + + +@pytest.mark.usefixtures("_chat_bridge_httpx") +@pytest.mark.parametrize(("sent_again", "logged_rejections"), _OUTCOMES_OF_THE_TURN_SENT_AGAIN) +def test_sync_messages_log_only_the_outcome_of_the_turn_sent_again( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + sent_again: httpx.Response, + logged_rejections: tuple[str, ...], +) -> None: + failure_log: Final = _FailureLog() + monkeypatch.setattr(litellm, "callbacks", [failure_log]) + respx_mock.post(url__regex=_RETRY_URL).mock(side_effect=[_RETRY_REJECTION, sent_again]) + + with pytest.raises(litellm.BadRequestError) if logged_rejections else nullcontext(): + litellm.anthropic.messages.create( + model="azure/gpt-5.4-nano", + api_base=_RETRY_API_BASE, + api_key="fake-key", + api_version="2025-04-01-preview", + max_tokens=64, + tools=_RETRY_TOOLS, + messages=_history_with_foreign_reasoning(), + ) + + assert failure_log.rejections == logged_rejections + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_chat_bridge_httpx") +@pytest.mark.parametrize( + ("messages", "rejection"), + [ + pytest.param( + [{"role": "user", "content": "Which city is colder?"}], + _RETRY_REJECTION, + id="no_encrypted_reasoning_to_drop", + ), + pytest.param(_history_with_foreign_reasoning(), _UNRELATED_REJECTION, id="unrelated_rejection"), + pytest.param( + _history_with_foreign_reasoning(), _UNRELATED_UNVERIFIED_REJECTION, id="unrelated_unverified_rejection" + ), + ], +) +async def test_messages_do_not_retry_other_rejections( + respx_mock: respx.MockRouter, messages: list[dict[str, object]], rejection: httpx.Response +) -> None: + route: Final = respx_mock.post(url__regex=_RETRY_URL).mock(return_value=rejection) + + with pytest.raises(litellm.BadRequestError): + await litellm.anthropic.messages.acreate( + model="azure/gpt-5.4-nano", + api_base=_RETRY_API_BASE, + api_key="fake-key", + api_version="2025-04-01-preview", + max_tokens=64, + tools=_RETRY_TOOLS, + messages=messages, + ) + + assert route.call_count == 1 + + +_STRIPPED_CHAT_MODEL: Final = "azure/gpt-5.6" +_STRIPPED_CHAT_API_BASE: Final = "https://stripped-reasoning.openai.azure.com" +_STRIPPED_CHAT_URL: Final = r".*stripped-reasoning\.openai\.azure\.com/openai/deployments/gpt-5\.6/chat/completions.*" + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_chat_bridge_httpx") +async def test_messages_do_not_retry_when_the_translated_request_already_dropped_the_reasoning( + respx_mock: respx.MockRouter, +) -> None: + route: Final = respx_mock.post(url__regex=_STRIPPED_CHAT_URL).mock(return_value=_RETRY_REJECTION) + + with pytest.raises(litellm.BadRequestError): + await litellm.anthropic.messages.acreate( + model=_STRIPPED_CHAT_MODEL, + api_base=_STRIPPED_CHAT_API_BASE, + api_key="fake-key", + max_tokens=64, + messages=_history_with_foreign_reasoning(), + ) + + assert route.call_count == 1 + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_chat_bridge_httpx") +async def test_openai_thinking_on_the_chat_adapter_keeps_reasoning_its_responses_route_replays( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "use_chat_completions_url_for_anthropic_messages", True) + route: Final = respx_mock.post(url__regex=r".*openai-thinking\.example/v1/responses").mock( + return_value=httpx.Response(200, json=_RETRY_REPLY) + ) + + await litellm.anthropic.messages.acreate( + model="openai/gpt-5-mini", + api_base="https://openai-thinking.example/v1", + api_key="fake-key", + max_tokens=2048, + thinking={"type": "enabled", "budget_tokens": 1024}, + messages=[ + {"role": "user", "content": "Which city is colder?"}, + { + "role": "assistant", + "content": [ + {"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAA_same_org")}, + {"type": "text", "text": "Oslo."}, + ], + }, + {"role": "user", "content": "And by how much?"}, + ], + ) + + assert route.call_count == 1 + assert [item["encrypted_content"] for item in _sent_reasoning_items(route.calls[0].request)] == ["gAAA_same_org"] + + +def _responses_sse(events: Sequence[Mapping[str, object]]) -> bytes: + return b"".join( + f"event: {event['type']}\ndata: {json.dumps({**event, 'sequence_number': number})}\n\n".encode() + for number, event in enumerate(events) + ) + + +def _reasoning_item(item_id: str, encrypted: str, summary: str = "") -> dict[str, object]: + return { + "type": "reasoning", + "id": item_id, + "summary": [{"type": "summary_text", "text": summary}] if summary else [], + "encrypted_content": encrypted, + } + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_chat_bridge_httpx") +async def test_streamed_consecutive_reasoning_items_each_become_their_own_block( + respx_mock: respx.MockRouter, +) -> None: + response: Final = {**_RETRY_REPLY, "id": "resp_two_reasoning", "status": "in_progress", "output": []} + message: Final = { + "type": "message", + "id": "msg_two_reasoning", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Oslo", "annotations": []}], + } + first: Final = _reasoning_item("rs_1", "gAAA_first") + second: Final = _reasoning_item("rs_2", "gAAA_second") + completed: Final = {**response, "status": "completed", "output": [first, second, message]} + events: Final = ( + {"type": "response.created", "response": response}, + {"type": "response.output_item.added", "output_index": 0, "item": _reasoning_item("rs_1", "")}, + {"type": "response.output_item.done", "output_index": 0, "item": first}, + {"type": "response.output_item.added", "output_index": 1, "item": _reasoning_item("rs_2", "")}, + {"type": "response.output_item.done", "output_index": 1, "item": second}, + {"type": "response.output_item.added", "output_index": 2, "item": {**message, "content": []}}, + { + "type": "response.output_text.delta", + "output_index": 2, + "item_id": "msg_two_reasoning", + "content_index": 0, + "delta": "Oslo", + }, + {"type": "response.output_item.done", "output_index": 2, "item": message}, + {"type": "response.completed", "response": completed}, + ) + respx_mock.post(url__regex=_RETRY_URL).mock( + return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, content=_responses_sse(events)) + ) + + stream: Final = await litellm.anthropic.messages.acreate( + model="azure/gpt-5.4-nano", + api_base=_RETRY_API_BASE, + api_key="fake-key", + api_version="2025-04-01-preview", + max_tokens=64, + tools=_RETRY_TOOLS, + stream=True, + messages=[{"role": "user", "content": "Which city is colder?"}], + ) + frames: Final = b"".join([chunk async for chunk in stream]).decode() + starts: Final = tuple( + json.loads(line.removeprefix("data: "))["content_block"] + for line in frames.splitlines() + if line.startswith("data: ") and '"content_block_start"' in line + ) + + assert [(block["type"], block.get("data")) for block in starts] == [ + ("redacted_thinking", encrypted_reasoning_signature("gAAA_first")), + ("redacted_thinking", encrypted_reasoning_signature("gAAA_second")), + ("text", None), + ] + + +def _streamed_blocks(frames: str) -> list[tuple[str, str, str | None]]: + events: Final = tuple( + json.loads(line.removeprefix("data: ")) for line in frames.splitlines() if line.startswith("data: ") + ) + deltas: Final = tuple(event for event in events if event["type"] == "content_block_delta") + + def block_deltas(index: int) -> tuple[Mapping[str, str], ...]: + return tuple(delta["delta"] for delta in deltas if delta["index"] == index) + + return [ + ( + event["content_block"]["type"], + "".join(delta.get("thinking", "") for delta in block_deltas(event["index"])), + next( + (delta["signature"] for delta in block_deltas(event["index"]) if delta["type"] == "signature_delta"), + None, + ), + ) + for event in events + if event["type"] == "content_block_start" + ] + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("_chat_bridge_httpx") +async def test_streamed_reasoning_summaries_stay_with_their_own_encrypted_item( + respx_mock: respx.MockRouter, +) -> None: + response: Final = {**_RETRY_REPLY, "id": "resp_two_summaries", "status": "in_progress", "output": []} + message: Final = { + "type": "message", + "id": "msg_two_summaries", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Oslo", "annotations": []}], + } + first: Final = _reasoning_item("rs_1", "gAAA_first", "Oslo is further north.") + second: Final = _reasoning_item("rs_2", "gAAA_second", "Check the forecast.") + completed: Final = {**response, "status": "completed", "output": [first, second, message]} + + def summary_delta(output_index: int, item_id: str, text: str) -> dict[str, object]: + return { + "type": "response.reasoning_summary_text.delta", + "output_index": output_index, + "item_id": item_id, + "summary_index": 0, + "delta": text, + } + + events: Final = ( + {"type": "response.created", "response": response}, + {"type": "response.output_item.added", "output_index": 0, "item": _reasoning_item("rs_1", "")}, + summary_delta(0, "rs_1", "Oslo is further north."), + {"type": "response.output_item.done", "output_index": 0, "item": first}, + {"type": "response.output_item.added", "output_index": 1, "item": _reasoning_item("rs_2", "")}, + summary_delta(1, "rs_2", "Check the forecast."), + {"type": "response.output_item.done", "output_index": 1, "item": second}, + {"type": "response.output_item.added", "output_index": 2, "item": {**message, "content": []}}, + { + "type": "response.output_text.delta", + "output_index": 2, + "item_id": "msg_two_summaries", + "content_index": 0, + "delta": "Oslo", + }, + {"type": "response.output_item.done", "output_index": 2, "item": message}, + {"type": "response.completed", "response": completed}, + ) + respx_mock.post(url__regex=_RETRY_URL).mock( + return_value=httpx.Response(200, headers={"content-type": "text/event-stream"}, content=_responses_sse(events)) + ) + + stream: Final = await litellm.anthropic.messages.acreate( + model="azure/gpt-5.4-nano", + api_base=_RETRY_API_BASE, + api_key="fake-key", + api_version="2025-04-01-preview", + max_tokens=64, + tools=_RETRY_TOOLS, + stream=True, + messages=[{"role": "user", "content": "Which city is colder?"}], + ) + frames: Final = b"".join([chunk async for chunk in stream]).decode() + + assert _streamed_blocks(frames) == [ + ("thinking", "Oslo is further north.", encrypted_reasoning_signature("gAAA_first")), + ("thinking", "Check the forecast.", encrypted_reasoning_signature("gAAA_second")), + ("text", "", None), + ] diff --git a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py index 6aec78bdb97..526ead24d26 100644 --- a/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py +++ b/tests/unit/llms/anthropic/pass_through/responses_adapters/test_responses_adapters_handler.py @@ -1,20 +1,26 @@ +import asyncio import datetime import json import os import sys -from collections.abc import Mapping -from typing import Final +from collections.abc import AsyncIterator, Mapping, Sequence +from contextlib import nullcontext +from typing import Final, cast from unittest.mock import AsyncMock, patch +import httpx import pytest import respx sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../../../.."))) import litellm +from litellm.integrations.custom_logger import CustomLogger +from litellm.litellm_core_utils.logging_worker import GLOBAL_LOGGING_WORKER from litellm.litellm_core_utils.prompt_templates.common_utils import ( TOOL_RESULT_IMAGE_BOUNDARY, TOOL_RESULT_IMAGE_PLACEHOLDER, + encrypted_reasoning_signature, ) from litellm.llms.anthropic.pass_through.responses_adapters.handler import ( LiteLLMMessagesToResponsesAPIHandler, @@ -22,6 +28,7 @@ from litellm.llms.anthropic.pass_through.responses_adapters.handler import ( ) from litellm.router import Router from litellm.types.integrations.custom_logger import CONVERTED_STREAM_KEYS, converted_stream_requested +from litellm.types.utils import StandardLoggingPayload MESSAGES = [{"role": "user", "content": "hello"}] CLAUDE_CODE_USER_ID = json.dumps({"device_id": "d" * 64, "account_uuid": "", "session_id": "session-abc"}) @@ -684,3 +691,371 @@ def test_drop_params_drops_the_mcp_toolset_tools_of_dropped_mcp_servers(monkeypa ) assert [tool["name"] for tool in responses_kwargs["tools"]] == ["get_weather"] + +TEXT_RESPONSE: Final = { + "id": "resp_replayed", + "object": "response", + "created_at": 1, + "status": "completed", + "model": "gpt-5.6-luna", + "output": [ + { + "type": "message", + "id": "msg_1", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": "done", "annotations": []}], + } + ], + "parallel_tool_calls": True, + "tool_choice": "auto", + "tools": [], + "usage": {"input_tokens": 3, "output_tokens": 4, "total_tokens": 7}, +} + +REJECTED_REPLAYED_REASONING: Final = { + "error": { + "message": "The encrypted content gAAA_foreign could not be verified.", + "type": "invalid_request_error", + "param": None, + "code": "invalid_encrypted_content", + } +} +UNVERIFIED_ORGANIZATION: Final = { + "error": {"message": "The organization could not be verified.", "type": "invalid_request_error"} +} + + +def _history(*, replayed_reasoning: bool) -> list[dict[str, object]]: + reasoning: Final = ( + ({"type": "redacted_thinking", "data": encrypted_reasoning_signature("gAAA_foreign")},) + if replayed_reasoning + else () + ) + return [ + {"role": "user", "content": "List the logs."}, + {"role": "assistant", "content": [*reasoning, {"type": "tool_use", "id": "call_1", "name": "Bash", "input": {}}]}, + {"role": "user", "content": [{"type": "tool_result", "tool_use_id": "call_1", "content": "app.log"}]}, + ] + + +def _replayed_reasoning_counts(upstream: respx.Route) -> list[int]: + return [ + sum(item.get("type") == "reasoning" for item in json.loads(call.request.content)["input"]) + for call in upstream.calls + ] + + +def _upstream(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, first: httpx.Response, stream: bool): + return _upstream_after(respx_mock, monkeypatch, (first,), stream) + + +def _upstream_after( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, failures: tuple[httpx.Response, ...], stream: bool +) -> respx.Route: + monkeypatch.setenv("OPENAI_API_KEY", "sk-replay-test") + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + served: Final = ( + httpx.Response(200, headers={"Content-Type": "text/event-stream"}, content=RESPONSES_SSE_BODY) + if stream + else httpx.Response(200, json=TEXT_RESPONSE) + ) + return respx_mock.post("https://api.openai.com/v1/responses").mock(side_effect=[*failures, served]) + + +@pytest.mark.asyncio +@pytest.mark.parametrize("stream", [False, True]) +async def test_openai_turn_is_sent_again_once_without_replayed_reasoning_the_backend_cannot_decrypt( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, stream: bool +) -> None: + upstream: Final = _upstream(respx_mock, monkeypatch, httpx.Response(400, json=REJECTED_REPLAYED_REASONING), stream) + + response: Final = await LiteLLMMessagesToResponsesAPIHandler.async_anthropic_messages_handler( + max_tokens=64, + messages=_history(replayed_reasoning=True), + model="openai/gpt-5.6-luna", + tools=[dict(BASH_TOOL)], + stream=stream, + custom_llm_provider="openai", + ) + events: Final = ( + [json.loads(chunk.decode().split("data: ", 1)[1]) async for chunk in cast(AsyncIterator[bytes], response)] + if stream + else [cast(Mapping[str, object], response)] + ) + + assert _replayed_reasoning_counts(upstream) == [1, 0] + assert events[-1]["type"] in ("message", "message_stop") + + +def test_sync_openai_turn_is_sent_again_once_without_replayed_reasoning_the_backend_cannot_decrypt( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + upstream: Final = _upstream(respx_mock, monkeypatch, httpx.Response(400, json=REJECTED_REPLAYED_REASONING), False) + + response: Final = LiteLLMMessagesToResponsesAPIHandler.anthropic_messages_handler( + max_tokens=64, + messages=_history(replayed_reasoning=True), + model="openai/gpt-5.6-luna", + tools=[dict(BASH_TOOL)], + custom_llm_provider="openai", + ) + + assert _replayed_reasoning_counts(upstream) == [1, 0] + assert cast(Mapping[str, object], response)["content"] == [{"type": "text", "text": "done"}] + + +@pytest.mark.asyncio +async def test_openai_turn_with_num_retries_sends_the_rejected_reasoning_only_once( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + upstream: Final = _upstream(respx_mock, monkeypatch, httpx.Response(400, json=REJECTED_REPLAYED_REASONING), False) + + await litellm.anthropic.messages.acreate( + max_tokens=64, + messages=_history(replayed_reasoning=True), + model="openai/gpt-5.6-luna", + tools=[dict(BASH_TOOL)], + num_retries=2, + ) + + assert _replayed_reasoning_counts(upstream) == [1, 0] + + +def test_sync_openai_turn_with_global_num_retries_sends_the_rejected_reasoning_once_and_keeps_the_setting( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + monkeypatch.setattr(litellm, "num_retries", 2) + upstream: Final = _upstream(respx_mock, monkeypatch, httpx.Response(400, json=REJECTED_REPLAYED_REASONING), False) + + litellm.anthropic.messages.create( + max_tokens=64, + messages=_history(replayed_reasoning=True), + model="openai/gpt-5.6-luna", + tools=[dict(BASH_TOOL)], + ) + + assert _replayed_reasoning_counts(upstream) == [1, 0] + assert litellm.num_retries == 2 + + +_SERVER_ERROR: Final = { + "error": {"message": "The server had an error while processing your request.", "type": "server_error"} +} + + +@pytest.mark.asyncio +async def test_openai_turn_retries_a_server_error_with_the_reasoning_and_drops_it_only_once_rejected( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + upstream: Final = _upstream_after( + respx_mock, + monkeypatch, + (httpx.Response(500, json=_SERVER_ERROR), httpx.Response(400, json=REJECTED_REPLAYED_REASONING)), + False, + ) + + response: Final = await litellm.anthropic.messages.acreate( + max_tokens=64, + messages=_history(replayed_reasoning=True), + model="openai/gpt-5.6-luna", + tools=[dict(BASH_TOOL)], + num_retries=2, + ) + + assert _replayed_reasoning_counts(upstream) == [1, 1, 0] + assert cast(Mapping[str, object], response)["content"] == [{"type": "text", "text": "done"}] + + +def test_sync_openai_turn_retries_a_server_error_with_the_reasoning_and_drops_it_only_once_rejected( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + upstream: Final = _upstream_after( + respx_mock, + monkeypatch, + (httpx.Response(500, json=_SERVER_ERROR), httpx.Response(400, json=REJECTED_REPLAYED_REASONING)), + False, + ) + + response: Final = litellm.anthropic.messages.create( + max_tokens=64, + messages=_history(replayed_reasoning=True), + model="openai/gpt-5.6-luna", + tools=[dict(BASH_TOOL)], + num_retries=2, + ) + + assert _replayed_reasoning_counts(upstream) == [1, 1, 0] + assert cast(Mapping[str, object], response)["content"] == [{"type": "text", "text": "done"}] + + +class _FailureLog(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.rejections: tuple[str, ...] = () + + def _record(self, kwargs: Mapping[str, object]) -> None: + error: Final = str(kwargs.get("exception")) + self.rejections = ( + *self.rejections, + next((phrase for phrase in ("encrypted content", "organization") if phrase in error), error), + ) + + def log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._record(kwargs) + + async def async_log_failure_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + self._record(kwargs) + + +_OUTCOMES_OF_THE_TURN_SENT_AGAIN: Final = ( + pytest.param(httpx.Response(200, json=TEXT_RESPONSE), (), id="served"), + pytest.param(httpx.Response(400, json=UNVERIFIED_ORGANIZATION), ("organization",), id="rejected"), +) + + +def _rejected_then(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch, sent_again: httpx.Response) -> None: + monkeypatch.setenv("OPENAI_API_KEY", "sk-replay-test") + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post("https://api.openai.com/v1/responses").mock( + side_effect=[httpx.Response(400, json=REJECTED_REPLAYED_REASONING), sent_again] + ) + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("sent_again", "logged_rejections"), _OUTCOMES_OF_THE_TURN_SENT_AGAIN) +async def test_openai_turn_logs_only_the_outcome_of_the_turn_sent_again( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + sent_again: httpx.Response, + logged_rejections: tuple[str, ...], +) -> None: + failure_log: Final = _FailureLog() + monkeypatch.setattr(litellm, "callbacks", [failure_log]) + _rejected_then(respx_mock, monkeypatch, sent_again) + + with pytest.raises(litellm.BadRequestError) if logged_rejections else nullcontext(): + await litellm.anthropic.messages.acreate( + max_tokens=64, + messages=_history(replayed_reasoning=True), + model="openai/gpt-5.6-luna", + tools=[dict(BASH_TOOL)], + ) + + assert failure_log.rejections == logged_rejections + + +@pytest.mark.parametrize(("sent_again", "logged_rejections"), _OUTCOMES_OF_THE_TURN_SENT_AGAIN) +def test_sync_openai_turn_logs_only_the_outcome_of_the_turn_sent_again( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + sent_again: httpx.Response, + logged_rejections: tuple[str, ...], +) -> None: + failure_log: Final = _FailureLog() + monkeypatch.setattr(litellm, "callbacks", [failure_log]) + _rejected_then(respx_mock, monkeypatch, sent_again) + + with pytest.raises(litellm.BadRequestError) if logged_rejections else nullcontext(): + litellm.anthropic.messages.create( + max_tokens=64, + messages=_history(replayed_reasoning=True), + model="openai/gpt-5.6-luna", + tools=[dict(BASH_TOOL)], + ) + + assert failure_log.rejections == logged_rejections + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + "replayed_reasoning, rejection", + [ + pytest.param(False, REJECTED_REPLAYED_REASONING, id="no_replayed_reasoning"), + pytest.param(True, UNVERIFIED_ORGANIZATION, id="unrelated_rejection"), + ], +) +async def test_openai_turn_is_not_sent_again_for_other_rejections( + respx_mock: respx.MockRouter, + monkeypatch: pytest.MonkeyPatch, + replayed_reasoning: bool, + rejection: Mapping[str, object], +) -> None: + upstream: Final = _upstream(respx_mock, monkeypatch, httpx.Response(400, json=rejection), False) + + with pytest.raises(litellm.BadRequestError): + await LiteLLMMessagesToResponsesAPIHandler.async_anthropic_messages_handler( + max_tokens=64, + messages=_history(replayed_reasoning=replayed_reasoning), + model="openai/gpt-5.6-luna", + tools=[dict(BASH_TOOL)], + custom_llm_provider="openai", + ) + + assert upstream.call_count == 1 + + +class _SuccessLog(CustomLogger): + def __init__(self) -> None: + super().__init__() + self.logged_reasoning_item_ids: tuple[tuple[object, ...], ...] = () + self.logged: Final = asyncio.Event() + + async def async_log_success_event( + self, kwargs: Mapping[str, object], response_obj: object, start_time: object, end_time: object + ) -> None: + payload: Final = cast(StandardLoggingPayload, kwargs["standard_logging_object"]) + logged: Final = cast(Mapping[str, Sequence[Mapping[str, Mapping[str, object]]]], payload["response"]) + message: Final = logged["choices"][0]["message"] + items: Final = cast(Sequence[Mapping[str, object]], message.get("reasoning_items") or ()) + self.logged_reasoning_item_ids = (*self.logged_reasoning_item_ids, tuple(item["id"] for item in items)) + self.logged.set() + + +def _reasoning_before_call(index: int) -> tuple[dict[str, object], dict[str, object]]: + return ( + {"type": "reasoning", "id": f"rs_{index}", "summary": [], "encrypted_content": f"gAAA_rs_{index}"}, + { + "type": "function_call", + "id": f"fc_{index}", + "status": "completed", + "arguments": "{}", + "call_id": f"call_{index}", + "name": "Bash", + }, + ) + + +@pytest.mark.asyncio +async def test_openai_turn_logs_every_reasoning_item_the_messages_client_was_served( + respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + success_log: Final = _SuccessLog() + monkeypatch.setattr(litellm, "callbacks", [success_log]) + monkeypatch.setenv("OPENAI_API_KEY", "sk-replay-test") + monkeypatch.setenv("DISABLE_AIOHTTP_TRANSPORT", "True") + litellm.in_memory_llm_clients_cache.flush_cache() + respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response( + 200, + json={**TEXT_RESPONSE, "output": [*_reasoning_before_call(1), *_reasoning_before_call(2)]}, + ) + ) + + reply: Final = await litellm.anthropic.messages.acreate( + max_tokens=64, + messages=[{"role": "user", "content": "List the logs twice."}], + model="openai/gpt-5.6-luna", + tools=[dict(BASH_TOOL)], + ) + GLOBAL_LOGGING_WORKER.start() + await asyncio.wait_for(success_log.logged.wait(), timeout=30.0) + + assert [block["type"] for block in reply["content"]].count("redacted_thinking") == 2 + assert success_log.logged_reasoning_item_ids == (("rs_1", "rs_2"),) diff --git a/tests/unit/llms/bedrock/chat/test_converse_transformation_nova_2.py b/tests/unit/llms/bedrock/chat/test_converse_transformation_nova_2.py index 58058a2e1d4..f29ad3b5844 100644 --- a/tests/unit/llms/bedrock/chat/test_converse_transformation_nova_2.py +++ b/tests/unit/llms/bedrock/chat/test_converse_transformation_nova_2.py @@ -637,8 +637,7 @@ class TestNova2MultiTurnMessageTranslation: assert "Simple addition" in texts assert "4." in texts - def test_should_keep_signed_thinking_blocks_as_reasoning_content(self): - """With a signature, reasoning is preserved as a reasoningContent block.""" + def test_should_drop_signed_thinking_blocks_from_another_model(self): bedrock_msgs = self._to_bedrock( [ {"role": "user", "content": "What is 2+2?"}, @@ -653,13 +652,9 @@ class TestNova2MultiTurnMessageTranslation: ] ) assistant = next(m for m in bedrock_msgs if m["role"] == "assistant") - rc_blocks = [b for b in assistant["content"] if "reasoningContent" in b] - assert len(rc_blocks) >= 1 - assert rc_blocks[0]["reasoningContent"]["reasoningText"]["text"] == "math" - assert rc_blocks[0]["reasoningContent"]["reasoningText"]["signature"] == "sig-1" + assert assistant["content"] == [{"text": "4."}] - def test_should_translate_inline_content_list_thinking_type(self): - """content=[{type:'thinking',...},{type:'text',...}] should also round-trip.""" + def test_should_drop_inline_signed_thinking_and_keep_the_answer(self): bedrock_msgs = self._to_bedrock( [ {"role": "user", "content": "Hi"}, @@ -674,11 +669,4 @@ class TestNova2MultiTurnMessageTranslation: ] ) assistant = next(m for m in bedrock_msgs if m["role"] == "assistant") - rc_blocks = [b for b in assistant["content"] if "reasoningContent" in b] - text_blocks = [ - b - for b in assistant["content"] - if "text" in b and "reasoningContent" not in b - ] - assert len(rc_blocks) >= 1 - assert any("Hello!" in b["text"] for b in text_blocks) + assert assistant["content"] == [{"text": "Hello!"}] diff --git a/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py index 84d3160e609..e418f6778ee 100644 --- a/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py +++ b/tests/unit/llms/bedrock/passthrough/test_bedrock_passthrough_transformation.py @@ -5,8 +5,9 @@ import struct import sys import types from binascii import crc32 -from collections.abc import Callable +from collections.abc import Callable, Mapping from datetime import datetime +from types import MappingProxyType from typing import Final from unittest.mock import patch @@ -17,6 +18,7 @@ import litellm from litellm.integrations.custom_logger import CustomLogger from litellm.litellm_core_utils.litellm_logging import Logging from litellm.llms.base_llm.passthrough.transformation import PassthroughStreamCollector +from litellm.llms.bedrock.common_utils import BedrockError from litellm.llms.bedrock.passthrough.transformation import ( BedrockPassthroughConfig, is_bedrock_streaming_endpoint, @@ -898,3 +900,197 @@ def test_stream_collector_logs_nothing_for_an_unrecognized_endpoint(): collector.add(_event_frame("messageStart", {"role": "assistant"})) assert collector.build_logged_response(_converse_stream_logging_obj()) is None + + +CONVERSE_BODY = {"messages": [{"role": "user", "content": [{"text": "hi"}]}]} +INVOKE_BODY = { + "anthropic_version": "bedrock-2023-05-31", + "max_tokens": 8, + "messages": [{"role": "user", "content": "hi"}], +} +DEPLOYMENT_GUARDRAIL: Final[Mapping[str, str]] = { + "guardrailIdentifier": "deployment-guardrail", + "guardrailVersion": "3", + "trace": "enabled", + "streamProcessingMode": "sync", +} +DEPLOYMENT_INVOKE_GUARDRAIL_HEADERS: Final[Mapping[str, str]] = { + "x-amzn-bedrock-guardrailidentifier": "deployment-guardrail", + "x-amzn-bedrock-guardrailversion": "3", + "x-amzn-bedrock-trace": "ENABLED", +} +CALLER_GUARDRAIL: Final[Mapping[str, str]] = { + "guardrailIdentifier": "caller-guardrail", + "guardrailVersion": "9", + "trace": "disabled", +} +CALLER_INVOKE_GUARDRAIL_HEADERS: Final[Mapping[str, str]] = { + "x-pass-X-Amzn-Bedrock-GuardrailIdentifier": "caller-guardrail", + "x-pass-X-Amzn-Bedrock-GuardrailVersion": "9", + "x-pass-X-Amzn-Bedrock-Trace": "DISABLED", +} + + +async def _sent_passthrough_request( + action: str, + body: dict[str, object], + deployment_params: Mapping[str, object] = MappingProxyType({}), + **kwargs: object, +) -> httpx.Request: + requests: Final[list[httpx.Request]] = [] + + def respond(request: httpx.Request) -> httpx.Response: + requests.append(request) + return httpx.Response(200, json={"ok": True}, request=request) + + router: Final = litellm.Router( + model_list=[ + { + "model_name": "bedrock-pt", + "litellm_params": { + "model": "bedrock/us.anthropic.claude-haiku-4-5-20251001-v1:0", + "aws_access_key_id": "AKIAROUTERKEY", + "aws_secret_access_key": "router-secret", + "aws_region_name": "us-west-2", + **deployment_params, + }, + } + ] + ) + await router.allm_passthrough_route( + model="bedrock-pt", + method="POST", + endpoint=f"/model/bedrock-pt/{action}", + data=body, + client=AsyncHTTPHandler(transport=httpx.MockTransport(respond)), + **kwargs, + ) + return requests[0] + + +def _signed_header_names(request: httpx.Request) -> frozenset[str]: + return frozenset(request.headers["authorization"].split("SignedHeaders=")[1].split(",")[0].split(";")) + + +def _invoke_guardrail_headers(request: httpx.Request) -> dict[str, str]: + return { + name: value + for name, value in request.headers.items() + if name.lower().startswith("x-amzn-bedrock-guardrail") or name.lower() == "x-amzn-bedrock-trace" + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("action", "sent_guardrail"), + [ + ("converse", {"guardrailIdentifier": "deployment-guardrail", "guardrailVersion": "3", "trace": "enabled"}), + ("converse-stream", DEPLOYMENT_GUARDRAIL), + ("converse/", {"guardrailIdentifier": "deployment-guardrail", "guardrailVersion": "3", "trace": "enabled"}), + ], +) +async def test_converse_passthrough_sends_the_deployment_guardrail_in_the_body( + action: str, sent_guardrail: Mapping[str, str] +) -> None: + request: Final = await _sent_passthrough_request( + action, CONVERSE_BODY, deployment_params={"guardrailConfig": DEPLOYMENT_GUARDRAIL} + ) + + assert json.loads(request.content) == {**CONVERSE_BODY, "guardrailConfig": sent_guardrail} + assert _invoke_guardrail_headers(request) == {} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +async def test_invoke_passthrough_sends_the_deployment_guardrail_as_signed_headers(action: str) -> None: + request: Final = await _sent_passthrough_request( + action, INVOKE_BODY, deployment_params={"guardrailConfig": DEPLOYMENT_GUARDRAIL} + ) + + assert _invoke_guardrail_headers(request) == DEPLOYMENT_INVOKE_GUARDRAIL_HEADERS + assert set(DEPLOYMENT_INVOKE_GUARDRAIL_HEADERS) <= _signed_header_names(request) + assert json.loads(request.content) == INVOKE_BODY + + +@pytest.mark.asyncio +@pytest.mark.parametrize( + ("action", "caller_guardrail"), + [("converse", CALLER_GUARDRAIL), ("converse-stream", {**CALLER_GUARDRAIL, "streamProcessingMode": "async"})], +) +async def test_deployment_guardrail_replaces_the_caller_converse_guardrail( + action: str, caller_guardrail: Mapping[str, str] +) -> None: + request: Final = await _sent_passthrough_request( + action, + {**CONVERSE_BODY, "guardrailConfig": caller_guardrail}, + deployment_params={"guardrailConfig": {"guardrailIdentifier": "deployment-guardrail"}}, + ) + + assert json.loads(request.content) == { + **CONVERSE_BODY, + "guardrailConfig": {"guardrailIdentifier": "deployment-guardrail"}, + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +async def test_deployment_guardrail_replaces_every_invoke_guardrail_header_the_caller_sent(action: str) -> None: + request: Final = await _sent_passthrough_request( + action, + INVOKE_BODY, + deployment_params={"guardrailConfig": {"guardrailIdentifier": "deployment-guardrail"}}, + request_headers=CALLER_INVOKE_GUARDRAIL_HEADERS, + ) + + assert _invoke_guardrail_headers(request) == {"x-amzn-bedrock-guardrailidentifier": "deployment-guardrail"} + assert json.loads(request.content) == INVOKE_BODY + + +@pytest.mark.asyncio +async def test_passthrough_without_a_deployment_guardrail_sends_the_caller_guardrail_as_is() -> None: + caller_body: Final = {**CONVERSE_BODY, "guardrailConfig": CALLER_GUARDRAIL} + + converse: Final = await _sent_passthrough_request("converse", caller_body) + invoke: Final = await _sent_passthrough_request( + "invoke", INVOKE_BODY, request_headers=CALLER_INVOKE_GUARDRAIL_HEADERS + ) + + assert json.loads(converse.content) == caller_body + assert _invoke_guardrail_headers(invoke) == { + "x-amzn-bedrock-guardrailidentifier": "caller-guardrail", + "x-amzn-bedrock-guardrailversion": "9", + "x-amzn-bedrock-trace": "DISABLED", + } + + +@pytest.mark.asyncio +@pytest.mark.parametrize("action", ["invoke", "invoke-with-response-stream"]) +@pytest.mark.parametrize("guardrail_config", ["deployment-guardrail", {"guardrailVersion": "3"}]) +async def test_invoke_passthrough_refuses_a_deployment_guardrail_it_cannot_send( + action: str, guardrail_config: object +) -> None: + with pytest.raises(BedrockError, match="guardrailConfig") as rejected: + await _sent_passthrough_request(action, INVOKE_BODY, deployment_params={"guardrailConfig": guardrail_config}) + + assert rejected.value.status_code == 400 + + +@pytest.mark.asyncio +async def test_converse_passthrough_forwards_a_non_object_deployment_guardrail_unchanged() -> None: + request: Final = await _sent_passthrough_request( + "converse", CONVERSE_BODY, deployment_params={"guardrailConfig": "deployment-guardrail"} + ) + + assert json.loads(request.content) == {**CONVERSE_BODY, "guardrailConfig": "deployment-guardrail"} + + +@pytest.mark.asyncio +async def test_passthrough_leaves_actions_without_a_guardrail_field_as_sent() -> None: + count_tokens_body: Final = {"input": {"converse": CONVERSE_BODY}} + + request: Final = await _sent_passthrough_request( + "count-tokens", count_tokens_body, deployment_params={"guardrailConfig": DEPLOYMENT_GUARDRAIL} + ) + + assert json.loads(request.content) == count_tokens_body + assert _invoke_guardrail_headers(request) == {} diff --git a/tests/unit/llms/openai/responses/test_openai_responses_transformation.py b/tests/unit/llms/openai/responses/test_openai_responses_transformation.py index 2fcb70ec40b..0a0f0d33543 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_transformation.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_transformation.py @@ -13,6 +13,8 @@ import litellm from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.llms.azure.responses.transformation import AzureOpenAIResponsesAPIConfig from litellm.llms.base_llm.responses.transformation import BaseResponsesAPIConfig +from litellm.llms.bedrock.responses.transformation import BedrockOpenAIResponsesConfig +from litellm.llms.bedrock_mantle.responses.transformation import BedrockMantleResponsesAPIConfig from litellm.llms.openai.responses.transformation import OpenAIResponsesAPIConfig from litellm.responses.litellm_completion_transformation.transformation import LiteLLMCompletionResponsesConfig from litellm.types.llms.openai import ( @@ -1298,7 +1300,15 @@ class TestOpenAIResponsesAPIConfig: ) return list(bridged.output) - @pytest.mark.parametrize("config", [OpenAIResponsesAPIConfig(), AzureOpenAIResponsesAPIConfig()]) + @pytest.mark.parametrize( + "config", + [ + OpenAIResponsesAPIConfig(), + AzureOpenAIResponsesAPIConfig(), + BedrockOpenAIResponsesConfig(), + BedrockMantleResponsesAPIConfig(), + ], + ) def test_claude_reasoning_minted_by_the_bridge_is_dropped_before_the_history_reaches_openai(self, config): saved_claude_turn = json.loads( json.dumps([item.model_dump() for item in self._claude_turn_bridged_to_responses_output()]) diff --git a/tests/unit/llms/vertex_ai/batches/test_handler.py b/tests/unit/llms/vertex_ai/batches/test_handler.py index 91a2adb9d19..2c246e065a8 100644 --- a/tests/unit/llms/vertex_ai/batches/test_handler.py +++ b/tests/unit/llms/vertex_ai/batches/test_handler.py @@ -30,6 +30,7 @@ from __future__ import annotations import asyncio import json +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import httpx @@ -41,9 +42,11 @@ from litellm.llms.vertex_ai.batches.handler import ( # noqa: E402 VertexAIBatchPrediction, ) from litellm.llms.vertex_ai.common_utils import VertexAIError # noqa: E402 +from litellm.proxy.utils import handle_exception_on_proxy # noqa: E402 from litellm.types.utils import LiteLLMBatch # noqa: E402 HMOD = "litellm.llms.vertex_ai.batches.handler" +_CLIENT_ERROR_429: Final = "Client error '429 Too Many Requests' for url 'https://x/batchPredictionJobs'" TOKEN = "ya29.fake-access-token" PROJECT = "my-project" LOCATION = "us-central1" @@ -370,19 +373,19 @@ def test_create_batch_sync_endpoint_without_deployed_model_raises_400(): client.post.assert_not_called() -def test_create_batch_sync_httpstatuserror_propagates(): - """``HTTPHandler.post`` raises for non-2xx via ``raise_for_status``; the - sync create path must surface that error, not swallow it.""" +def test_create_batch_sync_httpstatuserror_carries_vertex_status_body_and_headers(): h = _make_handler() client = MagicMock() request = httpx.Request("POST", "https://x/batchPredictionJobs") - err_response = httpx.Response(status_code=500, request=request, text="boom") + err_response: Final = httpx.Response( + status_code=429, request=request, text="boom", headers={"retry-after": "30"} + ) client.post.side_effect = httpx.HTTPStatusError( - "boom", request=request, response=err_response + _CLIENT_ERROR_429, request=request, response=err_response ) with patch(f"{HMOD}.get_httpx_client", return_value=client): - with pytest.raises(httpx.HTTPStatusError): + with pytest.raises(httpx.HTTPStatusError) as exc_info: h.create_batch( _is_async=False, create_batch_data=CREATE_DATA, @@ -393,6 +396,11 @@ def test_create_batch_sync_httpstatuserror_propagates(): timeout=600.0, max_retries=None, ) + assert exc_info.value.response.status_code == 429 + assert exc_info.value.response.headers["retry-after"] == "30" + assert handle_exception_on_proxy(exc_info.value).code == "429" + assert (getattr(exc_info.value, "message", None), getattr(exc_info.value, "text", None)) == ("boom", "boom") + assert str(exc_info.value).splitlines() == [_CLIENT_ERROR_429, "boom"] def test_create_batch_input_file_id_without_model_raises_400_before_post(): @@ -800,14 +808,16 @@ def test_cancel_batch_sync_httpstatuserror_logged_and_reraised(): client.get.assert_not_called() -def test_create_batch_async_httpstatuserror_logged_and_reraised(): +def test_create_batch_async_httpstatuserror_carries_vertex_status_body_and_headers(): h = _make_handler() async_client = MagicMock() request = httpx.Request("POST", "https://x/batchPredictionJobs") - err_response = httpx.Response(status_code=500, request=request, text="boom") + err_response: Final = httpx.Response( + status_code=429, request=request, text="boom", headers={"retry-after": "30"} + ) async_client.post = AsyncMock( side_effect=httpx.HTTPStatusError( - "boom", request=request, response=err_response + _CLIENT_ERROR_429, request=request, response=err_response ) ) @@ -825,8 +835,13 @@ def test_create_batch_async_httpstatuserror_logged_and_reraised(): timeout=600.0, max_retries=None, ) - with pytest.raises(httpx.HTTPStatusError): + with pytest.raises(httpx.HTTPStatusError) as exc_info: _run(coro) + assert exc_info.value.response.status_code == 429 + assert exc_info.value.response.headers["retry-after"] == "30" + assert handle_exception_on_proxy(exc_info.value).code == "429" + assert (getattr(exc_info.value, "message", None), getattr(exc_info.value, "text", None)) == ("boom", "boom") + assert str(exc_info.value).splitlines() == [_CLIENT_ERROR_429, "boom"] def test_async_retrieve_batch_non_200_raises(): @@ -1043,15 +1058,7 @@ async def test_vertex_list_batches(monkeypatch): @pytest.mark.asyncio -async def test_vertex_async_create_batch_logs_error_body_on_http_error(): - """ - When Vertex AI returns an HTTP error (e.g. 400), _async_create_batch should - re-raise httpx.HTTPStatusError (not swallow it) and log the response body. - - Before the fix the error body was lost because AsyncHTTPHandler.post() - calls raise_for_status() internally, raising before the handler's own - status-code check could log the body. - """ +async def test_vertex_async_create_batch_httpstatuserror_carries_vertex_error_body(): from litellm.llms.vertex_ai.batches.handler import VertexAIBatchPrediction handler = VertexAIBatchPrediction(gcs_bucket_name="test-bucket") @@ -1081,4 +1088,4 @@ async def test_vertex_async_create_batch_logs_error_body_on_http_error(): ) assert exc_info.value.response.status_code == 400 - assert "gemini-2.0-flash" in exc_info.value.response.text + assert "Do not support publisher model gemini-2.0-flash" in str(exc_info.value) diff --git a/tests/unit/llms/vertex_ai/batches/test_transformation.py b/tests/unit/llms/vertex_ai/batches/test_transformation.py index 064845a7266..1440a57a15f 100644 --- a/tests/unit/llms/vertex_ai/batches/test_transformation.py +++ b/tests/unit/llms/vertex_ai/batches/test_transformation.py @@ -21,6 +21,8 @@ import pytest from litellm.llms.vertex_ai.batches.transformation import ( # noqa: E402 VertexAIBatchTransformation, + is_vertex_anthropic_batch_model, + native_vertex_batch_row_stats, ) from litellm.llms.vertex_ai.common_utils import VertexAIError # noqa: E402 from litellm.types.utils import LiteLLMBatch # noqa: E402 @@ -493,3 +495,61 @@ def test_get_gcs_uri_prefix_keeps_passthrough_segment_so_output_lands_beside_inp T._get_gcs_uri_prefix_from_file(PASSTHROUGH_INPUT_FILE) == "gs://litellm-testing-bucket/litellm-vertex-files/passthrough/publishers/google/models/gemini-2.5-flash" ) + + +CLAUDE_INPUT_FILE = ( + "gs://litellm-testing-bucket/litellm-vertex-files/publishers/anthropic/models/claude-sonnet-4-6/uuid-1" +) + + +def test_claude_batch_job_targets_the_anthropic_publisher_model(): + job: Final = T.transform_openai_batch_request_to_vertex_ai_batch_request( + {"input_file_id": CLAUDE_INPUT_FILE}, vertex_project="p", vertex_location="us-east5" + ) + assert job["model"] == "publishers/anthropic/models/claude-sonnet-4-6" + + +@pytest.mark.parametrize( + ("model", "expected"), + [ + ("claude-sonnet-4-6", True), + ("vertex_ai/claude-opus-4-1@20250805", True), + ("publishers/anthropic/models/claude-haiku-4-5", True), + ("gemini-2.5-flash", False), + ("publishers/google/models/gemini-2.5-flash", False), + ], +) +def test_is_vertex_anthropic_batch_model(model: str, expected: bool): + assert is_vertex_anthropic_batch_model(model) is expected + + +def test_claude_output_row_is_costed_from_anthropic_usage_and_its_own_model(): + def cost_calculator(usage, model, custom_llm_provider=None, model_info=None, vertex_location=None) -> tuple[float, float]: + if model != "claude-sonnet-4-6": + return 0.0, 0.0 + return usage.prompt_tokens * 0.001, usage.completion_tokens * 0.002 + + stats: Final = native_vertex_batch_row_stats( + { + "custom_id": "r1", + "request": {"messages": [{"role": "user", "content": "hi"}], "anthropic_version": "vertex-2023-10-16"}, + "response": { + "id": "msg_vrtx_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": "hello"}], + "stop_reason": "end_turn", + "usage": {"input_tokens": 12, "output_tokens": 5}, + }, + }, + None, + model_info=None, + calculate_usage=lambda _: pytest.fail("Claude rows must not go through the Gemini usage parser"), + cost_calculator=cost_calculator, + ) + + assert stats is not None, "a successful Claude row must be costed, not counted as failed" + assert (stats.usage.prompt_tokens, stats.usage.completion_tokens, stats.total_tokens) == (12, 5, 17) + assert (stats.prompt_cost, stats.completion_cost) == pytest.approx((0.012, 0.010)) + assert (stats.prompt_cost, stats.completion_cost) == pytest.approx((12 * 0.001, 5 * 0.002)) diff --git a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py index dde3df6fcc0..c235e438847 100644 --- a/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py +++ b/tests/unit/llms/vertex_ai/files/test_vertex_ai_files_transformation.py @@ -5,7 +5,8 @@ Includes tests for Vertex AI batch output transformation to OpenAI format. import json import urllib.parse -from collections.abc import Mapping +from collections.abc import Mapping, Sequence +from typing import Final from types import MappingProxyType from urllib.parse import parse_qs, urlparse @@ -13,6 +14,8 @@ import httpx import pytest from unittest.mock import MagicMock +import litellm +from litellm.llms.base_llm.files.transformation import BaseFileUploadStream from litellm.llms.vertex_ai.files.transformation import ( VertexAIFilesConfig, _get_litellm_batch_custom_id_from_labels, @@ -165,6 +168,17 @@ class TestBatchObjectNaming: object_name = config._get_gcs_object_name_from_batch_jsonl([{"body": {"model": "gemini-2.5-flash"}}]) assert object_name.startswith("litellm-vertex-files/publishers/google/models/gemini-2.5-flash/") + @pytest.mark.parametrize( + ("model", "stored_model"), + [ + ("multimodalembedding@001", "publishers/google/models/multimodalembedding_001/"), + ("claude-haiku-4-5@20251001", "publishers/anthropic/models/claude-haiku-4-5@20251001/"), + ], + ) + def test_only_claude_models_keep_the_at_version_in_their_object_name(self, config, model, stored_model): + object_name: Final = config._get_gcs_object_name_from_batch_jsonl([{"body": {"model": model}}]) + assert object_name.startswith(f"litellm-vertex-files/{stored_model}") + def test_should_store_fine_tuned_endpoint_under_endpoints_path(self, config): """A numeric endpoint id must not be filed under publishers/google/models/gemini/, which the batch transformation later mangles into a nonexistent publisher model (LIT-6899).""" @@ -1948,3 +1962,420 @@ class TestVertexEmbeddingsBatchOutputTranslation: content = json.dumps(legacy_row).encode("utf-8") assert config._try_transform_vertex_batch_output_to_openai(content) == content + + +CLAUDE_DEPLOYMENT = {"gcs_bucket_name": "my-bucket", "model": "vertex_ai/claude-sonnet-4-6"} + + +def _claude_chat_entry(custom_id: str | None, content: str, **body: object) -> dict[str, object]: + entry: Final[dict[str, object]] = { + "method": "POST", + "url": "/v1/chat/completions", + "body": {"model": "claude-sonnet-4-6", "messages": [{"role": "user", "content": content}], **body}, + } + return entry if custom_id is None else {"custom_id": custom_id, **entry} + + +def _uploaded_rows( + config: VertexAIFilesConfig, litellm_params: Mapping[str, object], entries: Sequence[Mapping[str, object]] +) -> list: + upload: Final = config.transform_create_file_request( + model="", + create_file_data={ + "file": ("batch.jsonl", "\n".join(json.dumps(entry) for entry in entries).encode(), "application/jsonl"), + "purpose": "batch", + }, + optional_params={}, + litellm_params=dict(litellm_params), + ) + match upload: + case {"streaming_media_upload": {"body_stream": BaseFileUploadStream() as stream}}: + return [json.loads(line) for line in b"".join(stream.iter_bytes()).splitlines()] + case _: + pytest.fail(f"expected a streamed upload, got {upload!r}") + + +def _claude_output_row(custom_id: str, text: str, content: object = None) -> dict[str, object]: + return { + "custom_id": custom_id, + "status": "", + "processed_time": "2026-10-28T20:37:47.634+00:00", + "response": { + "id": "msg_vrtx_01", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-4-6", + "content": [{"type": "text", "text": text}] if content is None else content, + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 12, "output_tokens": 5}, + }, + } + + +class TestVertexClaudeBatch: + """Claude on Vertex batches only as `publishers/anthropic/models/` with Anthropic Messages rows + (https://github.com/BerriAI/litellm/issues/45671)""" + + @pytest.mark.parametrize("claude_model", ["claude-sonnet-4-6", "claude-haiku-4-5@20251001"]) + def test_claude_upload_creates_a_batch_job_for_the_anthropic_publisher_model(self, config, claude_model): + from litellm.llms.vertex_ai.batches.transformation import VertexAIBatchTransformation + + url: Final = config.get_complete_file_url( + api_base=None, + api_key=None, + model="", + optional_params={}, + litellm_params={**CLAUDE_DEPLOYMENT, "model": f"vertex_ai/{claude_model}"}, + data={"file": ("batch.jsonl", json.dumps(_claude_chat_entry("r1", "hi")).encode()), "purpose": "batch"}, + ) + object_name: Final = parse_qs(urlparse(url).query)["name"][0] + + job: Final = VertexAIBatchTransformation.transform_openai_batch_request_to_vertex_ai_batch_request( + {"input_file_id": f"gs://my-bucket/{object_name}"}, vertex_project="p", vertex_location="us-east5" + ) + + assert job["model"] == f"publishers/anthropic/models/{claude_model}", object_name + + def test_claude_body_model_without_deployment_is_filed_under_anthropic_publisher(self, config): + object_name: Final = config._get_gcs_object_name_from_batch_jsonl([{"body": {"model": "claude-haiku-4-5"}}]) + assert object_name.startswith("litellm-vertex-files/publishers/anthropic/models/claude-haiku-4-5/") + + def test_publisher_path_model_is_not_nested_under_google(self, config): + object_name: Final = config._get_gcs_object_name_from_batch_jsonl( + [{"body": {}}], deployment_model="vertex_ai/publishers/anthropic/models/claude-sonnet-4-6" + ) + assert object_name.startswith("litellm-vertex-files/publishers/anthropic/models/claude-sonnet-4-6/") + + def test_claude_rows_are_uploaded_in_anthropic_messages_shape(self, config): + rows: Final = _uploaded_rows( + config, + CLAUDE_DEPLOYMENT, + [_claude_chat_entry("request-1", "Hello!", max_tokens=50, temperature=0.2)], + ) + + assert rows == [ + { + "custom_id": "request-1", + "request": { + "messages": [{"role": "user", "content": [{"type": "text", "text": "Hello!"}]}], + "max_tokens": 50, + "temperature": 0.2, + "anthropic_version": "vertex-2023-10-16", + }, + } + ] + + def test_claude_system_message_becomes_the_anthropic_system_field(self, config): + entry: Final = _claude_chat_entry( + "r1", + "unused", + max_tokens=10, + messages=[{"role": "system", "content": "be brief"}, {"role": "user", "content": "Hello!"}], + ) + + (row,) = _uploaded_rows(config, CLAUDE_DEPLOYMENT, [entry]) + + assert row["request"]["system"] == [{"type": "text", "text": "be brief"}] + assert [message["role"] for message in row["request"]["messages"]] == ["user"] + + def test_claude_row_drops_an_unsupported_param_under_deployment_drop_params(self, config, monkeypatch): + monkeypatch.setitem(litellm.model_cost, "claude-zeta-9", {"supports_sampling_params": False}) + monkeypatch.setattr(litellm, "drop_params", False) + entry: Final = _claude_chat_entry("r1", "hi", max_tokens=10, top_p=0.5) + + (row,) = _uploaded_rows( + config, {**CLAUDE_DEPLOYMENT, "model": "vertex_ai/claude-zeta-9", "drop_params": True}, [entry] + ) + + assert row["request"]["max_tokens"] == 10 + assert "top_p" not in row["request"] + + def test_claude_row_drops_a_param_the_deployment_lists_in_additional_drop_params(self, config, monkeypatch): + monkeypatch.setitem(litellm.model_cost, "claude-zeta-9", {"supports_sampling_params": False}) + monkeypatch.setattr(litellm, "drop_params", False) + entry: Final = _claude_chat_entry("r1", "hi", max_tokens=10, top_p=0.5) + + (row,) = _uploaded_rows( + config, + {**CLAUDE_DEPLOYMENT, "model": "vertex_ai/claude-zeta-9", "additional_drop_params": ["top_p"]}, + [entry], + ) + + assert row["request"]["max_tokens"] == 10 + assert "top_p" not in row["request"] + + def test_claude_rows_take_deployment_params_as_defaults_the_row_overrides(self, config): + rows: Final = _uploaded_rows( + config, + {**CLAUDE_DEPLOYMENT, "max_tokens": 64, "temperature": 0.3}, + [_claude_chat_entry("r1", "hi"), _claude_chat_entry("r2", "hi", max_tokens=10)], + ) + + assert [(row["request"]["max_tokens"], row["request"]["temperature"]) for row in rows] == [(64, 0.3), (10, 0.3)] + + def test_claude_row_with_an_unsupported_param_is_refused_without_drop_params(self, config, monkeypatch): + from litellm.llms.vertex_ai.common_utils import VertexAIError + + monkeypatch.setitem(litellm.model_cost, "claude-zeta-9", {"supports_sampling_params": False}) + monkeypatch.setattr(litellm, "drop_params", False) + entry: Final = _claude_chat_entry("r1", "hi", max_tokens=10, top_p=0.5) + + with pytest.raises(VertexAIError, match=r"line 1 .*top_p=0.5.*drop_params") as exc_info: + _uploaded_rows(config, {**CLAUDE_DEPLOYMENT, "model": "vertex_ai/claude-zeta-9"}, [entry]) + assert exc_info.value.status_code == 400 + + def test_claude_deployment_wins_over_a_gemini_body_model(self, config): + entry: Final = _claude_chat_entry("r1", "hi", max_tokens=10, model="gemini-2.5-flash") + + (row,) = _uploaded_rows(config, CLAUDE_DEPLOYMENT, [entry]) + + assert row["request"]["anthropic_version"] == "vertex-2023-10-16" + assert "contents" not in row["request"] + + def test_gemini_deployment_still_uploads_generate_content_rows(self, config): + (row,) = _uploaded_rows( + config, + {"gcs_bucket_name": "my-bucket", "model": "vertex_ai/gemini-2.5-flash"}, + [_claude_chat_entry("r1", "hi")], + ) + assert row["request"]["contents"] == [{"role": "user", "parts": [{"text": "hi"}]}] + assert "custom_id" not in row + + def test_claude_embeddings_batch_line_is_rejected_with_400(self): + from litellm.llms.vertex_ai.common_utils import VertexAIError + + entry: Final = { + "custom_id": "e1", + "url": "/v1/embeddings", + "body": {"model": "claude-sonnet-4-6", "input": "hi"}, + } + with pytest.raises(VertexAIError) as exc_info: + _openai_batch_jsonl_entry_to_vertex_rows(entry, VertexAIFilesConfig()._map_openai_to_vertex_params) + assert exc_info.value.status_code == 400 + + def test_claude_output_rows_become_openai_batch_rows(self, config): + content: Final = b"\n".join( + json.dumps(row).encode() for row in (_claude_output_row("a", "first"), _claude_output_row("b", "second")) + ) + + rows: Final = [ + json.loads(line) for line in config._try_transform_vertex_batch_output_to_openai(content).splitlines() + ] + + assert [row["custom_id"] for row in rows] == ["a", "b"] + assert [row["error"] for row in rows] == [None, None] + bodies: Final = [row["response"]["body"] for row in rows] + assert [body["choices"][0]["message"]["content"] for body in bodies] == ["first", "second"] + assert [body["choices"][0]["finish_reason"] for body in bodies] == ["stop", "stop"] + assert [(body["usage"]["prompt_tokens"], body["usage"]["completion_tokens"]) for body in bodies] == [ + (12, 5), + (12, 5), + ] + assert {body["model"] for body in bodies} == {"claude-sonnet-4-6"} + + @pytest.mark.parametrize("failed_first", [False, True]) + def test_failed_claude_output_row_is_an_openai_error_row(self, config, failed_first): + failed_row: Final = { + "custom_id": "bad", + "status": 'Bad Request: {"type": "error", "error": {"type": "invalid_request_error", "message": "too long"}}', + "processed_time": "2026-10-28T20:37:48.413+00:00", + "response": {}, + } + rows_in: Final = (failed_row, _claude_output_row("ok", "fine")) + content: Final = b"\n".join(json.dumps(row).encode() for row in (rows_in if failed_first else rows_in[::-1])) + + rows_out: Final = { + row["custom_id"]: row + for row in ( + json.loads(line) for line in config._try_transform_vertex_batch_output_to_openai(content).splitlines() + ) + } + + assert rows_out["ok"]["response"]["body"]["choices"][0]["message"]["content"] == "fine" + assert rows_out["bad"]["response"] is None + assert rows_out["bad"]["error"]["code"] == "vertex_ai_error" + assert "too long" in rows_out["bad"]["error"]["message"] + + def test_claude_row_keeps_provider_only_params_and_leaves_out_litellm_ones(self, config): + (row,) = _uploaded_rows( + config, + CLAUDE_DEPLOYMENT, + [_claude_chat_entry("r1", "hi", max_tokens=10, top_k=5, drop_params=True, metadata={"team": "a"})], + ) + + assert row["request"]["top_k"] == 5 + assert not {"drop_params", "metadata", "extra_body"} & set(row["request"]) + + def test_claude_row_top_k_follows_drop_params_on_a_model_without_sampling_params(self, config, monkeypatch): + from litellm.llms.vertex_ai.common_utils import VertexAIError + + monkeypatch.setitem(litellm.model_cost, "claude-zeta-9", {"supports_sampling_params": False}) + monkeypatch.setattr(litellm, "drop_params", False) + deployment: Final = {**CLAUDE_DEPLOYMENT, "model": "vertex_ai/claude-zeta-9"} + entry: Final = _claude_chat_entry("r1", "hi", max_tokens=10, top_k=5) + + with pytest.raises(VertexAIError, match=r"line 1 .*top_k") as exc_info: + _uploaded_rows(config, deployment, [entry]) + (row,) = _uploaded_rows(config, {**deployment, "drop_params": True}, [entry]) + + assert exc_info.value.status_code == 400 + assert "top_k" not in row["request"] + + @pytest.mark.parametrize("param", [{"n": 2}, {"seed": 7}, {"presence_penalty": 0.5}, {"logprobs": True}]) + def test_claude_row_param_claude_does_not_support_follows_drop_params(self, config, monkeypatch, param): + from litellm.llms.vertex_ai.common_utils import VertexAIError + + monkeypatch.setattr(litellm, "drop_params", False) + (name,) = param + entry: Final = _claude_chat_entry("r1", "hi", max_tokens=10, **param) + + with pytest.raises(VertexAIError, match=rf"line 1 .*{name}") as exc_info: + _uploaded_rows(config, CLAUDE_DEPLOYMENT, [entry]) + (row,) = _uploaded_rows(config, {**CLAUDE_DEPLOYMENT, "drop_params": True}, [entry]) + + assert exc_info.value.status_code == 400 + assert name not in row["request"] + + def test_claude_row_refuses_a_tool_name_claude_would_rename(self, config): + from litellm.llms.vertex_ai.common_utils import VertexAIError + + def tool(name: str) -> dict[str, object]: + return {"type": "function", "function": {"name": name, "parameters": {"type": "object", "properties": {}}}} + + renamed: Final = _claude_chat_entry("r1", "hi", max_tokens=10, tools=[tool("get.weather")]) + kept: Final = _claude_chat_entry("r2", "hi", max_tokens=10, tools=[tool("get_weather")]) + + with pytest.raises(VertexAIError, match=r"line 1 .*get\.weather") as exc_info: + _uploaded_rows(config, CLAUDE_DEPLOYMENT, [renamed]) + (row,) = _uploaded_rows(config, CLAUDE_DEPLOYMENT, [kept]) + + assert exc_info.value.status_code == 400 + assert [t["name"] for t in row["request"]["tools"]] == ["get_weather"] + + def test_claude_row_refuses_a_user_tool_named_like_the_response_format_tool(self, config): + from litellm.llms.vertex_ai.common_utils import VertexAIError + + tool: Final = {"type": "function", "function": {"name": "json_tool_call", "parameters": {"type": "object"}}} + entry: Final = _claude_chat_entry("r1", "hi", max_tokens=10, tools=[tool]) + + with pytest.raises(VertexAIError, match=r"line 1 .*json_tool_call") as exc_info: + _uploaded_rows(config, CLAUDE_DEPLOYMENT, [entry]) + + assert exc_info.value.status_code == 400 + + def test_claude_row_leaves_out_stream(self, config): + entry: Final = _claude_chat_entry( + "r1", "hi", max_tokens=10, stream=True, stream_options={"include_usage": True} + ) + + (row,) = _uploaded_rows(config, CLAUDE_DEPLOYMENT, [entry]) + + assert "stream" not in row["request"] + assert "stream_options" not in row["request"] + + def test_claude_row_merges_extra_body_into_the_request(self, config): + (row,) = _uploaded_rows( + config, CLAUDE_DEPLOYMENT, [_claude_chat_entry("r1", "hi", max_tokens=10, extra_body={"top_k": 3})] + ) + + assert row["request"]["top_k"] == 3 + assert "extra_body" not in row["request"] + + def test_claude_response_format_row_round_trips_as_json_content(self, config): + schema: Final = {"type": "object", "properties": {"color": {"type": "string"}}, "required": ["color"]} + (row,) = _uploaded_rows( + config, + CLAUDE_DEPLOYMENT, + [ + _claude_chat_entry( + "r1", + "sky color?", + max_tokens=50, + response_format={"type": "json_schema", "json_schema": {"name": "sky", "schema": schema}}, + ) + ], + ) + assert "json_mode" not in row["request"] + (json_tool,) = row["request"]["tools"] + output_row: Final = _claude_output_row( + "r1", + "unused", + content=[{"type": "tool_use", "id": "toolu_1", "name": json_tool["name"], "input": {"color": "blue"}}], + ) + + (out,) = ( + json.loads(line) + for line in config._try_transform_vertex_batch_output_to_openai( + json.dumps(output_row).encode() + ).splitlines() + ) + + message: Final = out["response"]["body"]["choices"][0]["message"] + assert json.loads(message["content"]) == {"color": "blue"} + assert not message.get("tool_calls") + + def test_claude_output_keeps_a_user_tool_call_as_a_tool_call(self, config): + output_row: Final = _claude_output_row( + "r1", + "unused", + content=[{"type": "tool_use", "id": "toolu_1", "name": "get_weather", "input": {"city": "Paris"}}], + ) + + (out,) = ( + json.loads(line) + for line in config._try_transform_vertex_batch_output_to_openai( + json.dumps(output_row).encode() + ).splitlines() + ) + + (tool_call,) = out["response"]["body"]["choices"][0]["message"]["tool_calls"] + assert tool_call["function"]["name"] == "get_weather" + assert json.loads(tool_call["function"]["arguments"]) == {"city": "Paris"} + + def test_claude_input_file_read_back_is_returned_unchanged(self, config): + (row,) = _uploaded_rows(config, CLAUDE_DEPLOYMENT, [_claude_chat_entry("r1", "hi", max_tokens=10)]) + content: Final = json.dumps(row).encode() + + assert config._try_transform_vertex_batch_output_to_openai(content) == content + + def test_malformed_claude_message_fails_only_its_own_row(self, config): + malformed: Final = _claude_output_row("bad", "unused", content=5) + content: Final = json.dumps(_claude_output_row("ok", "fine")).encode() + b"\n" + json.dumps(malformed).encode() + + ok_row, bad_row = ( + json.loads(line) for line in config._try_transform_vertex_batch_output_to_openai(content).splitlines() + ) + + assert ok_row["response"]["body"]["choices"][0]["message"]["content"] == "fine" + assert bad_row["custom_id"] == "bad" + assert bad_row["response"] is None + assert bad_row["error"]["code"] == "transformation_error" + + def test_claude_line_without_custom_id_is_a_400_at_upload(self, config): + from litellm.llms.vertex_ai.common_utils import VertexAIError + + with pytest.raises(VertexAIError, match=r"line 2 .*custom_id") as exc_info: + _uploaded_rows( + config, CLAUDE_DEPLOYMENT, [_claude_chat_entry("r1", "hi", max_tokens=10), _claude_chat_entry(None, "hi")] + ) + + assert exc_info.value.status_code == 400 + + @pytest.mark.asyncio + async def test_streamed_claude_output_is_transformed(self, config): + async def chunks(): + yield json.dumps(_claude_output_row("a", "first")).encode() + b"\n" + + result: Final = await config.transform_file_content_stream( + stream_iterator=chunks(), + headers={}, + request_url="https://storage.googleapis.com/storage/v1/b/my-bucket/o/x?alt=media", + logging_obj=MagicMock(), + litellm_params={}, + ) + body: Final = b"".join([chunk async for chunk in result.stream_iterator]) + + (row,) = (json.loads(line) for line in body.splitlines()) + assert row["custom_id"] == "a" + assert row["response"]["body"]["choices"][0]["message"]["content"] == "first" diff --git a/tests/unit/proxy/auth/test_auth_utils.py b/tests/unit/proxy/auth/test_auth_utils.py index eb76901db45..a8b2ff4f6cd 100644 --- a/tests/unit/proxy/auth/test_auth_utils.py +++ b/tests/unit/proxy/auth/test_auth_utils.py @@ -414,18 +414,11 @@ class TestLogOnceIfBudgetReservationDisabled: def test_logs_info_only_once_when_enabled(self, caplog): with caplog.at_level(logging.INFO, logger="LiteLLM Proxy"): log_once_if_budget_reservation_disabled(disabled=False) - assert not any( - "disable_budget_reservation is enabled" in record.message - for record in caplog.records - ) + assert not any("disable_budget_reservation is enabled" in record.message for record in caplog.records) for _ in range(3): log_once_if_budget_reservation_disabled(disabled=True) - records = [ - record - for record in caplog.records - if "disable_budget_reservation is enabled" in record.message - ] + records = [record for record in caplog.records if "disable_budget_reservation is enabled" in record.message] assert len(records) == 1 assert records[0].levelno == logging.INFO @@ -753,19 +746,28 @@ def test_get_model_from_request_no_request_extracts_model(): ) -@pytest.mark.parametrize("provider,model", [ - ("laya", "english"), ("laya", "multilingual"), ("laya", "typed-decisions"), - ("bespoke", "nimble-latest"), ("bespoke", "bespokelabs/Bespoke-Nimble-9B"), -]) +@pytest.mark.parametrize( + "provider,model", + [ + ("laya", "english"), + ("laya", "multilingual"), + ("laya", "typed-decisions"), + ("bespoke", "nimble-latest"), + ("bespoke", "bespokelabs/Bespoke-Nimble-9B"), + ], +) @pytest.mark.parametrize("suffix", ["", "/"]) def test_oss_native_model_uses_the_classifier_permission_identity(provider: str, model: str, suffix: str) -> None: - assert get_model_from_request( - request_data={"model": model}, route=f"/{provider}/v1/systemone{suffix}" - ) == f"{provider}/{model}" + assert ( + get_model_from_request(request_data={"model": model}, route=f"/{provider}/v1/systemone{suffix}") + == f"{provider}/{model}" + ) @pytest.mark.parametrize("provider", ["laya", "bespoke"]) -@pytest.mark.parametrize("model", [None, "", "auto", "laya/english", "bespoke/nimble-latest", "unknown", ["english"], 7]) +@pytest.mark.parametrize( + "model", [None, "", "auto", "laya/english", "bespoke/nimble-latest", "unknown", ["english"], 7] +) def test_oss_native_model_cannot_implicitly_select_an_unauthorized_checkpoint(provider: str, model: object) -> None: with pytest.raises(HTTPException) as denied: get_model_from_request(request_data={"model": model}, route=f"/{provider}/v1/systemone") @@ -780,38 +782,49 @@ def test_laya_model_normalization_does_not_change_other_provider_routes() -> Non def _cache_prediction_router(): from litellm.router import Router - return Router(model_list=[ - { - "model_name": group, - "litellm_params": {"model": "anthropic/claude-sonnet-5", "api_key": "test-provider-key"}, - "model_info": {"id": deployment_id, "team_id": team_id}, - } - for group, deployment_id, team_id in ( - ("current-group", "current-id", None), ("candidate-group", "candidate-id", None), - ("own-group", "own-id", "prediction-team"), ("foreign-group", "foreign-id", "foreign-team"), - ) - ]) + return Router( + model_list=[ + { + "model_name": group, + "litellm_params": {"model": "anthropic/claude-sonnet-5", "api_key": "test-provider-key"}, + "model_info": {"id": deployment_id, "team_id": team_id}, + } + for group, deployment_id, team_id in ( + ("current-group", "current-id", None), + ("candidate-group", "candidate-id", None), + ("own-group", "own-id", "prediction-team"), + ("foreign-group", "foreign-id", "foreign-team"), + ) + ] + ) -@pytest.mark.parametrize("candidate,team_id,expected", [ - ("candidate-id", None, ["current-group", "candidate-group"]), - ("current-id", None, "current-group"), - ("missing-id", None, None), - ("candidate-group", None, None), - ("own-id", None, None), - ("own-id", "prediction-team", ["current-group", "own-group"]), - ("foreign-id", "prediction-team", None), -]) +@pytest.mark.parametrize( + "candidate,team_id,expected", + [ + ("candidate-id", None, ["current-group", "candidate-group"]), + ("current-id", None, "current-group"), + ("missing-id", None, None), + ("candidate-group", None, None), + ("own-id", None, None), + ("own-id", "prediction-team", ["current-group", "own-group"]), + ("foreign-id", "prediction-team", None), + ], +) def test_cache_prediction_auth_resolves_only_exact_deployment_ids(candidate, team_id, expected): - assert get_model_from_request( - request_data={ - "current_deployment_id": "current-id", "candidate_deployment_id": candidate, - "request": {"model": "caller-controlled-provider-model"}, - }, - route="/cost/predict-cache", - llm_router=_cache_prediction_router(), - team_id=team_id, - ) == expected + assert ( + get_model_from_request( + request_data={ + "current_deployment_id": "current-id", + "candidate_deployment_id": candidate, + "request": {"model": "caller-controlled-provider-model"}, + }, + route="/cost/predict-cache", + llm_router=_cache_prediction_router(), + team_id=team_id, + ) + == expected + ) def _cache_prediction_auth_app( @@ -834,18 +847,29 @@ def _cache_prediction_auth_app( router = _cache_prediction_router() allowed_models = ["current-group", "candidate-group", "own-group"] token = UserAPIKeyAuth( - api_key="test-proxy-key-hash", user_id="prediction-user", user_role=LitellmUserRoles.INTERNAL_USER, - models=allowed_models if key_models is None else key_models, team_id=team_id, + api_key="test-proxy-key-hash", + user_id="prediction-user", + user_role=LitellmUserRoles.INTERNAL_USER, + models=allowed_models if key_models is None else key_models, + team_id=team_id, team_models=allowed_models if team_models is None else team_models, - allowed_routes=allowed_routes, metadata=metadata or {}, + allowed_routes=allowed_routes, + metadata=metadata or {}, ) user = LiteLLM_UserTable( - user_id=token.user_id, user_role=LitellmUserRoles.INTERNAL_USER.value, models=user_models, + user_id=token.user_id, + user_role=LitellmUserRoles.INTERNAL_USER.value, + models=user_models, ) + async def authenticate(request, request_data, **_headers): await auth._enforce_key_and_fallback_model_access( - valid_token=token, request_data=request_data, route=request.url.path, request=request, - llm_model_list=router.get_model_list(), llm_router=router, + valid_token=token, + request_data=request_data, + route=request.url.path, + request=request, + llm_model_list=router.get_model_list(), + llm_router=router, ) return token @@ -878,22 +902,36 @@ def _cache_prediction_auth_app( def _cache_prediction_payload(candidate="candidate-id", current="current-id"): return { - "current_deployment_id": current, "candidate_deployment_id": candidate, - "request": {"messages": [{"role": "user", "content": [{ - "type": "text", "text": "Stable cached context", - "cache_control": {"type": "ephemeral"}, - }]}]}, + "current_deployment_id": current, + "candidate_deployment_id": candidate, + "request": { + "messages": [ + { + "role": "user", + "content": [ + { + "type": "text", + "text": "Stable cached context", + "cache_control": {"type": "ephemeral"}, + } + ], + } + ] + }, } @pytest.mark.asyncio -@pytest.mark.parametrize("allowed_routes,user_models,candidate,status_code", [ - (["/chat/completions"], ["current-group", "candidate-group"], "candidate-id", 403), - (["/cost/predict-cache"], ["current-group"], "candidate-id", 403), - (["/cost/*"], ["current-group", "candidate-group"], "candidate-id", 200), - (["/cost/predict-cache"], ["current-group"], "current-id", 200), - (["/cost/predict-cache"], ["current-group"], "missing-id", 404), -]) +@pytest.mark.parametrize( + "allowed_routes,user_models,candidate,status_code", + [ + (["/chat/completions"], ["current-group", "candidate-group"], "candidate-id", 403), + (["/cost/predict-cache"], ["current-group"], "candidate-id", 403), + (["/cost/*"], ["current-group", "candidate-group"], "candidate-id", 200), + (["/cost/predict-cache"], ["current-group"], "current-id", 200), + (["/cost/predict-cache"], ["current-group"], "missing-id", 404), + ], +) async def test_cache_prediction_authorizes_route_and_personal_models_before_provider_counts( monkeypatch, allowed_routes, user_models, candidate, status_code ): @@ -912,20 +950,27 @@ async def test_cache_prediction_authorizes_route_and_personal_models_before_prov @pytest.mark.asyncio @pytest.mark.parametrize("arm", ["current_deployment_id", "candidate_deployment_id"]) -@pytest.mark.parametrize("team_id,key_models,user_models,team_models", [ - (None, ["*"], ["*"], None), - (None, ["current-group", "candidate-group"], ["*"], None), - (None, ["*"], ["current-group", "candidate-group"], None), - ("prediction-team", ["*"], ["*"], ["current-group", "candidate-group"]), -]) +@pytest.mark.parametrize( + "team_id,key_models,user_models,team_models", + [ + (None, ["*"], ["*"], None), + (None, ["current-group", "candidate-group"], ["*"], None), + (None, ["*"], ["current-group", "candidate-group"], None), + ("prediction-team", ["*"], ["*"], ["current-group", "candidate-group"]), + ], +) async def test_cache_prediction_hides_foreign_and_missing_ids_before_model_authorization( monkeypatch, arm, team_id, key_models, user_models, team_models ): import httpx app, counts = _cache_prediction_auth_app( - monkeypatch, ["/cost/predict-cache"], user_models, - team_id=team_id, key_models=key_models, team_models=team_models, + monkeypatch, + ["/cost/predict-cache"], + user_models, + team_id=team_id, + key_models=key_models, + team_models=team_models, ) async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client: missing = await client.post("/cost/predict-cache", json={**_cache_prediction_payload(), arm: "missing-id"}) @@ -938,19 +983,26 @@ async def test_cache_prediction_hides_foreign_and_missing_ids_before_model_autho @pytest.mark.asyncio @pytest.mark.parametrize("arm", ["current_deployment_id", "candidate_deployment_id"]) -@pytest.mark.parametrize("key_models,team_models,status_code", [ - (["*"], ["*"], 200), - (["current-group", "candidate-group"], ["*"], 403), - (["*"], ["current-group", "candidate-group"], 403), -]) +@pytest.mark.parametrize( + "key_models,team_models,status_code", + [ + (["*"], ["*"], 200), + (["current-group", "candidate-group"], ["*"], 403), + (["*"], ["current-group", "candidate-group"], 403), + ], +) async def test_cache_prediction_checks_each_visible_team_deployment_model( monkeypatch, arm, key_models, team_models, status_code ): import httpx app, counts = _cache_prediction_auth_app( - monkeypatch, ["/cost/predict-cache"], ["*"], - team_id="prediction-team", key_models=key_models, team_models=team_models, + monkeypatch, + ["/cost/predict-cache"], + ["*"], + team_id="prediction-team", + key_models=key_models, + team_models=team_models, ) async with httpx.AsyncClient(transport=httpx.ASGITransport(app=app), base_url="http://test") as client: response = await client.post("/cost/predict-cache", json={**_cache_prediction_payload(), arm: "own-id"}) @@ -976,19 +1028,24 @@ async def test_cache_prediction_checks_each_visible_personal_deployment_model(mo @pytest.mark.asyncio -@pytest.mark.parametrize("header_tag,key_tags,limit,status_code,provider_calls", [ - ("limited", [], 1, 429, 1), - (None, ["limited"], 1, 429, 1), - ("limited", ["limited"], 4, 200, 4), - ("unlimited", [], 1, 200, 4), -]) +@pytest.mark.parametrize( + "header_tag,key_tags,limit,status_code,provider_calls", + [ + ("limited", [], 1, 429, 1), + (None, ["limited"], 1, 429, 1), + ("limited", ["limited"], 4, 200, 4), + ("unlimited", [], 1, 200, 4), + ], +) async def test_cache_prediction_preserves_authenticated_header_and_key_tag_rpm( monkeypatch, header_tag, key_tags, limit, status_code, provider_calls ): import httpx app, counts = _cache_prediction_auth_app( - monkeypatch, ["/cost/predict-cache"], ["current-group", "candidate-group"], + monkeypatch, + ["/cost/predict-cache"], + ["current-group", "candidate-group"], metadata={"tag_rpm_limit": {"limited": limit}, "tags": key_tags}, ) headers = {"x-litellm-tags": header_tag} if header_tag else {} @@ -1142,7 +1199,11 @@ def _azure_relay_router(): model_list=[ { "model_name": "gpt", - "litellm_params": {"model": "azure_ai/gpt-5.4-mini", "api_base": "https://a.services.ai.azure.com", "api_key": "k"}, + "litellm_params": { + "model": "azure_ai/gpt-5.4-mini", + "api_base": "https://a.services.ai.azure.com", + "api_key": "k", + }, }, { "model_name": "other-group", @@ -1364,8 +1425,14 @@ def test_get_model_from_request_resolves_video_id_model_with_router(): provider="vertex_ai", model_id="veo-3.1-generate-001", ) - llm_router = MagicMock() - llm_router.resolve_model_name_from_model_id.return_value = "gcp/google/veo-3.1-generate-001" + llm_router = Router( + model_list=[ + { + "model_name": "gcp/google/veo-3.1-generate-001", + "litellm_params": {"model": "vertex_ai/veo-3.1-generate-001"}, + } + ] + ) assert ( get_model_from_request( @@ -1375,7 +1442,148 @@ def test_get_model_from_request_resolves_video_id_model_with_router(): ) == "gcp/google/veo-3.1-generate-001" ) - llm_router.resolve_model_name_from_model_id.assert_called_once_with("veo-3.1-generate-001") + + +def test_get_model_from_request_checks_the_group_of_the_deployment_that_created_the_video() -> None: + from litellm.types.videos.utils import encode_video_id_with_provider + + llm_router: Final = Router( + model_list=[ + { + "model_name": name, + "litellm_params": {"model": "azure/sora-2", "api_base": f"https://{name}.example.com", "api_key": name}, + "model_info": {"id": f"dep-{name}"}, + } + for name in ("sora-team", "sora-shared") + ] + ) + video_id: Final = encode_video_id_with_provider( + video_id="video_raw", provider="azure", model_id="sora-2", deployment_id="dep-sora-shared" + ) + + assert ( + get_model_from_request( + request_data={"video_id": video_id}, + route="/v1/videos/{video_id}", + llm_router=llm_router, + ) + == "sora-shared" + ) + + +def test_get_model_from_request_checks_a_video_created_through_a_provider_wildcard_by_its_own_name() -> None: + from litellm.types.videos.utils import encode_video_id_with_provider + + llm_router: Final = Router( + model_list=[ + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*", "api_key": "wildcard-key"}, + "model_info": {"id": "dep-wildcard"}, + } + ] + ) + video_id: Final = encode_video_id_with_provider( + video_id="video_raw", provider="openai", model_id="sora-2", deployment_id="dep-wildcard" + ) + routed_ids: Final = [ + deployment["model_info"]["id"] for deployment in llm_router.pattern_router.get_deployments_by_pattern("sora-2") + ] + + assert routed_ids == ["dep-wildcard"] + assert ( + get_model_from_request( + request_data={"video_id": video_id}, + route="/v1/videos/{video_id}", + llm_router=llm_router, + ) + == "sora-2" + ) + + +def test_get_model_from_request_checks_a_wildcard_video_whose_bare_name_does_not_route_by_its_prefixed_name() -> None: + from litellm.types.videos.utils import encode_video_id_with_provider + + llm_router: Final = Router( + model_list=[ + { + "model_name": "gemini/*", + "litellm_params": {"model": "gemini/*", "api_key": "wildcard-key"}, + "model_info": {"id": "dep-wildcard"}, + } + ] + ) + video_id: Final = encode_video_id_with_provider( + video_id="operations/op_1", provider="gemini", model_id="unlisted-video-model", deployment_id="dep-wildcard" + ) + + assert ( + get_model_from_request( + request_data={"video_id": video_id}, + route="/v1/videos/{video_id}", + llm_router=llm_router, + ) + == "gemini/unlisted-video-model" + ) + + +def _sora_fallback_router() -> Router: + return Router( + model_list=[ + { + "model_name": name, + "litellm_params": {"model": "azure/sora-2", "api_base": f"https://{name}.example.com", "api_key": name}, + "model_info": {"id": f"dep-{name}"}, + } + for name in ("sora-pair", "sora-backup", "sora-last", "sora-other") + ], + fallbacks=[{"sora-pair": ["sora-backup"]}, {"sora-backup": ["sora-last"]}], + ) + + +@pytest.mark.parametrize( + ("general_settings", "requested_model", "creating_group", "expected"), + [ + ({"enforce_fallback_model_access": False}, "sora-pair", "sora-backup", "sora-pair"), + ({"enforce_fallback_model_access": False}, "sora-pair", "sora-last", "sora-pair"), + ({}, "sora-pair", "sora-backup", "sora-backup"), + ({"enforce_fallback_model_access": False}, "sora-other", "sora-backup", "sora-backup"), + ({"enforce_fallback_model_access": False}, None, "sora-backup", "sora-backup"), + ], + ids=[ + "opt-out-fallback", + "opt-out-chained-fallback", + "enforced", + "opt-out-no-fallback-relation", + "opt-out-no-requested-model", + ], +) +def test_get_model_from_request_checks_a_fallback_served_video_by_the_requested_group_only_under_the_opt_out( + monkeypatch: pytest.MonkeyPatch, + general_settings: dict[str, object], + requested_model: str | None, + creating_group: str, + expected: str, +) -> None: + from litellm.types.videos.utils import encode_video_id_with_provider + + monkeypatch.setattr("litellm.proxy.proxy_server.general_settings", general_settings) + video_id: Final = encode_video_id_with_provider( + video_id="video_raw", + provider="azure", + model_id="sora-2", + deployment_id=f"dep-{creating_group}", + requested_model=requested_model, + ) + + assert ( + get_model_from_request( + request_data={"video_id": video_id}, + route="/v1/videos/{video_id}", + llm_router=_sora_fallback_router(), + ) + == expected + ) _BATCH_DEPLOYMENT_ID = "8d0eaa7e6c6f54a425dfd0062cb6b0dc" @@ -4463,9 +4671,7 @@ class TestIsRequestBodySafeBlocksAwsIdentitySelectors: ), ], ) -def test_configurable_clientside_parameters( - allowed_param, input_value, should_return_true -): +def test_configurable_clientside_parameters(allowed_param, input_value, should_return_true): router = Router( model_list=[ { @@ -4499,9 +4705,7 @@ def test_get_customer_user_header_from_mapping_returns_customer_header_with_mixe def test_get_customer_user_header_from_mapping_no_customer_returns_none(): from litellm.proxy.auth.auth_utils import get_customer_user_header_from_mapping - mappings = [ - {"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"} - ] + mappings = [{"header_name": "X-OpenWebUI-User-Id", "litellm_user_role": "internal_user"}] result = get_customer_user_header_from_mapping(mappings) assert result is None @@ -4529,17 +4733,13 @@ def test_get_internal_user_header_from_mapping_returns_internal_header(): def test_get_internal_user_header_from_mapping_no_internal_returns_none(): from litellm.proxy.litellm_pre_call_utils import LiteLLMProxyRequestSetup - mappings = [ - {"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"} - ] + mappings = [{"header_name": "X-OpenWebUI-User-Email", "litellm_user_role": "customer"}] result = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping(mappings) assert result is None # Also support single mapping dict single_mapping = {"header_name": "X-Only-Customer", "litellm_user_role": "customer"} - result = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping( - single_mapping - ) + result = LiteLLMProxyRequestSetup.get_internal_user_header_from_mapping(single_mapping) assert result is None @@ -4561,9 +4761,7 @@ def test_get_internal_user_header_from_mapping_no_internal_returns_none(): def test_get_model_from_request(request_data, expected_model): from litellm.proxy.auth.auth_utils import get_model_from_request - request_data = { - "target_model_names": "gpt-3.5-turbo, gpt-4o-mini-general-deployment" - } + request_data = {"target_model_names": "gpt-3.5-turbo, gpt-4o-mini-general-deployment"} route = "/openai/deployments/gpt-3.5-turbo" model = get_model_from_request(request_data, "/v1/files") assert model == ["gpt-3.5-turbo", "gpt-4o-mini-general-deployment"] @@ -4625,9 +4823,7 @@ def test_get_model_from_request(request_data, expected_model): ), ], ) -def test_get_model_from_request_vertex_ai_passthrough( - request_data, route, expected_model -): +def test_get_model_from_request_vertex_ai_passthrough(request_data, route, expected_model): """Test that get_model_from_request correctly extracts Vertex AI model from URL""" from litellm.proxy.auth.auth_utils import get_model_from_request diff --git a/tests/unit/proxy/management_endpoints/test_ptu_consumption.py b/tests/unit/proxy/management_endpoints/test_ptu_consumption.py index a446afbd809..065ec93dcd0 100644 --- a/tests/unit/proxy/management_endpoints/test_ptu_consumption.py +++ b/tests/unit/proxy/management_endpoints/test_ptu_consumption.py @@ -232,14 +232,10 @@ def test_rows_keyed_by_a_name_a_wildcard_route_serves_are_sized_by_its_shared_de assert attached.results[0].breakdown.model_groups["ptu-chat"].metrics.ptu_hours == pytest.approx(1.0) -@pytest.mark.parametrize( - ("catch_all_routes", "expected_ptu_hours"), [(False, 1.0), (True, 0.0)], ids=["team-wildcard", "behind-catch-all"] -) -def test_a_teams_wildcard_reservation_reads_ptu_hours_only_while_the_router_serves_from_it( - monkeypatch, catch_all_routes, expected_ptu_hours -): - """The router tries the proxy-wide wildcards before a team's own, so a team's `ptu-*` - reservation behind an open `*` route never serves `ptu-chat`, and its row reads no PTU-hours.""" +@pytest.mark.parametrize("catch_all_routes", [False, True], ids=["team-wildcard", "beside-catch-all"]) +def test_a_teams_wildcard_reservation_reads_ptu_hours_even_beside_a_proxy_wide_catch_all(monkeypatch, catch_all_routes): + """The router tries a team's own wildcard before the proxy-wide ones, so a team's `ptu-*` + reservation serves `ptu-chat` whether or not an open `*` route exists, and its row reads its PTU-hours.""" monkeypatch.setenv("LITELLM_ENABLE_PTU_COST_ATTRIBUTION", "True") reservation: Final = { "model_name": "ptu-team-r", @@ -264,9 +260,7 @@ def test_a_teams_wildcard_reservation_reads_ptu_hours_only_while_the_router_serv attached: Final = with_ptu_consumption(_response(_day("2026-09-23", {"ptu-chat": one_hour})), router, "team-r") - assert attached.results[0].breakdown.model_groups["ptu-chat"].metrics.ptu_hours == pytest.approx( - expected_ptu_hours - ) + assert attached.results[0].breakdown.model_groups["ptu-chat"].metrics.ptu_hours == pytest.approx(1.0) def _mixed_ptu_router() -> Router: diff --git a/tests/unit/proxy/video_endpoints/test_endpoints.py b/tests/unit/proxy/video_endpoints/test_endpoints.py index 5052d75f9f9..9809d26826b 100644 --- a/tests/unit/proxy/video_endpoints/test_endpoints.py +++ b/tests/unit/proxy/video_endpoints/test_endpoints.py @@ -28,7 +28,8 @@ mock's call args), and a brand-new kwarg added to this layer surfaces as a failu from contextlib import ExitStack from dataclasses import dataclass -from typing import Any, Dict, Optional +from collections.abc import Awaitable, Callable +from typing import Any, Dict, Final, Optional from unittest.mock import AsyncMock, MagicMock, patch import orjson @@ -41,7 +42,9 @@ from litellm.proxy._types import UserAPIKeyAuth from litellm.proxy.common_request_processing import ProxyBaseLLMRequestProcessing from litellm.proxy.utils import ProxyLogging from litellm.router import Router +from litellm.types.videos.main import VideoObject from litellm.types.videos.utils import ( + decode_video_id_with_provider, encode_character_id_with_provider, encode_video_id_with_provider, ) @@ -70,6 +73,29 @@ RESOLVED_MODELS: Dict[str, str] = {VIDEO_MODEL_ID: "azure-sora"} SENTINEL = object() +def _video_with_empty_model_id( + hidden_params: dict[str, object] | None = None, +) -> VideoObject: + video: Final[VideoObject] = VideoObject( + id=encode_video_id_with_provider("video_new456", "azure", ""), + object="video", + status="queued", + ) + video.hidden_params = dict(hidden_params or {}) + return video + + +def _assert_id_resolves_to_deployment(response: object, *, deployment_id: str | None = None) -> None: + assert isinstance(response, VideoObject) + expected_deployment: Final = {} if deployment_id is None else {"deployment_id": deployment_id} + assert decode_video_id_with_provider(response.id) == { + "custom_llm_provider": "azure", + "model_id": VIDEO_MODEL_ID, + "video_id": "video_new456", + **expected_deployment, + } + + class FakeRequest: """Minimal stand-in. headers/query_params are read by the provider readers (mocked) and on the edit path the raw body is parsed for real via orjson.""" @@ -105,6 +131,9 @@ class Harness: assert self.base_process.call_count == 1 return dict(self.base_process.call_args.args[0].data) + def respond_with(self, response: object) -> None: + self.base_process.return_value = response # pyright: ignore[reportAny] # MagicMock subclasses Any in typeshed + def route_type(self) -> str: return self.base_process.call_args.kwargs["route_type"] @@ -118,6 +147,7 @@ def harness(): side_effect=lambda model_id, custom_llm_provider=None: RESOLVED_MODELS.get(model_id) ) router.resolve_model_name_from_model_id = resolve_model + router.get_deployment.return_value = None read_body = AsyncMock(return_value={}) batch_to_bytesio = AsyncMock(return_value=[b"filebytes"]) @@ -244,6 +274,36 @@ async def test_generation__input_reference_attached(harness): } +@pytest.mark.asyncio +async def test_generation__restamps_id_with_asked_for_model_over_router_deployment_id(harness: Harness) -> None: + harness.respond_with(_video_with_empty_model_id(hidden_params={"model_id": "deployment-uuid"})) + + resp: Final = await call_generation(harness, body={"model": VIDEO_MODEL_ID, "prompt": "x"}) + + _assert_id_resolves_to_deployment(resp, deployment_id="deployment-uuid") + + +@pytest.mark.asyncio +async def test_generation__restamps_id_with_request_model_when_no_hidden_model_id(harness: Harness) -> None: + harness.respond_with(_video_with_empty_model_id()) + + resp: Final = await call_generation(harness, body={"model": VIDEO_MODEL_ID, "prompt": "x"}) + + _assert_id_resolves_to_deployment(resp) + + +@pytest.mark.asyncio +async def test_generation__unroutable_request_model_falls_back_to_router_deployment_id(harness: Harness) -> None: + harness.respond_with(_video_with_empty_model_id(hidden_params={"model_id": "deployment-uuid"})) + + resp: Final = await call_generation(harness, body={"model": "team-public-sora", "prompt": "x"}) + + assert isinstance(resp, VideoObject) + decoded: Final = decode_video_id_with_provider(resp.id) + assert decoded.get("model_id") == "deployment-uuid" + assert decoded.get("deployment_id") == "deployment-uuid" + + @pytest.mark.asyncio async def test_generation__exception_routed_through_handler(harness): harness.base_process.side_effect = ValueError("provider boom") @@ -284,6 +344,24 @@ async def test_status__model_encoded_id_full_contract(harness): } +@pytest.mark.asyncio +async def test_status__restamps_response_id_with_decoded_model_id(harness: Harness) -> None: + harness.respond_with(_video_with_empty_model_id()) + + resp: Final = await call_status(harness, AZURE_VIDEO_ID) + + _assert_id_resolves_to_deployment(resp) + + +@pytest.mark.asyncio +async def test_status__id_without_a_deployment_takes_the_deployment_that_answered_the_poll(harness: Harness) -> None: + harness.respond_with(_video_with_empty_model_id(hidden_params={"model_id": "deployment-uuid"})) + + resp: Final = await call_status(harness, AZURE_VIDEO_ID) + + _assert_id_resolves_to_deployment(resp, deployment_id="deployment-uuid") + + @pytest.mark.asyncio async def test_status__plain_id_defaults_to_openai(harness): await call_status(harness, "video_plain") @@ -387,6 +465,30 @@ async def test_status_and_content__hand_the_served_member_to_the_request_process assert harness.base_process.call_args.args[0].served_model == "wan3-b" +@pytest.mark.asyncio +@pytest.mark.parametrize("call", [call_status, call_content]) +async def test_status_and_content__shared_provider_model_routes_to_the_creating_deployments_group( + harness: Harness, call: Callable[[Harness, str], Awaitable[object]] +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": name, + "litellm_params": {"model": "azure/sora-2", "api_base": f"https://{name}.example.com", "api_key": name}, + "model_info": {"id": f"dep-{name}"}, + } + for name in ("sora-team", "sora-shared") + ] + ) + harness.base_process.return_value = b"x" + video_id: Final = encode_video_id_with_provider("video_raw", "azure", "sora-2", deployment_id="dep-sora-shared") + + with patch.object(proxy_server, "llm_router", router): + await call(harness, video_id) + + assert harness.processor_data()["model"] == "sora-shared" + + # =========================================================================== # # POST /v1/videos/edits - video_edit # # =========================================================================== # @@ -422,6 +524,15 @@ async def test_edit__extracts_nested_video_id_full_contract(harness): } +@pytest.mark.asyncio +async def test_edit__restamps_new_video_id_with_source_model_id(harness: Harness) -> None: + harness.respond_with(_video_with_empty_model_id()) + + resp: Final = await call_edit(harness, body={"prompt": "brighter", "video": {"id": AZURE_VIDEO_ID}}) + + _assert_id_resolves_to_deployment(resp) + + @pytest.mark.asyncio async def test_edit__provider_from_body_data_for_plain_id(harness): """For a plain id, get_custom_provider_from_data (run for real) pulls the @@ -473,6 +584,23 @@ async def test_edit__json_string_video_reference_from_form_field(harness): assert harness.processor_data()["video_id"] == "video_plain" +@pytest.mark.asyncio +async def test_edit__uploaded_video_restamps_new_id_with_requested_model(harness: Harness) -> None: + import io + + upload: Final = StarletteUploadFile(file=io.BytesIO(b"rawmp4"), filename="clip.mp4") + harness.read_body.return_value = {"model": VIDEO_MODEL_ID, "prompt": "make it nighttime", "video": upload} + harness.respond_with(_video_with_empty_model_id()) + + resp: Final = await endpoints.video_edit( + request=FakeRequest(raw_body=b"multipart"), + fastapi_response=Response(), + user_api_key_dict=_user(), + ) + + _assert_id_resolves_to_deployment(resp) + + @pytest.mark.asyncio async def test_edit__uploaded_video_file_is_forwarded_not_dropped(harness): """A multipart-uploaded source video must be converted to bytes and attached @@ -565,6 +693,15 @@ async def test_remix__model_encoded_id_full_contract(harness): } +@pytest.mark.asyncio +async def test_remix__restamps_new_video_id_with_source_model_id(harness: Harness) -> None: + harness.respond_with(_video_with_empty_model_id()) + + resp: Final = await call_remix(harness, AZURE_VIDEO_ID, body={"prompt": "new colors"}) + + _assert_id_resolves_to_deployment(resp) + + @pytest.mark.asyncio async def test_remix__provider_from_body_data_not_request_body_reader(harness): """remix resolves the provider from data.get('custom_llm_provider'), never @@ -694,6 +831,20 @@ async def test_get_character__plain_id_defaults_openai_no_encode(harness): assert resp["id"] == "char_raw3" +@pytest.mark.asyncio +async def test_get_character__model_served_only_by_a_provider_wildcard_routes_through_it(harness: Harness) -> None: + router: Final = Router( + model_list=[{"model_name": "openai/*", "litellm_params": {"model": "openai/*", "api_key": "key-w"}}] + ) + harness.base_process.return_value = {"id": "char_raw4"} + character_id: Final = encode_character_id_with_provider("char_orig", "openai", "sora-2") + + with patch.object(proxy_server, "llm_router", router): + await call_get_character(harness, character_id) + + assert harness.processor_data()["model"] == "sora-2" + + # =========================================================================== # # POST /v1/videos/extensions - video_extension # # =========================================================================== # @@ -723,3 +874,12 @@ async def test_extension__extracts_nested_video_id_full_contract(harness): "custom_llm_provider": "azure", "model": "azure-sora", } + + +@pytest.mark.asyncio +async def test_extension__restamps_new_video_id_with_source_model_id(harness: Harness) -> None: + harness.respond_with(_video_with_empty_model_id()) + + resp: Final = await call_extension(harness, body={"prompt": "continue", "video": {"id": AZURE_VIDEO_ID}}) + + _assert_id_resolves_to_deployment(resp) diff --git a/tests/unit/proxy/video_endpoints/test_utils.py b/tests/unit/proxy/video_endpoints/test_utils.py index 9a2c208c075..01a2b7c2840 100644 --- a/tests/unit/proxy/video_endpoints/test_utils.py +++ b/tests/unit/proxy/video_endpoints/test_utils.py @@ -1,11 +1,12 @@ """ Pure-logic contract tests for litellm/proxy/video_endpoints/utils.py -Four helpers the video proxy endpoints lean on: +Five helpers the video proxy endpoints lean on: - extract_model_from_target_model_names: first model from a comma string / list - video_reference_to_id: normalize a video reference (dict / bare id / JSON string) to an id - get_custom_provider_from_data: provider precedence (top-level > extra_body) - encode_character_id_in_response: re-encode a response id in place + - encode_video_id_in_response: re-stamp a VideoObject id with the resolved model_id Every test asserts the exact result (or identity), so a mutation that flips a branch, drops a strip/filter, or changes precedence fails. The only collaborator @@ -13,19 +14,29 @@ is encode_character_id_with_provider, which runs for real; encoding assertions are checked by the genuine decode round-trip. """ +from typing import Final import pytest +from litellm.router import Router +from litellm.types.videos.main import DecodedVideoId from litellm.proxy.video_endpoints.utils import ( encode_character_id_in_response, + encode_video_id_in_response, extract_model_from_target_model_names, + fallback_requested_model, get_custom_provider_from_data, + routable_model_name, + video_model_group, video_reference_to_id, ) +from litellm.types.videos.main import VideoObject from litellm.types.videos.utils import ( decode_character_id_with_provider, + decode_video_id_with_provider, encode_character_id_with_provider, + encode_video_id_with_provider, ) # =========================================================================== # @@ -103,12 +114,7 @@ def test_provider__falsy_top_level_falls_through_to_extra_body(falsy): def test_provider__from_extra_body_dict(): - assert ( - get_custom_provider_from_data( - {"extra_body": {"custom_llm_provider": "bedrock"}} - ) - == "bedrock" - ) + assert get_custom_provider_from_data({"extra_body": {"custom_llm_provider": "bedrock"}}) == "bedrock" def test_provider__from_extra_body_json_string(): @@ -126,10 +132,7 @@ def test_provider__json_string_parsing_to_non_dict_is_none(): def test_provider__extra_body_provider_not_a_string_is_none(): - assert ( - get_custom_provider_from_data({"extra_body": {"custom_llm_provider": 123}}) - is None - ) + assert get_custom_provider_from_data({"extra_body": {"custom_llm_provider": 123}}) is None @pytest.mark.parametrize( @@ -161,9 +164,7 @@ def test_encode__dict_with_id_mutates_in_place_and_preserves_other_keys(): assert out is response # same dict, mutated in place assert out["object"] == "character" and out["name"] == "hero" - assert out["id"] == encode_character_id_with_provider( - "char_raw", "azure", "model-1" - ) + assert out["id"] == encode_character_id_with_provider("char_raw", "azure", "model-1") decoded = decode_character_id_with_provider(out["id"]) assert decoded["custom_llm_provider"] == "azure" assert decoded["model_id"] == "model-1" @@ -206,3 +207,387 @@ def test_encode__object_without_id_attr_returned_unchanged(): out = encode_character_id_in_response(resp, "azure", "model-1") assert out is resp assert not hasattr(resp, "id") + + +# =========================================================================== # +# encode_video_id_in_response +# =========================================================================== # + + +def _video(video_id: str, hidden_params: dict[str, object] | None = None) -> VideoObject: + video: Final[VideoObject] = VideoObject(id=video_id, object="video", status="queued") + video.hidden_params = dict(hidden_params or {}) + return video + + +def test_encode_video__fills_empty_model_id_and_deployment_from_hidden_params() -> None: + original_id: Final = encode_video_id_with_provider("video_raw", "openai", "") + video: Final = _video(original_id, hidden_params={"model_id": "deployment-1"}) + + out: Final = encode_video_id_in_response(video, fallback_provider=None, fallback_model_id=None) + + assert isinstance(out, VideoObject) + assert decode_video_id_with_provider(out.id) == { + "custom_llm_provider": "openai", + "model_id": "deployment-1", + "video_id": "video_raw", + "deployment_id": "deployment-1", + } + assert out is not video + assert video.id == original_id + assert out.hidden_params == {"model_id": "deployment-1"} + assert out.status == "queued" + + +def test_encode_video__asked_for_model_wins_over_hidden_deployment_id() -> None: + video: Final = _video( + encode_video_id_with_provider("video_raw", "openai", ""), hidden_params={"model_id": "deployment-1"} + ) + + out: Final = encode_video_id_in_response(video, fallback_provider=None, fallback_model_id="sora-alias") + + assert isinstance(out, VideoObject) + decoded: Final = decode_video_id_with_provider(out.id) + assert decoded.get("model_id") == "sora-alias" + assert decoded.get("deployment_id") == "deployment-1" + + +def test_encode_video__stamped_id_keeps_model_and_served_group_and_gains_routed_deployment() -> None: + original_id: Final = encode_video_id_with_provider("task_1", "volcengine", "seedance-group", "seedance-east") + video: Final = _video(original_id, hidden_params={"model_id": "deployment-1"}) + + out: Final = encode_video_id_in_response(video, fallback_provider=None, fallback_model_id="other") + + assert isinstance(out, VideoObject) + assert decode_video_id_with_provider(out.id) == { + "custom_llm_provider": "volcengine", + "model_id": "seedance-group", + "video_id": "task_1", + "served_model": "seedance-east", + "deployment_id": "deployment-1", + } + + +def test_encode_video__routed_deployment_replaces_the_one_in_the_id() -> None: + original_id: Final = encode_video_id_with_provider("video_raw", "openai", "sora", deployment_id="deployment-old") + video: Final = _video(original_id, hidden_params={"model_id": "deployment-new"}) + + out: Final = encode_video_id_in_response(video, fallback_provider=None, fallback_model_id=None) + + assert isinstance(out, VideoObject) + assert decode_video_id_with_provider(out.id).get("deployment_id") == "deployment-new" + + +@pytest.mark.parametrize( + ("polled_deployment", "answering_deployment", "returned_deployment"), + [ + (None, "deployment-routed", "deployment-routed"), + ("deployment-gone", "deployment-routed", "deployment-routed"), + ("deployment-creator", None, "deployment-creator"), + ], +) +def test_encode_video__a_poll_takes_the_deployment_that_answered_else_keeps_the_one_its_id_named( + polled_deployment: str | None, + answering_deployment: str | None, + returned_deployment: str, +) -> None: + polled_id: Final = encode_video_id_with_provider("video_raw", "openai", "sora", deployment_id=polled_deployment) + video: Final = _video( + encode_video_id_with_provider("video_raw", "openai", ""), + hidden_params={} if answering_deployment is None else {"model_id": answering_deployment}, + ) + + out: Final = encode_video_id_in_response( + video, + fallback_provider=None, + fallback_model_id="sora", + polled_video=decode_video_id_with_provider(polled_id), + ) + + assert isinstance(out, VideoObject) + assert out.id == encode_video_id_with_provider("video_raw", "openai", "sora", deployment_id=returned_deployment) + + +def test_encode_video__a_poll_keeps_the_group_a_fallback_create_was_serving() -> None: + polled_id: Final = encode_video_id_with_provider( + "video_raw", "azure", "sora-2", deployment_id="dep-backup", requested_model="sora-pair" + ) + video: Final = _video(encode_video_id_with_provider("video_raw", "azure", "sora-2")) + + out: Final = encode_video_id_in_response( + video, fallback_provider=None, fallback_model_id=None, polled_video=decode_video_id_with_provider(polled_id) + ) + + assert isinstance(out, VideoObject) + assert out.id == polled_id + + +def test_encode_video__a_remix_of_a_fallback_video_keeps_its_requested_group() -> None: + video: Final = _video( + encode_video_id_with_provider("video_new", "azure", "sora-2"), hidden_params={"model_id": "dep-backup"} + ) + + out: Final = encode_video_id_in_response( + video, fallback_provider=None, fallback_model_id=None, requested_model="sora-pair" + ) + + assert isinstance(out, VideoObject) + assert decode_video_id_with_provider(out.id).get("requested_model") == "sora-pair" + + +def test_encode_video__id_already_naming_the_routed_deployment_is_unchanged() -> None: + original_id: Final = encode_video_id_with_provider("video_raw", "openai", "sora", deployment_id="deployment-1") + video: Final = _video(original_id, hidden_params={"model_id": "deployment-1"}) + + assert encode_video_id_in_response(video, fallback_provider=None, fallback_model_id="other") is video + + +def test_encode_video__unrouted_response_keeps_the_deployment_in_its_id() -> None: + original_id: Final = encode_video_id_with_provider("video_raw", "openai", "", deployment_id="deployment-1") + + out: Final = encode_video_id_in_response(_video(original_id), fallback_provider=None, fallback_model_id="sora") + + assert isinstance(out, VideoObject) + decoded: Final = decode_video_id_with_provider(out.id) + assert decoded.get("model_id") == "sora" + assert decoded.get("deployment_id") == "deployment-1" + + +def test_encode_video__name_with_id_separator_is_never_written_into_the_id() -> None: + video: Final = _video( + encode_video_id_with_provider("video_raw", "openai", ""), hidden_params={"model_id": "deployment-1"} + ) + + out: Final = encode_video_id_in_response(video, fallback_provider=None, fallback_model_id="sora;model_id:evil") + + assert isinstance(out, VideoObject) + assert decode_video_id_with_provider(out.id) == { + "custom_llm_provider": "openai", + "model_id": "deployment-1", + "video_id": "video_raw", + "deployment_id": "deployment-1", + } + + +def test_encode_video__falls_back_to_given_model_id() -> None: + video: Final = _video(encode_video_id_with_provider("video_raw", "openai", "")) + + out: Final = encode_video_id_in_response(video, fallback_provider=None, fallback_model_id="sora-2") + + assert isinstance(out, VideoObject) + decoded: Final = decode_video_id_with_provider(out.id) + assert decoded.get("model_id") == "sora-2" + assert decoded.get("video_id") == "video_raw" + + +def test_encode_video__keeps_existing_model_id_over_fallback() -> None: + original_id: Final = encode_video_id_with_provider("video_raw", "azure", "deployment-1") + + out: Final = encode_video_id_in_response(_video(original_id), fallback_provider="openai", fallback_model_id="other") + + assert isinstance(out, VideoObject) + assert out.id == original_id + + +def test_encode_video__plain_id_encoded_with_fallback_provider() -> None: + video: Final = _video("video_raw", hidden_params={"model_id": "deployment-1"}) + + out: Final = encode_video_id_in_response(video, fallback_provider="openai", fallback_model_id=None) + + assert isinstance(out, VideoObject) + assert decode_video_id_with_provider(out.id) == { + "custom_llm_provider": "openai", + "model_id": "deployment-1", + "video_id": "video_raw", + "deployment_id": "deployment-1", + } + + +def test_encode_video__no_model_id_anywhere_unchanged() -> None: + video: Final = _video(encode_video_id_with_provider("video_raw", "openai", "")) + + out: Final = encode_video_id_in_response(video, fallback_provider=None, fallback_model_id=None) + + assert out is video + + +def test_encode_video__no_provider_anywhere_unchanged() -> None: + video: Final = _video("video_raw") + + out: Final = encode_video_id_in_response(video, fallback_provider=None, fallback_model_id="sora-2") + + assert out is video + + +@pytest.mark.parametrize("response", [b"video-bytes", {"id": "video_raw"}, None]) +def test_encode_video__non_video_object_passthrough(response: object) -> None: + assert encode_video_id_in_response(response, "openai", "sora-2") is response + + +def _router_with_groups(*names: str) -> Router: + return Router( + model_list=[ + {"model_name": name, "litellm_params": {"model": "openai/sora-2", "api_key": "group-key"}} for name in names + ] + ) + + +def test_routable_model_name__keeps_a_name_the_router_resolves() -> None: + assert routable_model_name("sora-alias", _router_with_groups("sora-alias"), "openai") == "sora-alias" + + +def test_routable_model_name__drops_a_name_the_router_cannot_route() -> None: + assert routable_model_name("team-public-name", _router_with_groups("sora-alias"), "openai") is None + + +def test_routable_model_name__keeps_the_name_without_a_router() -> None: + assert routable_model_name("sora-2", None, None) == "sora-2" + + +@pytest.mark.parametrize("model", [None, "", 7, "sora;model_id:evil"]) +def test_routable_model_name__never_returns_a_name_unsafe_for_the_id(model: object) -> None: + assert routable_model_name(model, _router_with_groups("sora;model_id:evil"), "openai") is None + + +def _two_groups_sharing_sora() -> Router: + return Router( + model_list=[ + { + "model_name": name, + "litellm_params": {"model": "azure/sora-2", "api_base": f"https://{name}.example.com", "api_key": name}, + "model_info": {"id": f"dep-{name}"}, + } + for name in ("sora-team", "sora-shared") + ] + ) + + +def _decoded(model_id: str, deployment_id: str | None) -> DecodedVideoId: + base: Final = DecodedVideoId(custom_llm_provider="azure", model_id=model_id, video_id="video_raw") + return base if deployment_id is None else {**base, "deployment_id": deployment_id} + + +@pytest.mark.parametrize("deployment_id", ["dep-sora-team", "dep-sora-shared"]) +def test_video_model_group__shared_provider_model_resolves_to_the_creating_deployments_group( + deployment_id: str, +) -> None: + router: Final = _two_groups_sharing_sora() + + assert video_model_group(_decoded("sora-2", deployment_id), router) == deployment_id.removeprefix("dep-") + + +@pytest.mark.parametrize("deployment_id", [None, "dep-gone"]) +def test_video_model_group__without_a_known_deployment_matches_plain_resolution(deployment_id: str | None) -> None: + router: Final = _two_groups_sharing_sora() + + assert video_model_group(_decoded("sora-2", deployment_id), router) == router.resolve_model_name_from_model_id( + "sora-2", "azure" + ) + + +def test_video_model_group__keeps_a_group_named_in_the_id_over_a_deployment_of_another_group() -> None: + assert video_model_group(_decoded("sora-team", "dep-sora-shared"), _two_groups_sharing_sora()) == "sora-team" + + +@pytest.mark.parametrize("deployment_id", ["dep-sora-team", "dep-sora-shared"]) +def test_video_model_group__keeps_an_alias_of_the_creating_deployments_group(deployment_id: str) -> None: + router: Final = Router( + model_list=[ + { + "model_name": name, + "litellm_params": {"model": f"openai/sora-{name}", "api_key": name}, + "model_info": {"id": f"dep-{name}"}, + } + for name in ("sora-team", "sora-shared") + ], + model_group_alias={"sora-alias": deployment_id.removeprefix("dep-")}, + ) + provider_model: Final = f"sora-{deployment_id.removeprefix('dep-')}" + + assert video_model_group(_decoded(provider_model, deployment_id), router) == "sora-alias" + + +def test_video_model_group__names_a_team_model_by_its_public_name() -> None: + router: Final = Router( + model_list=[ + { + "model_name": "model_name_team-1_internal", + "litellm_params": {"model": "azure/sora-2", "api_key": "team-key"}, + "model_info": {"id": "dep-team", "team_id": "team-1", "team_public_model_name": "team-sora"}, + } + ] + ) + + assert video_model_group(_decoded("sora-2", "dep-team"), router) == "team-sora" + + +def _fallback_router() -> Router: + return Router( + model_list=[ + { + "model_name": name, + "litellm_params": {"model": "azure/sora-2", "api_base": f"https://{name}.example.com", "api_key": name}, + "model_info": {"id": f"dep-{name}"}, + } + for name in ("sora-pair", "sora-backup", "sora-last") + ], + model_group_alias={"sora-alias": "sora-pair"}, + fallbacks=[{"sora-pair": ["sora-backup"]}, {"sora-backup": ["sora-last"]}], + ) + + +def _wildcard_router() -> Router: + return Router( + model_list=[ + { + "model_name": "gemini/*", + "litellm_params": {"model": "gemini/*", "api_key": "wildcard-key"}, + "model_info": {"id": "dep-wildcard"}, + } + ] + ) + + +@pytest.mark.parametrize( + ("requested_model", "answering_deployment", "expected"), + [ + ("sora-pair", "dep-sora-backup", "sora-pair"), + ("sora-pair", "dep-sora-last", "sora-pair"), + ("sora-backup", "dep-sora-pair", None), + ("sora-pair", "dep-sora-pair", None), + ("sora-alias", "dep-sora-pair", None), + ("sora-pair", None, None), + (None, "dep-sora-backup", None), + ], + ids=[ + "fallback-served", + "chained-fallback-served", + "served-by-a-group-it-never-falls-back-to", + "served-by-own-group", + "served-by-alias-target", + "unrouted", + "no-requested-model", + ], +) +def test_fallback_requested_model__names_the_asked_group_only_when_another_group_served_the_create( + requested_model: str | None, answering_deployment: str | None, expected: str | None +) -> None: + video: Final = _video( + encode_video_id_with_provider("video_raw", "azure", "sora-2"), + hidden_params={} if answering_deployment is None else {"model_id": answering_deployment}, + ) + + assert fallback_requested_model(requested_model, video, _fallback_router()) == expected + + +def test_fallback_requested_model__is_none_for_a_create_a_provider_wildcard_served() -> None: + video: Final = _video( + encode_video_id_with_provider("operations/op_1", "gemini", "veo-3.1-fast-generate-preview"), + hidden_params={"model_id": "dep-wildcard"}, + ) + + assert fallback_requested_model("gemini/veo-3.1-fast-generate-preview", video, _wildcard_router()) is None + + +def test_video_model_group__is_none_without_a_router() -> None: + assert video_model_group(_decoded("sora-2", "dep-sora-shared"), None) is None diff --git a/tests/unit/responses/mcp/test_tool_precheck.py b/tests/unit/responses/mcp/test_tool_precheck.py new file mode 100644 index 00000000000..832151ce2a6 --- /dev/null +++ b/tests/unit/responses/mcp/test_tool_precheck.py @@ -0,0 +1,318 @@ +import json +import sys +from collections import deque +from collections.abc import Callable, Iterable, Mapping, Sequence +from importlib.abc import MetaPathFinder +from importlib.machinery import ModuleSpec +from types import ModuleType +from typing import Final + +import httpx +import pytest +import respx + +import litellm +from litellm.responses.mcp.tool_precheck import may_name_litellm_mcp_gateway + +GATEWAY_REGISTRY_MODULE: Final = "litellm.proxy._experimental.mcp_server.mcp_server_manager" +GATEWAY_HANDLER_MODULES: Final = ( + "litellm.responses.mcp.litellm_proxy_mcp_handler", + "litellm.responses.mcp.chat_completions_handler", + "litellm.llms.anthropic.pass_through.messages.mcp_handler", +) + +FUNCTION_TOOL: Final = { + "type": "function", + "function": {"name": "noop", "description": "No operation", "parameters": {"type": "object", "properties": {}}}, +} +RESPONSES_FUNCTION_TOOL: Final = { + "type": "function", + "name": "noop", + "description": "No operation", + "parameters": {"type": "object", "properties": {}}, +} +ANTHROPIC_FUNCTION_TOOL: Final = { + "name": "noop", + "description": "No operation", + "input_schema": {"type": "object", "properties": {}}, +} +ANTHROPIC_MESSAGE_PAYLOAD: Final = { + "id": "msg_1", + "type": "message", + "role": "assistant", + "model": "claude-sonnet-5-5", + "content": [{"type": "text", "text": "ok"}], + "stop_reason": "end_turn", + "stop_sequence": None, + "usage": {"input_tokens": 1, "output_tokens": 1}, +} +GATEWAY_MCP_TOOL: Final = {"type": "mcp", "server_label": "litellm", "server_url": "litellm_proxy"} +GATEWAY_SERVER_MCP_TOOL: Final = {"type": "mcp", "server_label": "x", "server_url": "litellm_proxy/mcp/github"} +GATEWAY_PATH_MCP_TOOL: Final = {"type": "mcp", "server_label": "x", "server_url": "https://gw.example.com/mcp/github"} +REMOTE_MCP_TOOL: Final = { + "type": "mcp", + "server_label": "deepwiki", + "server_url": "https://mcp.deepwiki.com/sse", + "require_approval": "never", +} + +RESPONSES_PAYLOAD: Final = { + "id": "resp_1", + "object": "response", + "created_at": 1734366691, + "status": "completed", + "model": "gpt-4o-mini", + "output": [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Done.", "annotations": []}], + } + ], + "parallel_tool_calls": True, + "usage": { + "input_tokens": 10, + "output_tokens": 5, + "total_tokens": 15, + "output_tokens_details": {"reasoning_tokens": 0}, + }, + "tools": [], +} + + +@pytest.fixture +def gateway_handler_unavailable(monkeypatch: pytest.MonkeyPatch) -> None: + for module_name in GATEWAY_HANDLER_MODULES: + monkeypatch.setitem(sys.modules, module_name, None) + + +class _FailingImport(MetaPathFinder): + def __init__(self, module_name: str, error: ImportError) -> None: + self.module_name: Final = module_name + self.error: Final = error + + def find_spec( + self, fullname: str, path: Sequence[str] | None, target: ModuleType | None = None + ) -> ModuleSpec | None: + if fullname == self.module_name: + raise self.error + return None + + +@pytest.fixture( + params=[ + ModuleNotFoundError("No module named 'mcp'"), + ImportError("cannot import name 'CacheableResult' from 'mcp.types'"), + ModuleNotFoundError("No module named 'fastapi'"), + ], + ids=["mcp-missing", "mcp-1x", "proxy-extras-missing"], +) +def gateway_registry_unavailable( + request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch, gateway_handler_unavailable: None +) -> None: + monkeypatch.delitem(sys.modules, GATEWAY_REGISTRY_MODULE, raising=False) + monkeypatch.setattr(sys, "meta_path", [_FailingImport(GATEWAY_REGISTRY_MODULE, request.param), *sys.meta_path]) + + +class _ToolList(list[Mapping[str, object]]): + pass + + +@pytest.mark.parametrize( + "tools", + [[FUNCTION_TOOL], [REMOTE_MCP_TOOL], _ToolList([FUNCTION_TOOL]), (FUNCTION_TOOL, REMOTE_MCP_TOOL)], + ids=["function-tool", "remote-mcp-tool", "list-subclass", "tuple"], +) +@pytest.mark.usefixtures("gateway_handler_unavailable") +def test_completion_without_a_gateway_tool_does_not_import_the_gateway_handler(tools: Iterable[Mapping[str, object]]): + response: Final = litellm.completion( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + tools=tools, + mock_response="ok", + ) + + assert response.choices[0].message.content == "ok" + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("gateway_handler_unavailable") +async def test_acompletion_with_function_tool_does_not_import_the_gateway_handler(): + response: Final = await litellm.acompletion( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + tools=[FUNCTION_TOOL], + mock_response="ok", + ) + + assert response.choices[0].message.content == "ok" + + +@pytest.mark.parametrize( + "tool", + [GATEWAY_MCP_TOOL, GATEWAY_SERVER_MCP_TOOL, GATEWAY_PATH_MCP_TOOL], + ids=["litellm-proxy", "litellm-proxy-server", "gateway-path-url"], +) +@pytest.mark.usefixtures("gateway_handler_unavailable") +def test_completion_with_gateway_mcp_tool_still_reaches_the_handler(tool: Mapping[str, object]): + with pytest.raises(ModuleNotFoundError, match=r"litellm_proxy_mcp_handler|chat_completions_handler"): + litellm.completion( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + tools=[tool], + mock_response="ok", + ) + + +@pytest.mark.parametrize( + "tools", + [None, [RESPONSES_FUNCTION_TOOL], [REMOTE_MCP_TOOL]], + ids=["no-tools", "function-tool", "remote-mcp-tool"], +) +@pytest.mark.usefixtures("gateway_handler_unavailable") +def test_responses_without_a_gateway_tool_does_not_import_the_gateway_handler( + tools: list[Mapping[str, object]] | None, respx_mock: respx.MockRouter +): + upstream: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, json=RESPONSES_PAYLOAD) + ) + + response: Final = litellm.responses(model="openai/gpt-4o-mini", input="hi", tools=tools, api_key="sk-test") + + assert upstream.called + assert response.output[0].content[0].text == "Done." + + +@pytest.mark.usefixtures("gateway_handler_unavailable") +def test_responses_with_gateway_mcp_tool_still_reaches_the_handler(): + with pytest.raises(litellm.APIConnectionError, match=r"litellm_proxy_mcp_handler"): + litellm.responses(model="openai/gpt-4o-mini", input="hi", tools=[GATEWAY_MCP_TOOL], api_key="sk-test") + + +@pytest.mark.asyncio +@pytest.mark.usefixtures("gateway_handler_unavailable") +async def test_anthropic_messages_with_function_tool_does_not_import_the_gateway_handler( + monkeypatch: pytest.MonkeyPatch, respx_mock: respx.MockRouter +): + monkeypatch.setattr(litellm, "disable_aiohttp_transport", True) + upstream: Final = respx_mock.post("https://api.anthropic.com/v1/messages").mock( + return_value=httpx.Response(200, json=ANTHROPIC_MESSAGE_PAYLOAD) + ) + + response: Final = await litellm.anthropic_interface.messages.acreate( + model="anthropic/claude-sonnet-5-5", + messages=[{"role": "user", "content": "hi"}], + max_tokens=16, + tools=[ANTHROPIC_FUNCTION_TOOL], + api_key="sk-test", + ) + + assert upstream.called + assert response["content"][0]["text"] == "ok" + + +@pytest.mark.parametrize( + "tools, expected", + [ + (None, False), + ([], False), + ([FUNCTION_TOOL], False), + ([RESPONSES_FUNCTION_TOOL], False), + ([{"type": "web_search"}], False), + (["not-a-mapping"], False), + ([REMOTE_MCP_TOOL], False), + ([{"type": "mcp", "server_label": "x"}], False), + ([{"type": "mcp", "server_label": "x", "server_url": None}], False), + ([{"server_url": "litellm_proxy"}], False), + ([FUNCTION_TOOL, GATEWAY_MCP_TOOL], True), + ([GATEWAY_SERVER_MCP_TOOL], True), + ([GATEWAY_PATH_MCP_TOOL], True), + ([REMOTE_MCP_TOOL, GATEWAY_PATH_MCP_TOOL], True), + ], + ids=[ + "none", + "empty", + "chat-function", + "responses-function", + "web-search", + "non-mapping", + "remote-mcp", + "mcp-without-url", + "mcp-null-url", + "gateway-url-without-mcp-type", + "gateway-among-functions", + "gateway-server", + "gateway-path-url", + "gateway-path-after-remote-mcp", + ], +) +def test_may_name_litellm_mcp_gateway_matches_only_gateway_mcp_tools(tools: list[object] | None, expected: bool): + assert may_name_litellm_mcp_gateway(tools) is expected + + +@pytest.mark.parametrize( + "tools, expected", + [ + ([GATEWAY_PATH_MCP_TOOL], False), + ([FUNCTION_TOOL, GATEWAY_PATH_MCP_TOOL], False), + ([GATEWAY_MCP_TOOL], True), + ([GATEWAY_SERVER_MCP_TOOL, GATEWAY_PATH_MCP_TOOL], True), + ], + ids=["gateway-path-url", "gateway-path-among-functions", "litellm-proxy", "litellm-proxy-server-and-path-url"], +) +@pytest.mark.usefixtures("gateway_registry_unavailable") +def test_proxy_path_url_names_the_gateway_only_when_its_registry_loads(tools: list[object], expected: bool): + assert may_name_litellm_mcp_gateway(tools) is expected + + +@pytest.mark.usefixtures("gateway_registry_unavailable") +def test_responses_sends_a_proxy_path_mcp_url_to_the_provider_without_the_gateway_registry( + respx_mock: respx.MockRouter, +): + upstream: Final = respx_mock.post("https://api.openai.com/v1/responses").mock( + return_value=httpx.Response(200, json=RESPONSES_PAYLOAD) + ) + + response: Final = litellm.responses( + model="openai/gpt-4o-mini", input="hi", tools=[GATEWAY_PATH_MCP_TOOL], api_key="sk-test" + ) + + assert response.output[0].content[0].text == "Done." + assert json.loads(upstream.calls.last.request.content)["tools"] == [GATEWAY_PATH_MCP_TOOL] + + +@pytest.mark.usefixtures("gateway_registry_unavailable") +def test_completion_sends_a_proxy_path_mcp_url_to_the_provider_without_the_gateway_registry(): + response: Final = litellm.completion( + model="openai/gpt-4o-mini", + messages=[{"role": "user", "content": "hi"}], + tools=[GATEWAY_PATH_MCP_TOOL], + mock_response="ok", + ) + + assert response.choices[0].message.content == "ok" + + +class _ToolTuple(tuple[Mapping[str, object], ...]): + pass + + +@pytest.mark.parametrize( + "make_tools", + [ + lambda items: (item for item in items), + lambda items: iter(items), + lambda items: deque(items), + lambda items: _ToolList(items), + lambda items: _ToolTuple(items), + ], + ids=["generator", "iterator", "deque", "list-subclass", "tuple-subclass"], +) +@pytest.mark.usefixtures("gateway_handler_unavailable") +def test_precheck_reads_every_iterable_shape_without_the_gateway_handler( + make_tools: Callable[[list[Mapping[str, object]]], Iterable[Mapping[str, object]]], +): + assert may_name_litellm_mcp_gateway(make_tools([FUNCTION_TOOL, REMOTE_MCP_TOOL])) is False + assert may_name_litellm_mcp_gateway(make_tools([FUNCTION_TOOL, GATEWAY_MCP_TOOL])) is True diff --git a/tests/unit/responses/test_responses_prompt_management.py b/tests/unit/responses/test_responses_prompt_management.py index 1e8880424ea..9af92e7f9ca 100644 --- a/tests/unit/responses/test_responses_prompt_management.py +++ b/tests/unit/responses/test_responses_prompt_management.py @@ -70,10 +70,6 @@ def _patch_responses_dispatch(): import_module("litellm.responses.main").litellm, "get_llm_provider", side_effect=_provider_by_model, ), - patch.object( - import_module("litellm.responses.mcp.litellm_proxy_mcp_handler").LiteLLM_Proxy_MCP_Handler, "should_use_litellm_mcp_gateway", - return_value=False, - ), patch.object( import_module("litellm.responses.main").ProviderConfigManager, "get_provider_responses_api_config", return_value=None, @@ -158,7 +154,7 @@ class TestResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3]: + with patches[0], patches[1], patches[2]: import litellm litellm.responses( @@ -193,7 +189,7 @@ class TestResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3]: + with patches[0], patches[1], patches[2]: import litellm litellm.responses( @@ -216,7 +212,7 @@ class TestResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3]: + with patches[0], patches[1], patches[2]: import litellm litellm.responses( @@ -247,7 +243,7 @@ class TestResponsesAPIPromptManagement: logging_obj.model_call_details = {} patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3] as mock_handler: + with patches[0], patches[1], patches[2] as mock_handler: import litellm litellm.responses( @@ -273,7 +269,7 @@ class TestResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3] as mock_handler: + with patches[0], patches[1], patches[2] as mock_handler: import litellm litellm.responses( @@ -305,7 +301,7 @@ class TestResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3]: + with patches[0], patches[1], patches[2]: import litellm litellm.responses( @@ -328,7 +324,7 @@ class TestResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3] as mock_handler: + with patches[0], patches[1], patches[2] as mock_handler: import litellm litellm.responses( @@ -368,7 +364,7 @@ class TestResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3] as mock_handler: + with patches[0], patches[1], patches[2] as mock_handler: import litellm litellm.responses( @@ -399,8 +395,7 @@ class TestResponsesAPIPromptManagement: side_effect=_provider_by_model, ), patches[1], - patches[2], - patches[3] as mock_handler, + patches[2] as mock_handler, ): import litellm @@ -437,7 +432,7 @@ class TestAsyncResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3]: + with patches[0], patches[1], patches[2]: import litellm await litellm.aresponses( @@ -467,7 +462,7 @@ class TestAsyncResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3] as mock_handler: + with patches[0], patches[1], patches[2] as mock_handler: import litellm await litellm.aresponses( @@ -498,7 +493,7 @@ class TestAsyncResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3]: + with patches[0], patches[1], patches[2]: import litellm await litellm.aresponses( @@ -524,7 +519,7 @@ class TestAsyncResponsesAPIPromptManagement: ) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3] as mock_handler: + with patches[0], patches[1], patches[2] as mock_handler: import litellm await litellm.aresponses( @@ -657,7 +652,7 @@ async def test_async_guardrail_block_from_prompt_hook_reaches_caller_unwrapped() logging_obj.async_get_chat_completion_prompt = AsyncMock(side_effect=block) patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3], pytest.raises(HTTPException) as exc_info: + with patches[0], patches[1], patches[2], pytest.raises(HTTPException) as exc_info: await litellm.aresponses(input="Hi", model="gpt-4o", prompt_id="blocked", litellm_logging_obj=logging_obj) assert exc_info.value is block @@ -669,7 +664,7 @@ def test_sync_guardrail_block_from_prompt_hook_reaches_caller_unwrapped(): logging_obj.get_chat_completion_prompt.side_effect = block patches = _patch_responses_dispatch() - with patches[0], patches[1], patches[2], patches[3], pytest.raises(HTTPException) as exc_info: + with patches[0], patches[1], patches[2], pytest.raises(HTTPException) as exc_info: litellm.responses(input="Hi", model="gpt-4o", prompt_id="blocked", litellm_logging_obj=logging_obj) assert exc_info.value is block diff --git a/tests/unit/test_main.py b/tests/unit/test_main.py index be2eaafca22..a0677d88505 100644 --- a/tests/unit/test_main.py +++ b/tests/unit/test_main.py @@ -8,7 +8,7 @@ import logging import os import threading import urllib.parse -from collections.abc import Mapping +from collections.abc import Awaitable, Callable, Mapping from concurrent.futures import ThreadPoolExecutor from dataclasses import dataclass from datetime import datetime @@ -8124,6 +8124,19 @@ async def test_responses_with_retries(sync_mode): assert mock_responses.call_args.kwargs["max_retries"] == 0 +@pytest.mark.parametrize("retry_function", [acompletion_with_retries, aresponses_with_retries]) +@pytest.mark.asyncio +async def test_async_retry_helpers_stop_when_the_call_is_cancelled( + retry_function: Callable[..., Awaitable[object]], +) -> None: + upstream: Final = AsyncMock(side_effect=[asyncio.CancelledError(), MagicMock()]) + + with pytest.raises(asyncio.CancelledError): + await retry_function(model="gpt-4o", num_retries=3, original_function=upstream) + + assert upstream.await_count == 1 + + def test_azure_embedding_exceptions(): with pytest.raises(Exception, match="Mock error") as exc_info: litellm.embedding( diff --git a/tests/unit/test_router/test_router.py b/tests/unit/test_router/test_router.py index 8cde45e7603..fe9541cbe21 100644 --- a/tests/unit/test_router/test_router.py +++ b/tests/unit/test_router/test_router.py @@ -7371,7 +7371,8 @@ def test_get_wildcard_deployment_usable_by_team_prefers_the_team_pattern(): pytest.param(False, "rwild-chat", "team-r", ("team-r-wildcard",), id="team-pattern-when-no-proxy-wide-match"), pytest.param(False, "rwild-chat", "team-b", (), id="another-teams-pattern-never-serves"), pytest.param(False, "rwild-chat", None, (), id="no-team-no-team-pattern"), - pytest.param(True, "rwild-chat", "team-r", ("catch-all",), id="proxy-wide-pattern-beats-the-team-pattern"), + pytest.param(True, "rwild-chat", "team-r", ("team-r-wildcard",), id="team-pattern-beats-the-proxy-wide-pattern"), + pytest.param(True, "rwild-chat", "team-b", ("catch-all",), id="another-team-falls-through-to-proxy-wide"), pytest.param(True, "ptu-chat", "team-r", ("ptu-wildcard",), id="only-the-most-specific-proxy-wide-pattern"), ], ) @@ -7407,6 +7408,39 @@ def test_wildcard_route_deployments_is_the_one_pattern_routing_serves_from( assert tuple(Deployment.model_validate(d).model_info.id for d in early_deployments) == expected_ids +@pytest.mark.parametrize( + ("team_id", "expected_id"), + [ + pytest.param("team-a", "team-a-byok", id="team-key-uses-its-own-byok-wildcard"), + pytest.param("team-b", "shared-openrouter", id="team-without-byok-uses-the-shared-wildcard"), + pytest.param(None, "shared-openrouter", id="key-without-team-uses-the-shared-wildcard"), + ], +) +def test_team_byok_wildcard_is_served_before_a_proxy_wide_wildcard_of_the_same_name( + team_id: str | None, expected_id: str +): + """Regression: a proxy-wide openrouter/* shadowed a team's own openrouter/* BYOK deployment for that team's keys.""" + router: Final = litellm.Router( + model_list=[ + { + "model_name": "model_name_team-a_internal", + "litellm_params": {"model": "openrouter/*", "api_key": "sk-team-a"}, + "model_info": {"id": "team-a-byok", "team_id": "team-a", "team_public_model_name": "openrouter/*"}, + }, + { + "model_name": "openrouter/*", + "litellm_params": {"model": "openrouter/*", "api_key": "sk-shared"}, + "model_info": {"id": "shared-openrouter"}, + }, + ] + ) + requested: Final = "openrouter/openai/gpt-5-nano" + + served: Final = router.wildcard_route_deployments(requested, team_id) + + assert tuple(Deployment.model_validate(d).model_info.id for d in served) == (expected_id,) + + def test_get_deployment_credentials_with_provider_aws_bedrock_runtime_endpoint(): """ Test that get_deployment_credentials_with_provider correctly copies diff --git a/tests/unit/test_router/test_router_provider_endpoints.py b/tests/unit/test_router/test_router_provider_endpoints.py index 708322f4fec..9312204f471 100644 --- a/tests/unit/test_router/test_router_provider_endpoints.py +++ b/tests/unit/test_router/test_router_provider_endpoints.py @@ -16,6 +16,7 @@ from litellm.caching.dual_cache import DualCache from litellm.caching.in_memory_cache import InMemoryCache from litellm.constants import ROUTER_USAGE_COUNTED_TOKENS_METADATA_KEY from litellm.litellm_core_utils.streaming_handler import CustomStreamWrapper +from litellm.proxy._types import UserAPIKeyAuth from litellm.router_utils.router_callbacks.track_deployment_metrics import ( get_deployment_successes_for_current_minute, ) @@ -23,6 +24,7 @@ from litellm.types.caching import RedisPipelineIncrementOperation from litellm.types.llms.openai import HttpxBinaryResponseContent from litellm.types.router import RoutingStrategy from litellm.types.utils import ImageResponse, ModelResponse, RerankResponse +from litellm.types.videos.utils import encode_video_id_with_provider from openai import AsyncAzureOpenAI, AsyncOpenAI EVENT_TIMEOUT_SECONDS: Final = 5 @@ -730,3 +732,248 @@ async def test_router_moderation_endpoint_factory_invokes_default_model( assert route.call_count == 1 assert json.loads(route.calls[0].request.content) == {"input": "hello"} assert response.id == "modr-2" + + +def _sora_deployment( + model_name: str, api_key: str, deployment_id: str, access_groups: tuple[str, ...] = () +) -> dict[str, object]: + return { + "model_name": model_name, + "litellm_params": {"model": "openai/sora-2", "api_key": api_key}, + "model_info": {"id": deployment_id, "access_groups": list(access_groups)}, + } + + +def _two_account_sora_router() -> Router: + return Router( + model_list=[ + _sora_deployment("sora", "key-a", "dep-a"), + _sora_deployment("sora", "key-b", "dep-b"), + _sora_deployment("other", "key-x", "dep-x"), + ], + model_group_alias={"sora-alias": "sora"}, + ) + + +async def _poll_status_keys( + router: Router, + model: str, + deployment_id: str, + respx_mock: respx.MockRouter, + caller: UserAPIKeyAuth | None = None, +) -> set[str]: + route: Final = respx_mock.get("https://api.openai.com/v1/videos/video_raw").respond( + 200, json={"id": "video_raw", "object": "video", "status": "in_progress"} + ) + video_id: Final = encode_video_id_with_provider("video_raw", "openai", "sora", deployment_id=deployment_id) + metadata: Final = {} if caller is None else {"user_api_key_auth": caller} + for _ in range(10): + await router.avideo_status(model=model, video_id=video_id, litellm_metadata=dict(metadata)) + return {call.request.headers["authorization"] for call in route.calls} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ["sora", "sora-alias"]) +async def test_video_status_polls_the_deployment_that_created_the_video( + model: str, respx_mock: respx.MockRouter +) -> None: + keys: Final = await _poll_status_keys(_two_account_sora_router(), model, "dep-b", respx_mock) + + assert keys == {"Bearer key-b"} + + +def test_sync_video_status_polls_the_deployment_that_created_the_video(respx_mock: respx.MockRouter) -> None: + route: Final = respx_mock.get("https://api.openai.com/v1/videos/video_raw").respond( + 200, json={"id": "video_raw", "object": "video", "status": "in_progress"} + ) + router: Final = _two_account_sora_router() + video_id: Final = encode_video_id_with_provider("video_raw", "openai", "sora", deployment_id="dep-b") + + for _ in range(10): + router.video_status(model="sora", video_id=video_id) + + assert {call.request.headers["authorization"] for call in route.calls} == {"Bearer key-b"} + + +@pytest.mark.asyncio +async def test_video_status_on_a_wildcard_polls_the_deployment_that_created_the_video( + respx_mock: respx.MockRouter, +) -> None: + router: Final = Router( + model_list=[ + { + "model_name": "openai/*", + "litellm_params": {"model": "openai/*", "api_key": key}, + "model_info": {"id": dep}, + } + for dep, key in (("wc-a", "key-a"), ("wc-b", "key-b")) + ] + ) + + keys: Final = await _poll_status_keys(router, "openai/sora-2", "wc-b", respx_mock) + + assert keys == {"Bearer key-b"} + + +@pytest.mark.asyncio +async def test_video_status_never_pins_to_a_deployment_in_cooldown(respx_mock: respx.MockRouter) -> None: + router: Final = _two_account_sora_router() + router.cooldown_cache.add_deployment_to_cooldown( + model_id="dep-b", + original_exception=litellm.RateLimitError(message="upstream 429", llm_provider="openai", model="sora-2"), + exception_status=429, + cooldown_time=60, + ) + + keys: Final = await _poll_status_keys(router, "sora", "dep-b", respx_mock) + + assert keys == {"Bearer key-a"} + + +@pytest.mark.asyncio +async def test_video_status_on_a_routing_group_polls_the_member_deployment_that_created_the_video( + respx_mock: respx.MockRouter, +) -> None: + router: Final = Router( + model_list=[_sora_deployment("sora-a", "key-a", "dep-a"), _sora_deployment("sora-b", "key-b", "dep-b")], + routing_groups=[ + {"group_name": "sora-group", "models": ["sora-a", "sora-b"], "routing_strategy": "simple-shuffle"} + ], + ) + + keys: Final = await _poll_status_keys(router, "sora-group", "dep-b", respx_mock) + + assert keys == {"Bearer key-b"} + + +@pytest.mark.asyncio +async def test_video_status_never_pins_to_a_deployment_outside_the_requested_group( + respx_mock: respx.MockRouter, +) -> None: + keys: Final = await _poll_status_keys(_two_account_sora_router(), "sora", "dep-x", respx_mock) + + assert keys <= {"Bearer key-a", "Bearer key-b"} + + +@pytest.mark.asyncio +async def test_video_status_with_a_removed_deployment_routes_across_the_group(respx_mock: respx.MockRouter) -> None: + keys: Final = await _poll_status_keys(_two_account_sora_router(), "sora", "dep-gone", respx_mock) + + assert keys <= {"Bearer key-a", "Bearer key-b"} + assert keys + + +@pytest.mark.asyncio +@pytest.mark.parametrize(("key_access_group", "expected_key"), [("ag-a", "Bearer key-a"), ("ag-b", "Bearer key-b")]) +async def test_video_status_pins_only_to_a_deployment_in_the_callers_access_groups( + key_access_group: str, expected_key: str, respx_mock: respx.MockRouter +) -> None: + router: Final = Router( + model_list=[ + _sora_deployment("sora", "key-a", "dep-a", access_groups=("ag-a",)), + _sora_deployment("sora", "key-b", "dep-b", access_groups=("ag-b",)), + ] + ) + + keys: Final = await _poll_status_keys( + router, "sora", "dep-b", respx_mock, caller=UserAPIKeyAuth(models=[key_access_group]) + ) + + assert keys == {expected_key} + + +@pytest.mark.asyncio +@pytest.mark.parametrize("model", ["sora", "sora-alias"]) +async def test_pinned_video_status_logs_the_requested_model_group( + model: str, respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch +) -> None: + capture: Final = _RouterLoggingCapture("dep-b") + monkeypatch.setattr(litellm, "callbacks", [capture]) + + keys: Final = await _poll_status_keys(_two_account_sora_router(), model, "dep-b", respx_mock) + _, standard_logging_object = await capture.next_event() + + assert keys == {"Bearer key-b"} + assert isinstance(standard_logging_object, dict) + assert standard_logging_object["model_group"] == model + + +def _ordered_chat_router() -> Router: + return Router( + model_list=[ + { + "model_name": "chat-ordered", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": key, "order": order}, + "model_info": {"id": deployment_id}, + } + for deployment_id, key, order in (("dep-primary", "key-a", 1), ("dep-backup", "key-b", 2)) + ] + ) + + +@pytest.mark.asyncio +async def test_a_video_id_in_a_chat_body_never_overrides_deployment_order() -> None: + router: Final = _ordered_chat_router() + video_id: Final = encode_video_id_with_provider("video_raw", "openai", "chat-ordered", deployment_id="dep-backup") + + chats: Final = [ + await router.acompletion( + model="chat-ordered", messages=list(ROUTING_MESSAGES), mock_response="ok", video_id=video_id + ) + for _ in range(4) + ] + responses: Final = [ + await router.aresponses(model="chat-ordered", input="route this request", mock_response="ok", video_id=video_id) + for _ in range(4) + ] + + assert {chat._hidden_params["model_id"] for chat in chats} == {"dep-primary"} + assert {response._hidden_params["model_id"] for response in responses} == {"dep-primary"} + + +@pytest.mark.asyncio +async def test_video_lookup_404_never_cools_down_the_deployment_that_answered(respx_mock: respx.MockRouter) -> None: + respx_mock.get("https://api.openai.com/v1/videos/video_gone").respond( + 404, json={"error": {"code": "NotFound", "message": "Video not found"}} + ) + router: Final = Router( + model_list=[_sora_deployment("sora", "key-a", "dep-a"), _sora_deployment("sora", "key-b", "dep-b")], + num_retries=0, + allowed_fails=1, + ) + video_id: Final = encode_video_id_with_provider("video_gone", "openai", "sora", deployment_id="dep-b") + + for _ in range(3): + with pytest.raises(litellm.NotFoundError): + await router.avideo_status(model="sora", video_id=video_id) + + assert await router.cooldown_cache.async_get_active_cooldowns(["dep-a", "dep-b"], parent_otel_span=None) == [] + + +@pytest.mark.asyncio +async def test_chat_404_still_cools_down_the_deployment(respx_mock: respx.MockRouter) -> None: + respx_mock.post("https://api.openai.com/v1/chat/completions").respond( + 404, json={"error": {"code": "NotFound", "message": "Deployment not found"}} + ) + router: Final = Router( + model_list=[ + { + "model_name": "chat", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "key-a"}, + "model_info": {"id": "dep-chat"}, + }, + { + "model_name": "chat", + "litellm_params": {"model": "openai/gpt-4o-mini", "api_key": "key-b"}, + "model_info": {"id": "dep-chat-b"}, + }, + ], + num_retries=0, + allowed_fails=1, + ) + + for _ in range(3): + with contextlib.suppress(litellm.NotFoundError, litellm.RateLimitError): + await router.acompletion(model="chat", messages=list(ROUTING_MESSAGES)) + + assert await router.cooldown_cache.async_get_active_cooldowns(["dep-chat", "dep-chat-b"], parent_otel_span=None) diff --git a/tests/unit/test_video_generation.py b/tests/unit/test_video_generation.py index 5c4bc8ea652..df7e9c68b90 100644 --- a/tests/unit/test_video_generation.py +++ b/tests/unit/test_video_generation.py @@ -2,6 +2,7 @@ import asyncio import io import json import os +from typing import Final from unittest.mock import AsyncMock, MagicMock, patch import pytest @@ -1269,6 +1270,35 @@ def test_encode_video_id_with_provider_handles_azure_video_prefix(): assert encoded_twice == encoded_id # Should return the same encoded ID +@pytest.mark.parametrize( + "served_model, deployment_id, expected_extras", + [ + ("seedance-east", "dep-1", {"served_model": "seedance-east", "deployment_id": "dep-1"}), + (None, "dep-1", {"deployment_id": "dep-1"}), + ("seedance-east", None, {"served_model": "seedance-east"}), + (None, None, {}), + ], +) +def test_video_id_round_trips_served_model_and_deployment_id( + served_model: str | None, deployment_id: str | None, expected_extras: dict[str, str] +) -> None: + from litellm.types.videos.utils import ( + decode_video_id_with_provider, + encode_video_id_with_provider, + ) + + encoded: Final = encode_video_id_with_provider( + "task_1", "volcengine", "seedance-group", served_model=served_model, deployment_id=deployment_id + ) + + assert decode_video_id_with_provider(encoded) == { + "custom_llm_provider": "volcengine", + "model_id": "seedance-group", + "video_id": "task_1", + **expected_extras, + } + + class TestVideoListTransformation: """Tests for video list request/response transformation with provider ID encoding.""" diff --git a/uv.lock b/uv.lock index cc78423f359..0df2668e679 100644 --- a/uv.lock +++ b/uv.lock @@ -4594,7 +4594,7 @@ wheels = [ [[package]] name = "litellm" -version = "1.108.2" +version = "1.108.3" source = { editable = "." } dependencies = [ { name = "aiohttp" },