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

1"""Sync a LiteLLM provider into pi's models.json. 

2 

3pi ignores ANTHROPIC_BASE_URL/OPENAI_BASE_URL, so `lite pi` routes it through the 

4proxy by writing a provider entry instead. The key is stored as a $-reference so 

5the short-lived login token never lands on disk. 

6""" 

7 

8import json 

9import os 

10import tempfile 

11from collections.abc import Callable, Mapping 

12from dataclasses import dataclass 

13from enum import Enum 

14from pathlib import Path 

15from types import MappingProxyType 

16from typing import Annotated, Final 

17 

18import requests 

19from pydantic import BaseModel, ConfigDict, JsonValue, TypeAdapter, ValidationError, model_validator 

20from pydantic.types import StringConstraints 

21 

22PI_CONFIG_DIR_ENV: Final = "PI_CODING_AGENT_DIR" 

23PI_PROVIDER_NAME: Final = "litellm" 

24LITELLM_PROXY_API_KEY_ENV: Final = "LITELLM_PROXY_API_KEY" 

25_REJECTED_STATUSES: Final = frozenset((401, 403)) 

26 

27 

28class ListingFailure(str, Enum): 

29 """Why a proxy could not be listed, decided once where the HTTP outcome is classified. 

30 

31 `unreachable` means no response at all; the other kinds prove the proxy answered, so callers 

32 must not suggest checking whether it is running. 

33 """ 

34 

35 UNREACHABLE = "unreachable" 

36 REJECTED = "rejected" 

37 BAD_BODY = "bad_body" 

38 EMPTY = "empty" 

39 OTHER = "other" 

40 

41 

42@dataclass(frozen=True, slots=True) 

43class PiSyncError: 

44 message: str 

45 status: int | None = None 

46 kind: ListingFailure | None = None 

47 

48 

49@dataclass(frozen=True, slots=True) 

50class ModelLimits: 

51 context_window: int | None 

52 max_tokens: int | None 

53 

54 

55_NonEmptyString = Annotated[str, StringConstraints(min_length=1)] 

56 

57 

58class ListedModel(BaseModel): 

59 model_config = ConfigDict(frozen=True) 

60 

61 id: _NonEmptyString 

62 source_model: _NonEmptyString | None = None 

63 

64 

65class _ModelList(BaseModel): 

66 data: tuple[ListedModel, ...] 

67 

68 @model_validator(mode="after") 

69 def unique_id_mappings(self) -> "_ModelList": 

70 mappings: Final = frozenset((model.id, model.source_model or model.id) for model in self.data) 

71 if len(frozenset(model.id for model in self.data)) != len(mappings): 

72 raise ValueError("model ids must not map to multiple source models") 

73 return self 

74 

75 

76class _ModelGroup(BaseModel): 

77 model_group: str 

78 max_input_tokens: float | None = None 

79 max_output_tokens: float | None = None 

80 

81 

82class _ModelGroupList(BaseModel): 

83 data: tuple[_ModelGroup, ...] 

84 

85 

86def fetch_model_listing( 

87 base_url: str, 

88 api_key: str, 

89 *, 

90 get: Callable[..., requests.Response] = requests.get, 

91 headers: Mapping[str, str] = MappingProxyType({}), 

92) -> tuple[ListedModel, ...] | PiSyncError: 

93 url: Final = base_url.rstrip("/") + "/v1/models" 

94 try: 

95 resp: Final = get( 

96 url, 

97 headers={"Authorization": f"Bearer {api_key}", **headers}, # mutable-ok: requests headers require a dict 

98 timeout=10, 

99 ) 

100 except requests.RequestException as e: 

101 return PiSyncError(f"Could not list models from the proxy: {e}", kind=ListingFailure.UNREACHABLE) 

102 if resp.status_code != 200: 

103 return PiSyncError( 

104 f"The proxy returned HTTP {resp.status_code} for /v1/models; cannot list models.", 

105 resp.status_code, 

106 ListingFailure.REJECTED if resp.status_code in _REJECTED_STATUSES else ListingFailure.OTHER, 

107 ) 

108 try: 

109 listing: Final = _ModelList.model_validate(resp.json()) 

110 except (ValueError, ValidationError) as e: 

111 return PiSyncError(f"Unexpected /v1/models response from the proxy: {e}", kind=ListingFailure.BAD_BODY) 

112 models: Final = tuple(dict.fromkeys(listing.data)) 

113 if not models: 

114 return PiSyncError("The proxy returned no models for your key.", kind=ListingFailure.EMPTY) 

115 return models 

116 

117 

118def fetch_model_ids( 

119 base_url: str, 

120 api_key: str, 

121 *, 

122 get: Callable[..., requests.Response] = requests.get, 

123 headers: Mapping[str, str] = MappingProxyType({}), 

124) -> tuple[str, ...] | PiSyncError: 

125 listed: Final = fetch_model_listing(base_url, api_key, get=get, headers=headers) 

126 return listed if isinstance(listed, PiSyncError) else tuple(dict.fromkeys(model.id for model in listed)) 

127 

128 

129_NO_LIMITS: Final[Mapping[str, ModelLimits]] = MappingProxyType({}) 

130 

131 

132def fetch_model_limits( 

133 base_url: str, 

134 api_key: str, 

135 *, 

136 get: Callable[..., requests.Response] = requests.get, 

137) -> Mapping[str, ModelLimits]: 

138 """Best effort: pi falls back to its own defaults for models without limits, 

139 so an unavailable /model_group/info must not block the launch.""" 

