diff --git a/litellm/proxy/_lazy_openapi_snapshot.json b/litellm/proxy/_lazy_openapi_snapshot.json index 8331f748c6..46a514c087 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.json +++ b/litellm/proxy/_lazy_openapi_snapshot.json @@ -3572,7 +3572,7 @@ "/anthropic/{endpoint}": { "delete": { "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", - "operationId": "anthropic_proxy_route_anthropic__endpoint__put", + "operationId": "anthropic_proxy_route_anthropic__endpoint__delete", "parameters": [ { "in": "path", @@ -3616,7 +3616,7 @@ }, "get": { "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", - "operationId": "anthropic_proxy_route_anthropic__endpoint__put", + "operationId": "anthropic_proxy_route_anthropic__endpoint__get", "parameters": [ { "in": "path", @@ -3660,7 +3660,7 @@ }, "patch": { "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", - "operationId": "anthropic_proxy_route_anthropic__endpoint__put", + "operationId": "anthropic_proxy_route_anthropic__endpoint__patch", "parameters": [ { "in": "path", @@ -3704,7 +3704,7 @@ }, "post": { "description": "[Docs](https://docs.litellm.ai/docs/pass_through/anthropic_completion)", - "operationId": "anthropic_proxy_route_anthropic__endpoint__put", + "operationId": "anthropic_proxy_route_anthropic__endpoint__post", "parameters": [ { "in": "path", @@ -13260,7 +13260,7 @@ "/langfuse/{endpoint}": { "delete": { "description": "Call Langfuse via LiteLLM proxy. Works with Langfuse SDK.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/langfuse)", - "operationId": "langfuse_proxy_route_langfuse__endpoint__put", + "operationId": "langfuse_proxy_route_langfuse__endpoint__delete", "parameters": [ { "in": "path", @@ -13299,7 +13299,7 @@ }, "get": { "description": "Call Langfuse via LiteLLM proxy. Works with Langfuse SDK.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/langfuse)", - "operationId": "langfuse_proxy_route_langfuse__endpoint__put", + "operationId": "langfuse_proxy_route_langfuse__endpoint__get", "parameters": [ { "in": "path", @@ -13338,7 +13338,7 @@ }, "patch": { "description": "Call Langfuse via LiteLLM proxy. Works with Langfuse SDK.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/langfuse)", - "operationId": "langfuse_proxy_route_langfuse__endpoint__put", + "operationId": "langfuse_proxy_route_langfuse__endpoint__patch", "parameters": [ { "in": "path", @@ -13377,7 +13377,7 @@ }, "post": { "description": "Call Langfuse via LiteLLM proxy. Works with Langfuse SDK.\n\n[Docs](https://docs.litellm.ai/docs/pass_through/langfuse)", - "operationId": "langfuse_proxy_route_langfuse__endpoint__put", + "operationId": "langfuse_proxy_route_langfuse__endpoint__post", "parameters": [ { "in": "path", @@ -26883,7 +26883,7 @@ "/toolset/{toolset_name}/mcp": { "delete": { "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", - "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_put", + "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_delete", "parameters": [ { "in": "path", @@ -26922,7 +26922,7 @@ }, "get": { "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", - "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_put", + "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_get", "parameters": [ { "in": "path", @@ -26961,7 +26961,7 @@ }, "head": { "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", - "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_put", + "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_head", "parameters": [ { "in": "path", @@ -27000,7 +27000,7 @@ }, "options": { "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", - "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_put", + "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_options", "parameters": [ { "in": "path", @@ -27039,7 +27039,7 @@ }, "patch": { "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", - "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_put", + "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_patch", "parameters": [ { "in": "path", @@ -27078,7 +27078,7 @@ }, "post": { "description": "Namespace a toolset as its own MCP endpoint.\n\nConnecting to /toolset//mcp exposes exactly the tools defined in\nthe toolset. Access is enforced: non-admin API keys must have the toolset\nlisted in their object_permission.mcp_toolsets grant list, or the request\nwill be rejected with a 403.", - "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_put", + "operationId": "toolset_mcp_route_toolset__toolset_name__mcp_post", "parameters": [ { "in": "path", diff --git a/litellm/proxy/_lazy_openapi_snapshot.py b/litellm/proxy/_lazy_openapi_snapshot.py index 315f6a9742..309a0276aa 100644 --- a/litellm/proxy/_lazy_openapi_snapshot.py +++ b/litellm/proxy/_lazy_openapi_snapshot.py @@ -13,6 +13,7 @@ from pathlib import Path from typing import Dict, Optional SNAPSHOT_FILE = Path(__file__).parent / "_lazy_openapi_snapshot.json" +HTTP_METHODS = {"delete", "get", "head", "options", "patch", "post", "put"} def load_snapshot() -> Optional[Dict[str, Dict]]: @@ -25,6 +26,39 @@ def load_snapshot() -> Optional[Dict[str, Dict]]: return None +def _normalize_operation_ids(paths: Dict[str, Dict]) -> None: + """Make FastAPI-generated operation IDs stable for multi-method routes. + + FastAPI derives the default operation ID suffix from the first item in the + route's methods set. For routes registered with several HTTP methods, that + set iteration order can vary between processes, which makes the snapshot + drift even when no routes changed. + """ + for path_ops in paths.values(): + if not isinstance(path_ops, dict): + continue + + methods = {method for method in path_ops if method in HTTP_METHODS} + if not methods: + continue + + for method, operation in path_ops.items(): + if method not in HTTP_METHODS or not isinstance(operation, dict): + continue + + operation_id = operation.get("operationId") + if not isinstance(operation_id, str): + continue + + for suffix in methods: + suffix_token = f"_{suffix}" + if operation_id.endswith(suffix_token): + operation["operationId"] = ( + operation_id[: -len(suffix_token)] + f"_{method}" + ) + break + + def generate_snapshot() -> Dict[str, Dict]: import importlib @@ -52,13 +86,15 @@ def generate_snapshot() -> Dict[str, Dict]: if not feat_routes: continue full = get_openapi(title=app.title, version=app.version, routes=feat_routes) + paths = full.get("paths", {}) + _normalize_operation_ids(paths) # Group all of a feature's routes under one tag. - for path_ops in full.get("paths", {}).values(): + for path_ops in paths.values(): for op in path_ops.values(): if isinstance(op, dict): op["tags"] = [feat.name] fragments[feat.name] = { - "paths": full.get("paths", {}), + "paths": paths, "components": {"schemas": full.get("components", {}).get("schemas", {})}, } return fragments diff --git a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py index 6dd0288cb0..37be832d35 100644 --- a/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py +++ b/litellm/proxy/guardrails/guardrail_hooks/tool_permission.py @@ -225,10 +225,10 @@ class ToolPermissionGuardrail(CustomGuardrail): def _parse_tool_call_arguments( self, tool_call: ChatCompletionMessageToolCall - ) -> Dict[str, Any]: + ) -> tuple[Optional[Dict[str, Any]], Optional[str]]: arguments = getattr(tool_call.function, "arguments", None) if not arguments: - return {} + return None, "missing arguments" parsed_arguments: Any = {} try: @@ -236,22 +236,24 @@ class ToolPermissionGuardrail(CustomGuardrail): parsed_arguments = json.loads(arguments) elif isinstance(arguments, dict): parsed_arguments = arguments - except json.JSONDecodeError as exc: + else: + return None, "arguments must be a JSON object" + except (json.JSONDecodeError, TypeError) as exc: verbose_proxy_logger.warning( "Tool Permission Guardrail: Failed to decode arguments for tool %s: %s", tool_call.function.name, exc, ) - return {} + return None, "arguments could not be parsed" if isinstance(parsed_arguments, dict): - return parsed_arguments + return parsed_arguments, None verbose_proxy_logger.debug( - "Tool Permission Guardrail: Ignoring non-dict arguments for tool %s", + "Tool Permission Guardrail: Rejecting non-dict arguments for tool %s", tool_call.function.name, ) - return {} + return None, "arguments must be a JSON object" def _collect_argument_paths( self, @@ -331,10 +333,21 @@ class ToolPermissionGuardrail(CustomGuardrail): continue if rule.allowed_param_patterns and should_check_params: - arguments = self._parse_tool_call_arguments(tool_call) + arguments, parse_error = self._parse_tool_call_arguments(tool_call) + if parse_error: + default_message = f"Tool '{tool_identifier}' {parse_error} required by rule '{rule.id}'" + message = self.render_violation_message( + default=default_message, + context={"tool_name": tool_identifier, "rule_id": rule.id}, + ) + return False, rule.id, message if not arguments: - last_pattern_failure_msg = f"Tool '{tool_identifier}' is missing arguments required by rule '{rule.id}'" - continue + default_message = f"Tool '{tool_identifier}' is missing arguments required by rule '{rule.id}'" + message = self.render_violation_message( + default=default_message, + context={"tool_name": tool_identifier, "rule_id": rule.id}, + ) + return False, rule.id, message patterns_match, failure_message = self._patterns_match_for_rule( arguments=arguments, @@ -365,6 +378,33 @@ class ToolPermissionGuardrail(CustomGuardrail): ) return is_allowed, None, message + @staticmethod + def _get_mapping_value(item: Any, key: str) -> Any: + if isinstance(item, dict): + return item.get(key) + return getattr(item, key, None) + + @staticmethod + def _legacy_function_call_id(choice_index: int) -> str: + return f"legacy_function_call_{choice_index}" + + def _legacy_function_call_to_tool_call( + self, function_call: Any, choice_index: int + ) -> Optional[ChatCompletionMessageToolCall]: + if function_call is None: + return None + + function_name = self._get_mapping_value(function_call, "name") + arguments = self._get_mapping_value(function_call, "arguments") or "" + if not function_name: + return None + + return ChatCompletionMessageToolCall( + id=self._legacy_function_call_id(choice_index), + type="function", + function={"name": function_name, "arguments": arguments}, + ) + def _extract_tool_calls_from_response( self, response: ModelResponse ) -> List[ChatCompletionMessageToolCall]: @@ -379,13 +419,72 @@ class ToolPermissionGuardrail(CustomGuardrail): """ tool_calls = [] - for choice in response.choices: + for choice_index, choice in enumerate(response.choices): if isinstance(choice, Choices): for tool in choice.message.tool_calls or []: tool_calls.append(tool) + legacy_tool_call = self._legacy_function_call_to_tool_call( + getattr(choice.message, "function_call", None), choice_index + ) + if legacy_tool_call is not None: + tool_calls.append(legacy_tool_call) return tool_calls + def _get_request_tool_name(self, tool: Any) -> tuple[Optional[str], Optional[str]]: + tool_type = self._get_mapping_value(tool, "type") + if tool_type != "function": + return None, tool_type + + function = self._get_mapping_value(tool, "function") + tool_name = self._get_mapping_value(function, "name") + return tool_name, tool_type + + def _get_legacy_function_name(self, function: Any) -> Optional[str]: + return self._get_mapping_value(function, "name") + + def _get_named_tool_choice(self, data: dict) -> Optional[str]: + tool_choice = data.get("tool_choice") + if not tool_choice or tool_choice in ("auto", "none", "required"): + return None + if isinstance(tool_choice, str): + return tool_choice + if self._get_mapping_value(tool_choice, "type") != "function": + return None + return self._get_mapping_value( + self._get_mapping_value(tool_choice, "function"), "name" + ) + + def _get_named_function_call(self, data: dict) -> Optional[str]: + function_call = data.get("function_call") + if not function_call or function_call in ("auto", "none"): + return None + if isinstance(function_call, str): + return function_call + return self._get_mapping_value(function_call, "name") + + def _collect_request_tools(self, data: dict) -> List[tuple[str, Optional[str]]]: + request_tools: List[tuple[str, Optional[str]]] = [] + + for tool in data.get("tools") or []: + tool_name, tool_type = self._get_request_tool_name(tool) + if tool_name is not None: + request_tools.append((tool_name, tool_type)) + + for function in data.get("functions") or []: + function_name = self._get_legacy_function_name(function) + if function_name is not None: + request_tools.append((function_name, "function")) + + for forced_tool_name in ( + self._get_named_tool_choice(data), + self._get_named_function_call(data), + ): + if forced_tool_name is not None: + request_tools.append((forced_tool_name, "function")) + + return request_tools + def _modify_request_with_permission_errors( self, data: dict, @@ -410,19 +509,32 @@ class ToolPermissionGuardrail(CustomGuardrail): for tool_use in denied_tool_names: error_tool_names.add(tool_use) - # Modify the tools tools: Optional[List[ChatCompletionToolParam]] = data.get("tools") - if tools is None: - return data - - new_tools = [] - for tool in tools: - if tool["type"] != "function": - continue - tool_name: str = tool["function"]["name"] - if tool_name not in error_tool_names: + if tools is not None: + new_tools = [] + for tool in tools: + tool_name, tool_type = self._get_request_tool_name(tool) + if tool_type == "function" and tool_name in error_tool_names: + continue new_tools.append(tool) - data["tools"] = new_tools + data["tools"] = new_tools + + functions = data.get("functions") + if functions is not None: + data["functions"] = [ + function + for function in functions + if self._get_legacy_function_name(function) not in error_tool_names + ] + + named_tool_choice = self._get_named_tool_choice(data) + if named_tool_choice in error_tool_names: + data["tool_choice"] = "none" + + named_function_call = self._get_named_function_call(data) + if named_function_call in error_tool_names: + data["function_call"] = "none" + return data def _create_permission_error_result( @@ -472,7 +584,7 @@ class ToolPermissionGuardrail(CustomGuardrail): error_results[tool_use.id] = error_result # Modify the response content - for choice in response.choices: + for choice_index, choice in enumerate(response.choices): if isinstance(choice, Choices): filtered_tool_calls = [] error_messages = [] @@ -490,6 +602,15 @@ class ToolPermissionGuardrail(CustomGuardrail): filtered_tool_calls if filtered_tool_calls else None ) + legacy_tool_call = self._legacy_function_call_to_tool_call( + getattr(choice.message, "function_call", None), choice_index + ) + if legacy_tool_call is not None: + legacy_error_result = error_results.get(legacy_tool_call.id) + if legacy_error_result is not None: + choice.message.function_call = None + error_messages.append(legacy_error_result.content) + # Add error messages to content if error_messages: existing_content = choice.message.content @@ -519,21 +640,16 @@ class ToolPermissionGuardrail(CustomGuardrail): if self.should_run_guardrail(data=data, event_type=event_type) is not True: return data - new_tools: Optional[List[ChatCompletionToolParam]] = data.get("tools") - if new_tools is None: + new_tools = self._collect_request_tools(data) + if not new_tools: verbose_proxy_logger.warning( - "Tool Permission Guardrail: not running guardrail. No tools in data" + "Tool Permission Guardrail: not running guardrail. No tools or functions in data" ) return data # Check permissions for each tool denied_tool_names = [] - for tool in new_tools: - if tool["type"] != "function": - continue - tool_name: str = tool["function"]["name"] - tool_type: Optional[str] = tool.get("type") - + for tool_name, tool_type in new_tools: is_allowed, _, message = self._check_tool_permission(tool_name, tool_type) if not is_allowed and message is not None: diff --git a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py index 55d92e9141..716b4470d2 100644 --- a/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py +++ b/tests/test_litellm/proxy/guardrails/guardrail_hooks/test_tool_permission.py @@ -220,6 +220,27 @@ class TestToolPermissionGuardrail: assert tool_calls[0].id == "call_123" assert tool_calls[0].function.name == "Read" + def test_extract_tool_calls_legacy_function_call_format(self): + response = ModelResponse( + choices=[ + Choices( + message={ + "function_call": { + "name": "Read", + "arguments": '{"file_path": "/test/file.txt"}', + }, + } + ) + ] + ) + + tool_calls = self.guardrail._extract_tool_calls_from_response(response) + assert len(tool_calls) == 1 + assert isinstance(tool_calls[0], ChatCompletionMessageToolCall) + assert tool_calls[0].id == "legacy_function_call_0" + assert tool_calls[0].function.name == "Read" + assert tool_calls[0].function.arguments == '{"file_path": "/test/file.txt"}' + def test_extract_tool_calls_empty_response(self): response = ModelResponse(choices=[]) tool_calls = self.guardrail._extract_tool_calls_from_response(response) @@ -271,6 +292,31 @@ class TestToolPermissionGuardrail: data=data, user_api_key_dict=user_api_key_dict, response=response ) + @pytest.mark.asyncio + async def test_async_post_call_success_hook_with_denied_legacy_function_call_raises( + self, + ): + response = ModelResponse( + choices=[ + Choices( + message={ + "function_call": { + "name": "Read", + "arguments": "{}", + }, + } + ) + ] + ) + user_api_key_dict = UserAPIKeyAuth() + data = {"guardrails": ["test-tool-permission"]} + + with patch.object(self.guardrail, "should_run_guardrail", return_value=True): + with pytest.raises(GuardrailRaisedException): + await self.guardrail.async_post_call_success_hook( + data=data, user_api_key_dict=user_api_key_dict, response=response + ) + @pytest.mark.asyncio async def test_async_post_call_success_hook_param_patterns_allow(self): guardrail = ToolPermissionGuardrail( @@ -379,7 +425,9 @@ class TestToolPermissionGuardrail: assert "berri" in choice.message.content @pytest.mark.asyncio - async def test_async_post_call_success_hook_missing_arguments_default_allows(self): + async def test_async_post_call_success_hook_missing_arguments_blocks_param_rule( + self, + ): guardrail = ToolPermissionGuardrail( guardrail_name="mail-guardrail", rules=[ @@ -405,9 +453,52 @@ class TestToolPermissionGuardrail: data = {"guardrails": ["mail-guardrail"]} with patch.object(guardrail, "should_run_guardrail", return_value=True): - await guardrail.async_post_call_success_hook( - data=data, user_api_key_dict=user_api_key_dict, response=response - ) + with pytest.raises(GuardrailRaisedException): + await guardrail.async_post_call_success_hook( + data=data, user_api_key_dict=user_api_key_dict, response=response + ) + + @pytest.mark.asyncio + @pytest.mark.parametrize( + "arguments", + [ + "{not-json", + '["owner@berri.ai"]', + ], + ) + async def test_async_post_call_success_hook_malformed_arguments_blocks_param_rule( + self, arguments + ): + guardrail = ToolPermissionGuardrail( + guardrail_name="mail-guardrail", + rules=[ + { + "id": "deny_gmail", + "tool_name": r"^mail_mcp-send_email$", + "decision": "deny", + "allowed_param_patterns": {"to[]": r"^.+@gmail\.com$"}, + } + ], + default_action="allow", + on_disallowed_action="block", + ) + + tool_call = { + "function": { + "name": "mail_mcp-send_email", + "arguments": arguments, + }, + "type": "function", + } + response = ModelResponse(choices=[Choices(message={"tool_calls": [tool_call]})]) + user_api_key_dict = UserAPIKeyAuth() + data = {"guardrails": ["mail-guardrail"]} + + with patch.object(guardrail, "should_run_guardrail", return_value=True): + with pytest.raises(GuardrailRaisedException): + await guardrail.async_post_call_success_hook( + data=data, user_api_key_dict=user_api_key_dict, response=response + ) @pytest.mark.asyncio async def test_async_pre_call_hook_block_mode(self): @@ -430,6 +521,65 @@ class TestToolPermissionGuardrail: ) assert excinfo.value.status_code == 400 + @pytest.mark.asyncio + async def test_async_pre_call_hook_blocks_legacy_functions(self): + data = { + "functions": [ + {"name": "Bash", "description": "allowed"}, + {"name": "Read", "description": "denied"}, + ] + } + user_api_key_dict = UserAPIKeyAuth() + cache = DualCache(default_in_memory_ttl=1) + + with patch.object(self.guardrail, "should_run_guardrail", return_value=True): + with pytest.raises(HTTPException) as excinfo: + await self.guardrail.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="completion", + ) + assert excinfo.value.status_code == 400 + + @pytest.mark.asyncio + async def test_async_pre_call_hook_blocks_named_legacy_function_call(self): + data = { + "functions": [{"name": "Bash"}], + "function_call": {"name": "Read"}, + } + user_api_key_dict = UserAPIKeyAuth() + cache = DualCache(default_in_memory_ttl=1) + + with patch.object(self.guardrail, "should_run_guardrail", return_value=True): + with pytest.raises(HTTPException) as excinfo: + await self.guardrail.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="completion", + ) + assert excinfo.value.status_code == 400 + + @pytest.mark.asyncio + async def test_async_pre_call_hook_blocks_named_tool_choice(self): + data = { + "tools": [{"type": "function", "function": {"name": "Bash"}}], + "tool_choice": {"type": "function", "function": {"name": "Read"}}, + } + user_api_key_dict = UserAPIKeyAuth() + cache = DualCache(default_in_memory_ttl=1) + + with patch.object(self.guardrail, "should_run_guardrail", return_value=True): + with pytest.raises(HTTPException) as excinfo: + await self.guardrail.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="completion", + ) + assert excinfo.value.status_code == 400 + @pytest.mark.asyncio async def test_async_pre_call_hook_uses_custom_template(self): guardrail = ToolPermissionGuardrail( @@ -491,6 +641,41 @@ class TestToolPermissionGuardrail: assert "Bash" in tool_names assert "Read" not in tool_names + @pytest.mark.asyncio + async def test_async_pre_call_hook_rewrite_mode_filters_legacy_functions(self): + guardrail = ToolPermissionGuardrail( + guardrail_name="test-tool-permission", + rules=self.test_rules, + default_action="deny", + on_disallowed_action="rewrite", + ) + data = { + "functions": [ + {"name": "Bash", "description": "allowed"}, + {"name": "Read", "description": "denied"}, + ], + "function_call": {"name": "Read"}, + "tools": [ + {"type": "function", "function": {"name": "Bash"}}, + ], + "tool_choice": {"type": "function", "function": {"name": "Read"}}, + } + user_api_key_dict = UserAPIKeyAuth() + cache = DualCache(default_in_memory_ttl=1) + + with patch.object(guardrail, "should_run_guardrail", return_value=True): + new_data = await guardrail.async_pre_call_hook( + user_api_key_dict=user_api_key_dict, + cache=cache, + data=data, + call_type="completion", + ) + + assert isinstance(new_data, dict) + assert [function["name"] for function in new_data["functions"]] == ["Bash"] + assert new_data["function_call"] == "none" + assert new_data["tool_choice"] == "none" + def test_modify_response_with_permission_errors(self): # Setup a response with one tool_call tool_call = ChatCompletionMessageToolCall( @@ -522,6 +707,40 @@ class TestToolPermissionGuardrail: assert isinstance(choice.message.content, str) assert "Permission denied" in choice.message.content + def test_modify_response_with_permission_errors_filters_legacy_function_call(self): + response = ModelResponse( + choices=[ + Choices( + message={ + "function_call": { + "name": "Read", + "arguments": "{}", + }, + "content": "", + } + ) + ] + ) + tool_call = self.guardrail._extract_tool_calls_from_response(response)[0] + denied_tools = [ + ( + tool_call, + PermissionError( + tool_name="Read", + rule_id="deny_read", + message="Tool 'Read' denied by rule 'deny_read'", + ), + ) + ] + + self.guardrail._modify_response_with_permission_errors(response, denied_tools) + + choice = response.choices[0] + assert isinstance(choice, Choices) + assert choice.message.function_call is None + assert isinstance(choice.message.content, str) + assert "Permission denied" in choice.message.content + class TestToolPermissionGuardrailIntegration: """Integration tests for Tool Permission Guardrail""" diff --git a/tests/test_litellm/proxy/test_lazy_openapi_snapshot.py b/tests/test_litellm/proxy/test_lazy_openapi_snapshot.py new file mode 100644 index 0000000000..8bc39c93ee --- /dev/null +++ b/tests/test_litellm/proxy/test_lazy_openapi_snapshot.py @@ -0,0 +1,35 @@ +from litellm.proxy._lazy_openapi_snapshot import _normalize_operation_ids + + +def test_normalize_operation_ids_uses_each_http_method(): + paths = { + "/proxy/{endpoint}": { + "delete": {"operationId": "proxy_route_proxy__endpoint__put"}, + "get": {"operationId": "proxy_route_proxy__endpoint__put"}, + "post": {"operationId": "proxy_route_proxy__endpoint__put"}, + "put": {"operationId": "proxy_route_proxy__endpoint__put"}, + } + } + + _normalize_operation_ids(paths) + + operations = paths["/proxy/{endpoint}"] + assert operations["delete"]["operationId"] == "proxy_route_proxy__endpoint__delete" + assert operations["get"]["operationId"] == "proxy_route_proxy__endpoint__get" + assert operations["post"]["operationId"] == "proxy_route_proxy__endpoint__post" + assert operations["put"]["operationId"] == "proxy_route_proxy__endpoint__put" + + +def test_normalize_operation_ids_preserves_custom_ids(): + paths = { + "/proxy/{endpoint}": { + "get": {"operationId": "custom_operation"}, + "post": {"operationId": "custom_operation"}, + } + } + + _normalize_operation_ids(paths) + + operations = paths["/proxy/{endpoint}"] + assert operations["get"]["operationId"] == "custom_operation" + assert operations["post"]["operationId"] == "custom_operation"