Coverage for aipyapp/cli/command/cmd_llm.py: 0%

1from ... import T 

2from .base import CommandMode, Completable 

3from .base_parser import ParserCommand 

4from .utils import print_records 

5 

6class LLMCommand(ParserCommand): 

7 name = 'llm' 

8 description = T('LLM operations') 

9 modes = [CommandMode.MAIN, CommandMode.TASK] 

10 

11 def add_subcommands(self, subparsers): 

12 use_parser = subparsers.add_parser('use', help=T('Use a LLM provider')) 

13 use_parser.add_argument('provider', type=str, help=T('Provider name')) 

14 subparsers.add_parser('list', help=T('List LLM providers')) 

15 

16 def get_arg_values(self, arg, subcommand=None, partial_value=''): 

17 if subcommand == 'use' and arg.name == 'provider': 

18 ctx = self.manager.context 

19 return [Completable(client.name, str(client)) for client in ctx.tm.client_manager.clients.values()] 

20 return super().get_arg_values(arg, subcommand) 

21 

22 def cmd_list(self, args, ctx): 

23 rows = ctx.tm.list_llms() 

24 print_records(rows) 

25 

26 def cmd_use(self, args, ctx): 

27 if ctx.task: 

28 ret = ctx.task.use(args.provider) 

29 else: 

30 ret = ctx.tm.use(llm=args.provider) 

31 ctx.console.print(T('Success') if ret else T('Failed'), style='cyan' if ret else 'red') 

32 if ret: 

33 self.log.info(f'Use {args.provider} LLM provider') 

34 return ret 

35 

36 def cmd(self, args, ctx): 

37 self.cmd_list(args, ctx)