Coverage for src/appl/core/function.py: 94%

1from __future__ import annotations 

2 

3import inspect 

4import sys 

5import time 

6import traceback 

7from threading import Lock 

8from typing import Any, Callable, Literal, Optional 

9 

10from .compile import appl_compile 

11from .context import PromptContext 

12from .modifiers import Compositor 

13 

14 

15class PromptFunc: 

16 """A wrapper for an APPL function, can be called as a normal function. 

17 

18 The function contains a prompt context, which could be same as or 

19 copied from its caller function, or created from scratch, or resumed 

20 from the last run. 

21 """ 

22 

23 def __init__( 

24 self, 

25 func: Callable, 

26 ctx_method: str = "new", 

27 comp: Optional[Compositor] = None, 

28 default_return: Optional[Literal["prompt"]] = None, 

29 include_docstring: bool = False, 

30 new_ctx_func: Callable = PromptContext, 

31 # default_sep: Optional[str] = None, 

32 # ? set the default printer behavior for the prompt function? 

33 ): 

34 """Initialize the PromptFunc. 

35 

36 Args: 

37 func (Callable): the function being wrapped 

38 ctx_method (str): 

39 the method to deal with the child context, available methods includes: 

40 

41 - (default) "new" or "new_ctx": create a brand new context. 

42 - "copy" or "copy_ctx": 

43 copy from the parent's context, the change will not 

44 affect the parent's context. 

45 - "same" or "same_ctx": 

46 use the same context as the parent's, the change will 

47 affect the parent's context. 

48 - "resume" or "resume_ctx": 

49 resume its own context from the last run. 

50 For the first run, it will copy the parent's context. 

51 

52 comp (Compositor, optional): 

53 the default compositor to be used. Defaults to None. 

54 default_return (str, optional): 

55 The default return value, "prompt" means return the prompt within 

56 the function. Defaults to None. 

57 include_docstring (bool, optional): 

58 set to True to include the triple-quoted docstring in the prompt. 

59 Defaults to False. 

60 new_ctx_func (Callable, optional): 

61 the function to create a new context. Defaults to PromptContext. 

62 """ 

63 self._func = appl_compile(func) 

64 self._signature = inspect.signature(func) 

65 self._doc = func.__doc__ 

66 self._name = func.__name__ 

67 self._qualname = func.__qualname__ 

68 self._default_ctx_method = self._process_ctx_method(ctx_method) 

69 self._default_compositor = comp 

70 if default_return is not None and default_return != "prompt": 

71 raise NotImplementedError("Only support default_return='prompt' now.") 

72 self._default_return = default_return 

73 self._include_docstring = include_docstring 

74 self._new_ctx_func = new_ctx_func 

75 self._persist_ctx: Optional[PromptContext] = None 

76 self._reset_context_func: Optional[Callable[[], None]] = None 

77 # self._default_sep = default_sep 

78 

79 @property 

80 def compiled_func(self): 

81 """The compiled function.""" 

82 return self._func 

83 

84 def _process_ctx_method(self, ctx_method: str) -> str: 

85 available_methods = ["new", "copy", "same", "resume"] 

86 alias = [m + "_ctx" for m in available_methods] 

87 res = ctx_method 

88 if res in alias: 

89 res = available_methods[alias.index(ctx_method)] 

90 if res not in available_methods: 

91 raise ValueError(f"Unknown ctx_method: {ctx_method}") 

92 return res 

93 

94 def _run( 

95 self, 

96 parent_ctx: PromptContext, 

97 child_ctx: PromptContext, 

98 *args: Any, 

99 **kwargs: Any, 

100 ) -> Any: 

101 """Run the prompt function with desired context, deal with the exception.""" 

102 # if parent_ctx.is_outmost: 

103 # exc = None 

104 # try: 

105 # results = self._func(_ctx=child_ctx, *args, **kwargs) 

106 # except Exception: 

