Coverage for aipyapp/cli/cli_task.py: 0%

1#!/usr/bin/env python 

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

3from importlib.resources import read_text 

4 

5from rich.console import Console 

6from rich.text import Text 

7from prompt_toolkit import PromptSession 

8from prompt_toolkit.history import FileHistory 

9from prompt_toolkit.styles import Style 

10from prompt_toolkit.cursor_shapes import CursorShape 

11from prompt_toolkit.auto_suggest import AutoSuggestFromHistory 

12 

13from ..aipy import TaskManager 

14from .. import T, __version__, __respkg__ 

15from .command import CommandManager, TaskModeResult, CommandError, CommandResult 

16from ..display import DisplayManager 

17 

18STYLE_MAIN = { 

19 'completion-menu.completion': 'bg:green #ffffff', 

20 'completion-menu.completion.current': 'bg:#444444 #ffffff', 

21 'completion-menu.meta': 'bg:#000000 #999999', 

22 'completion-menu.meta.current': 'bg:#444444 #aaaaaa', 

23 'prompt': 'green', 

24 'bottom-toolbar': 'bg:#FFFFFF green' 

25} 

26 

27STYLE_AI = { 

28 'completion-menu.completion': 'bg:#008080 #ffffff', # 深蓝背景，白色文本 

29 'completion-menu.completion.current': 'bg:#005577 #ffffff', # 当前选中，亮蓝 

30 'completion-menu.meta': 'bg:#002244 #cccccc', # 补全项的 meta 信息 

31 'completion-menu.meta.current': 'bg:#005577 #eeeeee', # 当前选中的 meta 

32 'prompt': '#008080', 

33 'bottom-toolbar': "bg:#880000 #008080" 

34} 

35 

36class InteractiveConsole(): 

37 def __init__(self, tm, console, settings): 

38 self.tm = tm 

39 self.names = tm.client_manager.names 

40 self.history = FileHistory(str(settings['config_dir'] / ".history")) 

41 self.console = console 

42 self.settings = settings 

43 self.task = None 

44 self.style_main = Style.from_dict(STYLE_MAIN) 

45 self.style_task = Style.from_dict(STYLE_AI) 

46 self.command_manager = CommandManager(settings,tm, console) 

47 self.completer = self.command_manager 

48 self.session = PromptSession( 

49 history=self.history, 

50 completer=self.completer, 

51 auto_suggest=AutoSuggestFromHistory(), 

52 bottom_toolbar=self.get_bottom_toolbar, 

53 key_bindings=self.command_manager.create_key_bindings() 

54 ) 

55 

56 def get_main_status(self): 

57 status = self.tm.get_status() 

58 try: 

59 mcp_text = f" | MCP: {T('Enabled') if status['mcp_enabled'] else T('Disabled')}" 

60 except KeyError: 

61 mcp_text = "" 

62 return f"LLM: {status['llm']} | Role: {status['role']} | Display: {status['display']} | Tasks: {status['tasks']}{mcp_text}" 

63 

64 def get_task_status(self): 

65 if self.task: 

66 status = self.task.get_status() 

67 return f"LLM: {status['llm']} | Blocks: {status['blocks']} | Steps: {status['steps']}" 

68 return "" 

69 

70 def get_bottom_toolbar(self): 

71 if self.command_manager.is_task_mode(): 

72 status = self.get_task_status() 

73 text = f"[AI] {status}" 

74 else: 

75 status = self.get_main_status() 

76 text = f"[Main] {status}" 

77 return [('class:bottom-toolbar', text)] 

78 

79 def input_with_possible_multiline(self, prompt_text, task_mode=False): 

80 session = self.session 

81 style = self.style_task if task_mode else self.style_main 

82 cursor_shape = CursorShape.BEAM if not task_mode else CursorShape.BLOCK 

83 first_line = session.prompt([("class:prompt", prompt_text)], style=style, cursor=cursor_shape) 

84 if not first_line.endswith("\\"): 

85 return first_line 

86 # Multi-line input 

87 lines = [first_line.rstrip("\\")] 

88 while True: 

89 next_line = session.prompt([("class:prompt", "... ")], style=style) 

90 if next_line.endswith("\\"): 

