Coverage for src/dataknobs_llm/tools/base.py: 94%

1"""Base tool abstraction for LLM function calling. 

2 

3This module provides the base Tool class for implementing callable tools 

4that can be used with LLM function calling capabilities. 

5""" 

6 

7from abc import ABC, abstractmethod 

8from typing import Dict, Any 

9 

10 

11class Tool(ABC): 

12 """Abstract base class for LLM-callable tools. 

13 

14 A Tool represents a function that can be called by an LLM during generation. 

15 Each tool has a name, description, parameter schema, and execution logic. 

16 

17 Example: 

18 class CalculatorTool(Tool): 

19 def __init__(self): 

20 super().__init__( 

21 name="calculator", 

22 description="Performs basic arithmetic operations" 

23 ) 

24 

25 @property 

26 def schema(self) -> Dict[str, Any]: 

27 return { 

28 "type": "object", 

29 "properties": { 

30 "operation": { 

31 "type": "string", 

32 "enum": ["add", "subtract", "multiply", "divide"] 

33 }, 

34 "a": {"type": "number"}, 

35 "b": {"type": "number"} 

36 }, 

37 "required": ["operation", "a", "b"] 

38 } 

39 

40 async def execute(self, operation: str, a: float, b: float) -> float: 

41 if operation == "add": 

42 return a + b 

43 elif operation == "subtract": 

44 return a - b 

45 elif operation == "multiply": 

46 return a * b 

47 elif operation == "divide": 

48 return a / b 

49 else: 

50 raise ValueError(f"Unknown operation: {operation}") 

51 """ 

52 

53 def __init__( 

54 self, 

55 name: str, 

56 description: str, 

57 metadata: Dict[str, Any] | None = None 

58 ): 

59 """Initialize a tool. 

60 

61 Args: 

62 name: Unique identifier for the tool 

63 description: Human-readable description of what the tool does 

64 metadata: Optional metadata about the tool 

65 """ 

66 self.name = name 

67 self.description = description 

68 self.metadata = metadata or {} 

69 

70 @property 

71 @abstractmethod 

72 def schema(self) -> Dict[str, Any]: 

73 """Get JSON schema for tool parameters. 

74 

75 Returns a JSON Schema dictionary describing the parameters 

76 this tool accepts. The schema should follow the JSON Schema 

77 specification and is used by LLMs to understand how to call 

78 the tool. 

79 

80 Returns: 

81 JSON Schema dictionary for tool parameters 

82 

83 Example: 

84 { 

85 "type": "object", 

86 "properties": { 

87 "query": { 

88 "type": "string", 

89 "description": "The search query" 

90 }, 

91 "max_results": { 

92 "type": "integer", 

93 "description": "Maximum number of results", 

94 "default": 10 

95 } 

96 }, 

97 "required": ["query"] 

98 } 

99 """ 

100 pass 

101 

102 @abstractmethod 

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

104 """Execute the tool with given parameters. 

105 

106 This method performs the actual tool logic. Parameters are passed 

107 as keyword arguments matching the schema definition. 

108 

109 Args: 

110 **kwargs: Tool parameters matching the schema 

111 

112 Returns: 

113 Tool execution result (can be any JSON-serializable type) 

114 

115 Raises: 

116 Exception: If tool execution fails 

117 """ 

118 pass 

119 

120 def to_function_definition(self) -> Dict[str, Any]: 

121 """Convert tool to OpenAI function calling format. 

122 

123 Returns a dictionary in the format expected by OpenAI's function 

124 calling API. 

125 

126 Returns: 

127 Function definition dictionary 

128 

129 Example: 

130 { 

131 "name": "search_web", 

132 "description": "Search the web for information", 

133 "parameters": { 

134 "type": "object", 

135 "properties": { 

136 "query": {"type": "string"} 

137 }, 

138 "required": ["query"] 

139 } 

140 } 

141 """ 

142 return { 

143 "name": self.name, 

144 "description": self.description, 

145 "parameters": self.schema, 

146 } 

147 

148 def to_anthropic_tool_definition(self) -> Dict[str, Any]: 

149 """Convert tool to Anthropic tool format. 

150 

151 Returns a dictionary in the format expected by Anthropic's Claude API. 

152 

153 Returns: 

154 Tool definition dictionary 

155 """ 

156 return { 

157 "name": self.name, 

158 "description": self.description, 

159 "input_schema": self.schema, 

160 } 

161 

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

163 """Validate parameters against schema. 

164 

165 Optional method for parameter validation before execution. 

166 By default, assumes LLM provides valid parameters. 

167 

168 Args: 

169 **kwargs: Parameters to validate 

170 

171 Returns: 

172 True if valid, False otherwise 

173 """ 

174 # Basic validation - check required fields 

175 required = self.schema.get("required", []) 

176 return all(field in kwargs for field in required) 

177 

178 def __repr__(self) -> str: 

179 """String representation of tool.""" 

180 return f"Tool(name={self.name!r}, description={self.description!r})" 

181 

182 def __str__(self) -> str: 

183 """Human-readable string representation.""" 

184 return f"{self.name}: {self.description}"