Coverage for aipyapp/aipy/functions.py: 26%

1 import inspect

2 from typing import Any, Callable, Dict, Optional, List

3

4 from loguru import logger

5 from pydantic import create_model, ValidationError

6

7

8 class FunctionError(Exception):

9 """Base exception for function calls"""

10 pass

11

12 class FunctionNotFoundError(FunctionError):

13 """Function/method not found exception"""

14 pass

15

16 class ParameterValidationError(FunctionError):

17 """Parameter validation exception"""

18 pass

19

20 class FunctionManager:

21 """Function manager - manage all callable functions"""

22

23 def __init__(self):

24 """Initialize function manager"""

25 self.function_registry: Dict[str, Dict[str, Any]] = {}

26 self.logger = logger.bind(src=self.__class__.__name__)

27

28 def register_function(self, func: Callable, name: str = None) -> bool:

29 """Register function

30

31 Args:

32 func: Function to register

33 name: Custom function name, default is the function name

34

35 Returns:

36 Whether the function is registered successfully

37 """

38 func_name = name or func.__name__

39 sig = inspect.signature(func)

40

41 fields = {}

42 for param_name, param in sig.parameters.items():

43 annotation = param.annotation if param.annotation != inspect.Parameter.empty else Any

44 default = param.default if param.default != inspect.Parameter.empty else ...

45 fields[param_name] = (annotation, default)

46

47 ParamModel = create_model(f"{func_name}_Params", **fields)

48

49 self.function_registry[func_name] = {

50 "func": func,

51 "signature": str(sig),

52 "doc": inspect.getdoc(func) or "",

53 "param_model": ParamModel

54 }

55

56 self.logger.debug(f"Registered function: {func_name}")

57 return True

58

59 def register_functions(self, functions: Dict[str, Callable]) -> int:

60 """Batch register object instance methods

61

62 Args:

63 functions: Function name list

64 """

65 success_count = 0

66 for func_name, func in functions.items():

67 if self.register_function(func, func_name):

68 success_count += 1

69 self.logger.info(f"Registered {success_count}/{len(functions)} functions")

70 return success_count

71

72 def call(self, func_name: str, **kwargs) -> Any:

73 """Call the registered function

74

75 Args:

76 func_name: Function name

77 **kwargs: Function parameters

78

79 Returns:

80 Function execution result

81

82 Raises:

83 FunctionNotFoundError: Function not found

84 ParameterValidationError: Parameter validation failed

85 """

86 entry = self.function_registry.get(func_name)

87 if not entry:

88 self.logger.error(f"Function '{func_name}' not found")

89 raise FunctionNotFoundError(f"Function '{func_name}' not found")

90

91 fn = entry["func"]

92 ParamModel = entry.get("param_model")

93

94 try:

95 parsed = ParamModel(**kwargs)

96 self.logger.info(f"Calling function {func_name}, parameters: {parsed.model_dump()}")

97 result = fn(**parsed.model_dump())

98 return result

99

100 except ValidationError as e:

101 self.logger.error(f"Function '{func_name}' parameter validation failed: {e}")

102 raise ParameterValidationError(f"Parameter validation failed: {e}")

103 except Exception as e:

104 self.logger.error(f"Function '{func_name}' execution failed: {e}")

105 raise

106

107 def get_functions(self) -> Dict[str, Dict[str, str]]:

108 """Get all functions

109

110 Returns:

111 All functions

112 """

113 functions = {}

114 for name, meta in self.function_registry.items():

115 functions[name] = {

116 "signature": meta["signature"],

117 "docstring": meta["doc"],

118 }

119 return functions

120

121 def unregister_function(self, func_name: str) -> bool:

122 """Unregister function

123

124 Args:

125 func_name: Function name

126

127 Returns:

128 Whether the function is unregistered successfully

129 """

130 if func_name in self.function_registry:

131 del self.function_registry[func_name]

132 self.logger.debug(f"Unregistered function: {func_name}")

133 return True

134 return False

135

136 def clear_registry(self, group: Optional[str] = None):

137 """Clear the registry

138

139 Args:

140 group: Specify the group, None means clear all

141 """

142 if group is None:

143 self.function_registry.clear()

144 self.logger.info("Cleared all registered functions")

145 else:

146 to_remove = [

147 name for name, meta in self.function_registry.items()

148 if meta["group"] == group

149 ]

150 for name in to_remove:

151 del self.function_registry[name]

152 self.logger.info(f"Cleared functions in group {group}")

153

154 def get_registry(self) -> Dict[str, Dict[str, Any]]:

155 """Get the complete registry (read-only)

156

157 Returns:

158 Copy of the registry

159 """

160 return self.function_registry.copy()