107 # # if parent_ctx.is_outmost: # only print the trace in outmost appl function 

108 # traceback.print_exc() 

109 # print() # empty line 

110 # exc = sys.exc_info() 

111 

112 # if exc is not None: # the outmost appl function 

113 # _, e_instance, e_traceback = exc 

114 # # already printed above, clean up the traceback 

115 # if e_instance and e_traceback: 

116 # e_traceback.tb_next = None 

117 # raise e_instance.with_traceback(e_traceback) 

118 # else: # run normally, capture the exception at the outmost appl function 

119 # results = self._func(_ctx=child_ctx, *args, **kwargs) 

120 results = self._func(_ctx=child_ctx, *args, **kwargs) 

121 return results 

122 

123 def _call( 

124 self, *args: Any, _ctx: Optional[PromptContext] = None, **kwargs: Any 

125 ) -> Any: 

126 """Call the prompt function.""" 

127 parent_ctx = _ctx or self._new_ctx_func() 

128 is_class_method = kwargs.pop("_is_class_method", False) 

129 ctx_method = kwargs.pop("ctx_method", self._default_ctx_method) 

130 if ctx_method == "new": 

131 child_ctx = self._new_ctx_func() 

132 elif ctx_method == "copy": 

133 child_ctx = parent_ctx.copy() 

134 elif ctx_method == "same": 

135 child_ctx = parent_ctx.inherit() 

136 elif ctx_method == "resume": 

137 # NOTE: the resume method is not thread-safe 

138 # For standalone functions or class methods, they should de defined within the thread 

139 # So that the function or `cls` is thread-local 

140 # For object methods, the object need to be created within the thread 

141 # So that `self` is thread-local 

142 if is_class_method: 

143 self_or_cls = args[0] # is it guaranteed? need double check 

144 var_name = f"{self._name}_appl_ctx_" 

145 

146 def reset_context(): 

147 setattr(self_or_cls, var_name, None) 

148 

149 # try to retrieve the context from the class 

150 if (ctx := getattr(self_or_cls, var_name, None)) is None: 

151 # copy the parent context if not exist 

152 child_ctx = parent_ctx.copy() 

153 setattr(self_or_cls, var_name, child_ctx) 

154 else: 

155 # resume from the last run, but with clean locals 

156 child_ctx = ctx.inherit() 

157 else: 

158 

159 def reset_context(): 

160 self._persist_ctx = None 

161 

162 if self._persist_ctx is None: 

163 self._persist_ctx = parent_ctx 

164 

165 child_ctx = self._persist_ctx.inherit() 

166 

167 self._reset_context_func = reset_context 

168 else: 

169 raise ValueError(f"Unknown ctx_method: {ctx_method}") 

170 child_ctx.is_outmost = False 

171 child_ctx._func_name = self._name 

172 child_ctx._func_docstring = self._doc 

173 child_ctx._docstring_quote_count = self._func._docstring_quote_count 

174 child_ctx._include_docstring = self._include_docstring # set in the context 

175 

176 compositor: Optional[Compositor] = kwargs.pop( 

177 "compositor", self._default_compositor 

178 ) 

179 # push the compositor 

180 if compositor is not None: 

181 child_ctx.push_printer(compositor.push_args) 

182 

183 # run the function 

184 results = self._run(parent_ctx, child_ctx, *args, **kwargs) 

185 if results is None and self._default_return == "prompt": 

186 results = child_ctx.records # the default return result 

187 

188 # pop the compositor 

189 if compositor is not None: 

190 child_ctx.pop_printer() 

191 

192 if ctx_method == "same": 

193 parent_ctx.add_records(child_ctx.records, write_to_prompt=False) 

194 

195 return results 

196 

197 __call__ = _call 

198 

199 def __repr__(self): 

200 res = f"[PromptFunc]{self._name}{self._signature}" 

201 if self._doc is not None: 

202 res += f": {self._doc}" 

203 return res