diff --git a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py index dac5cdf36cb..8fe3effbe7b 100644 --- a/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py +++ b/litellm/llms/vertex_ai/gemini/vertex_and_google_ai_studio_gemini.py @@ -563,7 +563,19 @@ def _namespace_members_or_tool(tool: JsonValue) -> tuple[JsonValue, ...]: if not isinstance(tool, dict) or tool.get("type") != "namespace": return (tool,) members: Final = tool.get("tools") - return tuple(map(_as_chat_tool, members)) if isinstance(members, list) else () + return _members_named_once(tuple(map(_as_chat_tool, members))) if isinstance(members, list) else () + + +def _members_named_once( + members: tuple[JsonValue, ...], repeatable: frozenset[str] = frozenset() +) -> tuple[JsonValue, ...]: + names: Final = tuple(map(declared_function_name, members)) + first_index: Final = {name: index for index, name in reversed(tuple(enumerate(names)))} + return tuple( + member + for index, (member, name) in enumerate(zip(members, names, strict=True)) + if name is None or name in repeatable or first_index[name] == index + ) def _is_namespace(tool: JsonValue) -> bool: @@ -598,20 +610,31 @@ def _qualified_member(namespace: str, member: JsonValue) -> JsonValue: return {**member, "function": {**function, "name": qualified_member_name(namespace, name)}} -def _as_declared(tool: JsonValue, declarable: Counter[str | None]) -> tuple[JsonValue, ...]: +def _as_declared( + tool: JsonValue, declarable: Counter[str | None], top_level_names: frozenset[str] +) -> tuple[JsonValue, ...]: if _inlines_members(tool, declarable): return _namespace_members_or_tool(tool) namespace: Final = declared_function_name(tool) if not _is_namespace(tool) or namespace is None: return (tool,) - return tuple(_qualified_member(namespace, member) for member in _namespace_members_or_tool(tool)) + qualified: Final = (_qualified_member(namespace, member) for member in _namespace_members_or_tool(tool)) + return tuple(member for member in qualified if declared_function_name(member) not in top_level_names) + + +def _top_level_names_in(tools: tuple[JsonValue, ...]) -> frozenset[str]: + top_level_tools: Final = (tool for tool in tools if not _is_namespace(tool)) + return frozenset(name for name in map(declared_function_name, top_level_tools) if name is not None) def _with_namespace_members_inlined(tools: JsonValue, outside_declarations: JsonValue = None) -> JsonValue: if not isinstance(tools, list) or not any(map(_is_namespace, tools)): return tools - declarable: Final = Counter(_declarable_names_in((*tools, *_listed(outside_declarations)))) - return list(chain.from_iterable(_as_declared(tool, declarable) for tool in tools)) + candidates: Final = (*tools, *_listed(outside_declarations)) + declarable: Final = Counter(_declarable_names_in(candidates)) + top_level_names: Final = _top_level_names_in(candidates) + declared: Final = chain.from_iterable(_as_declared(tool, declarable, top_level_names) for tool in tools) + return list(_members_named_once(tuple(declared), top_level_names)) def _qualified_names_of(member_name: str, tool: JsonValue, declarable: Counter[str | None]) -> tuple[str, ...]: @@ -629,7 +652,7 @@ def _forced_function_names(name: str, tools: JsonValue) -> tuple[str, ...]: declarable: Final = Counter(_declarable_names_in(candidates)) named: Final = next((tool for tool in candidates if _is_namespace(tool) and _member_name_of(tool) == name), None) if named is not None: - declared: Final = map(declared_function_name, _as_declared(named, declarable)) + declared: Final = map(declared_function_name, _as_declared(named, declarable, frozenset())) return tuple(member for member in declared if member is not None) or (name,) qualified: Final = chain.from_iterable(_qualified_names_of(name, tool, declarable) for tool in candidates) return tuple(qualified) or (name,) diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index e93a7ceb6cb..4d3fe39b91d 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -1015,7 +1015,7 @@ class ToolPermissionGuardrail(CustomGuardrail): returned_members: Final = _grouped_values( chain( ((bare, (name, bare)) for name, bare in member_names if bare not in top_level_names), - ((name, (name, bare)) for name, bare in member_names), + ((name, (name, bare)) for name, bare in member_names if name not in top_level_names), ) ) return MappingProxyType( diff --git a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py index e33335aa431..1f7a457152d 100644 --- a/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py +++ b/tests/unit/llms/vertex_ai/gemini/test_vertex_and_google_ai_studio_gemini.py @@ -1836,12 +1836,8 @@ _TOP_LEVEL_GET_WEATHER = {"type": "function", "function": {k: v for k, v in _GET ), ([_TOP_LEVEL_GET_WEATHER, _WEATHER_NAMESPACE], ["get_weather", "weather__get_weather"]), ([_WEATHER_NAMESPACE, _TOP_LEVEL_GET_WEATHER], ["weather__get_weather", "get_weather"]), - ( - [{"type": "namespace", "name": "delete_db", "tools": [_GET_WEATHER, _GET_WEATHER]}], - ["delete_db__get_weather", "delete_db__get_weather"], - ), ], - ids=["namespaces_share_a_member", "top_level_first", "namespace_first", "namespace_repeats_a_member"], + ids=["namespaces_share_a_member", "top_level_first", "namespace_first"], ) def test_vertex_ai_map_tools_declares_qualified_member_names_when_another_tool_declares_a_member_name( tools: list[dict], declared: list[str] @@ -1851,6 +1847,55 @@ def test_vertex_ai_map_tools_declares_qualified_member_names_when_another_tool_d assert [declaration["name"] for declaration in mapped[0]["function_declarations"]] == declared +def test_vertex_ai_map_tools_declares_a_member_a_namespace_lists_twice_once_from_its_first_entry() -> None: + first = {**_GET_WEATHER, "description": "Current weather"} + repeat = {**_GET_WEATHER, "description": "Forecast"} + namespace = {"type": "namespace", "name": "weather", "tools": [first, repeat]} + + mapped = VertexGeminiConfig().map_function(value=[namespace], optional_params={}) + + declarations = mapped[0]["function_declarations"] + assert [(declaration["name"], declaration["description"]) for declaration in declarations] == [ + ("get_weather", "Current weather") + ] + + +@pytest.mark.parametrize( + ("tools", "declared"), + [ + ( + [ + _WEATHER_NAMESPACE, + _TOP_LEVEL_GET_WEATHER, + {"type": "function", "function": {"name": "weather__get_weather", "description": "Top level"}}, + ], + [("get_weather", None), ("weather__get_weather", "Top level")], + ), + ( + [ + {"type": "namespace", "name": "mcp__ops", "tools": [{"type": "function", "name": "exec"}]}, + {"type": "namespace", "name": "mcp", "tools": [{"type": "function", "name": "ops__exec"}]}, + {"type": "function", "function": {"name": "exec"}}, + {"type": "function", "function": {"name": "ops__exec"}}, + ], + [("mcp__ops__exec", None), ("exec", None), ("ops__exec", None)], + ), + ( + [_TOP_LEVEL_GET_WEATHER, _TOP_LEVEL_GET_WEATHER, _WEATHER_NAMESPACE], + [("get_weather", None), ("get_weather", None), ("weather__get_weather", None)], + ), + ], + ids=["top_level_function_has_the_qualified_name", "namespaces_build_one_qualified_name", "client_repeats_a_function"], +) +def test_vertex_ai_map_tools_never_repeats_a_generated_name( + tools: list[dict], declared: list[tuple[str, str | None]] +) -> None: + mapped = VertexGeminiConfig().map_function(value=tools, optional_params={}) + + declarations = mapped[0]["function_declarations"] + assert [(declaration["name"], declaration.get("description")) for declaration in declarations] == declared + + def test_vertex_ai_map_openai_params_declares_qualified_member_names_when_functions_repeat_a_member() -> None: mapped = VertexGeminiConfig().map_openai_params( non_default_params={ 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 a021625e49e..7d84c1ddb32 100644 --- a/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/unit/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -706,6 +706,37 @@ class TestToolPermissionGuardrail: assert [call.function.name for call in response.choices[0].message.tool_calls or ()] == kept + @pytest.mark.asyncio + async def test_post_call_judges_a_top_level_function_by_its_own_name_beside_a_member_qualified_the_same( + self, + ) -> None: + guardrail = ToolPermissionGuardrail( + guardrail_name="post-exec", + rules=[{"id": "deny_exec", "tool_name": r"^exec$", "decision": "deny"}], + default_action="allow", + on_disallowed_action="rewrite", + event_hook=GuardrailEventHooks.post_call, + ) + namespace = {"type": "namespace", "name": "ops", "tools": [{"type": "function", "name": "exec"}]} + top_level = {"type": "function", "function": {"name": "ops__exec"}} + function = Function(name="ops__exec", arguments="{}") + response = ModelResponse( + choices=[ + Choices( + finish_reason="tool_calls", + message={"tool_calls": [{"id": "call_ops", "type": "function", "function": function}]}, + ) + ] + ) + + await guardrail.async_post_call_success_hook( + data={"guardrails": ["post-exec"], "tools": [namespace, top_level]}, + user_api_key_dict=UserAPIKeyAuth(), + response=response, + ) + + assert [call.function.name for call in response.choices[0].message.tool_calls or ()] == ["ops__exec"] + @pytest.mark.asyncio async def test_async_post_call_success_hook_missing_arguments_blocks_param_rule( self,