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

1from collections.abc import Mapping 

2from typing import Final, Literal 

3 

4from pydantic import BaseModel, ConfigDict, Field, JsonValue, TypeAdapter 

5 

6TIER_NAMES: Final[tuple[str, ...]] = ("SIMPLE", "MEDIUM", "COMPLEX", "REASONING") 

7AUTOROUTER_MODEL_NAME: Final = "autorouter" 

8 

9 

10class ConfigGenerationError(Exception): 

11 """Raised when an AutorouteConfig references a model the discovery step didn't find.""" 

12 

13 

14class DiscoveredModel(BaseModel): 

15 model_config = ConfigDict(frozen=True) 

16 

17 name: str 

18 mode: str = "chat" 

19 

20 

21class _RawModelListing(BaseModel): 

22 model_config = ConfigDict(extra="ignore") 

23 

24 id: str 

25 # /v1/models attaches "mode" (sourced from the cost map) only for models it can resolve; 

26 # a model whose mode is unknown arrives without the field, so default it to chat rather 

27 # than dropping it, which keeps it selectable as a routing target in the wizard. 

28 mode: str = "chat" 

29 

30 

31_RAW_MODEL_LISTING_ADAPTER: Final = TypeAdapter(list[_RawModelListing]) 

32 

33 

34def parse_discovered_models(raw: list[JsonValue]) -> tuple[DiscoveredModel, ...]: 

35 """Validate a raw `/v1/models` response into typed models.""" 

36 parsed: Final = _RAW_MODEL_LISTING_ADAPTER.validate_python(raw) 

37 return tuple(DiscoveredModel(name=item.id, mode=item.mode) for item in parsed) 

38 

39 

40def chat_models(models: tuple[DiscoveredModel, ...]) -> tuple[DiscoveredModel, ...]: 

41 return tuple(m for m in models if m.mode == "chat") 

42 

43 

44def embedding_models(models: tuple[DiscoveredModel, ...]) -> tuple[DiscoveredModel, ...]: 

45 return tuple(m for m in models if m.mode == "embedding") 

46 

47 

48class HeuristicClassifier(BaseModel): 

49 model_config = ConfigDict(frozen=True) 

50 kind: Literal["heuristic"] = "heuristic" 

51 

52 

53class LLMClassifier(BaseModel): 

54 model_config = ConfigDict(frozen=True) 

55 kind: Literal["llm"] = "llm" 

56 model: str 

57 timeout_ms: int = 3000 

58 

59 

60ClassifierChoice = HeuristicClassifier | LLMClassifier 

61 

62 

63class NoSemanticMatching(BaseModel): 

64 model_config = ConfigDict(frozen=True) 

65 kind: Literal["none"] = "none" 

66 

67 

68class KeywordTierRule(BaseModel): 

69 model_config = ConfigDict(frozen=True) 

70 keywords: tuple[str, ...] 

71 tier: str 

72 

73 

74# Satisfies complexity_router's "semantic matching requires non-empty keyword_tier_rules" 

75# invariant with a sane starting point; the wizard lets the user override these per tier. 

76DEFAULT_KEYWORD_TIER_RULES: Final[tuple[KeywordTierRule, ...]] = ( 

77 KeywordTierRule(keywords=("hi", "hello", "thanks"), tier="SIMPLE"), 

78 KeywordTierRule(keywords=("explain", "how does"), tier="MEDIUM"), 

79 KeywordTierRule(keywords=("refactor", "implement", "debug"), tier="COMPLEX"), 

80 KeywordTierRule(keywords=("step by step", "think through", "prove"), tier="REASONING"), 

81) 

82 

83 

84class SemanticMatching(BaseModel): 

85 model_config = ConfigDict(frozen=True) 

86 kind: Literal["semantic"] = "semantic" 

87 embedding_model: str 

88 match_threshold: float = 0.5 

89 keyword_tier_rules: tuple[KeywordTierRule, ...] = DEFAULT_KEYWORD_TIER_RULES 