91 lines.append(next_line.rstrip("\\")) 

92 else: 

93 lines.append(next_line) 

94 break 

95 return "\n".join(lines) 

96 

97 def run_task(self, task, instruction, title=None): 

98 try: 

99 task.run(instruction, title=title) 

100 except (EOFError, KeyboardInterrupt): 

101 pass 

102 except Exception as e: 

103 self.console.print_exception() 

104 

105 def start_task_mode(self, task, instruction=None, title=None): 

106 if instruction: 

107 self.console.print(f"[AI] {T('Enter Ctrl+d or /done to end current task')}", style="dim color(240)") 

108 self.run_task(task, instruction, title=title) 

109 else: 

110 self.console.print(f"[AI] {T('Resuming task')}: {task.instruction[:32]}", style="dim color(240)") 

111 

112 while True: 

113 self.task = task 

114 self.command_manager.set_task_mode(task) 

115 try: 

116 user_input = self.input_with_possible_multiline(">>> ", task_mode=True).strip() 

117 if len(user_input) < 2: continue 

118 except (EOFError, KeyboardInterrupt): 

119 break 

120 

121 if user_input in ('/done', 'done'): 

122 break 

123 

124 if not user_input.startswith('/'): 

125 self.run_task(task, user_input) 

126 continue 

127 

128 try: 

129 self.command_manager.execute(user_input) 

130 except CommandError as e: 

131 self.console.print(f"[red]{e}[/red]") 

132 

133 try: 

134 task.done() 

135 except Exception as e: 

136 self.console.print_exception() 

137 self.task = None 

138 self.console.print(f"[{T('Exit AI mode')}]", style="dim") 

139 

140 def run(self): 

141 self.console.print(f"[Main] {T('Please enter an instruction or `/help` for more information')}", style="dim color(240)") 

142 tm = self.tm 

143 while True: 

144 self.command_manager.set_main_mode() 

145 try: 

146 user_input = self.input_with_possible_multiline(">> ").strip() 

147 if len(user_input) < 2: 

148 continue 

149 

150 if not user_input.startswith('/'): 

151 task = tm.new_task() 

152 self.start_task_mode(task, user_input) 

153 continue 

154 

155 try: 

156 ret = self.command_manager.execute(user_input) 

157 if isinstance(ret, CommandResult) and isinstance(ret.result, TaskModeResult): 

158 result = ret.result 

159 if result.instruction: 

160 task = tm.new_task() 

161 title = result.title 

162 else: 

163 task = result.task 

164 title = None 

165 self.start_task_mode(task, result.instruction, title=title) 

166 continue 

167 except CommandError as e: 

168 self.console.print(f"[red]{e}[/red]") 

169 except (EOFError, KeyboardInterrupt): 

170 break 

171 

172def get_logo_text(config_dir): 

173 path = config_dir / "logo.txt" 

174 if path.exists(): 

175 logo_text = path.read_text() 

176 else: 

177 logo_text = read_text(__respkg__, "logo.txt") 

178 return logo_text 

179 

180def main(settings): 

181 console = Console(record=True) 

182 console.print(f"🚀 Python use - AIPython ({__version__}) [[pink]https://aipy.app[/pink]]", style="bold green") 

183 logo_text = get_logo_text(settings['config_dir']) 

184 console.print(Text.from_ansi(logo_text)) 

185 

186 # 初始化显示效果管理器 

187 display_config = settings.get('display', {}) 

188 display_manager = DisplayManager(display_config, console=console) 

189 try: 

190 tm = TaskManager(settings, display_manager=display_manager) 

191 except Exception as e: 

192 console.print_exception() 

193 return 

194 

195 update = tm.get_update() 

196 if update and update.get('has_update'): 

197 console.print(f"[bold red]🔔 号外❗ {T('Update available')}: {update.get('latest_version')}") 

198 

199 if not tm.client_manager: 

200 console.print(f"[bold red]{T('No available LLM, please check the configuration file')}") 

201 return 

202 

203 cmd = settings.get('exec_cmd') 

204 if cmd: 

205 tm.new_task().run(cmd) 

206 return 

207 InteractiveConsole(tm, console, settings).run()