Coverage for aipyapp/aipy/runtime.py: 29%

1#! /usr/bin/env python3 

2# -*- coding: utf-8 -*- 

3import sys 

4from typing import Union, Any, Dict 

5from functools import wraps 

6import inspect 

7 

8from term_image.image import from_file, from_url 

9 

10from .functions import FunctionManager 

11from .. import T, TaskPlugin 

12from ..exec import PythonRuntime 

13from .blocks import CodeBlock 

14 

15def restore_output(func): 

16 @wraps(func) 

17 def wrapper(self, *args, **kwargs): 

18 old_stdout, old_stderr = sys.stdout, sys.stderr 

19 sys.stdout, sys.stderr = sys.__stdout__, sys.__stderr__ 

20 

21 try: 

22 return func(self, *args, **kwargs) 

23 finally: 

24 sys.stdout, sys.stderr = old_stdout, old_stderr 

25 return wrapper 

26 

27class CliPythonRuntime(PythonRuntime): 

28 def __init__(self, task): 

29 super().__init__(task.role.envs) 

30 self.gui = task.gui 

31 self.task = task 

32 self.display = task.display 

33 self._auto_install = task.settings.get('auto_install') 

34 self._auto_getenv = task.settings.get('auto_getenv') 

35 self.function_manager = FunctionManager() 

36 

37 def register_plugin(self, plugin: TaskPlugin): 

38 self.function_manager.register_functions(plugin.get_functions()) 

39 

40 @restore_output 

41 def install_packages(self, *packages: str) -> bool: 

42 """ 

43 Install third-party packages 

44 

45 Args: 

46 packages: The names of the packages to install 

47 

48 Returns: 

49 bool: True if the packages are installed successfully, False otherwise 

50 

51 Examples: 

52 >>> runtime.install_packages('requests', 'openai') 

53 True 

54 >>> runtime.install_packages('requests', 'openai') 

55 False 

56 """ 

57 message = f"LLM {T('Request to install third-party packages')}: {packages}" 

58 self.task.emit('runtime_message', message=message, status='warning') 

59 

60 if self.display: 

61 prompt = f"💬 {T('If you agree, please enter')} 'y'> " 

62 ok = self.display.confirm(prompt, auto=self._auto_install) 

63 else: 

64 ok = True 

65 

66 if ok: 

67 ret = self.ensure_packages(*packages) 

68 result_message = T("Package installation completed") if ret else T("Package installation failed") 

69 self.task.emit('runtime_message', message=result_message, status='success' if ret else 'error') 

70 return ret 

71 return False 

72 

73 @restore_output 

74 def get_env(self, name: str, default: str = None, *, desc: str = None) -> Union[str, None]: 

75 message = f"\n⚠️ LLM {T('Request to obtain environment variable {}, purpose', name)}: {desc}" 

76 self.task.emit('runtime_message', message=message) 

77 

78 try: 

79 value = self.envs[name][0] 

80 success_message = f"✅ {T('Environment variable {} exists, returned for code use', name)}" 

81 self.task.emit('runtime_message', message=success_message) 

82 except KeyError: 

83 if self._auto_getenv: 

84 auto_message = f"✅ {T('Auto confirm')}" 

85 self.task.emit('runtime_message', message=auto_message) 

86 value = None 

87 elif self.display: 

88 prompt = f"💬 {T('Environment variable {} not found, please enter', name)}: " 

89 value = self.display.input(prompt) 

90 value = value.strip() 

91 else: 

92 value = None 

93 if value: 

94 self.set_env(name, value, desc) 

95 return value or default 

96 

97 @restore_output 

98 def show_image(self, path: str = None, url: str = None) -> None: 

99 """ 

100 Display an image 

101 

102 Args: 

103 path: The path of the image 

104 url: The URL of the image 

105 """ 

106 self.task.emit('show_image', path=path, url=url) 

107 if not self.gui: 

108 image = from_file(path) if path else from_url(url) 

109 image.draw() 

110 

111 @restore_output 

112 def input(self, prompt: str) -> str: 

113 self.task.emit('runtime_input', prompt=prompt) 

114 if self.display: 

115 return self.display.input(prompt) 

116 return None 

117 

118 def get_block_by_name(self, block_name: str) -> Union[CodeBlock, None]: 

119 """ 

120 Get a code block by name 

121 

122 Args: 

123 block_name: The name of the code block 

124 

125 Returns: 

126 CodeBlock: The code block objector None if not found 

127 """ 

128 return self.task.code_blocks.get_block_by_name(block_name) 

129 

130 def call_function(self, name: str, **kwargs) -> Any: 

131 """ 

132 Call a registered function 

133 

134 Args: 

135 name: The name of the function to call 

136 **kwargs: The keyword arguments to pass to the function 

137 

138 Returns: 

139 Any: The result of the function call or raise an exception 

140 

141 Examples: 

142 >>> utils.call_function('get_env', name='PATH') 

143 '/usr/local/bin:/usr/bin:/bin:/usr/sbin:/sbin' 

144 >>> utils.call_function('get_env', name='PATH') 

145 None 

146 """ 

147 self.task.emit('call_function', funcname=name, kwargs=kwargs) 

148 try: 

149 result = self.function_manager.call(name, **kwargs) 

150 self.task.emit('call_function_result', funcname=name, kwargs=kwargs, result=result, success=True) 

151 return result 

152 except Exception as e: 

153 self.task.emit('call_function_result', funcname=name, kwargs=kwargs, result=None, success=False, error=str(e), exception=e) 

154 raise 

155 

156 def get_builtin_functions(self) -> Dict[str, Dict[str, str]]: 

157 """ 

158 根据函数签名和docstring，生成函数调用提示 

159 """ 

160 functions = {} 

161 

162 # 内置运行时函数 

163 builtin_names = ['set_state', 'get_block_state', 'set_persistent_state', 'get_persistent_state', 'install_packages', 'get_env', 'show_image', 'get_block_by_name', 'call_function'] 

164 for name in builtin_names: 

165 func_obj = getattr(self, name) 

166 docstring = func_obj.__doc__ 

167 signature = inspect.signature(func_obj) 

168 functions[name] = { 

169 'docstring': docstring, 

170 'signature': signature, 

171 } 

172 return functions 

173 

174 def get_plugin_functions(self): 

175 return self.function_manager.get_functions() 

176