diff --git a/basedpyright-code-budget.json b/basedpyright-code-budget.json index da788bf1ce3..9997e5dc0c6 100644 --- a/basedpyright-code-budget.json +++ b/basedpyright-code-budget.json @@ -123,7 +123,7 @@ "limit": 5 }, "reportUnnecessaryIsInstance": { - "limit": 826 + "limit": 823 }, "reportUntypedBaseClass": { "limit": 0 diff --git a/litellm/llms/openai/responses/guardrail_translation/handler.py b/litellm/llms/openai/responses/guardrail_translation/handler.py index 1530c154e93..caa20207a15 100644 --- a/litellm/llms/openai/responses/guardrail_translation/handler.py +++ b/litellm/llms/openai/responses/guardrail_translation/handler.py @@ -37,7 +37,7 @@ from typing import TYPE_CHECKING, Any, Final, Union, cast from openai.types.responses.response_function_tool_call import ResponseFunctionToolCall from openai.types.responses.tool_param import FunctionToolParam -from pydantic import BaseModel, TypeAdapter +from pydantic import BaseModel, TypeAdapter, ValidationError from typing_extensions import ReadOnly, TypedDict from litellm._logging import verbose_proxy_logger @@ -110,6 +110,88 @@ class ResponsesStreamChunk(TypedDict, total=False): content_index: ReadOnly[int] +_TOOL_TYPES_NOT_SENT_TO_GUARDRAIL: Final = frozenset( + {"web_search", "web_search_preview", "computer_use", "image_generation", "shell"} +) + + +_NAMESPACE_MEMBERS: Final = TypeAdapter(tuple[Mapping[str, object], ...]) + + +def _namespace_members(tool: Mapping[str, object]) -> tuple[Mapping[str, object], ...] | None: + """``None`` when the namespace has no member list and is flattened to one function named after it.""" + raw_members: Final = tool.get("tools") + if not isinstance(raw_members, Sequence) or isinstance(raw_members, (str, bytes)): + return None + try: + return _NAMESPACE_MEMBERS.validate_python(raw_members) + except ValidationError: + return () + + +def _is_function_member(member: Mapping[str, object]) -> bool: + return member.get("type") == "function" + + +def _qualified_member_name(namespace: str, member: Mapping[str, object]) -> str: + return f"{namespace}__{member.get('name') or ''}" + + +def _namespace_function_names(tool: Mapping[str, object]) -> tuple[str, ...]: + namespace: Final = str(tool.get("name") or "") + members: Final = _namespace_members(tool) + if members is None: + return (namespace,) + return tuple(_qualified_member_name(namespace, member) for member in members if _is_function_member(member)) + + +def _flattened_function_names(tools: Sequence[Mapping[str, object]]) -> tuple[str, ...]: + """Names the guardrail sees for ``tools`` once flattened to Chat Completions format.""" + top_level: Final = tuple( + str(tool.get("name") or "") for tool in tools if tool.get("type") in ("function", "custom") + ) + nested: Final = tuple( + name for tool in tools if tool.get("type") == "namespace" for name in _namespace_function_names(tool) + ) + return top_level + nested + + +def _merge_namespace_tool( + tool: Mapping[str, object], remapped_functions: Mapping[str, Mapping[str, object]] +) -> Mapping[str, object] | None: + namespace: Final = str(tool.get("name") or "") + members: Final = _namespace_members(tool) + if members is None: + return tool if namespace in remapped_functions else None + surviving: Final = tuple( + member + for member in members + if not _is_function_member(member) or _qualified_member_name(namespace, member) in remapped_functions + ) + if len(surviving) == len(members): + return tool + if not any(_is_function_member(member) for member in surviving): + return None + return {**tool, "tools": [*surviving]} # mutable-ok: request tools are JSON dicts sent to the provider + + +def _merge_original_tool( + tool: Mapping[str, object], + remapped_functions: Mapping[str, Mapping[str, object]], + remapped_passthrough: Sequence[Mapping[str, object]], +) -> Mapping[str, object] | None: + tool_type: Final = tool.get("type") + if tool_type in _TOOL_TYPES_NOT_SENT_TO_GUARDRAIL: + return tool + if tool_type == "function": + return remapped_functions.get(str(tool.get("name") or "")) + if tool_type == "custom": + return tool if str(tool.get("name") or "") in remapped_functions else None + if tool_type == "namespace": + return _merge_namespace_tool(tool, remapped_functions) + return tool if tool in remapped_passthrough else None + + def _next_stream_sequence_number(responses_so_far: Sequence[Any] | None) -> int: sequence_numbers: Final = ( item.get("sequence_number") if isinstance(item, dict) else getattr(item, "sequence_number", None) @@ -297,32 +379,37 @@ class OpenAIResponsesHandler(BaseTranslation): def _merge_tools_after_guardrail( self, - original_tools: list[dict[str, object]], + original_tools: Sequence[Mapping[str, object]], remapped: list[dict[str, object]], - ) -> list[dict[str, object]]: + ) -> Sequence[Mapping[str, object]]: """ - Merge remapped guardrailed tools with original tools that were not sent - to the guardrail (e.g. web_search, web_search_preview), preserving order. - Tools a guardrail appended (``remapped`` longer than ``original_tools``) - have no original slot and are kept so an injected tool is not dropped. + Rebuild the Responses tool list from ``original_tools`` and apply only the + guardrail's delta: tools it dropped are removed (namespace members + individually), tools it appended are kept at the end, and tools that never + reach the guardrail (web_search, computer_use, ...) stay untouched. + Namespace tools keep their wrapper and member names instead of the + ``namespace__tool`` functions the guardrail saw. """ if not original_tools: return remapped - result: Final[list[dict[str, object]]] = [] - j = 0 - for tool in original_tools: - if isinstance(tool, dict) and tool.get("type") in ( - "web_search", - "web_search_preview", - ): - result.append(tool) - else: - if j < len(remapped): - result.append(remapped[j]) - j += 1 - # Keep guardrail-appended tools that matched no original slot above. - result.extend(remapped[j:]) - return result + remapped_functions: Final = MappingProxyType( + {str(tool.get("name") or ""): tool for tool in remapped if tool.get("type") == "function"} + ) + remapped_passthrough: Final = tuple(tool for tool in remapped if tool.get("type") != "function") + original_function_names: Final = frozenset(_flattened_function_names(original_tools)) + kept: Final = tuple( + merged + for tool in original_tools + for merged in (_merge_original_tool(tool, remapped_functions, remapped_passthrough),) + if merged is not None + ) + appended: Final = tuple( + tool + for tool in remapped + if (tool.get("type") == "function" and str(tool.get("name") or "") not in original_function_names) + or (tool.get("type") != "function" and tool not in original_tools) + ) + return [*kept, *appended] # mutable-ok: request tools are JSON dicts sent to the provider def _apply_guardrailed_tools_to_data( self, diff --git a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py index 315b6948bd8..7e624f39c53 100644 --- a/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py +++ b/tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_handler.py @@ -1323,6 +1323,125 @@ class TestOpenAIResponsesHandlerToolInjection: assert "injected_tool" in names +class ToolDroppingGuardrail(CustomGuardrail): + """Guardrail that removes one flattened function tool by name.""" + + def __init__(self, drop: str, **kwargs): + super().__init__(**kwargs) + self.drop = drop + + async def apply_guardrail( + self, + inputs: GenericGuardrailAPIInputs, + request_data: dict, + input_type: Literal["request", "response"], + logging_obj: Optional[Any] = None, + ) -> GenericGuardrailAPIInputs: + return { + **inputs, + "tools": [t for t in inputs.get("tools") or [] if t["function"]["name"] != self.drop], + } + + +CODEX_NAMESPACE_TOOL = { + "type": "namespace", + "name": "mcp__confluence", + "description": "Tools from the confluence MCP server", + "tools": [ + { + "type": "function", + "name": "confluence_get_page", + "description": "Get a page", + "parameters": {"type": "object", "properties": {"page_id": {"type": "string"}}}, + }, + { + "type": "function", + "name": "confluence_search", + "description": "Search", + "parameters": {"type": "object", "properties": {"query": {"type": "string"}}}, + }, + ], +} + + +class TestOpenAIResponsesHandlerNamespaceTools: + """Regression for #39183: a guardrail on /v1/responses must not rewrite Codex + namespace MCP tools into flat ``namespace__tool`` functions.""" + + @pytest.mark.asyncio + async def test_namespace_tool_survives_tool_appending_guardrail(self): + handler = OpenAIResponsesHandler() + data = { + "input": [{"role": "user", "content": "fetch page 1", "type": "message"}], + "tools": [ + {"type": "function", "name": "shell", "parameters": {"type": "object", "properties": {}}}, + CODEX_NAMESPACE_TOOL, + ], + "model": "gpt-5.3-codex", + } + + result = await handler.process_input_messages(data, ToolAppendingGuardrail(guardrail_name="test")) + + assert [t["type"] for t in result["tools"]] == ["function", "namespace", "function"] + assert result["tools"][0]["name"] == "shell" + assert result["tools"][1] == CODEX_NAMESPACE_TOOL + assert result["tools"][2]["name"] == "injected_tool" + assert not any(t["name"].startswith("mcp__confluence__") for t in result["tools"]) + + @pytest.mark.asyncio + async def test_unchanged_guardrail_leaves_custom_mcp_and_memberless_namespace_tools_as_is(self): + handler = OpenAIResponsesHandler() + tools = [ + {"type": "custom", "name": "apply_patch", "description": "Apply a patch"}, + {"type": "mcp", "server_label": "docs", "server_url": "https://mcp.example.com"}, + {"type": "namespace", "name": "mcp__flat", "description": "Flattened to one function"}, + ] + data = { + "input": [{"role": "user", "content": "fetch page 1", "type": "message"}], + "tools": [dict(t) for t in tools], + "model": "gpt-5.3-codex", + } + + result = await handler.process_input_messages(data, MockPassThroughGuardrail(guardrail_name="test")) + + assert result["tools"] == tools + + @pytest.mark.asyncio + async def test_guardrail_can_still_drop_a_single_namespace_member(self): + handler = OpenAIResponsesHandler() + data = { + "input": [{"role": "user", "content": "fetch page 1", "type": "message"}], + "tools": [CODEX_NAMESPACE_TOOL], + "model": "gpt-5.3-codex", + } + + result = await handler.process_input_messages( + data, ToolDroppingGuardrail(drop="mcp__confluence__confluence_search", guardrail_name="test") + ) + + assert len(result["tools"]) == 1 + namespace_tool = result["tools"][0] + assert namespace_tool["type"] == "namespace" + assert namespace_tool["name"] == "mcp__confluence" + assert [t["name"] for t in namespace_tool["tools"]] == ["confluence_get_page"] + + @pytest.mark.asyncio + async def test_namespace_tool_removed_when_guardrail_drops_every_member(self): + handler = OpenAIResponsesHandler() + single_member = {**CODEX_NAMESPACE_TOOL, "tools": CODEX_NAMESPACE_TOOL["tools"][:1]} + data = { + "input": [{"role": "user", "content": "fetch page 1", "type": "message"}], + "tools": [{"type": "web_search_preview"}, single_member], + "model": "gpt-5.3-codex", + } + + result = await handler.process_input_messages( + data, ToolDroppingGuardrail(drop="mcp__confluence__confluence_get_page", guardrail_name="test") + ) + + assert result["tools"] == [{"type": "web_search_preview"}] + + class TestBuildBlockSseChunks: """build_block_sse_chunks turns a streaming ModifyResponseException into 200 SSE events""" diff --git a/type-discipline-budget.json b/type-discipline-budget.json index 52cb9628252..07bc874880f 100644 --- a/type-discipline-budget.json +++ b/type-discipline-budget.json @@ -1,9 +1,9 @@ { "LIT001": { - "limit": 22364 + "limit": 22352 }, "LIT002": { - "limit": 26777 + "limit": 26774 }, "LIT003": { "limit": 269 @@ -27,7 +27,7 @@ "limit": 0 }, "LIT010": { - "limit": 16507 + "limit": 16504 }, "LIT011": { "limit": 5535