Coverage for .venv/lib/python3.13/site-packages/litellm/proxy/client/cli/commands/autoroute/wizard.py: 0%

1import sys 

2from pathlib import Path 

3from typing import Final 

4 

5import click 

6import yaml 

7from InquirerPy import inquirer 

8from InquirerPy.base.control import Choice 

9from pydantic import JsonValue, TypeAdapter, ValidationError 

10 

11from .... import Client 

12from .config import ( 

13 DEFAULT_KEYWORD_TIER_RULES, 

14 TIER_NAMES, 

15 AutorouteConfig, 

16 ConfigGenerationError, 

17 DiscoveredModel, 

18 HeuristicClassifier, 

19 KeywordTierRule, 

20 LLMClassifier, 

21 NoSemanticMatching, 

22 SemanticMatching, 

23 build_generated_model_list, 

24 chat_models, 

25 embedding_models, 

26 master_key_from_config, 

27 parse_discovered_models, 

28 validate_config, 

29) 

30from .process import CONFIG_PATH, secure_create 

31 

32 

33def _is_interactive() -> bool: 

34 return sys.stdin.isatty() 

35 

36 

37def _fuzzy_pick(models: tuple[DiscoveredModel, ...], prompt_label: str, multiselect: bool) -> list[str]: 

38 """Type-to-filter picker over a (possibly huge) model pool, using InquirerPy's fzf-style fuzzy prompt. 

39 

40 A plain numbered table + typed index does not scale past a handful of models -- proxies with 

41 hundreds of model groups made that interaction unusable. This lets the user narrow the pool by 

42 typing a substring instead of scrolling/counting. 

43 

44 Assumes the caller already checked interactivity (run_configure_wizard does, once, up front) -- 

45 checking here too would check the wrong thing under test, where InquirerPy is driven through its 

46 own injected input/output rather than the real process stdin. 

47 """ 

48 choices: Final = [Choice(value=model.name, name=model.name) for model in models] 

49 toggle_hint: Final = "tab to toggle, " if multiselect else "" 

50 while True: 

51 result = inquirer.fuzzy( 

52 message=f"{prompt_label}: type to filter, {toggle_hint}enter to confirm", 

53 choices=choices, 

54 multiselect=multiselect, 

55 max_height="70%", 

56 ).execute() 

57 selected = result if multiselect else [result] 

58 if selected: 

59 return selected 

60 click.echo("Select at least one model.") 

61 

62 

63def _render_and_prompt_for_model(models: tuple[DiscoveredModel, ...], prompt_label: str) -> str: 

64 return _fuzzy_pick(models, prompt_label, multiselect=False)[0] 

65 

66 

67def _render_and_prompt_for_models(models: tuple[DiscoveredModel, ...], prompt_label: str) -> tuple[str, ...]: 

68 return tuple(_fuzzy_pick(models, prompt_label, multiselect=True)) 

69 

70 

71def _parse_keywords(raw: str) -> tuple[str, ...]: 

72 return tuple(keyword.strip() for keyword in raw.split(",") if keyword.strip()) 

73 

74 

75def _prompt_for_keyword_tier_rules() -> tuple[KeywordTierRule, ...]: 

76 """Let the user supply the semantic-matching keywords per tier, since matching those 

77 keywords against the request is the whole point of enabling it. Each prompt is prefilled 

78 with the built-in default, so pressing enter keeps it.""" 

79 click.echo("\nEnter example keywords/phrases per tier (comma-separated); press enter to keep the default:") 

80 defaults: Final = {rule.tier: rule.keywords for rule in DEFAULT_KEYWORD_TIER_RULES} 

81 

82 def _rule_for(tier: str) -> KeywordTierRule: 

83 default_keywords: Final = defaults.get(tier, ()) 

84 raw: Final = click.prompt(f" {tier} keywords", default=", ".join(default_keywords), show_default=True) 

85 return KeywordTierRule(keywords=_parse_keywords(raw) or default_keywords, tier=tier) 

86 

87 return tuple(_rule_for(tier) for tier in TIER_NAMES) 

88 

89 

90_RAW_CONFIG_ADAPTER: Final = TypeAdapter(dict[str, JsonValue]) 

91 

92 

