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

1 from collections import OrderedDict

2 import questionary

3 import requests

4

5 from loguru import logger

6

7 from .. import T

8 from ..config.llm import LLMConfig, PROVIDERS

9 from .trustoken import TrustToken

10

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

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

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