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

1#!/usr/bin/env python 

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

3 

4from rich.console import Console 

5from traitlets.config import Config 

6from IPython.terminal.prompts import ClassicPrompts, Token 

7from IPython.terminal.embed import embed, InteractiveShellEmbed 

8from IPython.core.magic import Magics, magics_class, line_cell_magic,line_magic, cell_magic, register_line_magic 

9 

10from ..aipy import TaskManager 

11from .. import T, __version__ 

12 

13class MainPrompt(ClassicPrompts): 

14 def in_prompt_tokens(self): 

15 return [(Token.Prompt, '>> ')] 

16 

17class TaskPrompt(ClassicPrompts): 

18 def in_prompt_tokens(self): 

19 return [(Token.Prompt, '>>> ')] 

20 

21def get_task_config(): 

22 c = Config() 

23 c.TerminalInteractiveShell.display_banner = False 

24 c.InteractiveShell.prompt_in1 = '>> ' 

25 c.InteractiveShell.prompt_in2 = '... ' 

26 c.InteractiveShell.prompt_out = '' 

27 c.InteractiveShell.banner1 = 'Use /ai "Task" to start a task' 

28 c.InteractiveShell.banner2 = 'For example: ai 获取Google网站标题\n' 

29 c.InteractiveShell.separate_in = '' 

30 c.InteractiveShell.separate_out = '' 

31 c.InteractiveShell.separate_out2 = '' 

32 c.InteractiveShell.prompts_class = TaskPrompt 

33 c.InteractiveShell.confirm_exit = False 

34 return c 

35 

36def get_main_config(): 

37 c = Config() 

38 c.TerminalInteractiveShell.display_banner = False 

39 c.InteractiveShell.prompt_in1 = '>> ' 

40 c.InteractiveShell.prompt_in2 = '... ' 

41 c.InteractiveShell.prompt_out = '' 

42 c.InteractiveShell.banner1 = 'Use /ai "Task" to start a task' 

43 c.InteractiveShell.banner2 = 'For example: ai 获取Google网站标题\n' 

44 c.InteractiveShell.separate_in = '' 

45 c.InteractiveShell.separate_out = '' 

46 c.InteractiveShell.separate_out2 = '' 

47 c.InteractiveShell.prompts_class = MainPrompt 

48 c.InteractiveShell.confirm_exit = False 

49 return c 

50 

51@magics_class 

52class AIMagics(Magics): 

53 def __init__(self, shell, ai): 

54 super().__init__(shell) 

55 self.ai = ai 

56 

57 @line_magic 

58 def task(self, line): 

59 task = self.ai.new_task() 

60 user_ns = {'task': task, 'settings': self.ai.settings} 

61 shell = InteractiveShellEmbed(user_ns=user_ns, config=get_task_config()) 

62 shell() 

63 

64 @line_magic 

65 def clear(self, _): 

66 self.ai.clear() 

67 

68 @line_magic 

69 def save(self, line): 

70 self.ai.save(line) 

71 

72 @line_cell_magic 

73 def ai(self, line, cell=None): 

74 print(line) 

75 print(cell) 

76 

77def main(settings): 

78 console = Console(record=True) 

79 console.print(f"[bold cyan]🚀 Python use - AIPython ({__version__}) [[green]https://aipy.app[/green]]") 

80 

81 try: 

82 ai = TaskManager(settings, console=console) 

83 except Exception as e: 

84 console.print_exception(e) 

85 return 

86 

87 update = ai.get_update(True) 

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

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

90 

91 if not ai.client_manager: 

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

93 return 

94 

95 names = ai.client_manager.names 

96 console.print(f"{T('Please use ai(task) to enter the task to be processed by AI (enter ai.use(llm) to switch to the following LLM:')}", style="green") 

97 console.print(f"[cyan]{T('Default')}: [green]{names['default']}，[cyan]{T('Enabled')}: [yellow]{' '.join(names['enabled'])}") 

98 

99 user_ns = {'AI': ai, 'settings': settings} 

100 shell = InteractiveShellEmbed(user_ns=user_ns, config=get_main_config()) 

101 shell.register_magics(AIMagics(shell, ai)) 

102 shell() 

103