90 

91 

92SemanticMatchingChoice = NoSemanticMatching | SemanticMatching 

93 

94 

95class AutorouteConfig(BaseModel): 

96 model_config = ConfigDict(frozen=True) 

97 

98 base_url: str 

99 api_key: str 

100 # Each tier maps to a pool of one or more models; complexity_router picks randomly among 

101 # them per request (or, in adaptive mode, learns which to prefer within the pool). 

102 tiers: dict[str, tuple[str, ...]] 

103 default_model: str 

104 classifier: ClassifierChoice = Field(default_factory=HeuristicClassifier) 

105 semantic_matching: SemanticMatchingChoice = Field(default_factory=NoSemanticMatching) 

106 adaptive: bool = False 

107 

108 

109def validate_config(config: AutorouteConfig, discovered: tuple[DiscoveredModel, ...]) -> None: 

110 """Raise ConfigGenerationError if config references a model discovery didn't return.""" 

111 chat_names: Final[frozenset[str]] = frozenset(m.name for m in chat_models(discovered)) 

112 embedding_names: Final[frozenset[str]] = frozenset(m.name for m in embedding_models(discovered)) 

113 

114 for tier, models in config.tiers.items(): 

115 for model in models: 

116 if model not in chat_names: 

117 raise ConfigGenerationError(f"Tier {tier} references unknown chat model '{model}'") 

118 

119 if config.default_model not in chat_names: 

120 raise ConfigGenerationError(f"default_model '{config.default_model}' is not a known chat model") 

121 

122 if isinstance(config.classifier, LLMClassifier) and config.classifier.model not in chat_names: 

123 raise ConfigGenerationError(f"classifier model '{config.classifier.model}' is not a known chat model") 

124 

125 if ( 

126 isinstance(config.semantic_matching, SemanticMatching) 

127 and config.semantic_matching.embedding_model not in embedding_names 

128 ): 

129 raise ConfigGenerationError( 

130 f"embedding model '{config.semantic_matching.embedding_model}' is not a known embedding model" 

131 ) 

132 

133 

134def _litellm_proxy_deployment(name: str, base_url: str, api_key: str) -> dict[str, JsonValue]: 

135 return { 

136 "model_name": name, 

137 "litellm_params": { 

138 "model": f"litellm_proxy/{name}", 

139 "api_base": base_url, 

140 "api_key": api_key, 

141 }, 

142 } 

143 

144 

145def build_generated_model_list(config: AutorouteConfig) -> list[JsonValue]: 

146 """Build the model_list for the ephemeral proxy's config.yaml. 

147 

148 Every real model referenced anywhere (tier targets, classifier, embedding) is deduplicated 

149 to exactly one `litellm_proxy/<name>` deployment forwarding to the customer's real proxy, 

150 plus one `auto_router/complexity_router` deployment tying the tiers together. 

151 """ 

152 referenced_names: Final = {model for models in config.tiers.values() for model in models} 

153 referenced_names.add(config.default_model) 

154 if isinstance(config.classifier, LLMClassifier): 

155 referenced_names.add(config.classifier.model) 

156 if isinstance(config.semantic_matching, SemanticMatching): 

157 referenced_names.add(config.semantic_matching.embedding_model) 

158 

159 proxy_deployments: Final = [ 

160 _litellm_proxy_deployment(name, config.base_url, config.api_key) for name in sorted(referenced_names) 

161 ] 

162 

163 complexity_router_config: Final[dict[str, JsonValue]] = { 

164 "tiers": {tier: list(models) for tier, models in config.tiers.items()}, 

165 "default_model": config.default_model, 

166 "return_raw_model_name": True, 

167 } 

168 if isinstance(config.classifier, LLMClassifier): 

169 complexity_router_config["classifier_type"] = "llm" 

