Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/guardrails/tool_name_extraction.py: 61%

1""" 

2Extract tool names from request body by route/call type. 

3 

4Used by auth (check_tools_allowlist) and ToolPolicyGuardrail so tool-format 

5knowledge lives in one place. Uses guardrail translation handlers where available, 

6with standalone extractors for generate_content and MCP. 

7""" 

8 

9from typing import Any, Final 

10 

11from litellm.litellm_core_utils.api_route_to_call_types import get_call_types_for_route 

12from litellm.llms import load_guardrail_translation_mappings 

13from litellm.types.utils import CallTypes 

14 

15# Call types that have no guardrail translation handler; we use standalone extractors 

16STANDALONE_EXTRACTORS: Final[dict[str, Any]] = {} 

17 

18 

19def _extract_generate_content_tool_names(data: dict) -> list[str]: 

20 """Google generateContent: tools[].functionDeclarations[].name""" 

21 names: Final[list[str]] = [] 

22 for tool in data.get("tools") or []: 22 ↛ 23line 22 didn't jump to line 23 because the loop on line 22 never started

23 if not isinstance(tool, dict): 

24 continue 

25 for decl in tool.get("functionDeclarations") or []: 

26 if isinstance(decl, dict) and decl.get("name"): 

27 names.append(str(decl["name"])) 

28 return names 

29 

30 

31def _extract_mcp_tool_names(data: dict) -> list[str]: 

32 """MCP call_tool: name or mcp_tool_name in body""" 

33 names: Final[list[str]] = [] 

34 name: Final = data.get("name") or data.get("mcp_tool_name") 

35 if name: 

36 names.append(str(name)) 

37 return names 

38 

39 

40def _register_standalone_extractors() -> None: 

41 if STANDALONE_EXTRACTORS: 

42 return 

43 STANDALONE_EXTRACTORS[CallTypes.generate_content.value] = _extract_generate_content_tool_names 

44 STANDALONE_EXTRACTORS[CallTypes.agenerate_content.value] = _extract_generate_content_tool_names 

45 STANDALONE_EXTRACTORS[CallTypes.call_mcp_tool.value] = _extract_mcp_tool_names 

46 

47 

48# Tool-capable call types (routes that can send tools in the request) 

49TOOL_CAPABLE_CALL_TYPES: Final = frozenset( 

50 { 

51 CallTypes.completion.value, 

52 CallTypes.acompletion.value, 

53 CallTypes.responses.value, 

54 CallTypes.aresponses.value, 

55 CallTypes.anthropic_messages.value, 

56 CallTypes.generate_content.value, 

57 CallTypes.agenerate_content.value, 

58 CallTypes.call_mcp_tool.value, 

59 } 

60) 

61 

62 

63def extract_request_tool_names(route: str, data: dict) -> list[str]: 

64 """ 

65 Extract tool names from the request body for the given route. 

66 Uses guardrail translation handlers when available, else standalone extractors 

67 for generate_content and MCP. Returns [] for non-tool-capable routes or when 

68 no tools are present. 

69 """ 

70 call_types: Final = get_call_types_for_route(route) 

71 if not call_types: 71 ↛ 72line 71 didn't jump to line 72 because the condition on line 71 was never true

72 return [] 

73 _register_standalone_extractors() 

74 mappings: Final = load_guardrail_translation_mappings() 

75 for call_type in call_types: 

76 if not isinstance(call_type, CallTypes): 76 ↛ 77line 76 didn't jump to line 77 because the condition on line 76 was never true

77 continue 

78 if call_type.value not in TOOL_CAPABLE_CALL_TYPES: 78 ↛ 79line 78 didn't jump to line 79 because the condition on line 78 was never true

79 continue 

80 if call_type.value in STANDALONE_EXTRACTORS: 

81 return STANDALONE_EXTRACTORS[call_type.value](data) 

82 handler_cls = mappings.get(call_type) 

83 if handler_cls is not None: 83 ↛ 75line 83 didn't jump to line 75 because the condition on line 83 was always true

84 names = handler_cls().extract_request_tool_names(data) 

85 if names: 85 ↛ 86line 85 didn't jump to line 86 because the condition on line 85 was never true

86 return names 

87 return []