Coverage for agentos/tools/function_calling.py: 43%

1""" 

2Function Calling Pipeline — Schema-validated tool invocation. 

3 

4Provides a complete function calling lifecycle: schema registration, LLM 

5tool_choice dispatch, argument validation, execution, and result formatting. 

6""" 

7 

8from __future__ import annotations 

9 

10import json 

11from collections.abc import Callable 

12from dataclasses import dataclass, field 

13from typing import Any 

14 

15import jsonschema 

16 

17 

18@dataclass 

19class ToolSchema: 

20 """OpenAI-compatible tool/function schema.""" 

21 

22 name: str 

23 description: str 

24 parameters: dict[str, Any] 

25 """JSON Schema for parameters.""" 

26 

27 required: list[str] = field(default_factory=list) 

28 """Required parameter names.""" 

29 

30 def to_openai(self) -> dict[str, Any]: 

31 """Convert to OpenAI function definition format.""" 

32 schema = { 

33 "type": self.parameters.get("type", "object"), 

34 "properties": self.parameters.get("properties", {}), 

35 } 

36 if self.required: 

37 schema["required"] = self.required 

38 return { 

39 "type": "function", 

40 "function": { 

41 "name": self.name, 

42 "description": self.description, 

43 "parameters": schema, 

44 }, 

45 } 

46 

47 def to_anthropic(self) -> dict[str, Any]: 

48 """Convert to Anthropic tool format.""" 

49 return { 

50 "name": self.name, 

51 "description": self.description, 

52 "input_schema": { 

53 "type": "object", 

54 "properties": self.parameters.get("properties", {}), 

55 "required": self.required, 

56 }, 

57 } 

58 

59 

60@dataclass 

61class ToolCall: 

62 """A parsed tool call from an LLM response.""" 

63 

64 id: str 

65 name: str 

66 arguments: dict[str, Any] 

67 

68 

69@dataclass 

70class ToolResult: 

71 """Result of executing a tool call.""" 

72 

73 call_id: str 

74 name: str 

75 success: bool 

76 output: Any = None 

77 error: str | None = None 

78 latency_ms: float = 0.0 

79 

80 

81class ToolRegistry: 

82 """ 

83 Registry of callable tools with schema validation. 

84 

85 Example:: 

86 

87 registry = ToolRegistry() 

88 registry.register( 

89 ToolSchema(name="get_weather", description="Get weather", parameters={ 

90 "type": "object", 

91 "properties": {"city": {"type": "string"}} 

92 }, required=["city"]), 

93 handler=lambda city: f"Weather in {city}: sunny" 

94 ) 

95 """ 

96 

97 def __init__(self): 

98 self._tools: dict[str, ToolSchema] = {} 

99 self._handlers: dict[str, Callable[..., Any]] = {} 

100 

101 def register( 

102 self, 

103 schema: ToolSchema, 

104 handler: Callable[..., Any], 

105 ) -> None: 

106 """Register a tool with its schema and handler function.""" 

107 name = schema.name 

108 if name in self._tools: 

109 raise ValueError(f"Tool '{name}' already registered") 

110 self._tools[name] = schema 

111 self._handlers[name] = handler 

112 

113 def unregister(self, name: str) -> None: 

114 """Remove a tool from the registry.""" 

115 self._tools.pop(name, None) 

116 self._handlers.pop(name, None) 

117 

118 def get_schema(self, name: str) -> ToolSchema | None: 

119 return self._tools.get(name) 

120 

121 def list_schemas(self) -> list[ToolSchema]: 

122 return list(self._tools.values()) 

123 

124 def to_openai_tools(self) -> list[dict[str, Any]]: 

125 """Export all tools as OpenAI function definitions.""" 

126 return [t.to_openai() for t in self._tools.values()] 

127 

128 def to_anthropic_tools(self) -> list[dict[str, Any]]: 

129 """Export all tools as Anthropic tool definitions.""" 

130 return [t.to_anthropic() for t in self._tools.values()] 

131 

132 def validate_arguments(self, name: str, arguments: dict) -> list[str]: 

133 """Validate arguments against tool schema. Returns list of errors.""" 

134 schema = self._tools.get(name) 

135 if schema is None: 

136 return [f"Unknown tool: {name}"] 

137 

138 errors: list[str] = [] 

139 

140 # Check required args 

141 for f in schema.required: 

142 if f not in arguments: 

143 errors.append(f"Missing required argument: {f}") 

144 

145 # JSON Schema validation 

146 try: 

147 jsonschema.validate(instance=arguments, schema=schema.parameters) 

148 except jsonschema.ValidationError as e: 

149 errors.append(f"Schema validation: {e.message}") 

150 

151 return errors 

152 

153 def execute(self, call: ToolCall) -> ToolResult: 

154 """ 

155 Validate and execute a tool call. 

156 

157 Args: 

158 call: Parsed tool call with name and arguments. 

159 

160 Returns: 

161 ToolResult with success/failure and output. 

162 """ 

163 import time 

164 

165 t0 = time.perf_counter() 

166 

167 errors = self.validate_arguments(call.name, call.arguments) 

168 if errors: 

169 return ToolResult( 

170 call_id=call.id, 

171 name=call.name, 

172 success=False, 

173 error="; ".join(errors), 

174 latency_ms=(time.perf_counter() - t0) * 1000, 

175 ) 

176 

177 handler = self._handlers.get(call.name) 

178 if handler is None: 

179 return ToolResult( 

180 call_id=call.id, 

181 name=call.name, 

182 success=False, 

183 error=f"No handler for tool: {call.name}", 

184 latency_ms=(time.perf_counter() - t0) * 1000, 

185 ) 

186 

187 try: 

188 output = handler(**call.arguments) 

189 return ToolResult( 

190 call_id=call.id, 

191 name=call.name, 

192 success=True, 

193 output=output, 

194 latency_ms=(time.perf_counter() - t0) * 1000, 

195 ) 

196 except Exception as exc: 

197 return ToolResult( 

198 call_id=call.id, 

199 name=call.name, 

200 success=False, 

201 error=f"{type(exc).__name__}: {exc}", 

202 latency_ms=(time.perf_counter() - t0) * 1000, 

203 ) 

204 

205 def execute_batch(self, calls: list[ToolCall]) -> list[ToolResult]: 

206 """Execute multiple tool calls. Independent calls run sequentially.""" 

207 return [self.execute(c) for c in calls] 

208 

209 def parse_tool_calls(self, raw_tool_calls: list[dict[str, Any]]) -> list[ToolCall]: 

210 """Parse raw LLM tool_call dicts into ToolCall objects.""" 

211 parsed: list[ToolCall] = [] 

212 for tc in raw_tool_calls: 

213 fn = tc.get("function", tc) 

214 args_raw = fn.get("arguments", "{}") 

215 if isinstance(args_raw, str): 

216 try: 

217 args = json.loads(args_raw) 

218 except json.JSONDecodeError: 

219 args = {} 

220 else: 

221 args = args_raw 

222 parsed.append( 

223 ToolCall( 

224 id=tc.get("id", ""), 

225 name=fn.get("name", ""), 

226 arguments=args, 

227 ) 

228 ) 

229 return parsed 

230 

231 @property 

232 def tool_count(self) -> int: 

233 return len(self._tools)