170 complexity_router_config["classifier_llm_config"] = { 

171 "model": config.classifier.model, 

172 "timeout_ms": config.classifier.timeout_ms, 

173 } 

174 if isinstance(config.semantic_matching, SemanticMatching): 

175 complexity_router_config["semantic_keyword_matching"] = True 

176 complexity_router_config["embedding_model"] = config.semantic_matching.embedding_model 

177 complexity_router_config["match_threshold"] = config.semantic_matching.match_threshold 

178 complexity_router_config["keyword_tier_rules"] = [ 

179 {"keywords": list(rule.keywords), "tier": rule.tier} for rule in config.semantic_matching.keyword_tier_rules 

180 ] 

181 if config.adaptive: 

182 complexity_router_config["adaptive"] = True 

183 

184 auto_router_litellm_params: Final[dict[str, JsonValue]] = { 

185 "model": "auto_router/complexity_router", 

186 "complexity_router_config": complexity_router_config, 

187 } 

188 # A bare "*" model_name looks like the obvious way to catch every request Claude Code 

189 # might send regardless of which model it thinks it's using, but Router's auto-router 

190 # registry is keyed by the literal requested model string (router.py:10711-10717), not 

191 # resolved through pattern/wildcard matching first -- so a "*" entry here would only ever 

192 # match a client that literally sends model="*", never an actual wildcard catch-all. Callers 

193 # instead need to make Claude Code request this "autorouter" name directly (see 

194 # ANTHROPIC_DEFAULT_*_MODEL in settings.py's merge_claude_settings_static_token). 

195 return [ 

196 *proxy_deployments, 

197 {"model_name": AUTOROUTER_MODEL_NAME, "litellm_params": auto_router_litellm_params}, 

198 ] 

199 

200 

201def build_generated_proxy_config(config: AutorouteConfig, master_key: str) -> dict[str, JsonValue]: 

202 """Full config.yaml content for the ephemeral proxy, including its own auth key. 

203 

204 master_key must live under general_settings, not litellm_settings -- the proxy server 

205 only ever reads general_settings.master_key (proxy_server.py:4530) to authenticate 

206 requests; a key placed under litellm_settings is silently ignored, leaving the proxy 

207 with no real auth at all. 

208 """ 

209 return { 

210 "model_list": build_generated_model_list(config), 

211 "general_settings": {"master_key": master_key}, 

212 } 

213 

214 

215def master_key_from_config(config: dict[str, JsonValue]) -> str | None: 

216 """The master key persisted in a generated config, or None when absent or blank. 

217 

218 Single definition of "this config already has a usable key", shared by `start` (reuse 

219 instead of minting) and the configure wizard (carry the key forward on rewrite) so the 

220 two sites can never disagree on what counts as one. Returned verbatim, never stripped: 

221 the proxy authenticates against the exact bytes under general_settings.master_key, so a 

222 normalized copy here would diverge from what the proxy expects. 

223 """ 

224 general_settings: Final = config.get("general_settings") 

225 if not isinstance(general_settings, Mapping): 

226 return None 

227 master_key: Final = general_settings.get("master_key") 

228 if isinstance(master_key, str) and master_key.strip(): 

229 return master_key 

230 return None 

231 

232 

233__all__ = [ 

234 "AUTOROUTER_MODEL_NAME", 

235 "DEFAULT_KEYWORD_TIER_RULES", 

236 "TIER_NAMES", 

237 "AutorouteConfig", 

238 "ClassifierChoice", 

239 "ConfigGenerationError", 

240 "DiscoveredModel", 

241 "HeuristicClassifier", 

242 "KeywordTierRule", 

243 "LLMClassifier", 

244 "NoSemanticMatching", 

245 "SemanticMatching", 

246 "SemanticMatchingChoice", 

247 "build_generated_model_list", 

248 "build_generated_proxy_config", 

249 "chat_models", 

250 "embedding_models", 

251 "master_key_from_config", 

252 "parse_discovered_models", 

253 "validate_config", 

254]