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

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

2

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

4 that can be used with LLM function calling capabilities.

5"""

6

7 from abc import ABC, abstractmethod

8 from typing import Dict, Any

9

10

11 class 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}"