140 url: Final = base_url.rstrip("/") + "/model_group/info" 

141 try: 

142 resp: Final = get( 

143 url, 

144 headers={"Authorization": f"Bearer {api_key}"}, # mutable-ok: requests headers require a dict 

145 timeout=10, 

146 ) 

147 if resp.status_code != 200: 

148 return _NO_LIMITS 

149 listing: Final = _ModelGroupList.model_validate(resp.json()) 

150 except (requests.RequestException, ValueError, ValidationError): 

151 return _NO_LIMITS 

152 return MappingProxyType( 

153 { 

154 group.model_group: ModelLimits( 

155 context_window=int(group.max_input_tokens) if group.max_input_tokens else None, 

156 max_tokens=int(group.max_output_tokens) if group.max_output_tokens else None, 

157 ) 

158 for group in listing.data 

159 } 

160 ) 

161 

162 

163def models_json_path(env: Mapping[str, str]) -> Path: 

164 override: Final = env.get(PI_CONFIG_DIR_ENV) 

165 root: Final = Path(override) if override else Path.home() / ".pi" / "agent" 

166 return root / "models.json" 

167 

168 

169def _model_entry( 

170 model_id: str, limits: Mapping[str, ModelLimits] 

171) -> dict[str, JsonValue]: # mutable-ok: JSON object is serialized 

172 limit: Final = limits.get(model_id) 

173 context: Final[dict[str, JsonValue]] = ( # mutable-ok: JSON field 

174 {"contextWindow": limit.context_window} if limit and limit.context_window else {} # mutable-ok: JSON field 

175 ) 

176 output: Final[dict[str, JsonValue]] = ( # mutable-ok: JSON field 

177 {"maxTokens": limit.max_tokens} if limit and limit.max_tokens else {} 

178 ) 

179 return {"id": model_id, **context, **output} # mutable-ok: JSON serialization requires a mutable object 

180 

181 

182def provider_block( 

183 base_url: str, 

184 model_ids: tuple[str, ...], 

185 limits: Mapping[str, ModelLimits] = _NO_LIMITS, 

186) -> dict[str, JsonValue]: # mutable-ok: JSON object is serialized 

187 """openai-completions is the one API shape every LiteLLM model serves. 

188 

189 Real contextWindow/maxTokens matter: pi otherwise assumes 128k/16384, which 

190 breaks compaction thresholds and over-asks models with smaller output caps. 

191 """ 

192 return { # mutable-ok: JSON serialization requires a mutable object 

193 "baseUrl": base_url.rstrip("/") + "/v1", 

194 "api": "openai-completions", 

195 "apiKey": f"${LITELLM_PROXY_API_KEY_ENV}", 

196 "models": [_model_entry(model_id, limits) for model_id in model_ids], # mutable-ok: JSON array 

197 } 

198 

199 

200_MODELS_FILE_ADAPTER: Final = TypeAdapter(dict[str, JsonValue]) 

201 

202 

203def sync_models_json( 

204 path: Path, 

205 base_url: str, 

206 model_ids: tuple[str, ...], 

207 limits: Mapping[str, ModelLimits] = _NO_LIMITS, 

208) -> PiSyncError | None: 

209 """Replace only the litellm provider entry, leaving the rest of the file intact.""" 

210 try: 

211 current: Final = _MODELS_FILE_ADAPTER.validate_json(path.read_text()) if path.exists() else {} 

212 except (OSError, ValidationError) as e: 

213 return PiSyncError(f"Could not read {path} as a JSON object: {e}. Fix or move the file, then retry.") 

214 existing_providers: Final = current.get("providers", {}) # mutable-ok: JSON object default 

215 if not isinstance(existing_providers, dict): 

216 return PiSyncError(f'"providers" in {path} is not an object; fix or move the file, then retry.') 

217 updated: Final = { # mutable-ok: JSON serialization requires a mutable object 

218 **current, 

219 "providers": { # mutable-ok: JSON serialization requires a mutable object 

220 **existing_providers, 

221 PI_PROVIDER_NAME: provider_block(base_url, model_ids, limits), 

222 }, 

223 } 

224 try: 

225 path.parent.mkdir(parents=True, exist_ok=True) 

226 except OSError as e: 

227 return PiSyncError(f"Could not write {path}: {e}") 

228 try: 

229 fd, tmp_name = tempfile.mkstemp(dir=path.parent, prefix=path.name + ".", suffix=".tmp") 

230 except OSError as e: 

231 return PiSyncError(f"Could not write {path}: {e}") 

232 try: 

233 with os.fdopen(fd, "w") as file: 

234 file.write(json.dumps(updated, indent=2) + "\n") 

235 os.replace(tmp_name, path) 

236 except OSError as e: 

237 try: 

238 os.unlink(tmp_name) 

239 except FileNotFoundError: 

240 pass 

241 return PiSyncError(f"Could not write {path}: {e}") 

242 return None 

243 

244 

245__all__ = ( 

246 "LITELLM_PROXY_API_KEY_ENV", 

247 "PI_CONFIG_DIR_ENV", 

248 "PI_PROVIDER_NAME", 

249 "ListedModel", 

250 "ListingFailure", 

251 "ModelLimits", 

252 "PiSyncError", 

253 "fetch_model_ids", 

254 "fetch_model_limits", 

255 "fetch_model_listing", 

256 "models_json_path", 

257 "provider_block", 

258 "sync_models_json", 

259)