Coverage for aipyapp/aipy/wizard.py: 0%

1from collections import OrderedDict 

2import questionary 

3import requests 

4 

5from loguru import logger 

6 

7from .. import T 

8from ..config.llm import LLMConfig, PROVIDERS 

9from .trustoken import TrustToken 

10 

11def get_models(providers, provider, api_key: str) -> list: 

12 """获取可用的模型列表""" 

13 provider_info = providers[provider] 

14 headers = { 

15 "Content-Type": "application/json" 

16 } 

17 

18 if provider == "Claude": 

19 headers["x-api-key"] = api_key 

20 headers["anthropic-version"] = "2023-06-01" 

21 else: 

22 headers["Authorization"] = f"Bearer {api_key}" 

23 

24 try: 

25 response = requests.get( 

26 f"{provider_info['api_base']}{provider_info['models_endpoint']}", 

27 headers=headers 

28 ) 

29 logger.info(f"获取模型列表: {response.text}") 

30 if response.status_code == 200: 

31 data = response.json() 

32 logger.info(f"获取模型列表成功: {data}") 

33 if provider in ["OpenAI", "DeepSeek", "xAI", "Claude"]: 

34 return [model["id"] for model in data["data"]] 

35 elif provider == "Gemini": 

36 return [model["name"] for model in data["models"]] 

37 return [] 

38 except Exception as e: 

39 logger.error(f"获取模型列表失败: {str(e)}") 

40 return [] # 如果API调用失败，返回空列表 

41 

42def select_provider(llm_config, default='Trustoken'): 

43 """选择提供商""" 

44 while True: 

45 name = questionary.select( 

46 T('Select LLM Provider'), 

47 choices=[ 

48 questionary.Choice(title=T('Trustoken is an intelligent API Key management service'), value='Trustoken', description=T('Recommended for beginners, easy to configure and feature-rich')), 

49 questionary.Choice(title=T('Other'), value='other') 

50 ], 

51 default=default 

52 ).unsafe_ask() 

53 if name == default: 

54 return default 

55 

56 names = [name for name in llm_config.providers.keys() if name != default] 

57 names.append('<--') 

58 name = questionary.select( 

59 T('Select other providers'), 

60 choices=names 

61 ).unsafe_ask() 

62 if name != '<--': 

63 return name 

64 

65 

66def config_llm(llm_config): 

67 """配置 LLM 提供商""" 

68 # 第一步：选择提供商 

69 name = select_provider(llm_config) 

70 if not name: 

71 return None 

72 provider = llm_config.providers[name] 

73 config = {"type": provider["type"]} 

74 

75 if name == 'Trustoken': 

76 def save_token(token): 

77 config['api_key'] = token 

78 

79 tt = TrustToken() 

80 if not tt.fetch_token(save_token): 

81 return None 

82 

83 config['model'] = 'auto' 

84 else: 

85 api_key = questionary.text( 

86 T('Enter your API key'), 

87 validate=lambda x: len(x) > 8 

88 ).unsafe_ask() 

89 config['api_key'] = api_key 

90 

91 # 获取可用模型列表 

92 available_models = get_models(llm_config.providers, name, api_key) 

93 if not available_models: 

94 logger.warning(T('Model list acquisition failed')) 

95 return None 

96 

97 # 第三步：选择模型 

98 model = questionary.select( 

99 T('Available Models'), 

100 choices=available_models 

101 ).unsafe_ask() 

102 config['model'] = model 

103 

104 # 保存配置 

105 config['enable'] = True 

106 current_config = llm_config.config 

107 current_config[name] = config 

108 llm_config.save_config(current_config) 

109 return current_config 

110