diff --git a/litellm/completion_extras/litellm_responses_transformation/transformation.py b/litellm/completion_extras/litellm_responses_transformation/transformation.py index da68b74cb6e..996818b6fa0 100644 --- a/litellm/completion_extras/litellm_responses_transformation/transformation.py +++ b/litellm/completion_extras/litellm_responses_transformation/transformation.py @@ -304,6 +304,136 @@ def _reasoning_items_from_output_items(output_items: Sequence[object]) -> tuple[ ) +_TOOL_SEARCH_ITEM_TYPES: Final = frozenset({"tool_search_call", "tool_search_output"}) +_TOOL_SEARCH_REPLAY_ITEM_TYPES: Final = frozenset({"reasoning", *_TOOL_SEARCH_ITEM_TYPES}) +_TOOL_SEARCH_ITEMS_KEY: Final = "tool_search_items" + + +def _json_field(value: object, key: str) -> object: + return cast(Mapping[str, object], value).get(key) if isinstance(value, dict) else None # cast-ok: untyped JSON + + +def _tool_search_replay_item(item: object) -> Mapping[str, object] | None: + raw: Final[object] = item.model_dump(exclude_unset=True) if isinstance(item, BaseModel) else item + if not isinstance(raw, dict): + return None + fields: Final = cast(dict[str, object], raw) # cast-ok: provider JSON item, keys forwarded verbatim + if fields.get("type") not in _TOOL_SEARCH_REPLAY_ITEM_TYPES: + return None + return {key: value for key, value in fields.items() if value is not None} + + +def _tool_search_items_from_output_items(output_items: object) -> tuple[Mapping[str, object], ...]: + if not isinstance(output_items, list): + return () + entries: Final = cast(Sequence[object], output_items) # cast-ok: each entry is re-validated + items: Final = tuple(item for item in (_tool_search_replay_item(entry) for entry in entries) if item is not None) + last_search: Final = max( + (index for index, item in enumerate(items) if item.get("type") in _TOOL_SEARCH_ITEM_TYPES), default=-1 + ) + return items[: last_search + 1] + + +def _tool_search_provider_fields( + items: Sequence[Mapping[str, object]], +) -> dict[str, object] | None: # mutable-ok: Message.provider_specific_fields is typed as a dict + return {_TOOL_SEARCH_ITEMS_KEY: [dict(item) for item in items]} if items else None + + +def _tool_search_input_items(msg: object) -> tuple[Mapping[str, object], ...]: + return _tool_search_items_from_output_items( + _json_field(_json_field(msg, "provider_specific_fields"), _TOOL_SEARCH_ITEMS_KEY) + ) + + +def carries_tool_search_items(provider_specific_fields: object) -> bool: + return bool(_json_field(provider_specific_fields, _TOOL_SEARCH_ITEMS_KEY)) + + +def _tool_definition_name(tool: object) -> tuple[str, ...]: + name: Final = _json_field(tool, "server_label" if _json_field(tool, "type") == "mcp" else "name") + return (name,) if isinstance(name, str) else () + + +def _namespace_members(tool: object) -> Sequence[object]: + members: Final = _json_field(tool, "tools") + if _json_field(tool, "type") != "namespace" or not isinstance(members, list): + return (tool,) + return cast(Sequence[object], members) # cast-ok: untyped JSON, each member re-checked + + +def _loaded_tools(item: Mapping[str, object]) -> Sequence[object]: + tools: Final = item.get("tools") + if item.get("type") != "tool_search_output" or not isinstance(tools, list): + return () + return cast(Sequence[object], tools) # cast-ok: untyped JSON, each tool re-checked + + +def _json_list(value: object) -> Sequence[object] | None: + return cast(Sequence[object], value) if isinstance(value, list) else None # cast-ok: untyped JSON + + +def _message_loaded_tools(message: object) -> Iterator[object]: + loaded: Final = chain.from_iterable(_loaded_tools(item) for item in _tool_search_input_items(message)) + return chain.from_iterable(_namespace_members(tool) for tool in loaded) + + +def replayed_tool_search_tools(request: object) -> tuple[object, ...]: + messages: Final = _json_list(_json_field(request, "messages")) or () + return tuple(chain.from_iterable(_message_loaded_tools(message) for message in messages)) + + +def replayed_tool_search_tool_names(request: object) -> tuple[str, ...]: + return tuple(chain.from_iterable(_tool_definition_name(tool) for tool in replayed_tool_search_tools(request))) + + +def _loads_any_tool(message: object, names: frozenset[str]) -> bool: + loaded_names: Final = chain.from_iterable(_tool_definition_name(tool) for tool in _message_loaded_tools(message)) + return any(name in names for name in loaded_names) + + +def _without_tool_search_items(message: object) -> object: + fields: Final = _json_field(message, "provider_specific_fields") + if not isinstance(fields, dict): + return message + kept: Final = { + key: value + for key, value in cast(Mapping[str, object], fields).items() # cast-ok: untyped JSON + if key != _TOOL_SEARCH_ITEMS_KEY + } + return {**cast(Mapping[str, object], message), "provider_specific_fields": kept} # cast-ok: dict checked above + + +def without_tool_search_items_loading( + request: object, names: frozenset[str] +) -> list[object] | None: # mutable-ok: request messages stay a list for every downstream handler + messages: Final = _json_list(_json_field(request, "messages")) + if messages is None: + return None + return [_without_tool_search_items(message) if _loads_any_tool(message, names) else message for message in messages] + + +def _effective_tools(request: object) -> Sequence[object]: + override: Final = _json_list(_json_field(_json_field(request, "extra_body"), "tools")) + return override if override is not None else (_json_list(_json_field(request, "tools")) or ()) + + +def _requests_tool_search(request: object) -> bool: + return any(_json_field(tool, "type") == "tool_search" for tool in _effective_tools(request)) + + +def _reasoning_ids(items: Sequence[Mapping[str, object]]) -> frozenset[str]: + ids: Final = (item.get("id") for item in items if item.get("type") == "reasoning") + return frozenset(item_id for item_id in ids if isinstance(item_id, str)) + + +def _assistant_replay_items(msg: "AllMessageValues", replay_tool_search: bool) -> tuple[Mapping[str, object], ...]: + searched: Final = _tool_search_input_items(msg) if replay_tool_search else () + searched_reasoning_ids: Final = _reasoning_ids(searched) + reasoning: Final = (item for item in _reasoning_input_items(msg) if item.get("id") not in searched_reasoning_ids) + return (*searched, *reasoning) + + def _as_chat_reasoning_items( reasoning_items: Sequence[_BuiltReasoningItem | ChatCompletionReasoningItem], ) -> list[ChatCompletionReasoningItem] | None: @@ -346,6 +476,7 @@ def _fold_choices_into_one( tool_calls: Sequence[Mapping[str, object]], trailing_reasoning_content: str | None, trailing_reasoning_item: _BuiltReasoningItem | None, + tool_search_items: Sequence[Mapping[str, object]], structured_output: bool = False, ) -> "Choices": """Clients that only read choices[0] must still see every text and tool call. @@ -373,6 +504,7 @@ def _fold_choices_into_one( tool_calls=list(tool_calls) or None, reasoning_content=" ".join(value for value in reasoning_contents if value) or None, reasoning_items=_as_chat_reasoning_items(reasoning_items), + provider_specific_fields=_tool_search_provider_fields(tool_search_items), ), finish_reason="tool_calls" if tool_calls else "stop", index=0, @@ -457,7 +589,9 @@ class _ChatToolCallDict(ChatCompletionToolCallChunk, total=False): provider_specific_fields: Mapping[str, object] -def tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _ChatToolCallDict: +def tool_call_dict_from_output_item( + item: Mapping[str, Any], index: int, *, include_namespace: bool = False +) -> _ChatToolCallDict: """Convert a ``function_call`` or ``custom_tool_call`` output item dict to a chat completions tool_call dict. Custom (grammar/freeform) tool calls carry their raw string payload in ``input`` rather than ``arguments``; both map to @@ -486,12 +620,24 @@ def tool_call_dict_from_output_item(item: Mapping[str, Any], index: int) -> _Cha provider_specific_fields = vars(raw_provider_fields) else: provider_specific_fields = None - if provider_specific_fields: - tool_call_dict["provider_specific_fields"] = provider_specific_fields - function_chunk["provider_specific_fields"] = provider_specific_fields + namespace: Final = item.get("namespace") if include_namespace else None + merged_provider_fields: Final = { + **(provider_specific_fields or {}), + **({"namespace": namespace} if isinstance(namespace, str) else {}), + } + if merged_provider_fields: + tool_call_dict["provider_specific_fields"] = merged_provider_fields + function_chunk["provider_specific_fields"] = merged_provider_fields return tool_call_dict +def _tool_call_namespace(tool_call: object) -> str | None: + top_level: Final = _json_field(_json_field(tool_call, "provider_specific_fields"), "namespace") + function_fields: Final = _json_field(_json_field(tool_call, "function"), "provider_specific_fields") + nested: Final = _json_field(function_fields, "namespace") + return top_level if isinstance(top_level, str) else nested if isinstance(nested, str) else None + + def _flat_responses_tool_choice(choice_type: str, name: str) -> ToolChoiceFunctionParam | ToolChoiceCustomParam: if choice_type == "custom": return ToolChoiceCustomParam(type="custom", name=name) @@ -641,9 +787,14 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): *, drop_params: bool = False, keep_prompt_cache_breakpoints: bool = False, + replay_tool_search: bool = False, + include_namespace: bool = False, ) -> tuple[list[object], str | None]: converted_input_items, instructions = self._convert_chat_completion_messages_to_responses_input( - messages, drop_params=drop_params + messages, + drop_params=drop_params, + replay_tool_search=replay_tool_search, + include_namespace=include_namespace, ) return ( converted_input_items @@ -653,7 +804,12 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ) def _convert_chat_completion_messages_to_responses_input( - self, messages: list["AllMessageValues"], *, drop_params: bool = False + self, + messages: list["AllMessageValues"], + *, + drop_params: bool = False, + replay_tool_search: bool = False, + include_namespace: bool = False, ) -> tuple[list[object], str | None]: input_items: Final[list[object]] = [] instructions: str | None = None @@ -735,7 +891,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ) ) elif role == "assistant" and tool_calls and isinstance(tool_calls, list): - input_items.extend(_reasoning_input_items(msg)) + input_items.extend(_assistant_replay_items(msg, replay_tool_search)) if content: input_items.append( { @@ -758,6 +914,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): input_tool_call["name"] = function["name"] if "arguments" in function: input_tool_call["arguments"] = function["arguments"] + if include_namespace and ( + namespace := _tool_call_namespace(cast(object, tool_call)) # cast-ok: untyped + ): + input_tool_call["namespace"] = namespace input_items.append(input_tool_call) elif isinstance(custom, dict): input_items.append( @@ -772,7 +932,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): raise ValueError(f"tool call not supported: {tool_call}") elif content is not None: if role == "assistant": - input_items.extend(_reasoning_input_items(msg)) + input_items.extend(_assistant_replay_items(msg, replay_tool_search)) input_items.append( { "type": "message", @@ -783,7 +943,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): } ) elif role == "assistant": - input_items.extend(_reasoning_input_items(msg)) + input_items.extend(_assistant_replay_items(msg, replay_tool_search)) return input_items, instructions @@ -890,6 +1050,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): litellm_logging_obj: "LiteLLMLoggingObj", client: object | None = None, ) -> dict: + optional_params = self._extract_extra_body_params(optional_params) provider: Final = _served_provider(litellm_params) base_model: Final = litellm_params.get("base_model") supports_prompt_cache_breakpoint: Final = supports_openai_prompt_cache_breakpoint(model, provider) or ( @@ -901,6 +1062,8 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): messages, drop_params=bool(litellm_params.get("drop_params") or litellm.drop_params), keep_prompt_cache_breakpoints=supports_prompt_cache_breakpoint, + replay_tool_search=_requests_tool_search(cast(object, optional_params)), # cast-ok: upcast + include_namespace=True, ) # OpenAI's Responses API rejects an empty input. For a system-only # request, carry the system message as a system-role input item instead @@ -925,8 +1088,6 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): ) instructions: Final = None if is_system_only_request else converted_instructions - optional_params = self._extract_extra_body_params(optional_params) - # Build responses API request using the reverse transformation logic responses_api_request: Final = ResponsesAPIOptionalRequestParams() @@ -1100,7 +1261,9 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): # Tool calls accumulate into the single trailing tool_calls choice # like the typed branches above; a choice per call would hide every # call after choices[0] from chat clients - accumulated_tool_calls.append(tool_call_dict_from_output_item(raw_item, tool_call_index)) + accumulated_tool_calls.append( + tool_call_dict_from_output_item(raw_item, tool_call_index, include_namespace=True) + ) tool_call_index += 1 elif handle_raw_dict_callback is not None: choice, index = handle_raw_dict_callback(item=raw_item, index=index) @@ -1109,22 +1272,31 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): else: pass # don't fail request if item in list is not supported + tool_search_items: Final = _tool_search_items_from_output_items(output_items) if not choices and not accumulated_tool_calls: - return choices + if not tool_search_items: + return choices + return [ + LiteLLMResponsesTransformationHandler._build_empty_choice( + output_items, "stop", tool_search_items=tool_search_items + ) + ] return [ _fold_choices_into_one( choices=choices, tool_calls=accumulated_tool_calls, trailing_reasoning_content=reasoning_content, trailing_reasoning_item=pending_reasoning_item, + tool_search_items=tool_search_items, structured_output=structured_output, ) ] @staticmethod - def _build_empty_incomplete_choice( + def _build_empty_choice( output_items: Sequence[object], - finish_reason: Literal["length", "content_filter"], + finish_reason: Literal["stop", "length", "content_filter"], + tool_search_items: Sequence[Mapping[str, object]] = (), ) -> "Choices": from litellm.types.utils import Choices, Message @@ -1136,9 +1308,10 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): if summary_block.get("text") ) message: Final = Message( - content="", + content=None if tool_search_items else "", reasoning_content=reasoning_content if reasoning_content else None, reasoning_items=_as_chat_reasoning_items(reasoning_items), + provider_specific_fields=_tool_search_provider_fields(tool_search_items), ) return Choices(message=message, finish_reason=finish_reason, index=0) @@ -1265,7 +1438,7 @@ class LiteLLMResponsesTransformationHandler(CompletionTransformationBridge): raw_response.incomplete_details.reason if raw_response.incomplete_details is not None else None ) if len(choices) == 0: - choices.append(self._build_empty_incomplete_choice(output_items, incomplete_finish_reason)) + choices.append(self._build_empty_choice(output_items, incomplete_finish_reason)) else: for choice in choices: choice.finish_reason = incomplete_finish_reason @@ -1802,7 +1975,9 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): # New output item added output_item = parsed_chunk.get("item", {}) if output_item.get("type") in ("function_call", "custom_tool_call"): - converted: Final = tool_call_dict_from_output_item(output_item, parsed_chunk.get("output_index", 0)) + converted: Final = tool_call_dict_from_output_item( + output_item, parsed_chunk.get("output_index", 0), include_namespace=True + ) provider_specific_fields: Final = converted.get("provider_specific_fields") function_chunk: Final = ChatCompletionToolCallFunctionChunk( @@ -1815,14 +1990,13 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): tool_call_index = OpenAiResponsesToChatCompletionStreamIterator._sequential_tool_call_index( tool_call_index_map, parsed_chunk.get("output_index", 0) ) - tool_call_chunk: Final = ChatCompletionToolCallChunk( + tool_call_chunk: Final = _ChatToolCallDict( id=converted["id"], index=tool_call_index, type="function", function=function_chunk, + **({"provider_specific_fields": provider_specific_fields} if provider_specific_fields else {}), ) - if provider_specific_fields: - tool_call_chunk.provider_specific_fields = provider_specific_fields return ModelResponseStream( choices=[ @@ -1878,7 +2052,7 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): delta=Delta( tool_calls=( tool_call_dict_from_output_item( - output_item, parsed_chunk.get("output_index", 0) + output_item, parsed_chunk.get("output_index", 0), include_namespace=True ), ) ), @@ -1973,6 +2147,9 @@ class OpenAiResponsesToChatCompletionStreamIterator(BaseModelResponseIterator): delta=Delta( content="", reasoning_items=terminal_reasoning_items_typed, + provider_specific_fields=_tool_search_provider_fields( + _tool_search_items_from_output_items(cast(object, output_items)) # cast-ok: untyped + ), ), finish_reason=finish_reason, ) diff --git a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py index ad33c47fc51..63bbd6dc904 100644 --- a/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py +++ b/litellm/litellm_core_utils/llm_response_utils/convert_dict_to_response.py @@ -18,6 +18,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( ) from litellm.types.llms.databricks import DatabricksTool from litellm.types.llms.openai import ( + ChatCompletionReasoningItem, ChatCompletionThinkingBlock, ImageURLListItem, OpenAIModerationResponse, @@ -156,6 +157,21 @@ def _clear_later_replay_slice_metadata(choice: StreamingChoices) -> None: del choice.enhancements +def _first_replay_delta( + message: Mapping[str, object], reasoning_items: Sequence[ChatCompletionReasoningItem] | None +) -> Delta: + return Delta( + content=message.get("content"), + role=message["role"], + function_call=message.get("function_call"), + tool_calls=message.get("tool_calls"), + reasoning_items=None if reasoning_items is None else list(reasoning_items), + provider_specific_fields=_cached_provider_specific_fields( + _CACHED_MESSAGE_PROVIDER_FIELDS.validate_python(message) + ), + ) + + def _invalid_choices_message(response_object: Mapping[str, object]) -> str: raw_keys: Final = list(response_object.keys()) if "choices" not in response_object: @@ -217,15 +233,7 @@ async def convert_to_streaming_response_async( t["index"] = index pydantic_tool_calls.append(ChatCompletionDeltaToolCall(**t)) choice["message"]["tool_calls"] = pydantic_tool_calls - delta = Delta( - content=choice["message"].get("content", None), - role=choice["message"]["role"], - function_call=choice["message"].get("function_call", None), - tool_calls=choice["message"].get("tool_calls", None), - provider_specific_fields=_cached_provider_specific_fields( - _CACHED_MESSAGE_PROVIDER_FIELDS.validate_python(choice["message"]) - ), - ) + delta = _first_replay_delta(choice["message"], choice["message"].get("reasoning_items")) finish_reason = choice.get("finish_reason", None) if finish_reason is None: @@ -754,6 +762,7 @@ def convert_to_model_response_object( provider_specific_fields=provider_specific_fields, reasoning_content=reasoning_content, thinking_blocks=thinking_blocks, + reasoning_items=choice["message"].get("reasoning_items", None), annotations=choice["message"].get("annotations", None), images=_normalize_images_for_message(choice["message"].get("images", None)), ) diff --git a/litellm/litellm_core_utils/prompt_templates/common_utils.py b/litellm/litellm_core_utils/prompt_templates/common_utils.py index ce22bb776ab..710187ac663 100644 --- a/litellm/litellm_core_utils/prompt_templates/common_utils.py +++ b/litellm/litellm_core_utils/prompt_templates/common_utils.py @@ -1807,18 +1807,62 @@ def cache_control_as_prompt_cache_breakpoint(block: _MarkedT) -> _MarkedT: ) -LITELLM_INTERNAL_MESSAGE_FIELDS: Final = frozenset({"thinking_blocks", "reasoning_content", "provider_specific_fields"}) +LITELLM_RESPONSE_ONLY_MESSAGE_FIELDS: Final = frozenset( + {"thinking_blocks", "reasoning_items", "provider_specific_fields"} +) +LITELLM_INTERNAL_MESSAGE_FIELDS: Final = LITELLM_RESPONSE_ONLY_MESSAGE_FIELDS | {"reasoning_content"} + + +def _tool_call_without_provider_specific_fields(tool_call: object) -> object: + if not isinstance(tool_call, dict): + return tool_call + fields: Final = cast(dict[str, object], tool_call) # cast-ok: untyped JSON tool call, keys forwarded verbatim + function_fields: Final = ( + cast("dict[str, object]", fields.get("function")) # cast-ok: untyped JSON function chunk + if isinstance(fields.get("function"), dict) + else None + ) + if "provider_specific_fields" not in fields and ( + function_fields is None or "provider_specific_fields" not in function_fields + ): + return cast("dict[str, object]", tool_call) # cast-ok: untyped JSON tool call, no strip needed + return { + key: ( + {fkey: fvalue for fkey, fvalue in function_fields.items() if fkey != "provider_specific_fields"} + if key == "function" and function_fields is not None + else value + ) + for key, value in fields.items() + if key != "provider_specific_fields" + } + + +def _message_without_fields(message: AllMessageValues, dropped: frozenset[str]) -> AllMessageValues: + kept: Final = {key: value for key, value in message.items() if key not in dropped} + tool_calls: Final = message.get("tool_calls") + if not isinstance(tool_calls, list): + return cast(AllMessageValues, kept) # cast-ok: same TypedDict minus dropped keys + return cast( # cast-ok: same TypedDict minus dropped keys, tool calls minus their provider_specific_fields + AllMessageValues, + { + **kept, + "tool_calls": [ + _tool_call_without_provider_specific_fields(tool_call) + for tool_call in cast(list[object], tool_calls) # cast-ok: untyped JSON tool calls + ], + }, + ) + + +def strip_litellm_response_only_fields(message: AllMessageValues) -> AllMessageValues: + """Drop the fields litellm attaches to a returned message (e.g. Anthropic thinking blocks or Responses + API items) that no OpenAI-compatible endpoint accepts back, so a reply can be sent to a strict one.""" + return _message_without_fields(message, LITELLM_RESPONSE_ONLY_MESSAGE_FIELDS) def strip_litellm_internal_message_fields(message: AllMessageValues) -> AllMessageValues: - """Drop the fields litellm attaches to assistant messages (e.g. when translating Anthropic thinking - blocks) that OpenAI-compatible endpoints with strict schemas reject as extra inputs.""" - if LITELLM_INTERNAL_MESSAGE_FIELDS.isdisjoint(message): - return message - return cast( # cast-ok: same TypedDict minus internal keys - AllMessageValues, - {key: value for key, value in message.items() if key not in LITELLM_INTERNAL_MESSAGE_FIELDS}, - ) + """``strip_litellm_response_only_fields`` plus ``reasoning_content``, for endpoints that reject it too.""" + return _message_without_fields(message, LITELLM_INTERNAL_MESSAGE_FIELDS) def filter_value_from_dict(dictionary: dict, key: str, depth: int = 0) -> Any: diff --git a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py index 9dde1432fae..6c18128a608 100644 --- a/litellm/litellm_core_utils/streaming_chunk_builder_utils.py +++ b/litellm/litellm_core_utils/streaming_chunk_builder_utils.py @@ -2,7 +2,7 @@ import base64 import time from collections.abc import Callable, Iterator, Mapping, Sequence from dataclasses import dataclass, replace -from itertools import groupby +from itertools import chain, groupby from types import MappingProxyType from typing import TYPE_CHECKING, Final, Optional, TypeAlias, TypedDict, Union, cast @@ -1330,10 +1330,29 @@ def _stamp_streaming_usage_cost(usage: Usage, response: ModelResponse, logging_o _set_hidden_response_cost(response, float(computed_cost)) +def _field_of(value: object, name: str) -> object: + if isinstance(value, dict): + return cast(Mapping[str, object], value).get(name) # cast-ok: chunk dicts are untyped JSON + return getattr(value, name, None) + + +def _delta_reasoning_items(delta: object) -> Sequence[object]: + items: Final = _field_of(delta, "reasoning_items") + return cast(Sequence[object], items) if isinstance(items, list) else () # cast-ok: items forwarded verbatim + + +def _chunk_reasoning_items(chunk: object) -> Sequence[object]: + choices: Final = _field_of(chunk, "choices") + if not isinstance(choices, list) or not choices: + return () + return _delta_reasoning_items(_field_of(cast(Sequence[object], choices)[0], "delta")) # cast-ok: untyped chunk + + _NON_TEXT_DELTA_FIELDS: Final = ( "tool_calls", "function_call", "reasoning_content", + "reasoning_items", "thinking_blocks", "annotations", "audio", @@ -1510,6 +1529,11 @@ def stream_chunk_builder( reasoning_chunks ) + typed_chunks: Final = cast(Sequence[object], chunks) # cast-ok: chunks are untyped upstream + reasoning_items: Final = tuple(chain.from_iterable(_chunk_reasoning_items(chunk) for chunk in typed_chunks)) + if reasoning_items: + response["choices"][0]["message"]["reasoning_items"] = list(reasoning_items) + annotation_chunks: Final = [ chunk for chunk in chunks diff --git a/litellm/llms/fireworks_ai/chat/transformation.py b/litellm/llms/fireworks_ai/chat/transformation.py index bd7c0d11b7d..d0551c22441 100644 --- a/litellm/llms/fireworks_ai/chat/transformation.py +++ b/litellm/llms/fireworks_ai/chat/transformation.py @@ -13,7 +13,10 @@ from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLogging from litellm.litellm_core_utils.llm_response_utils.get_headers import ( get_response_headers, ) -from litellm.litellm_core_utils.prompt_templates.common_utils import unpack_legacy_defs +from litellm.litellm_core_utils.prompt_templates.common_utils import ( + strip_litellm_response_only_fields, + unpack_legacy_defs, +) from litellm.secret_managers.main import get_secret_str from litellm.types.llms.openai import ( AllMessageValues, @@ -520,15 +523,8 @@ class FireworksAIConfig(FireworksAIMixin, OpenAIGPTConfig): llm_provider="fireworks_ai", ) filter_value_from_dict(cast(dict, message), "cache_control") - # Remove fields not permitted by FireworksAI (additionalProperties: false - # on their ChatMessage schema) that may cause: - # "Extra inputs are not permitted, field: 'messages[n].'" - if isinstance(message, dict): - m = cast(dict, message) - m.pop("provider_specific_fields", None) - m.pop("thinking_blocks", None) - return messages + return [strip_litellm_response_only_fields(message) for message in messages] # Cached index of fireworks_ai/* entries from litellm.model_cost. Building # this index requires a full scan of model_cost (tens of thousands of diff --git a/litellm/llms/groq/chat/transformation.py b/litellm/llms/groq/chat/transformation.py index ebe3319b149..6980a25e2eb 100644 --- a/litellm/llms/groq/chat/transformation.py +++ b/litellm/llms/groq/chat/transformation.py @@ -11,6 +11,7 @@ from pydantic import BaseModel, TypeAdapter, ValidationError import litellm from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj +from litellm.litellm_core_utils.prompt_templates.common_utils import strip_litellm_response_only_fields from litellm.llms.openai.chat.gpt_transformation import ( OpenAIChatCompletionStreamingHandler, ) @@ -152,7 +153,7 @@ class GroqChatConfig(OpenAILikeChatConfig): for k, v in _message.items(): if v is not None: new_message[k] = v - messages[idx] = new_message + messages[idx] = strip_litellm_response_only_fields(new_message) if is_async: return super()._transform_messages(messages=messages, model=model, is_async=True) diff --git a/litellm/llms/mistral/chat/transformation.py b/litellm/llms/mistral/chat/transformation.py index 0d45482c101..a40970ca429 100644 --- a/litellm/llms/mistral/chat/transformation.py +++ b/litellm/llms/mistral/chat/transformation.py @@ -15,6 +15,7 @@ from litellm._logging import verbose_logger from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj from litellm.litellm_core_utils.prompt_templates.common_utils import ( handle_messages_with_content_list_to_str_conversion, + strip_litellm_internal_message_fields, strip_none_values_from_message, ) from litellm.llms.openai.chat.gpt_transformation import ( @@ -424,17 +425,14 @@ class MistralConfig(OpenAIGPTConfig): @classmethod def _strip_output_only_fields(cls, message: AllMessageValues) -> AllMessageValues: """ - ``reasoning_content`` and ``thinking_blocks`` are output-only fields that - LiteLLM attaches to assistant responses. Mistral's input schema forbids + LiteLLM attaches output-only fields (``reasoning_content``, ``thinking_blocks``, + ``provider_specific_fields``) to assistant responses. Mistral's input schema forbids unknown fields, so replaying them verbatim in a follow-up turn triggers a 422 ``extra_forbidden``. Drop them before the request is sent. """ if message["role"] != "assistant": return message - return cast( - AllMessageValues, - {k: v for k, v in message.items() if k not in ("reasoning_content", "thinking_blocks")}, - ) + return strip_litellm_internal_message_fields(message) @classmethod def _handle_name_in_message(cls, message: AllMessageValues) -> AllMessageValues: diff --git a/litellm/llms/openai/chat/guardrail_translation/handler.py b/litellm/llms/openai/chat/guardrail_translation/handler.py index a12f5f164e4..bf8b42efe4c 100644 --- a/litellm/llms/openai/chat/guardrail_translation/handler.py +++ b/litellm/llms/openai/chat/guardrail_translation/handler.py @@ -25,6 +25,9 @@ from typing_extensions import NotRequired, ReadOnly, TypedDict import litellm from litellm._logging import verbose_proxy_logger +from litellm.completion_extras.litellm_responses_transformation.transformation import ( + replayed_tool_search_tool_names, +) from litellm.litellm_core_utils.streaming_chunk_builder_utils import stream_chunk_builder from litellm.llms.base_llm.guardrail_translation.base_translation import ( BaseTranslation, @@ -256,7 +259,8 @@ class OpenAIChatCompletionsHandler(BaseTranslation): return None if images else "no scannable content" def extract_request_tool_names(self, data: dict) -> list[str]: - """Extract tool names from OpenAI chat completions request (tools[].function.name, functions[].name).""" + """Extract tool names from OpenAI chat completions request (tools[].function.name, functions[].name, + and the tools a replayed tool_search_output loads).""" names: Final[list[str]] = [] for tool in data.get("tools") or []: if isinstance(tool, dict) and tool.get("type") == "function": @@ -266,7 +270,7 @@ class OpenAIChatCompletionsHandler(BaseTranslation): for fn in data.get("functions") or []: if isinstance(fn, dict) and fn.get("name"): names.append(str(fn["name"])) - return names + return [*names, *replayed_tool_search_tool_names(cast(object, data))] # cast-ok: upcast, read field by field def _extract_inputs( self, diff --git a/litellm/llms/together_ai/chat/transformation.py b/litellm/llms/together_ai/chat/transformation.py index 948bd3c8e14..a0eb0c75d35 100644 --- a/litellm/llms/together_ai/chat/transformation.py +++ b/litellm/llms/together_ai/chat/transformation.py @@ -9,7 +9,6 @@ from types import MappingProxyType from typing import ( Final, Literal, - cast, # noqa: TID251 # rebuilding a TypedDict minus keys has no checked spelling overload, ) @@ -18,6 +17,7 @@ from typing_extensions import ReadOnly, TypedDict import litellm from litellm._logging import verbose_logger from litellm.exceptions import UnsupportedParamsError +from litellm.litellm_core_utils.prompt_templates.common_utils import strip_litellm_response_only_fields from litellm.router_utils.reasoning_effort_capability import declared_reasoning_efforts_for_model from litellm.types.llms.openai import AllMessageValues from litellm.utils import supports_function_calling, supports_reasoning, supports_response_schema @@ -25,7 +25,6 @@ from litellm.utils import supports_function_calling, supports_reasoning, support from ...openai.chat.gpt_transformation import OpenAIGPTConfig TOOL_CALLING_PARAMS: Final = ("tools", "tool_choice", "function_call") -LITELLM_INTERNAL_ASSISTANT_FIELDS: Final = frozenset({"thinking_blocks", "provider_specific_fields"}) FUNCTION_CALLING_DOCS_URL: Final = "https://docs.together.ai/docs/function-calling" STRUCTURED_OUTPUTS_DOCS_URL: Final = "https://docs.together.ai/docs/inference/chat/structured-outputs" @@ -174,12 +173,9 @@ def _drop_response_format(passed_params: Container[str], model: str, drop_params def _without_litellm_internal_fields(message: AllMessageValues) -> AllMessageValues: - if message["role"] != "assistant" or LITELLM_INTERNAL_ASSISTANT_FIELDS.isdisjoint(message): + if message["role"] != "assistant": return message - return cast( # cast-ok: rebuilding the same TypedDict minus internal keys loses the narrowed type - "AllMessageValues", - {key: value for key, value in message.items() if key not in LITELLM_INTERNAL_ASSISTANT_FIELDS}, - ) + return strip_litellm_response_only_fields(message) class TogetherAIChatConfig(OpenAIGPTConfig): diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index 226718fc406..94e48762101 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -4,11 +4,16 @@ from collections.abc import AsyncGenerator, AsyncIterable, Mapping, Sequence from typing import Any, Final, Literal, TypedDict from fastapi import HTTPException +from pydantic import TypeAdapter from typing_extensions import ReadOnly, Required from litellm import ChatCompletionToolParam from litellm._logging import verbose_proxy_logger from litellm.caching.dual_cache import DualCache +from litellm.completion_extras.litellm_responses_transformation.transformation import ( + replayed_tool_search_tools, + without_tool_search_items_loading, +) from litellm.exceptions import GuardrailRaisedException from litellm.integrations.custom_guardrail import ( CustomGuardrail, @@ -40,6 +45,7 @@ from litellm.types.utils import ( ) GUARDRAIL_NAME: Final = "tool_permission" +_REQUEST_DATA: Final = TypeAdapter(object) def _object_mapping(value: object) -> Mapping[str, object] | None: @@ -52,6 +58,12 @@ def _object_list(value: object) -> Sequence[object] | None: return value if isinstance(value, list) else None +def _function_tool_name(tool: object) -> str | None: + fields: Final = _object_mapping(tool) + name: Final = fields.get("name") if fields is not None and fields.get("type") == "function" else None + return name if isinstance(name, str) else None + + class _ToolPermissionRuleFields(TypedDict, total=False): """The config-file shape a :class:`ToolPermissionRule` is built from.""" @@ -653,6 +665,9 @@ class ToolPermissionGuardrail(CustomGuardrail): if forced_tool_name is not None: request_tools.append((forced_tool_name, "function")) + replayed: Final = map(_function_tool_name, replayed_tool_search_tools(_REQUEST_DATA.validate_python(data))) + request_tools.extend((name, "function") for name in replayed if name is not None) + return request_tools def _modify_request_with_permission_errors( @@ -701,6 +716,12 @@ class ToolPermissionGuardrail(CustomGuardrail): if named_function_call in error_tool_names: data["function_call"] = "none" + messages: Final = without_tool_search_items_loading( + _REQUEST_DATA.validate_python(data), frozenset(denied_tool_names) + ) + if messages is not None: + data["messages"] = messages # rebind-ok: rewrite mode edits the request in place + return data def _create_permission_error_result( diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py index f54e4bc30f7..59001c8bc8d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_policy/tool_policy_guardrail.py @@ -23,9 +23,11 @@ Configuration in proxy config YAML: from typing import TYPE_CHECKING, Final, Literal, Optional from fastapi import HTTPException +from pydantic import TypeAdapter from typing_extensions import TypedDict, Unpack from litellm._logging import verbose_proxy_logger +from litellm.completion_extras.litellm_responses_transformation.transformation import replayed_tool_search_tool_names from litellm.integrations.custom_guardrail import ( CustomGuardrail, log_guardrail_information, @@ -38,6 +40,7 @@ if TYPE_CHECKING: from litellm.litellm_core_utils.litellm_logging import Logging as LiteLLMLoggingObj GUARDRAIL_NAME: Final = "tool_policy" +_REQUEST_DATA: Final = TypeAdapter(object) class _CustomGuardrailOptions(TypedDict, total=False, extra_items=object): @@ -133,15 +136,18 @@ class ToolPolicyGuardrail(CustomGuardrail): """ if input_type == "request": tools: Final = inputs.get("tools") or [] - tool_names = [ + function_tool_names: Final = [ t["function"]["name"] for t in tools if isinstance(t, dict) and isinstance(t.get("function"), dict) and t["function"].get("name") ] - if not tool_names: - route: Final = _get_request_route_from_data(request_data) - if route: - tool_names = extract_request_tool_names(route, request_data) + route: Final = _get_request_route_from_data(request_data) + replayed_tool_names: Final = replayed_tool_search_tool_names(_REQUEST_DATA.validate_python(request_data)) + tool_names = ( + extract_request_tool_names(route, request_data) + if route and not function_tool_names + else [*function_tool_names, *replayed_tool_names] + ) else: tool_calls: Final = inputs.get("tool_calls") or [] tool_names = [] diff --git a/litellm/responses/litellm_completion_transformation/transformation.py b/litellm/responses/litellm_completion_transformation/transformation.py index ca738ce22b3..7f5d143e98e 100644 --- a/litellm/responses/litellm_completion_transformation/transformation.py +++ b/litellm/responses/litellm_completion_transformation/transformation.py @@ -2868,13 +2868,18 @@ class LiteLLMCompletionResponsesConfig: else _attribute_fields(provider_fields) ) + namespace: Final = getattr(tool_call_item, "namespace", None) + merged_provider_fields: Final = { + **(provider_specific_fields if isinstance(provider_specific_fields, dict) else {}), + **({"namespace": namespace} if isinstance(namespace, str) else {}), + } function_dict: Final[dict[str, object]] = { "name": item.name, "arguments": item.arguments, } - if provider_specific_fields: - function_dict["provider_specific_fields"] = provider_specific_fields + if merged_provider_fields: + function_dict["provider_specific_fields"] = merged_provider_fields tool_call_dict: Final[dict[str, object]] = { "id": LiteLLMCompletionResponsesConfig.tool_call_id_from_responses_item( @@ -2886,8 +2891,8 @@ class LiteLLMCompletionResponsesConfig: "index": index, } - if provider_specific_fields: - tool_call_dict["provider_specific_fields"] = provider_specific_fields + if merged_provider_fields: + tool_call_dict["provider_specific_fields"] = merged_provider_fields return tool_call_dict diff --git a/litellm/tool_loop.py b/litellm/tool_loop.py index 76339e8d210..aaa0d62eea7 100644 --- a/litellm/tool_loop.py +++ b/litellm/tool_loop.py @@ -46,19 +46,24 @@ def _expect_model_response(response: object) -> ModelResponse: def _assistant_tool_call(tool_call: ChatCompletionMessageToolCall) -> ChatCompletionAssistantToolCall: - return ChatCompletionAssistantToolCall( - id=tool_call.id, - type="function", - function=ChatCompletionToolCallFunctionChunk( - name=tool_call.function.name, arguments=tool_call.function.arguments - ), + provider_fields: Final[object] = getattr(tool_call, "provider_specific_fields", None) + return cast( # cast-ok: dict literal with provider_specific_fields spread in only when set + "ChatCompletionAssistantToolCall", + { + "id": tool_call.id, + "type": "function", + "function": ChatCompletionToolCallFunctionChunk( + name=tool_call.function.name, arguments=tool_call.function.arguments + ), + **({"provider_specific_fields": provider_fields} if provider_fields else {}), + }, ) def _assistant_message( - message: Message, tool_calls: tuple[ChatCompletionMessageToolCall, ...], replays_server_tool_results: bool + message: Message, tool_calls: tuple[ChatCompletionMessageToolCall, ...], replays_provider_specific_fields: bool ) -> ChatCompletionAssistantMessage: - return cast( # cast-ok: dict literal with thinking_blocks, reasoning_items and provider_specific_fields spread in only when set + return cast( # cast-ok: dict literal with the optional message fields spread in only when set "ChatCompletionAssistantMessage", { "role": "assistant", @@ -71,7 +76,9 @@ def _assistant_message( ("reasoning_items", getattr(message, "reasoning_items", None)), ( "provider_specific_fields", - getattr(message, "provider_specific_fields", None) if replays_server_tool_results else None, + getattr(message, "provider_specific_fields", None) + if replays_provider_specific_fields + else None, ), ) if value is not None @@ -105,6 +112,18 @@ def _client_tool_calls( ) +def _replays_provider_specific_fields( + message: Message, + tool_calls: tuple[ChatCompletionMessageToolCall, ...], + client_tool_calls: tuple[ChatCompletionMessageToolCall, ...], +) -> bool: + from litellm.completion_extras.litellm_responses_transformation.transformation import carries_tool_search_items + + return client_tool_calls != tool_calls or carries_tool_search_items( + getattr(message, "provider_specific_fields", None) + ) + + def run_tool_loop( *, model: str, @@ -148,7 +167,9 @@ def run_tool_loop( tool_results = tuple(execute_tool(tool_call) for tool_call in client_tool_calls) history = ( *history, - _assistant_message(message, tool_calls, client_tool_calls != tool_calls), + _assistant_message( + message, tool_calls, _replays_provider_specific_fields(message, tool_calls, client_tool_calls) + ), *tool_results, ) raise ToolLoopMaxRoundsExceeded(max_rounds) @@ -197,7 +218,9 @@ async def arun_tool_loop( tool_results = tuple([await execute_tool(tool_call) for tool_call in client_tool_calls]) history = ( *history, - _assistant_message(message, tool_calls, client_tool_calls != tool_calls), + _assistant_message( + message, tool_calls, _replays_provider_specific_fields(message, tool_calls, client_tool_calls) + ), *tool_results, ) raise ToolLoopMaxRoundsExceeded(max_rounds) diff --git a/tests/integration/_support/tool_search.py b/tests/integration/_support/tool_search.py new file mode 100644 index 00000000000..8230fdc00cd --- /dev/null +++ b/tests/integration/_support/tool_search.py @@ -0,0 +1,340 @@ +from __future__ import annotations + +import json +import uuid +from collections.abc import Callable, Mapping, Sequence +from dataclasses import dataclass +from functools import reduce +from itertools import chain +from typing import Final + +from integration._support import prompt_cache_breakpoint as pcb +from integration._support import responses_vendor as rv +from integration._support.client import string_value +from integration._support.database import read_rows +from integration._support.openai_wire import answering_model_discovery +from integration._support.wire import Reply, Request +from pydantic import JsonValue + + +WEATHER: Final = "get_weather" +TIME: Final = "get_time" +ONLY_SEARCH: Final = "only-search" +TWO_SEARCHES: Final = "two-searches" +REJECT_REPLAY: Final = "reject-replay" +REJECTED: Final = "scripted: the replayed tool_search_output was rejected" +TOOL_SEARCH: Final[dict[str, JsonValue]] = {"type": "tool_search"} +PARAMETERS: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], +} +REPLAY_TYPES: Final = frozenset({"reasoning", "tool_search_call", "tool_search_output", "function_call"}) + + +def function(name: str) -> dict[str, JsonValue]: + return {"name": name, "description": f"Look up {name}", "parameters": PARAMETERS} + + +def deferred(name: str) -> dict[str, JsonValue]: + return {"type": "function", "function": function(name), "defer_loading": True} + + +def chat_tools(*names: str) -> list[JsonValue]: + return [TOOL_SEARCH, *(deferred(name) for name in names)] + + +def loaded(name: str) -> dict[str, JsonValue]: + return {"type": "function", **function(name), "strict": False, "defer_loading": True} + + +def namespace(name: str, *members: str) -> dict[str, JsonValue]: + return {"type": "namespace", "name": name, "description": f"{name} tools", "tools": [loaded(m) for m in members]} + + +def reasoning(marker: str, label: str) -> dict[str, JsonValue]: + return { + "id": f"rs_{label}_{marker}", + "type": "reasoning", + "summary": [], + "encrypted_content": f"enc-{label}-{marker}", + } + + +def search_call(marker: str, label: str = "a") -> dict[str, JsonValue]: + return { + "id": f"tsc_{label}_{marker}", + "type": "tool_search_call", + "arguments": {"query": f"tools {label}"}, + "execution": "server", + "status": "completed", + } + + +def search_output(marker: str, tools: Sequence[JsonValue], label: str = "a") -> dict[str, JsonValue]: + return { + "id": f"tso_{label}_{marker}", + "type": "tool_search_output", + "execution": "server", + "status": "completed", + "tools": list(tools), + } + + +def call_id(marker: str, name: str = WEATHER) -> str: + return f"call_{name}_{marker}" + + +def replayed_call(marker: str, name: str = WEATHER) -> dict[str, JsonValue]: + return { + "type": "function_call", + "call_id": call_id(marker, name), + "name": name, + "arguments": '{"city": "Paris"}', + "namespace": name, + } + + +def call_item(marker: str, name: str = WEATHER) -> dict[str, JsonValue]: + return {**replayed_call(marker, name), "id": f"fc_{name}_{marker}", "status": "completed"} + + +def searched(marker: str, label: str, name: str) -> tuple[dict[str, JsonValue], ...]: + return (reasoning(marker, label), search_call(marker, label), search_output(marker, [loaded(name)], label)) + + +def turn_one_output(marker: str) -> tuple[dict[str, JsonValue], ...]: + return (*searched(marker, "a", WEATHER), reasoning(marker, "b"), call_item(marker)) + + +def turn_two_output(marker: str) -> tuple[dict[str, JsonValue], ...]: + return (*searched(marker, "c", TIME), call_item(marker, TIME)) + + +def only_search_output(marker: str) -> tuple[dict[str, JsonValue], ...]: + return searched(marker, "a", WEATHER) + + +def on_the_wire(item: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return {**item, "created_by": None} if item.get("type") == "tool_search_call" else dict(item) + + +def _message_item(identity: str, text: str) -> dict[str, JsonValue]: + return { + "id": f"msg_{identity}", + "type": "message", + "role": "assistant", + "status": "completed", + "content": [{"type": "output_text", "text": text, "annotations": []}], + } + + +def _response(identity: str, model: str, output: Sequence[JsonValue], status: str) -> dict[str, JsonValue]: + return { + "id": identity, + "object": "response", + "created_at": 1, + "status": status, + "model": model, + "output": list(output), + "usage": rv.USAGE, + } + + +def _item_events(index: int, item: Mapping[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]: + added: Final[dict[str, JsonValue]] = {"type": "response.output_item.added", "output_index": index, "item": {**item}} + done: Final[dict[str, JsonValue]] = {"type": "response.output_item.done", "output_index": index, "item": {**item}} + if item.get("type") == "function_call": + return ( + {**added, "item": {**item, "arguments": "", "status": "in_progress"}}, + { + "type": "response.function_call_arguments.delta", + "output_index": index, + "item_id": item["id"], + "delta": item["arguments"], + }, + done, + ) + if item.get("type") == "message": + contents: Final = rv.ITEMS.validate_python(item["content"]) + return ( + {**added, "item": {**item, "content": [], "status": "in_progress"}}, + { + "type": "response.output_text.delta", + "output_index": index, + "item_id": item["id"], + "content_index": 0, + "delta": contents[0]["text"], + }, + done, + ) + return (added, done) + + +def stream_chunks(identity: str, model: str, output: Sequence[Mapping[str, JsonValue]]) -> tuple[bytes, ...]: + created: Final[dict[str, JsonValue]] = { + "type": "response.created", + "response": _response(identity, model, [], "in_progress"), + } + items: Final = tuple(chain.from_iterable(_item_events(index, item) for index, item in enumerate(output))) + completed: Final[dict[str, JsonValue]] = { + "type": "response.completed", + "response": _response(identity, model, [dict(item) for item in output], "completed"), + } + events: Final = (created, *items, completed) + return tuple(rv.sse({**event, "sequence_number": number}) for number, event in enumerate(events)) + + +def reply( + identity: str, model: str, output: Sequence[Mapping[str, JsonValue]], *, stream: bool, pause: float = 0 +) -> Reply: + if stream: + return Reply( + content_type="text/event-stream", + chunks=stream_chunks(identity, model, output), + pause_between_chunks=pause, + ) + return Reply(body=json.dumps(_response(identity, model, [dict(item) for item in output], "completed")).encode()) + + +def _count(items: Sequence[Mapping[str, JsonValue]], kind: str) -> int: + return sum(1 for item in items if item.get("type") == kind) + + +def _scripted_output(text: str, marker: str, items: Sequence[Mapping[str, JsonValue]]) -> tuple[JsonValue, ...]: + results: Final = _count(items, "function_call_output") + if TWO_SEARCHES in text and results == 1: + return tuple(on_the_wire(item) for item in turn_two_output(marker)) + if results or (ONLY_SEARCH in text and _count(items, "tool_search_output")): + return (_message_item(f"{marker}-{results}", rv.answer(marker)),) + chosen: Final = only_search_output(marker) if ONLY_SEARCH in text else turn_one_output(marker) + return tuple(on_the_wire(item) for item in chosen) + + +def scripted_with_pause(pause: float) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + body: Final = pcb.body_of(request) + text: Final = request.body.decode() + marker: Final = rv.newest_marker(text) or uuid.uuid4().hex + items: Final = pcb.input_items(request) + if REJECT_REPLAY in text and _count(items, "tool_search_output"): + return rv.error(400, REJECTED, "invalid_request_error") + return reply( + f"resp_{marker}-{uuid.uuid4().hex}", + string_value(body["model"]), + [rv.JSON_OBJECT.validate_python(item) for item in _scripted_output(text, marker, items)], + stream=body.get("stream") is True, + pause=pause, + ) + + return answering_model_discovery(respond) + + +scripted: Final = scripted_with_pause(0) + + +def replayed(request: Request) -> list[dict[str, JsonValue]]: + return [item for item in pcb.input_items(request) if item.get("type") in REPLAY_TYPES] + + +def first_message(text: str) -> dict[str, JsonValue]: + choices: Final = rv.ITEMS.validate_python(rv.JSON_OBJECT.validate_json(text)["choices"]) + return rv.JSON_OBJECT.validate_python(choices[0]["message"]) + + +def provider_fields(value: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + fields: Final = value.get("provider_specific_fields") + return rv.JSON_OBJECT.validate_python(fields) if isinstance(fields, dict) else {} + + +def tool_calls(message: Mapping[str, JsonValue]) -> list[dict[str, JsonValue]]: + return rv.ITEMS.validate_python(message.get("tool_calls") or []) + + +def namespaces(message: Mapping[str, JsonValue]) -> list[JsonValue]: + return [provider_fields(call).get("namespace") for call in tool_calls(message)] + + +def chunks_of(text: str) -> tuple[dict[str, JsonValue], ...]: + return tuple( + rv.JSON_OBJECT.validate_json(line.removeprefix("data: ")) + for line in text.splitlines() + if line.startswith("data: {") + ) + + +def _deltas(chunks: Sequence[Mapping[str, JsonValue]]) -> tuple[dict[str, JsonValue], ...]: + choices: Final = tuple(rv.ITEMS.validate_python(chunk.get("choices") or []) for chunk in chunks) + return tuple(rv.JSON_OBJECT.validate_python(found[0].get("delta") or {}) for found in choices if found) + + +def _merged(left: Mapping[str, JsonValue], right: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return {**left, **{key: value for key, value in right.items() if value is not None}} + + +@dataclass(frozen=True, slots=True) +class _CallParts: + index: int + id: str | None + name: str | None + arguments: str + fields: Mapping[str, JsonValue] + + +def _call_parts(delta: Mapping[str, JsonValue]) -> tuple[_CallParts, ...]: + calls: Final = rv.ITEMS.validate_python(delta.get("tool_calls") or []) + functions: Final = tuple(rv.JSON_OBJECT.validate_python(call.get("function") or {}) for call in calls) + return tuple( + _CallParts( + index=int(str(call.get("index") or 0)), + id=call["id"] if isinstance(call.get("id"), str) else None, + name=fn["name"] if isinstance(fn.get("name"), str) else None, + arguments=str(fn.get("arguments") or ""), + fields=_merged(provider_fields(fn), provider_fields(call)), + ) + for call, fn in zip(calls, functions, strict=True) + ) + + +def _rebuilt_call(parts: Sequence[_CallParts]) -> dict[str, JsonValue]: + fields: Final = reduce(_merged, (part.fields for part in parts), {}) + ids: Final = tuple(part.id for part in parts if part.id) + names: Final = tuple(part.name for part in parts if part.name) + return { + "id": ids[0], + "type": "function", + "function": {"name": names[0], "arguments": "".join(part.arguments for part in parts)}, + **({"provider_specific_fields": fields} if fields else {}), + } + + +def assistant_from_stream(chunks: Sequence[Mapping[str, JsonValue]]) -> dict[str, JsonValue]: + deltas: Final = _deltas(chunks) + parts: Final = tuple(chain.from_iterable(_call_parts(delta) for delta in deltas)) + indexes: Final = sorted({part.index for part in parts}) + calls: Final[list[JsonValue]] = [ + _rebuilt_call([part for part in parts if part.index == index]) for index in indexes + ] + fields: Final = reduce(_merged, (provider_fields(delta) for delta in deltas), {}) + reasoning_items: Final[list[JsonValue]] = list( + chain.from_iterable(rv.ITEMS.validate_python(delta.get("reasoning_items") or []) for delta in deltas) + ) + content: Final = "".join(str(delta.get("content") or "") for delta in deltas) + return { + "role": "assistant", + "content": content or None, + **({"tool_calls": calls} if calls else {}), + **({"provider_specific_fields": fields} if fields else {}), + **({"reasoning_items": reasoning_items} if reasoning_items else {}), + } + + +def tool_result(marker: str, name: str = WEATHER) -> dict[str, JsonValue]: + return {"role": "tool", "tool_call_id": call_id(marker, name), "content": "sunny"} + + +def spend_rows_of(model: str, response_id: str) -> list[dict[str, JsonValue]]: + rows: Final = read_rows( + 'SELECT request_id, status, cache_hit FROM "LiteLLM_SpendLogs" WHERE model_group = %s', (model,) + ) + return [row for row in rows if rv.same_response(string_value(row["request_id"]), response_id)] diff --git a/tests/integration/authorization/test_tool_search_replay_allowlist.py b/tests/integration/authorization/test_tool_search_replay_allowlist.py new file mode 100644 index 00000000000..afc8047a5e8 --- /dev/null +++ b/tests/integration/authorization/test_tool_search_replay_allowlist.py @@ -0,0 +1,166 @@ +import uuid +from collections.abc import Sequence +from typing import Final + +import httpx +import pytest +from integration._support import prompt_cache_breakpoint as pcb +from integration._support import responses_vendor as rv +from integration._support import tool_search as ts +from integration._support.client import Gateway, Scenario +from integration._support.openai_wire import answering_model_discovery, chat_reply +from integration._support.wire import Reply, Request, Wire, wire_server +from pydantic import JsonValue + +from tests.integration.caching.response_cache_case import await_entry, redis_store + +pytestmark: Final = pytest.mark.timeout(120) + + +def _shell() -> str: + return f"run_shell_{uuid.uuid4().hex[:12]}" + + +def _replay_body(model: str, marker: str, loads: Sequence[JsonValue], **fields: JsonValue) -> dict[str, JsonValue]: + call: Final[dict[str, JsonValue]] = { + "id": ts.call_id(marker), + "type": "function", + "function": {"name": ts.WEATHER, "arguments": '{"city": "Paris"}'}, + "provider_specific_fields": {"namespace": ts.WEATHER}, + } + items: Final[list[JsonValue]] = [ts.search_call(marker), ts.search_output(marker, loads)] + return { + "model": model, + "messages": [ + {"role": "user", "content": pcb.prompt(marker)}, + { + "role": "assistant", + "content": None, + "tool_calls": [call], + "provider_specific_fields": {"tool_search_items": items}, + }, + ts.tool_result(marker), + ], + "tools": ts.chat_tools(ts.WEATHER), + **pcb.NO_CACHE, + **fields, + } + + +def _bridge(scenario: Scenario, wire: Wire, **parameters: JsonValue) -> str: + return scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1", **parameters) + + +def _posts(wire: Wire, marker: str) -> int: + return len(pcb.with_marker(pcb.drained_posts(wire), marker)) + + +def _assert_denied(response: httpx.Response, wire: Wire, marker: str, tool: str) -> None: + assert response.status_code == 403, response.text + error: Final = rv.JSON_OBJECT.validate_python(rv.JSON_OBJECT.validate_json(response.text)["error"]) + assert error["type"] == "tool_access_denied", response.text + assert tool in str(error["message"]) and ts.WEATHER not in str(error["message"]), response.text + assert _posts(wire, marker) == 0, "a denied replay reached the upstream" + + +def _assert_served(response: httpx.Response, wire: Wire, marker: str) -> None: + assert response.status_code == 200, response.text + assert ts.first_message(response.text)["content"] == rv.answer(marker), response.text + assert _posts(wire, marker) == 1 + + +def test_key_without_an_allowlist_still_replays_any_loaded_tool(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(ts.scripted) as wire, gateway.scenario() as scenario: + model: Final = _bridge(scenario, wire) + key: Final = scenario.key(models=[model]) + body: Final = _replay_body(model, marker, [ts.loaded(ts.WEATHER), ts.loaded(_shell())]) + _assert_served(gateway.request("POST", "/v1/chat/completions", body, key=key), wire, marker) + + +@pytest.mark.parametrize("placement", ["tools", "extra_body", "deployment", "absent"]) +def test_allowlisted_key_cannot_replay_a_search_that_loaded_another_tool(gateway: Gateway, placement: str) -> None: + marker: Final = uuid.uuid4().hex + shell: Final = _shell() + with wire_server(ts.scripted) as wire, gateway.scenario() as scenario: + model: Final = _bridge(scenario, wire, **({"tools": [ts.TOOL_SEARCH]} if placement == "deployment" else {})) + key: Final = scenario.key(models=[model], metadata={"allowed_tools": [ts.WEATHER]}) + loads: Final[list[JsonValue]] = [ts.loaded(ts.WEATHER), ts.loaded(shell)] + overrides: Final[dict[str, dict[str, JsonValue]]] = { + "tools": {}, + "extra_body": { + "tools": [ts.deferred(ts.WEATHER)], + "extra_body": {"tools": [ts.TOOL_SEARCH, ts.loaded(ts.WEATHER)]}, + }, + "deployment": {"tools": [ts.deferred(ts.WEATHER)]}, + "absent": {"tools": [{"type": "function", "function": ts.function(ts.WEATHER)}]}, + } + body: Final = {**_replay_body(model, marker, loads), **overrides[placement]} + _assert_denied(gateway.request("POST", "/v1/chat/completions", body, key=key), wire, marker, shell) + + +def test_allowlisted_key_cannot_replay_a_denied_namespace_member(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + shell: Final = _shell() + with wire_server(ts.scripted) as wire, gateway.scenario() as scenario: + model: Final = _bridge(scenario, wire) + key: Final = scenario.key(models=[model], metadata={"allowed_tools": [ts.WEATHER]}) + body: Final = _replay_body(model, marker, [ts.namespace("ops", ts.WEATHER, shell)]) + _assert_denied(gateway.request("POST", "/v1/chat/completions", body, key=key), wire, marker, shell) + + +def _chat_upstream(request: Request) -> Reply: + marker: Final = rv.newest_marker(request.body.decode()) + return chat_reply(f"chatcmpl-{marker}", "gpt-4o-mini", rv.answer(marker), stream=False) + + +def test_plain_chat_deployment_also_counts_replayed_search_loads(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + shell: Final = _shell() + with wire_server(answering_model_discovery(_chat_upstream)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(api_base=f"{wire.url}/v1") + key: Final = scenario.key(models=[model], metadata={"allowed_tools": [ts.WEATHER]}) + body: Final = _replay_body(model, marker, [ts.loaded(shell)]) + _assert_denied(gateway.request("POST", "/v1/chat/completions", body, key=key), wire, marker, shell) + + +def test_team_allowlist_applies_and_a_key_allowlist_overrides_it(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + shell: Final = _shell() + with wire_server(ts.scripted) as wire, gateway.scenario() as scenario: + model: Final = _bridge(scenario, wire) + team: Final = scenario.team(models=[model], metadata={"allowed_tools": [ts.WEATHER]}) + team_key: Final = scenario.key(team_id=team, models=[model]) + wider_key: Final = scenario.key(team_id=team, models=[model], metadata={"allowed_tools": [ts.WEATHER, shell]}) + body: Final = _replay_body(model, marker, [ts.loaded(ts.WEATHER), ts.loaded(shell)]) + _assert_denied(gateway.request("POST", "/v1/chat/completions", body, key=team_key), wire, marker, shell) + _assert_served(gateway.request("POST", "/v1/chat/completions", body, key=wider_key), wire, marker) + + +def test_allowed_replay_is_served_and_its_cache_hit_still_checks_the_allowlist(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + shell: Final = _shell() + with wire_server(ts.scripted) as wire, gateway.scenario() as scenario: + model: Final = _bridge(scenario, wire) + allowed: Final = scenario.key(models=[model], metadata={"allowed_tools": [ts.WEATHER]}) + cached: Final = { + key: value + for key, value in _replay_body(model, marker, [ts.loaded(ts.WEATHER)]).items() + if key != "cache" + } + first: Final = gateway.request("POST", "/v1/chat/completions", cached, key=allowed) + _assert_served(first, wire, marker) + await_entry(redis_store(), str(rv.JSON_OBJECT.validate_json(first.text)["id"])) + hit: Final = gateway.request("POST", "/v1/chat/completions", cached, key=allowed) + assert hit.status_code == 200 and hit.headers.get("x-litellm-cache-key"), (hit.text, dict(hit.headers)) + assert ts.first_message(hit.text)["content"] == rv.answer(marker), hit.text + shell_body: Final = { + key: value + for key, value in _replay_body(model, marker, [ts.loaded(shell)]).items() + if key != "cache" + } + populated: Final = gateway.request("POST", "/v1/chat/completions", shell_body) + assert populated.status_code == 200, populated.text + await_entry(redis_store(), str(rv.JSON_OBJECT.validate_json(populated.text)["id"])) + assert _posts(wire, marker) == 1 + _assert_denied(gateway.request("POST", "/v1/chat/completions", shell_body, key=allowed), wire, marker, shell) diff --git a/tests/integration/authorization/test_tool_search_replay_guardrails.py b/tests/integration/authorization/test_tool_search_replay_guardrails.py new file mode 100644 index 00000000000..14f647ecf94 --- /dev/null +++ b/tests/integration/authorization/test_tool_search_replay_guardrails.py @@ -0,0 +1,276 @@ +import json +import uuid +from collections.abc import Iterator +from dataclasses import dataclass +from pathlib import Path +from typing import Final + +import httpx +import pytest +from integration._support import prompt_cache_breakpoint as pcb +from integration._support import responses_vendor as rv +from integration._support import tool_search as ts +from integration._support.client import Gateway +from integration._support.database import write_rows +from integration._support.process import graceful_stop_seconds, owned_proxy +from integration._support.wire import Wire, wire_server +from pydantic import JsonValue + +from tests.integration.caching.response_cache_case import await_entry, redis_store + +pytestmark: Final = pytest.mark.timeout(2 * graceful_stop_seconds() + 120) + +_POLICY: Final = "integration-tool-policy" +_BLOCK: Final = "tool-permission-block" +_REWRITE: Final = "tool-permission-rewrite" + + +@dataclass(frozen=True, slots=True) +class _Guarded: + proxy: Gateway + wire: Wire + model: str + shell: str + blocked: str + + +def _permission(name: str, shell: str, action: str) -> dict[str, JsonValue]: + return { + "guardrail_name": name, + "litellm_params": { + "guardrail": "tool_permission", + "mode": "pre_call", + "default_on": False, + "default_action": "allow", + "on_disallowed_action": action, + "rules": [{"id": "deny_shell", "tool_name": f"^{shell}$", "decision": "deny"}], + }, + } + + +def _config(directory: Path, model: str, upstream: str, shell: str) -> Path: + path: Final = directory / "tool_search_guardrails.yaml" + path.write_text( + json.dumps( + { + "model_list": [ + { + "model_name": model, + "litellm_params": {"model": pcb.MODEL, "api_base": f"{upstream}/v1", "api_key": "sk-upstream"}, + } + ], + "general_settings": { + "master_key": "os.environ/LITELLM_MASTER_KEY", + "database_url": "os.environ/DATABASE_URL", + "store_model_in_db": True, + }, + "router_settings": {"disable_cooldowns": True}, + "litellm_settings": { + "cache": True, + "cache_params": {"type": "redis", "host": "os.environ/REDIS_HOST", "port": "os.environ/REDIS_PORT"}, + }, + "guardrails": [ + { + "guardrail_name": _POLICY, + "litellm_params": {"guardrail": "tool_policy", "mode": "pre_call", "default_on": True}, + }, + _permission(_BLOCK, shell, "block"), + _permission(_REWRITE, shell, "rewrite"), + ], + } + ) + ) + return path + + +@pytest.fixture(scope="module") +def guarded(session_gateway: Gateway, tmp_path_factory: pytest.TempPathFactory) -> Iterator[_Guarded]: + suffix: Final = uuid.uuid4().hex[:12] + shell: Final = f"run_shell_{suffix}" + blocked: Final = f"wipe_disk_{suffix}" + model: Final = f"tool-search-guarded-{suffix}" + policy: Final = session_gateway.request( + "POST", "/v1/tool/policy", {"tool_name": blocked, "input_policy": "blocked"} + ) + assert policy.status_code == 200, policy.text + directory: Final = tmp_path_factory.mktemp("tool_search_guardrails") + try: + with ( + wire_server(ts.scripted) as wire, + owned_proxy(session_gateway, directory, {}, config=_config(directory, model, wire.url, shell), workers=2) + as proxy, + ): + yield _Guarded(proxy, wire, model, shell, blocked) + finally: + write_rows('DELETE FROM "LiteLLM_ToolTable" WHERE tool_name = %s', (blocked,)) + + +def _call(marker: str, name: str) -> dict[str, JsonValue]: + return { + "id": ts.call_id(marker, name), + "type": "function", + "function": {"name": name, "arguments": '{"city": "Paris"}'}, + "provider_specific_fields": {"namespace": name}, + } + + +def _searched_turn(marker: str, name: str, label: str, loads: list[JsonValue]) -> list[JsonValue]: + items: Final[list[JsonValue]] = [ts.search_call(marker, label), ts.search_output(marker, loads, label)] + return [ + { + "role": "assistant", + "content": None, + "tool_calls": [_call(marker, name)], + "provider_specific_fields": {"tool_search_items": items}, + }, + ts.tool_result(marker, name), + ] + + +def _two_turn_body(guarded: _Guarded, marker: str, **fields: JsonValue) -> dict[str, JsonValue]: + return { + "model": guarded.model, + "messages": [ + {"role": "user", "content": pcb.prompt(marker)}, + *_searched_turn(marker, ts.WEATHER, "a", [ts.loaded(ts.WEATHER)]), + *_searched_turn(marker, ts.TIME, "b", [ts.loaded(ts.TIME), ts.loaded(guarded.shell)]), + ], + "tools": ts.chat_tools(ts.WEATHER, ts.TIME), + **fields, + } + + +def _posts(wire: Wire, marker: str) -> tuple[bytes, ...]: + return tuple(request.body for request in pcb.with_marker(pcb.drained_posts(wire), marker)) + + +def _assert_refused(response: httpx.Response, wire: Wire, marker: str, message: str, tool: str) -> None: + assert response.status_code == 400, response.text + assert message in response.text and tool in response.text, response.text + assert _posts(wire, marker) == (), "a refused replay reached the upstream" + + +def _one_search_body( + guarded: _Guarded, marker: str, loads: list[JsonValue], **fields: JsonValue +) -> dict[str, JsonValue]: + return { + "model": guarded.model, + "messages": [{"role": "user", "content": pcb.prompt(marker)}, *_searched_turn(marker, ts.WEATHER, "a", loads)], + "tools": ts.chat_tools(ts.WEATHER), + **fields, + } + + +def _blocked_load_body(guarded: _Guarded, marker: str) -> dict[str, JsonValue]: + return _one_search_body(guarded, marker, [ts.loaded(ts.WEATHER), ts.loaded(guarded.blocked)]) + + +def _cache(guarded: _Guarded, marker: str, body: dict[str, JsonValue], key: str | None = None) -> None: + populated: Final = guarded.proxy.request("POST", "/v1/chat/completions", body, key=key) + assert populated.status_code == 200, populated.text + await_entry(redis_store(), str(rv.JSON_OBJECT.validate_json(populated.text)["id"])) + hit: Final = guarded.proxy.request("POST", "/v1/chat/completions", body, key=key) + assert hit.status_code == 200 and hit.headers.get("x-litellm-cache-key"), (hit.text, dict(hit.headers)) + assert ts.first_message(hit.text)["content"] == rv.answer(marker), hit.text + assert len(_posts(guarded.wire, marker)) == 1 + + +def _assert_policy_refuses_the_replay(guarded: _Guarded) -> None: + marker: Final = uuid.uuid4().hex + refused: Final = guarded.proxy.request("POST", "/v1/chat/completions", _blocked_load_body(guarded, marker)) + _assert_refused(refused, guarded.wire, marker, "Violated tool policy", guarded.blocked) + detail: Final = rv.JSON_OBJECT.validate_python(rv.JSON_OBJECT.validate_json(refused.text)["error"]) + assert rv.JSON_OBJECT.validate_python(detail["provider_specific_fields"])["blocked_tools"] == [guarded.blocked] + + +def test_tool_policy_blocks_a_replayed_search_that_loaded_a_blocked_tool(guarded: _Guarded) -> None: + control_marker: Final = uuid.uuid4().hex + control: Final = guarded.proxy.request( + "POST", + "/v1/chat/completions", + { + "model": guarded.model, + "messages": [{"role": "user", "content": pcb.prompt(control_marker)}], + "tools": [{"type": "function", "function": ts.function(guarded.blocked)}], + }, + ) + _assert_refused(control, guarded.wire, control_marker, "Violated tool policy", guarded.blocked) + _assert_policy_refuses_the_replay(guarded) + _assert_policy_refuses_the_replay(guarded) + + +def test_tool_permission_block_mode_refuses_a_replayed_search_that_loaded_a_denied_tool(guarded: _Guarded) -> None: + marker: Final = uuid.uuid4().hex + response: Final = guarded.proxy.request( + "POST", "/v1/chat/completions", _two_turn_body(guarded, marker, guardrails=[_BLOCK]) + ) + _assert_refused(response, guarded.wire, marker, "Violated guardrail policy", guarded.shell) + + +def test_tool_permission_rewrite_mode_drops_only_the_search_that_loaded_the_denied_tool(guarded: _Guarded) -> None: + marker: Final = uuid.uuid4().hex + response: Final = guarded.proxy.request( + "POST", "/v1/chat/completions", _two_turn_body(guarded, marker, guardrails=[_REWRITE]) + ) + assert response.status_code == 200, response.text + assert ts.first_message(response.text)["content"] == rv.answer(marker), response.text + (posted,) = pcb.with_marker(pcb.drained_posts(guarded.wire), marker) + assert ts.replayed(posted) == [ + ts.search_call(marker, "a"), + ts.search_output(marker, [ts.loaded(ts.WEATHER)], "a"), + ts.replayed_call(marker, ts.WEATHER), + ts.replayed_call(marker, ts.TIME), + ] + assert guarded.shell not in posted.body.decode() + + +def test_guardrails_leave_an_allowed_replay_untouched(guarded: _Guarded) -> None: + marker: Final = uuid.uuid4().hex + body: Final = _one_search_body(guarded, marker, [ts.loaded(ts.WEATHER)], guardrails=[_BLOCK]) + response: Final = guarded.proxy.request("POST", "/v1/chat/completions", body) + assert response.status_code == 200, response.text + (posted,) = pcb.with_marker(pcb.drained_posts(guarded.wire), marker) + assert ts.replayed(posted) == [ + ts.search_call(marker, "a"), + ts.search_output(marker, [ts.loaded(ts.WEATHER)], "a"), + ts.replayed_call(marker, ts.WEATHER), + ] + + +def test_guardrails_serve_an_allowed_replay_from_the_cache(guarded: _Guarded) -> None: + marker: Final = uuid.uuid4().hex + _cache(guarded, marker, _one_search_body(guarded, marker, [ts.loaded(ts.WEATHER)], guardrails=[_BLOCK])) + + +def test_tool_policy_refuses_a_replay_whose_answer_an_exempt_key_cached(guarded: _Guarded) -> None: + marker: Final = uuid.uuid4().hex + body: Final = _blocked_load_body(guarded, marker) + with guarded.proxy.scenario() as scenario: + exempt: Final = scenario.key(models=[guarded.model], metadata={"disable_global_guardrails": True}) + _cache(guarded, marker, body, key=exempt) + refused: Final = guarded.proxy.request("POST", "/v1/chat/completions", body) + _assert_refused(refused, guarded.wire, marker, "Violated tool policy", guarded.blocked) + + +def test_tool_permission_block_mode_refuses_a_replay_whose_answer_is_cached(guarded: _Guarded) -> None: + marker: Final = uuid.uuid4().hex + _cache(guarded, marker, _two_turn_body(guarded, marker)) + response: Final = guarded.proxy.request( + "POST", "/v1/chat/completions", _two_turn_body(guarded, marker, guardrails=[_BLOCK]) + ) + _assert_refused(response, guarded.wire, marker, "Violated guardrail policy", guarded.shell) + + +def test_tool_permission_rewrite_mode_is_not_served_the_cached_unrewritten_answer(guarded: _Guarded) -> None: + marker: Final = uuid.uuid4().hex + _cache(guarded, marker, _two_turn_body(guarded, marker)) + rewritten: Final = _two_turn_body(guarded, marker, guardrails=[_REWRITE]) + first: Final = guarded.proxy.request("POST", "/v1/chat/completions", rewritten) + assert first.status_code == 200, first.text + (posted,) = pcb.with_marker(pcb.drained_posts(guarded.wire), marker) + assert guarded.shell not in posted.body.decode() + await_entry(redis_store(), str(rv.JSON_OBJECT.validate_json(first.text)["id"])) + hit: Final = guarded.proxy.request("POST", "/v1/chat/completions", rewritten) + assert hit.status_code == 200 and hit.headers.get("x-litellm-cache-key"), (hit.text, dict(hit.headers)) + assert ts.first_message(hit.text)["content"] == rv.answer(marker), hit.text + assert _posts(guarded.wire, marker) == () diff --git a/tests/integration/providers/test_responses_bridge_tool_search_chaos.py b/tests/integration/providers/test_responses_bridge_tool_search_chaos.py index 2a3b98c3659..bec1f5799eb 100644 --- a/tests/integration/providers/test_responses_bridge_tool_search_chaos.py +++ b/tests/integration/providers/test_responses_bridge_tool_search_chaos.py @@ -19,9 +19,12 @@ import pytest import yaml from _pytest.mark.structures import ParameterSet from integration._support import claude_code as cc +from integration._support import prompt_cache_breakpoint as pcb from integration._support import responses_vendor as rv +from integration._support import tool_search as ts from integration._support.anthropic_thinking import BEDROCK_MODEL, EVENT_STREAM, Event, aws_chunks, sse_chunks -from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment +from integration._support.client import Gateway, Scenario, eventually, gateway_from_environment, string_value +from integration._support.openai_wire import answering_model_discovery from integration._support.database import read_rows from integration._support.process import OwnedProxy, graceful_stop_seconds, owned_proxy_process from integration._support.wire import Reply, Request, Wire, wire_server @@ -1008,3 +1011,311 @@ def test_a_global_drop_params_drops_a_malformed_defer_loading_and_serves_the_sea ("get_weather", True), ("lookup_order", None), ], tools + + +_BRIDGE_OUTAGE: Final = "scripted: upstream outage" + +BridgeEndpoint: TypeAlias = Literal["chat", "messages", "responses"] + + +@dataclass(frozen=True, slots=True) +class _BridgeCall: + endpoint: BridgeEndpoint + marker: str + stream: bool + replay: bool + + +@dataclass(frozen=True, slots=True) +class _BridgeServed: + call: _BridgeCall + status: int + text: str + + +_BRIDGE_ENDPOINTS: Final[tuple[BridgeEndpoint, ...]] = ("chat", "chat", "messages", "responses") + + +def _bridge_calls( + count: int, endpoint: Callable[[int], BridgeEndpoint], stream: Callable[[int], bool] +) -> tuple[_BridgeCall, ...]: + return tuple( + _BridgeCall(endpoint(index), uuid.uuid4().hex, stream(index), endpoint(index) != "messages" and index % 3 != 0) + for index in range(count) + ) + + +def _bridge_chat_messages(call: _BridgeCall) -> list[JsonValue]: + user: Final[dict[str, JsonValue]] = {"role": "user", "content": pcb.prompt(call.marker)} + if not call.replay: + return [user] + items: Final[list[JsonValue]] = [ + ts.search_call(call.marker), + ts.search_output(call.marker, [ts.loaded(ts.WEATHER)]), + ] + tool_call: Final[dict[str, JsonValue]] = { + "id": ts.call_id(call.marker), + "type": "function", + "function": {"name": ts.WEATHER, "arguments": '{"city": "Paris"}'}, + "provider_specific_fields": {"namespace": ts.WEATHER}, + } + return [ + user, + { + "role": "assistant", + "content": None, + "tool_calls": [tool_call], + "provider_specific_fields": {"tool_search_items": items}, + }, + ts.tool_result(call.marker), + ] + + +def _bridge_responses_input(call: _BridgeCall) -> list[JsonValue]: + user: Final[dict[str, JsonValue]] = {"type": "message", "role": "user", "content": pcb.prompt(call.marker)} + if not call.replay: + return [user] + return [ + user, + *ts.searched(call.marker, "a", ts.WEATHER), + ts.replayed_call(call.marker), + {"type": "function_call_output", "call_id": ts.call_id(call.marker), "output": "sunny"}, + ] + + +def _bridge_expected_replay(call: _BridgeCall) -> list[dict[str, JsonValue]]: + if not call.replay: + return [] + match call.endpoint: + case "chat": + return [ + ts.search_call(call.marker), + ts.search_output(call.marker, [ts.loaded(ts.WEATHER)]), + ts.replayed_call(call.marker), + ] + case "responses": + return [*ts.searched(call.marker, "a", ts.WEATHER), ts.replayed_call(call.marker)] + case "messages": + return [] + + +def _bridge_request(call: _BridgeCall, model: str) -> tuple[str, dict[str, JsonValue]]: + common: Final[dict[str, JsonValue]] = {"model": model, "stream": call.stream, "num_retries": 0, **pcb.NO_CACHE} + match call.endpoint: + case "chat": + return "/v1/chat/completions", { + **common, + "messages": _bridge_chat_messages(call), + "tools": ts.chat_tools(ts.WEATHER), + } + case "messages": + return "/v1/messages", { + **common, + "max_tokens": 64, + "messages": [{"role": "user", "content": pcb.prompt(call.marker)}], + "tools": [{"name": ts.WEATHER, "description": f"Look up {ts.WEATHER}", "input_schema": ts.PARAMETERS}], + } + case "responses": + return "/v1/responses", { + **common, + "input": _bridge_responses_input(call), + "tools": [ts.TOOL_SEARCH, ts.loaded(ts.WEATHER)], + } + + +async def _bridge_send(client: httpx.AsyncClient, key: str, model: str, call: _BridgeCall) -> _BridgeServed: + route, body = _bridge_request(call, model) + headers: Final = {"Authorization": f"Bearer {key}", "anthropic-version": "2023-06-01"} + async with client.stream("POST", route, json=body, headers=headers) as response: + raw: Final = await response.aread() + return _BridgeServed(call, response.status_code, raw.decode()) + + +async def _bridge_burst( + gateway: Gateway, model: str, calls: tuple[_BridgeCall, ...], *, tolerate_transport_errors: bool = False +) -> tuple[_BridgeServed, ...]: + async with httpx.AsyncClient(base_url=str(gateway.client.base_url), timeout=120, trust_env=False) as client: + results: Final = await asyncio.gather( + *(_bridge_send(client, gateway.key, model, call) for call in calls), + return_exceptions=tolerate_transport_errors, + ) + unexpected: Final = [ + result + for result in results + if isinstance(result, BaseException) and not isinstance(result, httpx.TransportError) + ] + assert unexpected == [], unexpected + return tuple(result for result in results if isinstance(result, _BridgeServed)) + + +def _bridge_assert_chat_answered(served: _BridgeServed) -> None: + assistant: Final = ( + ts.assistant_from_stream(ts.chunks_of(served.text)) if served.call.stream else ts.first_message(served.text) + ) + if served.call.replay: + assert assistant["content"] == rv.answer(served.call.marker), served.text + return + assert ts.provider_fields(assistant).get("tool_search_items") == [ + *ts.searched(served.call.marker, "a", ts.WEATHER) + ], served.text + assert ts.namespaces(assistant) == [ts.WEATHER], served.text + + +def _bridge_assert_answered(served: _BridgeServed) -> None: + assert served.status == 200, served.text + match served.call.endpoint: + case "chat": + _bridge_assert_chat_answered(served) + case "messages": + assert '"tool_use"' in served.text and ts.call_id(served.call.marker) in served.text, served.text + assert "tool_search" not in served.text, served.text + case "responses": + expected: Final = rv.answer(served.call.marker) if served.call.replay else f"tso_a_{served.call.marker}" + assert expected in served.text, served.text + + +def _bridge_assert_upstream_saw_each_call_once(received: tuple[Request, ...], calls: tuple[_BridgeCall, ...]) -> None: + posts: Final = tuple(request for request in received if request.method == "POST") + assert sorted(rv.newest_marker(request.body.decode()) or "" for request in posts) == sorted( + call.marker for call in calls + ), [request.body for request in posts] + expected: Final = {call.marker: _bridge_expected_replay(call) for call in calls} + assert [ts.replayed(request) for request in posts] == [ + expected[rv.newest_marker(request.body.decode()) or ""] for request in posts + ], [request.body for request in posts] + + +def _bridge_caller_id(served: _BridgeServed) -> str: + if not served.call.stream: + return string_value(rv.JSON_OBJECT.validate_json(served.text)["id"]) + first: Final = ts.chunks_of(served.text)[0] + match served.call.endpoint: + case "chat": + return string_value(first["id"]) + case "messages": + return string_value(rv.JSON_OBJECT.validate_python(first["message"])["id"]) + case "responses": + return string_value(rv.JSON_OBJECT.validate_python(first["response"])["id"]) + + +def _bridge_assert_one_success_row(model: str, served: _BridgeServed) -> None: + caller_id: Final = _bridge_caller_id(served) + rows: Final = eventually(lambda: ts.spend_rows_of(model, caller_id), lambda found: len(found) >= 1, seconds=60) + assert [row["status"] for row in rows] == ["success"], rows + + +def _bridge_outage_for(down: frozenset[str]) -> Callable[[Request], Reply]: + def respond(request: Request) -> Reply: + if rv.newest_marker(request.body.decode()) in down: + return rv.error(503, _BRIDGE_OUTAGE, "server_error") + return ts.scripted(request) + + return answering_model_discovery(respond) + + +async def test_upstream_outage_mid_burst_fails_only_its_calls_and_every_replay_lands_once(gateway: Gateway) -> None: + calls: Final = _bridge_calls(24, lambda index: _BRIDGE_ENDPOINTS[index % 4], lambda index: (index // 4) % 2 == 0) + down: Final = frozenset(call.marker for index, call in enumerate(calls) if index % 5 == 1) + with wire_server(_bridge_outage_for(down)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1") + served: Final = await _bridge_burst(gateway, model, calls) + assert len(served) == len(calls) + for item in served: + if item.call.marker in down: + assert item.status == 503 and _BRIDGE_OUTAGE in item.text, item.text + assert rv.answer(item.call.marker) not in item.text, item.text + else: + _bridge_assert_answered(item) + recovery: Final = _BridgeCall("chat", uuid.uuid4().hex, stream=True, replay=True) + (recovered,) = await _bridge_burst(gateway, model, (recovery,)) + _bridge_assert_answered(recovered) + _bridge_assert_upstream_saw_each_call_once(wire.drain(), (*calls, recovery)) + for item in (*served, recovered): + if item.call.marker not in down: + _bridge_assert_one_success_row(model, item) + + +async def test_slow_upstream_streams_keep_items_intact_without_a_duplicate_upstream_call(gateway: Gateway) -> None: + calls: Final = _bridge_calls(12, lambda index: _BRIDGE_ENDPOINTS[index % 4], lambda index: (index // 4) % 2 == 0) + with wire_server(ts.scripted_with_pause(0.3)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1") + served: Final = await _bridge_burst(gateway, model, calls) + assert len(served) == len(calls) + for item in served: + _bridge_assert_answered(item) + _bridge_assert_upstream_saw_each_call_once(wire.drain(), calls) + for item in served: + _bridge_assert_one_success_row(model, item) + + +def _bridge_chaos_config(wire: Wire, directory: Path) -> Path: + template: Final = Path("tests/integration/proxy_config.yaml").read_text() + base: Final = rv.JSON_OBJECT.validate_python(yaml.safe_load(template)) + config: Final = { + **base, + "model_list": [ + { + "model_name": _CONFIG_MODEL, + "litellm_params": {"model": pcb.MODEL, "api_base": f"{wire.url}/v1", "api_key": "sk-upstream"}, + } + ], + } + path: Final = directory / "responses-bridge-tool-search-chaos.yaml" + path.write_text(yaml.safe_dump(config)) + return path + + +def _bridge_open_upstream_connections(pid: int, upstream: str) -> int: + port: Final = urlsplit(upstream).port + return sum( + 1 + for connection in psutil.Process(pid).net_connections(kind="tcp") + if connection.status == psutil.CONN_ESTABLISHED and connection.raddr and connection.raddr.port == port + ) + + +@pytest.mark.timeout(2 * graceful_stop_seconds() + 120) +async def test_worker_kill_mid_replay_burst_leaves_the_sibling_replaying_search_items( + gateway: Gateway, tmp_path: Path +) -> None: + calls: Final = _bridge_calls(20, lambda index: _BRIDGE_ENDPOINTS[index % 4], lambda _: False) + release: Final = threading.Event() + held_markers: Final[SimpleQueue[str]] = SimpleQueue() + + def held(request: Request) -> Reply: + held_markers.put(rv.newest_marker(request.body.decode()) or "") + assert release.wait(timeout=90), "The burst was never released" + return ts.scripted(request) + + with wire_server(answering_model_discovery(held)) as wire: + with owned_proxy_process( + gateway, tmp_path, {}, config=_bridge_chaos_config(wire, tmp_path), workers=2 + ) as owned: + workers: Final = eventually( + lambda: pcb.started_worker_pids(owned.log), + lambda pids: len(pids) == 2, + seconds=graceful_stop_seconds() + 60, + ) + burst: Final = asyncio.create_task( + _bridge_burst(owned.gateway, _CONFIG_MODEL, calls, tolerate_transport_errors=True) + ) + await asyncio.to_thread(eventually, held_markers.qsize, lambda size: size == len(calls), 60) + held_by: Final = MappingProxyType( + {pid: _bridge_open_upstream_connections(pid, wire.url) for pid in workers} + ) + assert sum(held_by.values()) == len(calls), held_by + victim_pid, survivor_pid = sorted(workers, key=held_by.__getitem__) + with owned.replacing_worker(victim_pid) as kill: + kill.freeze() + await asyncio.to_thread(kill) + release.set() + served: Final = await burst + assert len(served) == held_by[survivor_pid], (held_by, len(served)) + for item in served: + _bridge_assert_answered(item) + follow_up: Final = _BridgeCall("chat", uuid.uuid4().hex, stream=True, replay=True) + (answered,) = await _bridge_burst(owned.gateway, _CONFIG_MODEL, (follow_up,)) + _bridge_assert_answered(answered) + _bridge_assert_upstream_saw_each_call_once(wire.drain(), (*calls, follow_up)) + for item in (*served, answered): + _bridge_assert_one_success_row(_CONFIG_MODEL, item) diff --git a/tests/integration/providers/test_responses_bridge_tool_search_replay.py b/tests/integration/providers/test_responses_bridge_tool_search_replay.py new file mode 100644 index 00000000000..0f240cc2d2a --- /dev/null +++ b/tests/integration/providers/test_responses_bridge_tool_search_replay.py @@ -0,0 +1,496 @@ +import asyncio +import uuid +from collections.abc import Iterable, Mapping, Sequence +from itertools import chain +from typing import Final, Literal, TypeAlias, cast + +import httpx +import pytest +from anthropic import Anthropic +from integration._support import prompt_cache_breakpoint as pcb +from integration._support import responses_vendor as rv +from integration._support import tool_search as ts +from integration._support.client import Gateway, Scenario, eventually, string_value +from integration._support.wire import Wire, wire_server +from openai import AsyncOpenAI, OpenAI +from openai.types.chat import ChatCompletionMessageParam, ChatCompletionToolUnionParam +from pydantic import JsonValue + +from tests.integration.caching.response_cache_case import await_entry, redis_store + +pytestmark: Final = pytest.mark.timeout(120) + +Client: TypeAlias = Literal["httpx", "openai_sync", "openai_async"] +_ANTHROPIC_TOOL: Final[dict[str, JsonValue]] = { + "name": ts.WEATHER, + "description": "Look up get_weather", + "input_schema": ts.PARAMETERS, +} + + +def _base_url(gateway: Gateway) -> str: + return str(gateway.client.base_url).rstrip("/") + + +def _sdk_body(body: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + return {key: value for key, value in body.items() if key not in ("model", "messages", "tools", "stream")} + + +def _sync_sdk(gateway: Gateway, body: Mapping[str, JsonValue]) -> tuple[object, ...]: + messages: Final = cast(Iterable[ChatCompletionMessageParam], body["messages"]) + tools: Final = cast(Iterable[ChatCompletionToolUnionParam], body["tools"]) + model: Final = string_value(body["model"]) + with OpenAI(api_key=gateway.key, base_url=f"{_base_url(gateway)}/v1", max_retries=0) as client: + if body.get("stream") is True: + chunks: Final = client.chat.completions.create( + model=model, messages=messages, tools=tools, stream=True, extra_body=_sdk_body(body) + ) + return tuple(chunk.model_dump() for chunk in chunks) + completion: Final = client.chat.completions.create( + model=model, messages=messages, tools=tools, stream=False, extra_body=_sdk_body(body) + ) + return (completion.model_dump(),) + + +async def _async_sdk(gateway: Gateway, body: Mapping[str, JsonValue]) -> tuple[object, ...]: + messages: Final = cast(Iterable[ChatCompletionMessageParam], body["messages"]) + tools: Final = cast(Iterable[ChatCompletionToolUnionParam], body["tools"]) + model: Final = string_value(body["model"]) + async with AsyncOpenAI(api_key=gateway.key, base_url=f"{_base_url(gateway)}/v1", max_retries=0) as client: + if body.get("stream") is True: + chunks: Final = await client.chat.completions.create( + model=model, messages=messages, tools=tools, stream=True, extra_body=_sdk_body(body) + ) + return tuple([chunk.model_dump() async for chunk in chunks]) + completion: Final = await client.chat.completions.create( + model=model, messages=messages, tools=tools, stream=False, extra_body=_sdk_body(body) + ) + return (completion.model_dump(),) + + +def _send(gateway: Gateway, client: Client, body: Mapping[str, JsonValue]) -> tuple[dict[str, JsonValue], ...]: + match client: + case "httpx": + response: Final = gateway.request("POST", "/v1/chat/completions", dict(body)) + assert response.status_code == 200, response.text + if body.get("stream") is True: + return ts.chunks_of(response.text) + return (rv.JSON_OBJECT.validate_json(response.text),) + case "openai_sync": + return tuple(rv.JSON_OBJECT.validate_python(item) for item in _sync_sdk(gateway, body)) + case "openai_async": + return tuple(rv.JSON_OBJECT.validate_python(item) for item in asyncio.run(_async_sdk(gateway, body))) + + +def _assistant(served: Sequence[Mapping[str, JsonValue]], stream: bool) -> dict[str, JsonValue]: + if stream: + return ts.assistant_from_stream(served) + choices: Final = rv.ITEMS.validate_python(served[0]["choices"]) + return rv.JSON_OBJECT.validate_python(choices[0]["message"]) + + +def _sent_message(message: Mapping[str, JsonValue]) -> dict[str, JsonValue]: + kept: Final = ("role", "content", "tool_calls", "provider_specific_fields", "reasoning_items") + return {key: value for key, value in message.items() if key in kept and value is not None} + + +def _search_items(marker: str) -> list[JsonValue]: + return [*ts.searched(marker, "a", ts.WEATHER)] + + +def _replay(marker: str) -> list[dict[str, JsonValue]]: + return [*ts.searched(marker, "a", ts.WEATHER), ts.reasoning(marker, "b"), ts.replayed_call(marker)] + + +def _body( + model: str, messages: Sequence[JsonValue], *, stream: bool = False, **extra: JsonValue +) -> dict[str, JsonValue]: + return { + "model": model, + "messages": list(messages), + "tools": ts.chat_tools(ts.WEATHER), + "stream": stream, + **pcb.NO_CACHE, + **extra, + } + + +def _user(marker: str, *words: str) -> dict[str, JsonValue]: + return {"role": "user", "content": " ".join((pcb.prompt(marker), *words))} + + +def _two_turns( + gateway: Gateway, wire: Wire, model: str, client: Client, stream: bool, marker: str +) -> tuple[dict[str, JsonValue], list[dict[str, JsonValue]], tuple[str, str]]: + user: Final = _user(marker) + first: Final = _send(gateway, client, _body(model, [user], stream=stream)) + assistant: Final = _assistant(first, stream) + pcb.posted(wire, marker) + history: Final = [user, _sent_message(assistant), ts.tool_result(marker)] + second: Final = _send(gateway, client, _body(model, history, stream=stream)) + answered: Final = _assistant(second, stream) + assert answered["content"] == rv.answer(marker), answered + identities: Final = (string_value(first[0]["id"]), string_value(second[0]["id"])) + return assistant, ts.replayed(pcb.posted(wire, marker)), identities + + +def _assert_turn_one(assistant: Mapping[str, JsonValue], marker: str) -> None: + assert ts.provider_fields(assistant).get("tool_search_items") == _search_items(marker), assistant + assert [call["id"] for call in ts.tool_calls(assistant)] == [ts.call_id(marker)], assistant + assert ts.namespaces(assistant) == [ts.WEATHER], assistant + + +def _assert_spend(model: str, response_ids: Sequence[str]) -> None: + rows: Final = eventually( + lambda: tuple(ts.spend_rows_of(model, response_id) for response_id in response_ids), + lambda found: all(found), + seconds=70, + ) + assert [[row["status"] for row in found] for found in rows] == [["success"]] * len(response_ids), rows + + +@pytest.mark.parametrize("client", ["httpx", "openai_sync", "openai_async"]) +def test_chat_bridge_returns_and_replays_tool_search_items_non_streaming(gateway: Gateway, client: Client) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(ts.scripted) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1") + assistant, replayed, response_ids = _two_turns(gateway, wire, model, client, False, marker) + _assert_turn_one(assistant, marker) + assert replayed == _replay(marker) + _assert_spend(model, response_ids) + + +@pytest.mark.parametrize("client", ["httpx", "openai_async"]) +def test_chat_bridge_streams_tool_search_items_and_replays_the_rebuilt_message( + gateway: Gateway, client: Client +) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(ts.scripted) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1") + assistant, replayed, response_ids = _two_turns(gateway, wire, model, client, True, marker) + _assert_turn_one(assistant, marker) + assert replayed == _replay(marker) + _assert_spend(model, response_ids) + + +def _history(marker: str, **assistant_fields: JsonValue) -> list[JsonValue]: + call: Final[dict[str, JsonValue]] = { + "id": ts.call_id(marker), + "type": "function", + "function": {"name": ts.WEATHER, "arguments": '{"city": "Paris"}'}, + } + return [ + _user(marker), + {"role": "assistant", "content": None, "tool_calls": [call], **assistant_fields}, + ts.tool_result(marker), + ] + + +def _replayed_history(marker: str) -> list[JsonValue]: + return _history( + marker, + tool_calls=[ + { + "id": ts.call_id(marker), + "type": "function", + "function": {"name": ts.WEATHER, "arguments": '{"city": "Paris"}'}, + "provider_specific_fields": {"namespace": ts.WEATHER}, + } + ], + provider_specific_fields={"tool_search_items": _search_items(marker)}, + reasoning_items=[ts.reasoning(marker, "a"), ts.reasoning(marker, "b")], + ) + + +def _replayed_upstream( + gateway: Gateway, + scenario: Scenario, + wire: Wire, + marker: str, + body: Mapping[str, JsonValue], + **model_params: JsonValue, +) -> list[dict[str, JsonValue]]: + model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1", **model_params) + response: Final = gateway.request("POST", "/v1/chat/completions", {**body, "model": model}) + assert response.status_code == 200, response.text + assert ts.first_message(response.text)["content"] == rv.answer(marker), response.text + return ts.replayed(pcb.posted(wire, marker)) + + +def test_turn_without_tool_search_strips_search_items_but_keeps_namespace_and_reasoning(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(ts.scripted) as wire, gateway.scenario() as scenario: + body: Final = { + **_body("", _replayed_history(marker)), + "tools": [{"type": "function", "function": ts.function(ts.WEATHER)}], + } + assert _replayed_upstream(gateway, scenario, wire, marker, body) == [ + ts.reasoning(marker, "a"), + ts.reasoning(marker, "b"), + ts.replayed_call(marker), + ] + + +def test_tool_search_only_in_extra_body_tools_still_replays(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(ts.scripted) as wire, gateway.scenario() as scenario: + body: Final = { + **_body("", _replayed_history(marker)), + "tools": [ts.deferred(ts.WEATHER)], + "extra_body": {"tools": [ts.TOOL_SEARCH, ts.loaded(ts.WEATHER)]}, + } + assert _replayed_upstream(gateway, scenario, wire, marker, body) == _replay(marker) + + +def test_extra_body_tools_without_tool_search_win_over_tools_and_skip_the_replay(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(ts.scripted) as wire, gateway.scenario() as scenario: + body: Final = { + **_body("", _replayed_history(marker)), + "extra_body": {"tools": [{"type": "function", **ts.function(ts.WEATHER)}]}, + } + assert _replayed_upstream(gateway, scenario, wire, marker, body) == [ + ts.reasoning(marker, "a"), + ts.reasoning(marker, "b"), + ts.replayed_call(marker), + ] + + +def test_deployment_tool_search_replays_items_the_request_did_not_ask_for(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(ts.scripted) as wire, gateway.scenario() as scenario: + body: Final = {**_body("", _replayed_history(marker)), "tools": [ts.deferred(ts.WEATHER)]} + replayed: Final = _replayed_upstream(gateway, scenario, wire, marker, body, tools=[ts.TOOL_SEARCH]) + assert replayed == _replay(marker) + + +_MALFORMED_ITEMS: Final[dict[str, JsonValue]] = { + "int": 7, + "string": "tool_search_call", + "empty-object": {}, + "empty-list": [], + "null": None, + "junk-list": [1, "x", None, {"type": "message", "content": "hi"}, {"no": "type"}], + "5kb-string": "x" * 5120, +} + + +@pytest.mark.parametrize("items", list(_MALFORMED_ITEMS.values()), ids=list(_MALFORMED_ITEMS)) +def test_malformed_tool_search_items_are_ignored_without_a_server_error(gateway: Gateway, items: JsonValue) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(ts.scripted) as wire, gateway.scenario() as scenario: + body: Final = _body("", _history(marker, provider_specific_fields={"tool_search_items": items})) + assert _replayed_upstream(gateway, scenario, wire, marker, body) == [ + {k: v for k, v in ts.replayed_call(marker).items() if k != "namespace"} + ] + + +def test_search_output_with_non_list_tools_and_items_repeated_on_two_turns_replay_verbatim(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + odd_output: Final = {**ts.search_output(marker, []), "tools": "not-a-list"} + items: Final[list[JsonValue]] = [ts.search_call(marker), odd_output] + first_call: Final = { + "id": ts.call_id(marker), + "type": "function", + "function": {"name": ts.WEATHER, "arguments": "{}"}, + } + messages: Final[list[JsonValue]] = [ + _user(marker), + {"role": "assistant", "content": "one", "provider_specific_fields": {"tool_search_items": items}}, + {"role": "user", "content": "again"}, + { + "role": "assistant", + "content": None, + "tool_calls": [first_call], + "provider_specific_fields": {"tool_search_items": items}, + }, + ts.tool_result(marker), + ] + with wire_server(ts.scripted) as wire, gateway.scenario() as scenario: + replayed: Final = _replayed_upstream(gateway, scenario, wire, marker, _body("", messages)) + assert replayed == [ + ts.search_call(marker), + odd_output, + ts.search_call(marker), + odd_output, + {"type": "function_call", "call_id": ts.call_id(marker), "name": ts.WEATHER, "arguments": "{}"}, + ] + + +_NAMESPACES: Final[dict[str, JsonValue]] = {"int": 3, "list": ["get_weather"], "empty": "", "5kb": "n" * 5120} + + +@pytest.mark.parametrize("value", list(_NAMESPACES.values()), ids=list(_NAMESPACES)) +def test_malformed_namespace_is_dropped_and_a_long_one_is_forwarded(gateway: Gateway, value: JsonValue) -> None: + marker: Final = uuid.uuid4().hex + call: Final[dict[str, JsonValue]] = { + "id": ts.call_id(marker), + "type": "function", + "function": {"name": ts.WEATHER, "arguments": '{"city": "Paris"}'}, + "provider_specific_fields": {"namespace": value}, + } + with wire_server(ts.scripted) as wire, gateway.scenario() as scenario: + body: Final = _body("", _history(marker, tool_calls=[call])) + bare: Final = {k: v for k, v in ts.replayed_call(marker).items() if k != "namespace"} + expected: Final = {**bare, "namespace": value} if isinstance(value, str) and value else bare + assert _replayed_upstream(gateway, scenario, wire, marker, body) == [expected] + + +def test_upstream_rejecting_the_replay_reaches_the_caller(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(ts.scripted) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1") + messages: Final = [_user(marker, ts.REJECT_REPLAY), *_replayed_history(marker)[1:]] + response: Final = gateway.request("POST", "/v1/chat/completions", {**_body(model, messages)}) + assert response.status_code == 400, response.text + assert ts.REJECTED in response.text, response.text + assert ts.replayed(pcb.posted(wire, marker)) == _replay(marker) + healthy: Final = gateway.request("POST", "/v1/chat/completions", _body(model, [_user(uuid.uuid4().hex)])) + assert healthy.status_code == 200, healthy.text + + +@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) +def test_turn_that_only_searched_returns_the_items_and_continues_the_search(gateway: Gateway, stream: bool) -> None: + marker: Final = uuid.uuid4().hex + user: Final = _user(marker, ts.ONLY_SEARCH) + with wire_server(ts.scripted) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1") + first: Final = _send(gateway, "httpx", _body(model, [user], stream=stream)) + assistant: Final = _assistant(first, stream) + choices: Final = chain.from_iterable(rv.ITEMS.validate_python(chunk["choices"]) for chunk in first) + finish: Final = [choice["finish_reason"] for choice in choices if choice.get("finish_reason")] + assert finish == ["stop"], first + assert assistant.get("content") in (None, ""), assistant + assert ts.provider_fields(assistant).get("tool_search_items") == _search_items(marker), assistant + assert "tool_calls" not in assistant or not assistant["tool_calls"], assistant + pcb.posted(wire, marker) + second: Final = _send(gateway, "httpx", _body(model, [user, _sent_message(assistant)], stream=stream)) + assert _assistant(second, stream)["content"] == rv.answer(marker), second + assert ts.replayed(pcb.posted(wire, marker)) == _search_items(marker) + + +def test_two_searches_in_two_turns_replay_in_order(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + user: Final = _user(marker, ts.TWO_SEARCHES) + with wire_server(ts.scripted) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1") + tools: Final = ts.chat_tools(ts.WEATHER, ts.TIME) + first: Final = _assistant(_send(gateway, "httpx", {**_body(model, [user]), "tools": tools}), False) + pcb.posted(wire, marker) + history: Final[list[JsonValue]] = [user, _sent_message(first), ts.tool_result(marker)] + second: Final = _assistant(_send(gateway, "httpx", {**_body(model, history), "tools": tools}), False) + assert ts.provider_fields(second).get("tool_search_items") == [*ts.searched(marker, "c", ts.TIME)], second + assert ts.namespaces(second) == [ts.TIME], second + pcb.posted(wire, marker) + final: Final = [*history, _sent_message(second), ts.tool_result(marker, ts.TIME)] + third: Final = _assistant(_send(gateway, "httpx", {**_body(model, final), "tools": tools}), False) + assert third["content"] == rv.answer(marker), third + assert ts.replayed(pcb.posted(wire, marker)) == [ + *_replay(marker), + *ts.searched(marker, "c", ts.TIME), + ts.replayed_call(marker, ts.TIME), + ] + + +def _cached_twin( + gateway: Gateway, body: Mapping[str, JsonValue], stream: bool +) -> tuple[httpx.Response, httpx.Response]: + first: Final = gateway.request("POST", "/v1/chat/completions", dict(body)) + assert first.status_code == 200, first.text + first_frame: Final = ts.chunks_of(first.text)[0] if stream else rv.JSON_OBJECT.validate_json(first.text) + first_id: Final = string_value(first_frame["id"]) + await_entry(redis_store(), first_id) + second: Final = gateway.request("POST", "/v1/chat/completions", dict(body)) + assert second.status_code == 200, second.text + assert second.headers.get("x-litellm-cache-key"), dict(second.headers) + return first, second + + +@pytest.mark.parametrize("stream", [False, True], ids=["non-streaming", "streaming"]) +def test_cache_hit_turn_carries_the_items_and_replays_them(gateway: Gateway, stream: bool) -> None: + marker: Final = uuid.uuid4().hex + user: Final = _user(marker) + with wire_server(ts.scripted) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1") + cached_body: Final = { + key: value for key, value in _body(model, [user], stream=stream).items() if key != "cache" + } + first, hit = _cached_twin(gateway, cached_body, stream) + pcb.posted(wire, marker) + served: Final = ts.chunks_of(hit.text) if stream else (rv.JSON_OBJECT.validate_json(hit.text),) + assistant: Final = _assistant(served, stream) + _assert_turn_one(assistant, marker) + second: Final = _send(gateway, "httpx", _body(model, [user, _sent_message(assistant), ts.tool_result(marker)])) + assert _assistant(second, False)["content"] == rv.answer(marker), second + assert ts.replayed(pcb.posted(wire, marker)) == _replay(marker) + first_frames: Final = ts.chunks_of(first.text) if stream else (rv.JSON_OBJECT.validate_json(first.text),) + assert served[0]["id"] == first_frames[0]["id"] + + +def test_messages_endpoint_sees_the_same_tool_use_on_both_sides(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(ts.scripted) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1") + with Anthropic(api_key=gateway.key, base_url=_base_url(gateway), max_retries=0) as client: + message: Final = client.messages.create( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": pcb.prompt(marker)}], + tools=[_ANTHROPIC_TOOL], # pyright: ignore[reportArgumentType] # JSON tool definition + ) + pcb.posted(wire, marker) + uses: Final = [block for block in message.content if block.type == "tool_use"] + assert [(block.name, block.input) for block in uses] == [(ts.WEATHER, {"city": "Paris"})], message + assert "tool_search" not in message.model_dump_json(), message.model_dump_json() + with client.messages.stream( + model=model, + max_tokens=64, + messages=[{"role": "user", "content": pcb.prompt(marker)}], + tools=[_ANTHROPIC_TOOL], # pyright: ignore[reportArgumentType] # JSON tool definition + ) as stream: + streamed: Final = stream.get_final_message() + pcb.posted(wire, marker) + streamed_uses: Final = [block for block in streamed.content if block.type == "tool_use"] + assert [(block.name, block.input) for block in streamed_uses] == [(ts.WEATHER, {"city": "Paris"})] + assert "tool_search" not in streamed.model_dump_json(), streamed.model_dump_json() + + +def test_messages_endpoint_turn_that_only_searched_is_answered(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + with wire_server(ts.scripted) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1") + response: Final = gateway.request( + "POST", + "/v1/messages", + { + "model": model, + "max_tokens": 64, + "messages": [{"role": "user", "content": f"{pcb.prompt(marker)} {ts.ONLY_SEARCH}"}], + "tools": [_ANTHROPIC_TOOL], + }, + ) + assert response.status_code == 200, response.text + body: Final = rv.JSON_OBJECT.validate_json(response.text) + assert body["stop_reason"] == "end_turn", body + assert [block for block in rv.ITEMS.validate_python(body["content"]) if block.get("type") == "tool_use"] == [] + pcb.posted(wire, marker) + + +def test_native_responses_forwards_tool_search_items_and_namespace_verbatim(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + history: Final[list[JsonValue]] = [ + {"type": "message", "role": "user", "content": pcb.prompt(marker)}, + *ts.searched(marker, "a", ts.WEATHER), + ts.replayed_call(marker), + {"type": "function_call_output", "call_id": ts.call_id(marker), "output": "sunny"}, + ] + with wire_server(ts.scripted) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1") + response: Final = gateway.request( + "POST", + "/v1/responses", + {"model": model, "input": history, "tools": [ts.TOOL_SEARCH, ts.loaded(ts.WEATHER)], **pcb.NO_CACHE}, + ) + assert response.status_code == 200, response.text + assert rv.answer(marker) in response.text, response.text + assert pcb.input_items(pcb.posted(wire, marker)) == history diff --git a/tests/integration/providers/test_responses_bridge_tool_search_wire.py b/tests/integration/providers/test_responses_bridge_tool_search_wire.py index b00ba3edfd6..03cd4ce9dc4 100644 --- a/tests/integration/providers/test_responses_bridge_tool_search_wire.py +++ b/tests/integration/providers/test_responses_bridge_tool_search_wire.py @@ -19,8 +19,9 @@ from integration._support import claude_code as cc from integration._support import prompt_cache_breakpoint as pcb from integration._support import responses_vendor as rv from integration._support.anthropic_thinking import BEDROCK_MODEL, EVENT_STREAM, MODEL, Event, aws_chunks, sse_chunks -from integration._support.client import Gateway, Scenario, eventually, object_value +from integration._support.client import Gateway, Scenario, eventually, object_value, string_value from integration._support.database import read_rows +from integration._support.openai_wire import answering_model_discovery, responses_reply from integration._support.vertex import service_account_json from integration._support.wire import Reply, Request, Wire, wire_server from integration.providers.test_bedrock_auth_wire import RESPONSE as CONVERSE_RESPONSE @@ -1324,3 +1325,113 @@ def test_chat_bridge_forwards_defer_loading_so_tool_search_has_a_deferred_tool( _defer_wire_function("get_time"), _defer_wire_function("get_news", defer_loading=False), ] + + +_OURS_PARAMETERS: Final[dict[str, JsonValue]] = { + "type": "object", + "properties": {"city": {"type": "string"}}, + "required": ["city"], +} +_OURS_TOOL_SEARCH: Final[dict[str, JsonValue]] = {"type": "tool_search"} + + +def _ours_function(name: str, **extra: JsonValue) -> dict[str, JsonValue]: + return {"name": name, "description": f"Look up {name}", "parameters": _OURS_PARAMETERS, **extra} + + +def _ours_wire_function(name: str, **extra: JsonValue) -> dict[str, JsonValue]: + return {"type": "function", "strict": None, **_ours_function(name), **extra} + + +_OURS_CALL_ID: Final = "call_weather" +_OURS_TOOL_SEARCH_CALL: Final[dict[str, JsonValue]] = { + "id": "tsc_weather", + "type": "tool_search_call", + "arguments": {"query": "weather"}, + "execution": "server", + "status": "completed", +} +_OURS_TOOL_SEARCH_OUTPUT: Final[dict[str, JsonValue]] = { + "id": "tso_weather", + "type": "tool_search_output", + "execution": "server", + "status": "completed", + "tools": [_ours_wire_function("get_weather", defer_loading=True)], +} +_OURS_WEATHER_CALL: Final[dict[str, JsonValue]] = { + "type": "function_call", + "call_id": _OURS_CALL_ID, + "name": "get_weather", + "arguments": '{"city": "Paris"}', + "namespace": "get_weather", +} + + +def _ours_searches_then_answers(request: Request) -> Reply: + body: Final = pcb.body_of(request) + marker: Final = rv.newest_marker(request.body.decode()) + identity: Final = f"resp_{marker}-{uuid.uuid4().hex}" + if any(item.get("type") == "function_call_output" for item in pcb.input_items(request)): + return responses_reply(identity, string_value(body["model"]), rv.answer(marker), stream=False) + return Reply( + body=json.dumps( + { + "id": identity, + "object": "response", + "created_at": 1, + "status": "completed", + "model": body["model"], + "output": [ + _OURS_TOOL_SEARCH_CALL, + _OURS_TOOL_SEARCH_OUTPUT, + {**_OURS_WEATHER_CALL, "id": "fc_weather", "status": "completed"}, + ], + "usage": rv.USAGE, + } + ).encode() + ) + + +def test_chat_bridge_replays_tool_search_items_on_the_next_turn(gateway: Gateway) -> None: + marker: Final = uuid.uuid4().hex + tools: Final = [ + _OURS_TOOL_SEARCH, + {"type": "function", "function": _ours_function("get_weather"), "defer_loading": True}, + ] + user_turn: Final[dict[str, JsonValue]] = {"role": "user", "content": pcb.prompt(marker)} + with wire_server(answering_model_discovery(_ours_searches_then_answers)) as wire, gateway.scenario() as scenario: + model: Final = scenario.model(model=pcb.MODEL, api_base=f"{wire.url}/v1") + first: Final = gateway.request( + "POST", "/v1/chat/completions", {"model": model, "messages": [user_turn], "tools": tools, **pcb.NO_CACHE} + ) + assert first.status_code == 200, first.text + assistant_turn: Final = rv.JSON_OBJECT.validate_python( + rv.ITEMS.validate_python(rv.JSON_OBJECT.validate_json(first.text)["choices"])[0]["message"] + ) + assert rv.JSON_OBJECT.validate_python(assistant_turn["provider_specific_fields"])["tool_search_items"] == [ + _OURS_TOOL_SEARCH_CALL, + _OURS_TOOL_SEARCH_OUTPUT, + ], first.text + pcb.posted(wire, marker) + + second: Final = gateway.request( + "POST", + "/v1/chat/completions", + { + "model": model, + "messages": [ + user_turn, + assistant_turn, + {"role": "tool", "tool_call_id": _OURS_CALL_ID, "content": "sunny"}, + ], + "tools": tools, + **pcb.NO_CACHE, + }, + ) + assert second.status_code == 200, second.text + replayed: Final = [ + item + for item in pcb.input_items(pcb.posted(wire, marker)) + if item["type"] not in ("message", "function_call_output") + ] + assert replayed == [_OURS_TOOL_SEARCH_CALL, _OURS_TOOL_SEARCH_OUTPUT, _OURS_WEATHER_CALL] diff --git a/tests/integration/sdk/test_responses_bridge_tool_search_sdk.py b/tests/integration/sdk/test_responses_bridge_tool_search_sdk.py new file mode 100644 index 00000000000..5119b23c303 --- /dev/null +++ b/tests/integration/sdk/test_responses_bridge_tool_search_sdk.py @@ -0,0 +1,201 @@ +import asyncio +import uuid +from collections.abc import Sequence +from typing import Final, cast + +import pytest +from integration._support import prompt_cache_breakpoint as pcb +from integration._support import responses_vendor as rv +from integration._support import tool_search as ts +from integration._support.wire import Wire, wire_server +from pydantic import JsonValue + +import litellm +from litellm import Router +from litellm.types.llms.openai import AllMessageValues, ChatCompletionToolMessage, ChatCompletionToolParam +from litellm.types.utils import ChatCompletionMessageToolCall, Choices, Message, ModelResponse + +pytestmark: Final = pytest.mark.timeout(120) + +_API_KEY: Final = "sk-scripted-provider" + + +def _tools(*tools: JsonValue) -> list[ChatCompletionToolParam]: + return cast(list[ChatCompletionToolParam], list(tools)) + + +def _messages(*messages: object) -> list[AllMessageValues]: + return cast(list[AllMessageValues], list(messages)) + + +def _user(marker: str, *words: str) -> dict[str, JsonValue]: + return {"role": "user", "content": " ".join((pcb.prompt(marker), *words))} + + +def _replay(marker: str) -> list[dict[str, JsonValue]]: + return [*ts.searched(marker, "a", ts.WEATHER), ts.reasoning(marker, "b"), ts.replayed_call(marker)] + + +def _message(response: object) -> Message: + assert isinstance(response, ModelResponse), type(response) + choice: Final = response.choices[0] + assert isinstance(choice, Choices), type(choice) + return choice.message + + +def _assert_turn_one(message: Message, marker: str) -> None: + fields: Final = rv.JSON_OBJECT.validate_python(message.provider_specific_fields or {}) + assert fields.get("tool_search_items") == [*ts.searched(marker, "a", ts.WEATHER)], message + calls: Final = message.tool_calls or [] + call_fields: Final = [ + rv.JSON_OBJECT.validate_python(getattr(call, "provider_specific_fields", None) or {}) for call in calls + ] + assert call_fields == [{"namespace": ts.WEATHER}], message + + +def _kwargs(wire: Wire) -> dict[str, str]: + return {"api_base": f"{wire.url}/v1", "api_key": _API_KEY} + + +def test_completion_returns_and_replays_tool_search_items() -> None: + marker: Final = uuid.uuid4().hex + tools: Final = _tools(*ts.chat_tools(ts.WEATHER)) + with wire_server(ts.scripted) as wire: + user: Final = _user(marker) + first: Final = _message( + litellm.completion(model=pcb.MODEL, messages=_messages(user), tools=tools, **_kwargs(wire)) + ) + _assert_turn_one(first, marker) + pcb.posted(wire, marker) + second: Final = _message( + litellm.completion( + model=pcb.MODEL, + messages=_messages(user, first, ts.tool_result(marker)), + tools=tools, + **_kwargs(wire), + ) + ) + assert second.content == rv.answer(marker), second + assert ts.replayed(pcb.posted(wire, marker)) == _replay(marker) + + +def test_streamed_completion_rebuilt_by_stream_chunk_builder_replays_tool_search_items() -> None: + marker: Final = uuid.uuid4().hex + tools: Final = _tools(*ts.chat_tools(ts.WEATHER)) + + async def run(wire: Wire) -> tuple[Message, Message]: + user: Final = _user(marker) + stream: Final = await litellm.acompletion( + model=pcb.MODEL, messages=_messages(user), tools=tools, stream=True, **_kwargs(wire) + ) + chunks: Final = [chunk async for chunk in stream] # pyright: ignore[reportGeneralTypeIssues] # stream wrapper + built: Final = _message(litellm.stream_chunk_builder(chunks, messages=_messages(user))) + pcb.posted(wire, marker) + answered: Final = _message( + await litellm.acompletion( + model=pcb.MODEL, + messages=_messages(user, built.model_dump(), ts.tool_result(marker)), + tools=tools, + **_kwargs(wire), + ) + ) + return built, answered + + with wire_server(ts.scripted) as wire: + built, answered = asyncio.run(run(wire)) + _assert_turn_one(built, marker) + assert [item["id"] for item in built.reasoning_items or []] == [f"rs_a_{marker}", f"rs_b_{marker}"], built + assert answered.content == rv.answer(marker), answered + assert ts.replayed(pcb.posted(wire, marker)) == _replay(marker) + + +def test_streamed_turn_without_tool_search_replays_its_reasoning() -> None: + marker: Final = uuid.uuid4().hex + tools: Final = _tools({"type": "function", "function": ts.function(ts.WEATHER)}) + with wire_server(ts.scripted) as wire: + user: Final = _user(marker) + stream: Final = litellm.completion( + model=pcb.MODEL, messages=_messages(user), tools=tools, stream=True, **_kwargs(wire) + ) + chunks: Final = list(stream) # pyright: ignore[reportArgumentType] # stream wrapper + built: Final = _message(litellm.stream_chunk_builder(chunks, messages=_messages(user))) + pcb.posted(wire, marker) + answered: Final = _message( + litellm.completion( + model=pcb.MODEL, + messages=_messages(user, built.model_dump(), ts.tool_result(marker)), + tools=tools, + **_kwargs(wire), + ) + ) + assert answered.content == rv.answer(marker), answered + assert ts.replayed(pcb.posted(wire, marker)) == [ + ts.reasoning(marker, "a"), + ts.reasoning(marker, "b"), + ts.replayed_call(marker), + ] + + +def test_router_completion_replays_tool_search_items() -> None: + marker: Final = uuid.uuid4().hex + tools: Final = _tools(*ts.chat_tools(ts.WEATHER)) + with wire_server(ts.scripted) as wire: + router: Final = Router( + model_list=[{"model_name": "bridged", "litellm_params": {"model": pcb.MODEL, **_kwargs(wire)}}], + num_retries=0, + ) + user: Final = _user(marker) + first: Final = _message(router.completion(model="bridged", messages=_messages(user), tools=tools)) + _assert_turn_one(first, marker) + pcb.posted(wire, marker) + second: Final = _message( + router.completion(model="bridged", messages=_messages(user, first, ts.tool_result(marker)), tools=tools) + ) + assert second.content == rv.answer(marker), second + assert ts.replayed(pcb.posted(wire, marker)) == _replay(marker) + + +def _sunny(tool_call: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage: + return {"role": "tool", "tool_call_id": tool_call.id, "content": "sunny"} + + +async def _async_sunny(tool_call: ChatCompletionMessageToolCall) -> ChatCompletionToolMessage: + return _sunny(tool_call) + + +def _posted_in_order(wire: Wire, marker: str) -> Sequence[list[dict[str, JsonValue]]]: + return [ts.replayed(request) for request in pcb.with_marker(pcb.drained_posts(wire), marker)] + + +def test_run_tool_loop_keeps_tool_search_items_and_namespace_between_rounds() -> None: + marker: Final = uuid.uuid4().hex + with wire_server(ts.scripted) as wire: + answer: Final = litellm.run_tool_loop( + model=pcb.MODEL, + messages=_messages(_user(marker)), + tools=_tools(*ts.chat_tools(ts.WEATHER)), + execute_tool=_sunny, + **_kwargs(wire), + ) + assert answer == rv.answer(marker) + assert _posted_in_order(wire, marker) == [[], _replay(marker)] + + +def test_async_run_tool_loop_keeps_two_searches_in_order() -> None: + marker: Final = uuid.uuid4().hex + with wire_server(ts.scripted) as wire: + answer: Final = asyncio.run( + litellm.arun_tool_loop( + model=pcb.MODEL, + messages=_messages(_user(marker, ts.TWO_SEARCHES)), + tools=_tools(*ts.chat_tools(ts.WEATHER, ts.TIME)), + execute_tool=_async_sunny, + **_kwargs(wire), + ) + ) + assert answer == rv.answer(marker) + assert _posted_in_order(wire, marker) == [ + [], + _replay(marker), + [*_replay(marker), *ts.searched(marker, "c", ts.TIME), ts.replayed_call(marker, ts.TIME)], + ] 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 056ee1bd717..e5c515f1a09 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 @@ -331,6 +331,187 @@ def test_chunk_parser_function_call_added_produces_tool_use(): assert choice.finish_reason is None +def test_streamed_function_call_namespace_replays_on_next_turn(): + from unittest.mock import Mock + + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + OpenAiResponsesToChatCompletionStreamIterator, + ) + from litellm.types.llms.openai import ResponsesAPIStreamEvents + + iterator = OpenAiResponsesToChatCompletionStreamIterator(streaming_response=None, sync_stream=True) + streamed = iterator.chunk_parser( + { + "type": ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, + "item": { + "type": "function_call", + "name": "get_weather", + "call_id": "call-ns", + "arguments": '{"city": "Paris"}', + "namespace": "get_weather", + }, + } + ) + tool_call = streamed.choices[0].delta.tool_calls[0].model_dump(exclude_none=True) + + request = LiteLLMResponsesTransformationHandler().transform_request( + model="gpt-5.5", + messages=[ + {"role": "user", "content": "weather in Paris"}, + {"role": "assistant", "content": None, "tool_calls": [tool_call]}, + {"role": "tool", "tool_call_id": "call-ns", "content": "sunny"}, + ], + optional_params={}, + litellm_params={}, + headers={}, + litellm_logging_obj=Mock(), + ) + + function_calls = [item for item in request["input"] if item.get("type") == "function_call"] + assert [call["namespace"] for call in function_calls] == ["get_weather"] + + +def test_function_level_namespace_replays_when_tool_call_level_is_missing(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + + tool_call = { + "id": "call-ns", + "type": "function", + "function": { + "name": "get_weather", + "arguments": '{"city": "Paris"}', + "provider_specific_fields": {"namespace": "get_weather"}, + }, + } + + input_items, _ = LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api( + [ + {"role": "user", "content": "weather in Paris"}, + {"role": "assistant", "content": None, "tool_calls": [tool_call]}, + {"role": "tool", "tool_call_id": "call-ns", "content": "sunny"}, + ], + include_namespace=True, + ) + + function_calls = [item for item in input_items if item.get("type") == "function_call"] + assert [call["namespace"] for call in function_calls] == ["get_weather"] + + +def test_shared_converter_leaves_tool_search_items_and_namespace_out_by_default(): + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + ) + + tool_call = { + "id": "call-ns", + "type": "function", + "function": {"name": "get_weather", "arguments": "{}"}, + "provider_specific_fields": {"namespace": "weather"}, + } + search_items = [ + {"type": "tool_search_call", "id": "tsc_1", "arguments": {"paths": ["weather"]}}, + {"type": "tool_search_output", "id": "tso_1", "tools": [{"type": "function", "name": "get_weather"}]}, + ] + + input_items, _ = LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api( + [ + {"role": "user", "content": "weather in Paris"}, + { + "role": "assistant", + "content": None, + "tool_calls": [tool_call], + "provider_specific_fields": {"tool_search_items": search_items}, + }, + {"role": "tool", "tool_call_id": "call-ns", "content": "sunny"}, + ] + ) + + assert input_items == [ + {"type": "message", "role": "user", "content": [{"type": "input_text", "text": "weather in Paris"}]}, + {"type": "function_call", "call_id": "call-ns", "name": "get_weather", "arguments": "{}"}, + {"type": "function_call_output", "call_id": "call-ns", "output": [{"type": "input_text", "text": "sunny"}]}, + ] + + +def test_streamed_tool_search_turn_replays_with_its_reasoning_item(): + from unittest.mock import Mock + + import litellm + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + LiteLLMResponsesTransformationHandler, + OpenAiResponsesToChatCompletionStreamIterator, + ) + from litellm.types.llms.openai import ResponsesAPIStreamEvents + + function_call = { + "type": "function_call", + "id": "fc_1", + "name": "get_weather", + "call_id": "call-ts", + "arguments": '{"city": "Paris"}', + "namespace": "get_weather", + "status": "completed", + } + tool_search_call = { + "type": "tool_search_call", + "id": "tsc_1", + "arguments": {"paths": ["get_weather"]}, + "call_id": None, + "created_by": None, + "execution": "server", + "status": "completed", + } + tool_search_output = { + "type": "tool_search_output", + "id": "tso_1", + "call_id": None, + "execution": "server", + "status": "completed", + "tools": [{"type": "function", "name": "get_weather", "parameters": {}, "defer_loading": True}], + } + reasoning = {"type": "reasoning", "id": "rs_1", "encrypted_content": "enc", "summary": []} + later_reasoning = {"type": "reasoning", "id": "rs_2", "encrypted_content": "enc_2", "summary": []} + iterator = OpenAiResponsesToChatCompletionStreamIterator(streaming_response=None, sync_stream=True) + chunks = [ + iterator.chunk_parser({"type": ResponsesAPIStreamEvents.OUTPUT_ITEM_ADDED, "item": function_call}), + iterator.chunk_parser( + { + "type": "response.completed", + "response": { + "id": "resp_1", + "status": "completed", + "output": [reasoning, tool_search_call, tool_search_output, later_reasoning, function_call], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + }, + } + ), + ] + user = {"role": "user", "content": "weather in Paris"} + message = litellm.stream_chunk_builder(chunks, messages=[user]).choices[0].message + + request = LiteLLMResponsesTransformationHandler().transform_request( + model="gpt-5.5", + messages=[user, message.model_dump(), {"role": "tool", "tool_call_id": "call-ts", "content": "sunny"}], + optional_params={"tools": [{"type": "tool_search"}]}, + litellm_params={}, + headers={}, + litellm_logging_obj=Mock(), + ) + + assert [(item["type"], item.get("id")) for item in request["input"][1:-1]] == [ + ("reasoning", "rs_1"), + ("tool_search_call", "tsc_1"), + ("tool_search_output", "tso_1"), + ("reasoning", "rs_2"), + ("function_call", None), + ] + assert request["input"][2] == {key: value for key, value in tool_search_call.items() if value is not None} + assert request["input"][5]["namespace"] == "get_weather" + + def test_transform_response_with_reasoning_and_output(): """Test transform_response handles ResponsesAPIResponse with reasoning items and output messages.""" from unittest.mock import Mock @@ -3857,6 +4038,207 @@ def test_transform_response_incomplete_reasoning_only_returns_empty_length_choic assert result.usage.completion_tokens_details.reasoning_tokens == 16 +@pytest.mark.parametrize( + ("status", "incomplete_reason", "finish_reason"), + [("completed", None, "stop"), ("incomplete", "max_output_tokens", "length")], +) +def test_search_only_turn_keeps_tool_search_items_for_the_next_request(status, incomplete_reason, finish_reason): + search_items = [ + {"type": "tool_search_call", "id": "tsc_1", "arguments": {"paths": ["get_weather"]}, "status": "completed"}, + {"type": "tool_search_output", "id": "tso_1", "status": "completed", "tools": [{"name": "get_weather"}]}, + ] + raw_response = _make_incomplete_responses_api_response( + incomplete_reason, + [{"type": "reasoning", "id": "rs_raw", "summary": []}, _make_reasoning_only_output_item(), *search_items], + status=status, + ) + + choice = _call_transform_response(LiteLLMResponsesTransformationHandler(), raw_response).choices[0] + + stored = choice.message.provider_specific_fields["tool_search_items"] + assert [item["id"] for item in stored] == ["rs_raw", "rs_incomplete", "tsc_1", "tso_1"] + assert stored[2:] == search_items + assert [item["id"] for item in choice.message.reasoning_items] == ["rs_raw", "rs_incomplete"] + assert choice.finish_reason == finish_reason + replayed, _ = LiteLLMResponsesTransformationHandler().convert_chat_completion_messages_to_responses_api( + [choice.message.model_dump()], replay_tool_search=True + ) + assert [item["id"] for item in replayed] == ["rs_raw", "rs_incomplete", "tsc_1", "tso_1"] + + +_CLIENT_SEARCH_CALL: Final = { + "type": "tool_search_call", + "id": "tsc_client", + "arguments": {"query": "weather"}, + "call_id": "call_search", + "execution": "client", + "status": "completed", +} + + +def _non_streamed_client_search_turn() -> "ModelResponse": + raw_response = _make_incomplete_responses_api_response(None, [_CLIENT_SEARCH_CALL], status="completed") + return _call_transform_response(LiteLLMResponsesTransformationHandler(), raw_response) + + +def _streamed_client_search_turn() -> "ModelResponse": + from litellm.completion_extras.litellm_responses_transformation.transformation import ( + OpenAiResponsesToChatCompletionStreamIterator, + ) + + iterator: Final = OpenAiResponsesToChatCompletionStreamIterator(streaming_response=None, sync_stream=True) + completed: Final = iterator.chunk_parser( + { + "type": "response.completed", + "response": { + "id": "resp_client_search", + "status": "completed", + "output": [_CLIENT_SEARCH_CALL], + "usage": {"input_tokens": 1, "output_tokens": 1, "total_tokens": 2}, + }, + } + ) + built: Final = litellm.stream_chunk_builder([completed], messages=[{"role": "user", "content": "weather in Paris"}]) + assert isinstance(built, litellm.ModelResponse) + return built + + +@pytest.mark.parametrize( + "client_search_turn", + [_non_streamed_client_search_turn, _streamed_client_search_turn], + ids=["non-streaming", "streaming"], +) +def test_client_search_turn_returns_its_call_and_replays_with_the_client_output(client_search_turn): + choice: Final = client_search_turn().choices[0] + client_output: Final = { + "type": "tool_search_output", + "call_id": "call_search", + "execution": "client", + "status": "completed", + "tools": [{"type": "function", "name": "get_weather", "parameters": {}, "defer_loading": True}], + } + + assert choice.finish_reason == "stop" + assert choice.message.provider_specific_fields["tool_search_items"] == [_CLIENT_SEARCH_CALL] + answered: Final = { + **choice.message.model_dump(), + "provider_specific_fields": {"tool_search_items": [_CLIENT_SEARCH_CALL, client_output]}, + } + request: Final = LiteLLMResponsesTransformationHandler().transform_request( + model="gpt-5.5", + messages=[{"role": "user", "content": "weather in Paris"}, answered], + optional_params={"tools": [{"type": "tool_search", "execution": "client"}]}, + litellm_params={}, + headers={}, + litellm_logging_obj=Mock(), + ) + assert request["input"][1:3] == [_CLIENT_SEARCH_CALL, client_output] + + +def test_reasoning_after_a_tool_search_replays_after_the_search_it_followed(): + from unittest.mock import Mock + + from openai.types.responses import ResponseFunctionToolCall + + search_items = [ + {"type": "tool_search_call", "id": "tsc_1", "arguments": {"paths": ["get_weather"]}, "status": "completed"}, + {"type": "tool_search_output", "id": "tso_1", "status": "completed", "tools": [{"name": "get_weather"}]}, + ] + function_calls = [ + ResponseFunctionToolCall( + type="function_call", id=f"fc_{city}", call_id=f"call_{city}", name="get_weather", arguments="{}" + ) + for city in ("paris", "berlin") + ] + raw_response = _make_incomplete_responses_api_response( + None, + [ + {"type": "reasoning", "id": "rs_search", "summary": [], "encrypted_content": "enc_1"}, + *search_items, + {"type": "reasoning", "id": "rs_call", "summary": [], "encrypted_content": "enc_2"}, + *function_calls, + ], + status="completed", + ) + message = _call_transform_response(LiteLLMResponsesTransformationHandler(), raw_response).choices[0].message + + request = LiteLLMResponsesTransformationHandler().transform_request( + model="gpt-5.5", + messages=[ + {"role": "user", "content": "weather in Paris and Berlin"}, + message.model_dump(), + {"role": "tool", "tool_call_id": "call_paris", "content": "sunny"}, + {"role": "tool", "tool_call_id": "call_berlin", "content": "rainy"}, + ], + optional_params={"tools": [{"type": "tool_search"}]}, + litellm_params={}, + headers={}, + litellm_logging_obj=Mock(), + ) + + assert [(item["type"], item.get("id", item.get("call_id"))) for item in request["input"][1:]] == [ + ("reasoning", "rs_search"), + ("tool_search_call", "tsc_1"), + ("tool_search_output", "tso_1"), + ("reasoning", "rs_call"), + ("function_call", "call_paris"), + ("function_call", "call_berlin"), + ("function_call_output", "call_paris"), + ("function_call_output", "call_berlin"), + ] + + +@pytest.mark.parametrize( + ("tools", "replayed_types"), + [ + ( + [{"type": "tool_search"}], + ["message", "tool_search_call", "tool_search_output", "function_call", "function_call_output"], + ), + ( + [{"type": "function", "function": {"name": "get_weather", "parameters": {"type": "object"}}}], + ["message", "function_call", "function_call_output"], + ), + ], + ids=["tool_search_requested", "tool_search_not_requested"], +) +def test_tool_search_items_replay_only_to_a_request_that_carries_tool_search(tools, replayed_types): + from unittest.mock import Mock + + messages: Final = [ + {"role": "user", "content": "weather in Paris"}, + { + "role": "assistant", + "content": None, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"Paris"}'}, + } + ], + "provider_specific_fields": { + "tool_search_items": [ + {"type": "tool_search_call", "id": "tsc_1", "arguments": {"paths": ["get_weather"]}}, + {"type": "tool_search_output", "id": "tso_1", "tools": [{"type": "function", "name": "get_weather"}]}, + ] + }, + }, + {"role": "tool", "tool_call_id": "call_1", "content": "sunny"}, + ] + + request: Final = LiteLLMResponsesTransformationHandler().transform_request( + model="gpt-5-mini", + messages=messages, + optional_params={"tools": tools}, + litellm_params={}, + headers={}, + litellm_logging_obj=Mock(), + ) + + assert [item["type"] for item in request["input"]] == replayed_types + + def test_transform_response_incomplete_content_filter_maps_finish_reason(): handler = LiteLLMResponsesTransformationHandler() raw_response = _make_incomplete_responses_api_response("content_filter", [_make_reasoning_only_output_item()]) diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py index 8fdc6c8ba47..f918d4cb79e 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_dict_to_response.py @@ -1485,6 +1485,39 @@ def test_convert_to_model_response_object_preserves_provider_specific_fields_fro == web_search_results ) +def test_convert_to_model_response_object_keeps_reasoning_items_from_a_cached_response(): + reasoning_items: Final = [ + {"id": "rs_a", "type": "reasoning", "summary": [], "encrypted_content": "enc-a"}, + {"id": "rs_b", "type": "reasoning", "summary": [], "encrypted_content": "enc-b"}, + ] + result: Final = convert_to_model_response_object( + model_response_object=ModelResponse(), + response_object={ + "id": "chatcmpl-cached", + "object": "chat.completion", + "created": 1728933352, + "model": "gpt-5.6", + "choices": [ + { + "index": 0, + "message": {"role": "assistant", "content": None, "reasoning_items": reasoning_items}, + "finish_reason": "tool_calls", + } + ], + }, + stream=False, + start_time=datetime.now(), + end_time=datetime.now(), + hidden_params=None, + _response_headers=None, + convert_tool_call_to_json_mode=False, + ) + assert isinstance(result, ModelResponse) + choice: Final = result.choices[0] + assert isinstance(choice, Choices) + assert choice.message.reasoning_items == reasoning_items + + def test_convert_to_model_response_object_provider_specific_fields_merges_extra_keys(): """ Test that provider_specific_fields from the response are merged with diff --git a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py index e9d96f057ae..6b017f9ff79 100644 --- a/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py +++ b/tests/unit/litellm_core_utils/llm_response_utils/test_convert_to_streaming_response.py @@ -9,6 +9,8 @@ replay yields multiple word-shaped chunks instead of a single one-shot content frame, restoring per-token cadence on cache hits. """ +from typing import Final + import pytest import litellm @@ -218,6 +220,37 @@ async def test_async_tool_calls_and_function_call_only_on_first_chunk(): assert c.choices[0].delta.function_call is None +@pytest.mark.asyncio +async def test_async_reasoning_items_and_provider_specific_fields_only_on_first_chunk(): + reasoning_items: Final = [{"id": "rs_1", "type": "reasoning", "summary": [], "encrypted_content": "enc-1"}] + provider_specific_fields: Final = { + "tool_search_items": [{"id": "tsc_1", "type": "tool_search_call", "arguments": {"query": "weather"}}] + } + payload: Final = { + **_async_payload(), + "choices": [ + { + "index": 0, + "message": { + "role": "assistant", + "content": "cached content slices here", + "reasoning_items": reasoning_items, + "provider_specific_fields": provider_specific_fields, + }, + "finish_reason": "stop", + } + ], + } + chunks: Final = await _collect_async(payload) + assert len(chunks) > 1 + first_delta: Final = chunks[0].choices[0].delta + assert getattr(first_delta, "reasoning_items", None) == reasoning_items + assert getattr(first_delta, "provider_specific_fields", None) == provider_specific_fields + for c in chunks[1:]: + assert getattr(c.choices[0].delta, "reasoning_items", None) is None + assert getattr(c.choices[0].delta, "provider_specific_fields", None) is None + + @pytest.mark.asyncio async def test_async_logprobs_only_on_first_chunk(): payload = _async_payload(content="cached token logprobs") diff --git a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py index f06620954e8..1fda656060a 100644 --- a/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py +++ b/tests/unit/litellm_core_utils/prompt_templates/test_litellm_core_utils_prompt_templates_common_utils.py @@ -4,7 +4,7 @@ import functools import json import os import sys -from typing import Final +from typing import Final, cast import pytest @@ -31,6 +31,7 @@ from litellm.litellm_core_utils.prompt_templates.common_utils import ( salvage_concatenated_tool_arguments, split_concatenated_json_objects, strip_encrypted_reasoning_from_messages, + strip_litellm_response_only_fields, system_messages_first, update_messages_with_model_file_ids, with_prompt_cache_breakpoint, @@ -2334,3 +2335,30 @@ def test_native_redacted_thinking_data_only_unwraps_the_converse_base64_layer(da @pytest.mark.parametrize("data", [_NATIVE_REDACTED_DATA, _CONVERSE_REDACTED_DATA], ids=["native", "legacy-wrapped"]) def test_converse_redacted_content_wraps_claude_native_data_exactly_once(data: str) -> None: assert converse_redacted_content(data) == _CONVERSE_REDACTED_DATA + + +def test_strip_litellm_response_only_fields_drops_nested_tool_call_provider_fields() -> None: + message: Final = { + "role": "assistant", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "provider_specific_fields": {"namespace": "ns"}, + "function": { + "name": "get_weather", + "arguments": '{"city": "Paris"}', + "provider_specific_fields": {"namespace": "ns"}, + }, + } + ], + } + + stripped: Final = strip_litellm_response_only_fields(message) + + tool_call: Final = cast("dict[str, object]", stripped["tool_calls"][0]) + assert "provider_specific_fields" not in tool_call + function: Final = cast("dict[str, object]", tool_call["function"]) + assert "provider_specific_fields" not in function + assert function["name"] == "get_weather" + assert function["arguments"] == '{"city": "Paris"}' diff --git a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py index 3f73eace59f..f8cd0416b1c 100644 --- a/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py +++ b/tests/unit/litellm_core_utils/test_streaming_chunk_builder_utils.py @@ -2138,3 +2138,45 @@ def test_stream_chunk_builder_joins_openai_audio_output_pieces() -> None: "expires_at": 1767225600, "transcript": "Yes.", } + + +def test_stream_chunk_builder_keeps_reasoning_items_on_text_only_replies(): + from litellm.types.utils import Delta, ModelResponseStream, StreamingChoices + + reasoning_items: Final[list[dict[str, object]]] = [ + {"type": "reasoning", "id": "rs_1", "summary": [], "encrypted_content": "enc_1"} + ] + chunks: Final = [ + ModelResponseStream( + id="chatcmpl-rs", + created=1, + model="gpt-5.5", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason=None, + index=0, + delta=Delta(content="Hello", role="assistant"), + ) + ], + ), + ModelResponseStream( + id="chatcmpl-rs", + created=1, + model="gpt-5.5", + object="chat.completion.chunk", + choices=[ + StreamingChoices( + finish_reason="stop", + index=0, + delta=Delta(content=" world", role="assistant", reasoning_items=reasoning_items), + ) + ], + ), + ] + + response: Final = stream_chunk_builder(chunks=list(chunks)) + + assert response is not None + assert response.choices[0].message.content == "Hello world" + assert response.choices[0].message.reasoning_items == reasoning_items diff --git a/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py b/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py index 16e2221dc84..04b791cf289 100644 --- a/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py +++ b/tests/unit/llms/fireworks_ai/chat/test_fireworks_ai_chat_transformation.py @@ -494,6 +494,46 @@ def test_transform_messages_helper_strips_thinking_blocks_but_keeps_reasoning_co assert out[1]["content"] == "I can help." +def test_tool_search_reply_sent_to_fireworks_drops_litellm_response_fields(): + tool_search_reply = { + "role": "assistant", + "content": "", + "reasoning_items": [{"type": "reasoning", "id": "rs_1", "summary": []}], + "provider_specific_fields": {"tool_search_items": [{"type": "tool_search_call", "id": "ts_1"}]}, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + "provider_specific_fields": {"namespace": "weather"}, + } + ], + } + request = FireworksAIConfig().transform_request( + model="accounts/fireworks/models/gpt-oss-120b", + messages=[ + {"role": "user", "content": "Weather in Paris?"}, + tool_search_reply, + {"role": "tool", "tool_call_id": "call_1", "content": "sunny"}, + ], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert request["messages"][1] == { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + } + ], + } + + # ----------------------------------------------------------------------------- # Regression tests for legacy / OpenAPI $ref defs in tool parameters. # diff --git a/tests/unit/llms/groq/chat/test_groq_chat_transformation.py b/tests/unit/llms/groq/chat/test_groq_chat_transformation.py index 26a6c3972bf..95611e344dc 100644 --- a/tests/unit/llms/groq/chat/test_groq_chat_transformation.py +++ b/tests/unit/llms/groq/chat/test_groq_chat_transformation.py @@ -483,3 +483,43 @@ class TestGroqReasoning: parsed_chunk.choices[0].delta.tool_calls[0]["function"]["name"] == "test_function" ) + + +def test_tool_search_reply_sent_to_groq_drops_litellm_response_fields(): + tool_search_reply = { + "role": "assistant", + "content": "", + "reasoning_items": [{"type": "reasoning", "id": "rs_1", "summary": []}], + "provider_specific_fields": {"tool_search_items": [{"type": "tool_search_call", "id": "ts_1"}]}, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + "provider_specific_fields": {"namespace": "weather"}, + } + ], + } + request = GroqChatConfig().transform_request( + model="openai/gpt-oss-120b", + messages=[ + {"role": "user", "content": "Weather in Paris?"}, + tool_search_reply, + {"role": "tool", "tool_call_id": "call_1", "content": "sunny"}, + ], + optional_params={}, + litellm_params={}, + headers={}, + ) + + assert request["messages"][1] == { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + } + ], + } diff --git a/tests/unit/llms/mistral/test_mistral_chat_transformation.py b/tests/unit/llms/mistral/test_mistral_chat_transformation.py index 23b58c789d2..2e88bccd49b 100644 --- a/tests/unit/llms/mistral/test_mistral_chat_transformation.py +++ b/tests/unit/llms/mistral/test_mistral_chat_transformation.py @@ -830,6 +830,47 @@ class TestMistralStripsOutputOnlyFields: assert assistant_message["content"] == "Follow-up" assert assistant_message["role"] == "assistant" + def test_tool_search_reply_drops_litellm_response_fields(self): + tool_search_reply = { + "role": "assistant", + "content": "", + "reasoning_items": [{"type": "reasoning", "id": "rs_1", "summary": []}], + "provider_specific_fields": {"tool_search_items": [{"type": "tool_search_call", "id": "ts_1"}]}, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + "provider_specific_fields": {"namespace": "weather"}, + } + ], + } + messages = cast( + List[AllMessageValues], + [ + {"role": "user", "content": "Weather in Paris?"}, + tool_search_reply, + {"role": "tool", "tool_call_id": "call_1", "content": "sunny"}, + ], + ) + + result = cast( + List[AllMessageValues], + MistralConfig().transform_messages(messages=messages, model="mistral-medium-3-5"), + ) + + assert result[1] == { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + } + ], + } + def test_non_assistant_messages_are_untouched(self): messages = cast( List[AllMessageValues], diff --git a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py index f0894573df8..fdf3fd62eff 100644 --- a/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/unit/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -94,6 +94,22 @@ class RecordingMaskingGuardrail(MockGuardrail): return await super().apply_guardrail(inputs, request_data, input_type, logging_obj) +class ToolCallRecordingGuardrail(CustomGuardrail): + def __init__(self, guardrail_name: str) -> None: + super().__init__(guardrail_name=guardrail_name) + self.seen_tool_calls: list[object] = [] + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict[str, object], + input_type: Literal["request", "response"], + logging_obj: LiteLLMLoggingObj | None = None, + ) -> GenericGuardrailAPIInputs: + self.seen_tool_calls.extend(inputs.get("tool_calls", [])) + return inputs + + class LastTextDroppingGuardrail(CustomGuardrail): async def apply_guardrail( self, @@ -1173,6 +1189,39 @@ class TestOpenAIResponsesHandlerToolCallExtraction: assert exc_info.value.status_code == 400 assert "Response blocked by guardrail" in str(exc_info.value.detail) + @pytest.mark.asyncio + async def test_namespaced_function_call_reaches_output_guardrails_in_the_unchanged_tool_call_shape(self): + guardrail: Final = ToolCallRecordingGuardrail(guardrail_name="recorder") + response: Final = ResponsesAPIResponse( + id="resp_namespaced", + created_at=1764901066, + model="gpt-5.5", + object="response", + status="completed", + output=[ + ResponseFunctionToolCall( + arguments='{"city":"Paris"}', + call_id="call_namespaced", + name="get_weather", + namespace="weather_tools", + type="function_call", + id="fc_namespaced", + status="completed", + ) + ], + ) + + await OpenAIResponsesHandler().process_output_response(response, guardrail) + + assert guardrail.seen_tool_calls == [ + { + "id": "call_namespaced", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city":"Paris"}'}, + "index": 0, + } + ] + def test_extract_mixed_content_with_text_and_tool_calls(self): """Test extracting both text and tool calls from response""" handler = OpenAIResponsesHandler() diff --git a/tests/unit/llms/together_ai/chat/test_together_ai_chat_transformation.py b/tests/unit/llms/together_ai/chat/test_together_ai_chat_transformation.py index 7a8027d5379..608264e83c3 100644 --- a/tests/unit/llms/together_ai/chat/test_together_ai_chat_transformation.py +++ b/tests/unit/llms/together_ai/chat/test_together_ai_chat_transformation.py @@ -488,6 +488,46 @@ def test_transform_request_keeps_reasoning_content_strips_internal_fields(): _assert_internal_fields_stripped_reasoning_kept(request["messages"]) +def test_tool_search_reply_sent_to_together_drops_litellm_response_fields(): + tool_search_reply = { + "role": "assistant", + "content": "", + "reasoning_items": [{"type": "reasoning", "id": "rs_1", "summary": []}], + "provider_specific_fields": {"tool_search_items": [{"type": "tool_search_call", "id": "ts_1"}]}, + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + "provider_specific_fields": {"namespace": "weather"}, + } + ], + } + request = TogetherAIChatConfig().transform_request( + model="openai/gpt-oss-120b", + messages=[ + {"role": "user", "content": "Weather in Paris?"}, + tool_search_reply, + {"role": "tool", "tool_call_id": "call_1", "content": "sunny"}, + ], + optional_params={}, + litellm_params={"custom_llm_provider": "together_ai"}, + headers={}, + ) + + assert request["messages"][1] == { + "role": "assistant", + "content": "", + "tool_calls": [ + { + "id": "call_1", + "type": "function", + "function": {"name": "get_weather", "arguments": '{"city": "Paris"}'}, + } + ], + } + + async def test_async_transform_request_keeps_reasoning_content_strips_internal_fields(): request = await TogetherAIChatConfig().async_transform_request( model=REASONING_MODEL, diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py index 427fa43ffd5..4e19c7ed328 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -750,6 +750,78 @@ class TestToolPermissionGuardrail: assert new_data["function_call"] == "none" assert new_data["tool_choice"] == "none" + @staticmethod + def _tool_search_turn(*loaded_names: str) -> dict: + return { + "role": "assistant", + "content": None, + "provider_specific_fields": { + "tool_search_items": [ + {"type": "tool_search_call", "call_id": "ts_1", "execution": "client"}, + { + "type": "tool_search_output", + "call_id": "ts_1", + "tools": [{"type": "function", "name": name} for name in loaded_names], + }, + ] + }, + } + + def _replayed_run_shell_request(self) -> dict: + return { + "tools": [{"type": "function", "function": {"name": "get_weather"}}, {"type": "tool_search"}], + "messages": [ + {"role": "user", "content": "check the weather, then clean up the disk"}, + self._tool_search_turn("get_weather"), + self._tool_search_turn("run_shell"), + ], + } + + @staticmethod + def _deny_run_shell(on_disallowed_action: str) -> ToolPermissionGuardrail: + return ToolPermissionGuardrail( + guardrail_name="test-tool-permission", + rules=[{"id": "deny_shell", "tool_name": r"^run_shell$", "decision": "deny"}], + default_action="allow", + on_disallowed_action=on_disallowed_action, + ) + + @pytest.mark.asyncio + async def test_async_pre_call_hook_blocks_tool_loaded_by_replayed_tool_search(self): + guardrail = self._deny_run_shell("block") + + with patch.object(guardrail, "should_run_guardrail", return_value=True): + with pytest.raises(HTTPException) as excinfo: + await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(default_in_memory_ttl=1), + data=self._replayed_run_shell_request(), + call_type="completion", + ) + + assert excinfo.value.status_code == 400 + assert "run_shell" in json.dumps(excinfo.value.detail) + + @pytest.mark.asyncio + async def test_async_pre_call_hook_rewrite_drops_replayed_tool_search_loading_denied_tool(self): + guardrail = self._deny_run_shell("rewrite") + data = self._replayed_run_shell_request() + + with patch.object(guardrail, "should_run_guardrail", return_value=True): + new_data = await guardrail.async_pre_call_hook( + user_api_key_dict=UserAPIKeyAuth(), + cache=DualCache(default_in_memory_ttl=1), + data=data, + call_type="completion", + ) + + assert isinstance(new_data, dict) + user_turn, weather_search, shell_search = new_data["messages"] + assert user_turn == {"role": "user", "content": "check the weather, then clean up the disk"} + assert weather_search == self._tool_search_turn("get_weather") + assert shell_search == {"role": "assistant", "content": None, "provider_specific_fields": {}} + assert [tool.get("function", {}).get("name") for tool in new_data["tools"]] == ["get_weather", None] + @pytest.mark.asyncio async def test_async_post_call_streaming_iterator_hook_plain_text_yields_chunks( self, diff --git a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py index 9113ac5015f..f40465e8d7b 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_policy_guardrail.py @@ -186,3 +186,39 @@ async def test_response_tool_calls_as_objects(guardrail): await guardrail.apply_guardrail( inputs=inputs, request_data={}, input_type="response" ) + + +@pytest.mark.asyncio +async def test_blocked_tool_loaded_by_a_replayed_tool_search_raises_next_to_allowed_tools(guardrail): + request_data = { + "tools": [{"type": "function", "function": {"name": "get_weather"}}, {"type": "tool_search"}], + "messages": [ + {"role": "user", "content": "clean up the disk"}, + { + "role": "assistant", + "content": None, + "provider_specific_fields": { + "tool_search_items": [ + {"type": "tool_search_call", "call_id": "ts_1", "execution": "client"}, + { + "type": "tool_search_output", + "call_id": "ts_1", + "tools": [{"type": "function", "name": "run_shell", "parameters": {"type": "object"}}], + }, + ] + }, + }, + ], + } + with patch( + "litellm.proxy.db.tool_registry_writer.get_tool_policy_registry", + return_value=_registry_mock({"get_weather": "untrusted", "run_shell": "blocked"}), + ): + with pytest.raises(HTTPException) as exc_info: + await guardrail.apply_guardrail( + inputs=_tool_request_inputs(["get_weather"]), + request_data=request_data, + input_type="request", + ) + assert exc_info.value.status_code == 400 + assert exc_info.value.detail["blocked_tools"] == ["run_shell"] diff --git a/tests/unit/proxy/test_tools_allowlist_enforcement.py b/tests/unit/proxy/test_tools_allowlist_enforcement.py index 3e9c7c14b95..e855ef47dd5 100644 --- a/tests/unit/proxy/test_tools_allowlist_enforcement.py +++ b/tests/unit/proxy/test_tools_allowlist_enforcement.py @@ -227,6 +227,89 @@ class TestCheckToolsAllowlist: assert exc_info.value.type == ProxyErrorTypes.tool_access_denied assert "restricted_tool" in str(exc_info.value.message) + @pytest.mark.parametrize( + "loaded_tool", + [ + {"type": "function", "name": "run_shell", "parameters": {"type": "object"}}, + { + "type": "namespace", + "name": "ops", + "tools": [{"type": "function", "name": "run_shell", "parameters": {"type": "object"}}], + }, + ], + ids=["function", "namespace_member"], + ) + @pytest.mark.asyncio + async def test_disallowed_tool_loaded_by_a_replayed_tool_search_raises_on_chat_route(self, loaded_tool): + token = _token(metadata={"allowed_tools": ["get_weather"]}) + body = { + "tools": [{"type": "function", "function": {"name": "get_weather"}}, {"type": "tool_search"}], + "messages": [ + {"role": "user", "content": "clean up the disk"}, + { + "role": "assistant", + "content": None, + "provider_specific_fields": { + "tool_search_items": [ + {"type": "tool_search_call", "call_id": "ts_1", "execution": "client"}, + { + "type": "tool_search_output", + "call_id": "ts_1", + "tools": [{"type": "function", "name": "get_weather"}, loaded_tool], + }, + ] + }, + }, + ], + } + with pytest.raises(ProxyException) as exc_info: + await check_tools_allowlist( + request_body=body, + valid_token=token, + team_object=None, + route="/v1/chat/completions", + ) + assert exc_info.value.type == ProxyErrorTypes.tool_access_denied + assert "run_shell" in str(exc_info.value.message) + assert "get_weather" not in str(exc_info.value.message) + + @pytest.mark.parametrize( + "search_tools", + [{"extra_body": {"tools": [{"type": "tool_search"}]}}, {}], + ids=["extra_body_tool_search", "deployment_adds_tool_search"], + ) + @pytest.mark.asyncio + async def test_disallowed_tool_loaded_by_a_replayed_tool_search_raises_without_tool_search_in_tools( + self, search_tools + ): + token = _token(metadata={"allowed_tools": ["get_weather"]}) + body = { + "tools": [{"type": "function", "function": {"name": "get_weather"}}], + **search_tools, + "messages": [ + {"role": "user", "content": "clean up the disk"}, + { + "role": "assistant", + "content": None, + "provider_specific_fields": { + "tool_search_items": [ + {"type": "tool_search_call", "call_id": "ts_1", "execution": "client"}, + {"type": "tool_search_output", "call_id": "ts_1", "tools": [{"name": "run_shell"}]}, + ] + }, + }, + ], + } + with pytest.raises(ProxyException) as exc_info: + await check_tools_allowlist( + request_body=body, + valid_token=token, + team_object=None, + route="/v1/chat/completions", + ) + assert exc_info.value.type == ProxyErrorTypes.tool_access_denied + assert "run_shell" in str(exc_info.value.message) + @pytest.mark.asyncio async def test_team_allowlist_used_when_key_empty(self): token = _token( diff --git a/tests/unit/test_tool_loop.py b/tests/unit/test_tool_loop.py index 5c14e644d43..e3ebfe7258f 100644 --- a/tests/unit/test_tool_loop.py +++ b/tests/unit/test_tool_loop.py @@ -361,6 +361,66 @@ def test_responses_bridge_replays_reasoning_items_across_rounds(respx_mock: resp assert reasoning_index < function_call_index, f"reasoning item must precede function_call: {item_types}" +def test_responses_bridge_replays_tool_search_items_across_rounds(respx_mock: respx.MockRouter) -> None: + round_one: Final = _responses_payload( + "resp_1", + [ + {"type": "reasoning", "id": "rs_1", "summary": [], "encrypted_content": "enc_1"}, + {"type": "tool_search_call", "id": "tsc_1", "execution": "server", "arguments": {"paths": ["weather"]}}, + { + "type": "tool_search_output", + "id": "tso_1", + "execution": "server", + "tools": [{"type": "namespace", "name": "weather", "tools": [{"type": "function", "name": "get_weather"}]}], + }, + { + "type": "function_call", + "id": "fc_1", + "call_id": "call_1", + "name": "get_weather", + "namespace": "weather", + "arguments": '{"city": "Paris"}', + "status": "completed", + }, + ], + ) + round_two: Final = _responses_payload( + "resp_2", + [ + { + "type": "message", + "id": "msg_1", + "status": "completed", + "role": "assistant", + "content": [{"type": "output_text", "text": "Paris is 72F", "annotations": []}], + } + ], + ) + respx_mock.post("https://api.openai.com/v1/responses").mock( + side_effect=[httpx.Response(200, json=round_one), httpx.Response(200, json=round_two)] + ) + + litellm.run_tool_loop( + model="openai/responses/gpt-5.5", + messages=[{"role": "user", "content": "weather in Paris?"}], + tools=({"type": "tool_search"}, *WEATHER_TOOLS), + execute_tool=_tool_result, + api_key="sk-test", + ) + + second_input: Final = _request_bodies(respx_mock)[1]["input"] + replayed: Final = [ + (item["type"], item.get("id"), item.get("namespace")) for item in second_input if item["type"] != "message" + ] + assert replayed == [ + ("reasoning", "rs_1", None), + ("tool_search_call", "tsc_1", None), + ("tool_search_output", "tso_1", None), + ("function_call", None, "weather"), + ("function_call_output", None, None), + ] + + async def test_arun_tool_loop_two_rounds(respx_mock: respx.MockRouter, monkeypatch: pytest.MonkeyPatch) -> None: from litellm.llms.custom_httpx.http_handler import AsyncHTTPHandler