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

1#!/usr/bin/env python

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

3

4 from rich.console import Console

5 from traitlets.config import Config

6 from IPython.terminal.prompts import ClassicPrompts, Token

7 from IPython.terminal.embed import embed, InteractiveShellEmbed

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

9

10 from ..aipy import TaskManager

11 from .. import T, __version__

12

13 class MainPrompt(ClassicPrompts):

14 def in_prompt_tokens(self):

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

16

17 class TaskPrompt(ClassicPrompts):

18 def in_prompt_tokens(self):

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

20

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

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

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

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