fix(responses): keep namespace custom tools through guardrail merges and Mantle params identity

litellm/llms/bedrock_mantle/responses/transformation.py

@@ -237,10 +237,15 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
         )
         hoisted: Final = hoist_additional_tools(input, params.get("tools"))
         normalized_input: Final = self._normalize_codex_input_items(hoisted.input)
+        request_params: Final = (
+            self._params_with_hoisted_tools(params, hoisted)
+            if hoisted.hoisted
+            else response_api_optional_request_params
+        )
         return super().transform_responses_api_request(
             model=model,
             input=normalized_input,
-            response_api_optional_request_params=self._params_with_hoisted_tools(params, hoisted),
+            response_api_optional_request_params=request_params,
             litellm_params=litellm_params,
             headers=headers,
         )
@@ -249,8 +254,6 @@ class BedrockMantleResponsesAPIConfig(BedrockMantleAuthMixin, OpenAIResponsesAPI
     def _params_with_hoisted_tools(
         cls, params: Mapping[str, object], hoisted: HoistedAdditionalTools
     ) -> dict[str, object]:
-        if not hoisted.hoisted:
-            return dict(params)
         supported_tools: Final = cls._filter_unsupported_tools(list(hoisted.tools))
         if supported_tools:
             return {**params, "tools": supported_tools}

litellm/llms/openai/responses/guardrail_translation/tool_merge.py

@@ -6,8 +6,10 @@ from typing import Final, TypeAlias
 from pydantic import BaseModel, TypeAdapter, ValidationError
 
 from litellm._logging import verbose_logger
+from litellm.responses.litellm_completion_transformation.custom_tools import custom_tool_grammar_suffix
 from litellm.responses.litellm_completion_transformation.transformation import (
     NAMESPACE_DESCRIPTION_SEPARATOR,
+    NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS,
     LiteLLMCompletionResponsesConfig,
 )
 
@@ -34,8 +36,8 @@ def _validated_tools(values: Iterable[object]) -> tuple[Tool, ...]:
     return tuple(tool for tool in validated if tool is not None)
 
 
-def _is_function(tool: Tool) -> bool:
-    return tool.get("type") == "function"
+def _has_chat_tool(member: Tool) -> bool:
+    return member.get("type") in NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS
 
 
 def _chat_tool_key(tool: Tool) -> str:
@@ -67,18 +69,19 @@ def _function_fields(tool: Tool) -> Tool:
     return function if function is not None else MappingProxyType({})
 
 
-def _without_namespace_prefix(key: str, value: object, prefix: str) -> object:
-    if key != "description" or not isinstance(value, str) or not value.startswith(prefix):
+def _member_description(key: str, value: object, prefix: str, suffix: str) -> object:
+    if key != "description" or not isinstance(value, str):
         return value
-    return value[len(prefix) :]
+    return value.removeprefix(prefix).removesuffix(suffix)
 
 
 def _rebuilt_member(member: Tool, flattened: Tool, guardrailed: Tool, namespace_description: str) -> Tool:
     flattened_function: Final = _function_fields(flattened)
     prefix: Final = f"{namespace_description}{NAMESPACE_DESCRIPTION_SEPARATOR}" if namespace_description else ""
+    suffix: Final = custom_tool_grammar_suffix(member.get("format")) if member.get("type") == "custom" else ""
     changed_function: Final = MappingProxyType(
         {
-            key: _without_namespace_prefix(key, value, prefix)
+            key: _member_description(key, value, prefix, suffix)
             for key, value in _function_fields(guardrailed).items()
             if flattened_function.get(key) != value
         }
@@ -93,8 +96,8 @@ def _rebuilt_member(member: Tool, flattened: Tool, guardrailed: Tool, namespace_
     return {**member, **changed_extras, **changed_function}  # mutable-ok: json.dumps rejects MappingProxyType
 
 
-def _rebuilt_function_members(
-    function_members: Sequence[Tool],
+def _rebuilt_flattened_members(
+    flattened_members: Sequence[Tool],
     flattened_group: Sequence[Tool],
     group_keys: Sequence[IndexedKey],
     guardrailed_by_key: Mapping[IndexedKey, Tool],
@@ -106,7 +109,7 @@ def _rebuilt_function_members(
         else member
         if guardrailed_by_key[key] == flattened
         else _rebuilt_member(member, flattened, guardrailed_by_key[key], namespace_description)
-        for member, flattened, key in zip(function_members, flattened_group, group_keys)
+        for member, flattened, key in zip(flattened_members, flattened_group, group_keys)
     )
 
 
@@ -118,9 +121,9 @@ def _rebuilt_namespace(
     guardrailed_by_key: Mapping[IndexedKey, Tool],
 ) -> tuple[Tool, ...]:
     namespace_description: Final = str(original.get("description") or "")
-    rebuilt_functions: Final = iter(
-        _rebuilt_function_members(
-            tuple(member for member in members if _is_function(member)),
+    rebuilt_flattened: Final = iter(
+        _rebuilt_flattened_members(
+            tuple(member for member in members if _has_chat_tool(member)),
             flattened_group,
             group_keys,
             guardrailed_by_key,
@@ -129,7 +132,7 @@ def _rebuilt_namespace(
     )
     rebuilt_members: Final = tuple(
         rebuilt
-        for rebuilt in (next(rebuilt_functions) if _is_function(member) else member for member in members)
+        for rebuilt in (next(rebuilt_flattened) if _has_chat_tool(member) else member for member in members)
         if rebuilt is not None
     )
     if not rebuilt_members:
@@ -149,7 +152,7 @@ def _merged_original(
     if guardrailed_group == tuple(flattened_group):
         return (original,)
     members: Final = _namespace_members(original) if original.get("type") == "namespace" else ()
-    if members and sum(map(_is_function, members)) == len(flattened_group):
+    if members and sum(map(_has_chat_tool, members)) == len(flattened_group):
         return _rebuilt_namespace(original, members, flattened_group, group_keys, guardrailed_by_key)
     if not guardrailed_group:
         return ()

litellm/responses/litellm_completion_transformation/custom_tools.py

@@ -45,21 +45,32 @@ class _ToolNameFields(BaseModel):
     tools: tuple[object, ...] = ()
 
 
-def _custom_tool_names_of(tool: object) -> tuple[str, ...]:
+def _tool_name_fields_of(tool: object) -> _ToolNameFields | None:
     try:
-        parsed: Final = _ToolNameFields.model_validate(tool)
+        return _ToolNameFields.model_validate(tool)
     except ValidationError:
+        return None
+
+
+def _custom_tool_name_of(tool: object) -> str | None:
+    parsed: Final = _tool_name_fields_of(tool)
+    if parsed is None or parsed.type != "custom" or not parsed.name:
+        return None
+    return parsed.name
+
+
+def _nested_tools_of(tool: object) -> tuple[object, ...]:
+    parsed: Final = _tool_name_fields_of(tool)
+    if parsed is None or parsed.type != "namespace":
         return ()
-    if parsed.type == "custom":
-        return (parsed.name,) if parsed.name else ()
-    if parsed.type != "namespace":
-        return ()
-    return tuple(name for nested_tool in parsed.tools for name in _custom_tool_names_of(nested_tool))
+    return parsed.tools
 
 
 def extract_custom_tool_names(tools: Sequence[object] | None) -> set[str]:
-    """Extract names of tools defined as ``type: "custom"``, at the top level or inside a ``namespace`` tool."""
-    return {name for tool in tools or () for name in _custom_tool_names_of(tool)}
+    """Extract names of ``type: "custom"`` tools, at the top level or one level inside a ``namespace`` tool."""
+    top_level: Final = tuple(tools or ())
+    nested: Final = tuple(nested_tool for tool in top_level for nested_tool in _nested_tools_of(tool))
+    return {name for tool in (*top_level, *nested) if (name := _custom_tool_name_of(tool)) is not None}
 
 
 def is_custom_tool_call(tool_name: str, custom_tool_names: set[str]) -> bool:
@@ -155,7 +166,7 @@ def validated_allowed_callers(value: object) -> list[str] | None:
         raise ValueError("allowed_callers must be a list of strings") from exc
 
 
-def _grammar_suffix(fmt: object) -> str:
+def custom_tool_grammar_suffix(fmt: object) -> str:
     try:
         parsed: Final = _CustomToolFormat.model_validate(fmt)
     except ValidationError:
@@ -179,7 +190,9 @@ def convert_custom_tool_to_function_tool(tool: Mapping[str, object]) -> ChatComp
     raw_name: Final = tool.get("name")
     name: Final = raw_name if isinstance(raw_name, str) else ""
     raw_description: Final = tool.get("description")
-    description = (raw_description if isinstance(raw_description, str) else "") + _grammar_suffix(tool.get("format"))
+    description: Final = (raw_description if isinstance(raw_description, str) else "") + custom_tool_grammar_suffix(
+        tool.get("format")
+    )
     allowed_callers: Final = validated_allowed_callers(tool.get("allowed_callers"))
     function_chunk: Final = ChatCompletionToolParamFunctionChunk(
         name=name,

litellm/responses/litellm_completion_transformation/transformation.py

@@ -109,6 +109,7 @@ NamespaceTool: TypeAlias = Mapping[str, object]
 ResponseTools: TypeAlias = Sequence[Mapping[str, object]] | None
 ChatToolParam: TypeAlias = ChatCompletionToolParam | OpenAIMcpServerTool
 NAMESPACE_DESCRIPTION_SEPARATOR: Final = "\n\n"
+NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS: Final = frozenset({"function", "custom"})
 
 
 @dataclass(frozen=True, slots=True)
@@ -1891,7 +1892,7 @@ class LiteLLMCompletionResponsesConfig:
         nested: bool,
     ) -> ChatCompletionToolParam | None:
         tool_type: Final = namespace_tool.get("type")
-        if nested and tool_type not in ("function", "custom"):
+        if nested and tool_type not in NAMESPACE_MEMBER_TYPES_WITH_CHAT_TOOLS:
             return None
 
         raw_description: Final = str(namespace_tool.get("description") or "")

tests/test_litellm/llms/bedrock_mantle/test_bedrock_mantle_responses_transformation.py

@@ -773,6 +773,12 @@ class TestBedrockMantleCodexAdditionalTools:
         assert body["input"] == codex_agentic_items
         assert "tools" not in body
 
+    def test_input_without_additional_tools_sanitizes_tools_on_the_caller_params_object(self):
+        params = {"tools": [{"type": "function", "name": "wait", "parameters": '{"type": "object"}'}]}
+        body = self._transform(input=[self._USER_MESSAGE], params=params)
+        assert body["tools"][0]["parameters"] == {"type": "object"}
+        assert params["tools"][0]["parameters"] == {"type": "object"}
+
     def test_malformed_additional_tools_item_without_tools_list_is_stripped(self):
         body = self._transform(
             input=[

tests/test_litellm/llms/openai/responses/test_openai_responses_guardrail_tool_merge.py

@@ -133,7 +133,20 @@ def test_namespace_keeps_a_non_function_member_when_a_function_member_is_edited(
     assert merged[0]["tools"][1] == custom_member
 
 
-def test_namespace_keeps_its_non_function_members_when_every_function_member_is_dropped():
+def test_namespace_keeps_its_custom_member_when_every_function_member_is_dropped():
+    custom_member = {"type": "custom", "name": "grep", "description": "Grep", "format": {"type": "text"}}
+    original = [
+        {"type": "namespace", "name": "ns", "description": "NS", "tools": [_function("read"), custom_member]},
+        _function("a"),
+    ]
+    groups = _groups(original)
+
+    merged = merge_guardrailed_tools(original, groups, [groups[0][1], groups[1][0]])
+
+    assert list(merged) == [{"type": "namespace", "name": "ns", "description": "NS", "tools": [custom_member]}, _function("a")]
+
+
+def test_namespace_custom_member_is_dropped_when_the_guardrail_drops_its_chat_form():
     custom_member = {"type": "custom", "name": "grep", "description": "Grep", "format": {"type": "text"}}
     original = [
         {"type": "namespace", "name": "ns", "description": "NS", "tools": [_function("read"), custom_member]},
@@ -143,7 +156,22 @@ def test_namespace_keeps_its_non_function_members_when_every_function_member_is_
 
     merged = merge_guardrailed_tools(original, groups, [groups[1][0]])
 
-    assert list(merged) == [{"type": "namespace", "name": "ns", "description": "NS", "tools": [custom_member]}, _function("a")]
+    assert list(merged) == [_function("a")]
+
+
+def test_custom_member_description_edit_lands_without_the_namespace_prefix_or_grammar_block():
+    grammar = {"type": "grammar", "syntax": "lark", "definition": "start: X"}
+    custom_member = {"type": "custom", "name": "exec", "description": "Run a command", "format": grammar}
+    original = [{"type": "namespace", "name": "shell", "description": "Shell", "tools": [custom_member]}]
+    groups = _groups(original)
+    assert groups[0][0]["function"]["description"] == "Shell\n\nRun a command\n\nFormat:\n```lark\nstart: X\n```"
+    edited = copy.deepcopy(_flat(groups))
+    edited[0]["function"]["description"] = "Shell\n\nRun a command (guarded)\n\nFormat:\n```lark\nstart: X\n```"
+
+    merged = merge_guardrailed_tools(original, groups, edited)
+
+    guarded_member = {**custom_member, "description": "Run a command (guarded)"}
+    assert list(merged) == [{"type": "namespace", "name": "shell", "description": "Shell", "tools": [guarded_member]}]
 
 
 def test_member_extras_edited_by_the_guardrail_land_on_that_member():