Coverage for src / local_deep_research / advanced_search_system / tools / base_tool.py: 96%

1""" 

2Base class for all agent-compatible tools. 

3Defines the common interface and shared functionality for different tools. 

4""" 

5 

6from loguru import logger 

7from abc import ABC, abstractmethod 

8from typing import Any, Dict 

9 

10 

11# Safe type mapping for parameter validation (prevents eval() RCE) 

12TYPE_MAP = { 

13 "str": str, 

14 "int": int, 

15 "float": float, 

16 "bool": bool, 

17 "list": list, 

18 "dict": dict, 

19 "tuple": tuple, 

20 "set": set, 

21} 

22 

23 

24class BaseTool(ABC): 

25 """Abstract base class for all agent-compatible tools.""" 

26 

27 def __init__(self, name: str, description: str): 

28 """ 

29 Initialize the tool. 

30 

31 Args: 

32 name: The name of the tool 

33 description: A description of what the tool does 

34 """ 

35 self.name = name 

36 self.description = description 

37 self.parameters: Dict[str, Dict[str, Any]] = {} 

38 

39 @abstractmethod 

40 def execute(self, **kwargs) -> Any: 

41 """ 

42 Execute the tool with the given parameters. 

43 

44 Args: 

45 **kwargs: Tool-specific parameters 

46 

47 Returns: 

48 Any: The result of the tool execution 

49 """ 

50 pass 

51 

52 def get_schema(self) -> Dict[str, Any]: 

53 """ 

54 Get the JSON schema for the tool's parameters. 

55 

56 Returns: 

57 Dict[str, Any]: The JSON schema 

58 """ 

59 return { 

60 "name": self.name, 

61 "description": self.description, 

62 "parameters": self.parameters, 

63 } 

64 

65 def validate_parameters(self, **kwargs) -> bool: 

66 """ 

67 Validate the provided parameters against the tool's schema. 

68 

69 Args: 

70 **kwargs: Parameters to validate 

71 

72 Returns: 

73 bool: True if parameters are valid, False otherwise 

74 """ 

75 for param_name, param_schema in self.parameters.items(): 

76 if param_name not in kwargs: 

77 if param_schema.get("required", False): 

78 logger.error(f"Missing required parameter: {param_name}") 

79 return False 

80 continue 

81 

82 param_value = kwargs[param_name] 

83 param_type = param_schema.get("type") 

84 

85 # Use safe type mapping instead of eval() to prevent RCE 

86 if param_type: 86 ↛ 95line 86 didn't jump to line 95 because the condition on line 86 was always true

87 expected_type = TYPE_MAP.get(param_type) 

88 if expected_type and not isinstance(param_value, expected_type): 

89 logger.error( 

90 f"Invalid type for parameter {param_name}: " 

91 f"expected {param_type}, got {type(param_value).__name__}" 

92 ) 

93 return False 

94 

95 if ( 

96 "enum" in param_schema 

97 and param_value not in param_schema["enum"] 

98 ): 

99 logger.error(f"Invalid value for parameter {param_name}") 

100 return False 

101 

102 return True 

103 

104 def _log_execution(self, **kwargs) -> None: 

105 """ 

106 Log tool execution details. 

107 

108 Args: 

109 **kwargs: Parameters used in execution 

110 """ 

111 # Don't log kwargs directly as it may contain sensitive data 

112 logger.info(f"Executing tool {self.name} with {len(kwargs)} parameters") 

113 

114 def _log_result(self, result: Any) -> None: 

115 """ 

116 Log tool execution result. 

117 

118 Args: 

119 result: The result of the tool execution 

120 """ 

121 logger.info( 

122 f"Tool {self.name} execution completed with result: {result}" 

123 )