93def _load_persisted_master_key(config_path: Path) -> str | None: 

94 """The master key from an existing generated config, so a rewrite carries it forward. 

95 

96 Lenient on a missing or corrupt file: configure is the regeneration path, so it must 

97 succeed from any prior state; a key that cannot be read is simply not carried and `start` 

98 mints a fresh one. 

99 """ 

100 if not config_path.exists(): 

101 return None 

102 try: 

103 raw: Final = _RAW_CONFIG_ADAPTER.validate_python(yaml.safe_load(config_path.read_text())) 

104 except (OSError, UnicodeDecodeError, yaml.YAMLError, ValidationError): 

105 return None 

106 return master_key_from_config(raw) 

107 

108 

109def run_configure_wizard(ctx: click.Context) -> Path: 

110 """Discover the caller's accessible models, walk them through tier assignment, write config.""" 

111 base_url: Final = ctx.obj["base_url"] 

112 api_key: Final = ctx.obj["api_key"] 

113 client: Final = Client(base_url=base_url, api_key=api_key) 

114 

115 raw_models: Final = client.models.list() 

116 if not isinstance(raw_models, list): 

117 raise click.ClickException( 

118 f"Unexpected response from /v1/models: expected a list, got {type(raw_models).__name__}" 

119 ) 

120 discovered: Final = parse_discovered_models(raw_models) 

121 chat_pool: Final = chat_models(discovered) 

122 embedding_pool: Final = embedding_models(discovered) 

123 

124 if not chat_pool: 

125 raise click.ClickException("Your key has no chat-capable models available on this proxy.") 

126 

127 if not _is_interactive(): 

128 raise click.ClickException("`lite autoroute configure` requires an interactive terminal.") 

129 

130 click.echo("Assign model(s) to each complexity tier (from what your key can access):") 

131 tiers: Final = {tier: _render_and_prompt_for_models(chat_pool, tier) for tier in TIER_NAMES} 

132 default_model: Final = tiers["MEDIUM"][0] 

133 

134 classifier = HeuristicClassifier() 

135 if click.confirm("\nUse an LLM classifier instead of the free heuristic scorer?", default=False): 

136 classifier_model: Final = _render_and_prompt_for_model(chat_pool, "LLM classifier") 

137 classifier = LLMClassifier(model=classifier_model) 

138 

139 semantic_matching = NoSemanticMatching() 

140 if embedding_pool and click.confirm("\nEnable semantic keyword matching?", default=False): 

141 embedding_model: Final = _render_and_prompt_for_model(embedding_pool, "semantic embeddings") 

142 keyword_tier_rules: Final = _prompt_for_keyword_tier_rules() 

143 semantic_matching = SemanticMatching(embedding_model=embedding_model, keyword_tier_rules=keyword_tier_rules) 

144 

145 adaptive: Final = click.confirm("\nEnable adaptive (bandit) selection on top of tiering?", default=False) 

146 

147 config: Final = AutorouteConfig( 

148 base_url=base_url, 

149 api_key=api_key, 

150 tiers=tiers, 

151 default_model=default_model, 

152 classifier=classifier, 

153 semantic_matching=semantic_matching, 

154 adaptive=adaptive, 

155 ) 

156 try: 

157 validate_config(config, discovered) 

158 except ConfigGenerationError as e: 

159 raise click.ClickException(str(e)) 

160 

161 model_list: Final = build_generated_model_list(config) 

162 persisted_master_key: Final = _load_persisted_master_key(CONFIG_PATH) 

163 generated: Final[dict[str, JsonValue]] = ( 

164 {"model_list": model_list, "general_settings": {"master_key": persisted_master_key}} 

165 if persisted_master_key is not None 

166 else {"model_list": model_list} 

167 ) 

168 CONFIG_PATH.parent.mkdir(parents=True, exist_ok=True) 

169 with secure_create(CONFIG_PATH) as f: 

170 yaml.safe_dump(generated, f, sort_keys=False) 

171 

172 click.echo(f"\nWrote {CONFIG_PATH}") 

173 for tier, models in tiers.items(): 

174 click.echo(f" {tier}: {', '.join(models)}") 

175 return CONFIG_PATH 

176 

177 

178__all__ = ["run_configure_wizard"]