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

1#! /usr/bin/env python3

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

3 import sys

4 from typing import Union, Any, Dict

5 from functools import wraps

6 import inspect

7

8 from term_image.image import from_file, from_url

9

10 from .functions import FunctionManager

11 from .. import T, TaskPlugin

12 from ..exec import PythonRuntime

13 from .blocks import CodeBlock

14

15 def 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

27 class 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