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

1#!/usr/bin/env python

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

3 from importlib.resources import read_text

4

5 from rich.console import Console

6 from rich.text import Text

7 from prompt_toolkit import PromptSession

8 from prompt_toolkit.history import FileHistory

9 from prompt_toolkit.styles import Style

10 from prompt_toolkit.cursor_shapes import CursorShape

11 from prompt_toolkit.auto_suggest import AutoSuggestFromHistory

12

13 from ..aipy import TaskManager

14 from .. import T, __version__, __respkg__

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

16 from ..display import DisplayManager

17

18 STYLE_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

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

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

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

180 def 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()