Coverage for src / osiris_cli / model_fetcher.py: 0%

139 statements  

« prev     ^ index     » next       coverage.py v7.13.0, created at 2025-12-27 17:41 +0200

1#!/usr/bin/env python3 

2""" 

3Dynamic Model Fetcher for Osiris CLI 

4Fetches available models from each provider's API in real-time 

5""" 

6 

7import httpx 

8from typing import List, Dict, Optional 

9from rich.console import Console 

10import os 

11 

12console = Console() 

13 

14 

15class ModelFetcher: 

16 """Fetch available models from provider APIs""" 

17 

18 def __init__(self): 

19 self.cache = {} 

20 

21 async def fetch_openrouter_models(self) -> List[Dict]: 

22 """Fetch all models from OpenRouter API""" 

23 try: 

24 async with httpx.AsyncClient(timeout=10.0) as client: 

25 response = await client.get("https://openrouter.ai/api/v1/models") 

26 response.raise_for_status() 

27 data = response.json() 

28 

29 models = [] 

30 for model in data.get("data", []): 

31 models.append({ 

32 "id": model["id"], 

33 "name": model.get("name", model["id"]), 

34 "context_length": model.get("context_length", 0), 

35 "pricing": model.get("pricing", {}), 

36 "description": model.get("description", ""), 

37 }) 

38 

39 return models 

40 except Exception as e: 

41 console.print(f"[yellow]Warning: Could not fetch OpenRouter models: {e}[/yellow]") 

42 return self._get_fallback_openrouter_models() 

43 

44 async def fetch_openai_models(self, api_key: str) -> List[Dict]: 

45 """Fetch all models from OpenAI API""" 

46 try: 

47 async with httpx.AsyncClient(timeout=10.0) as client: 

48 response = await client.get( 

49 "https://api.openai.com/v1/models", 

50 headers={"Authorization": f"Bearer {api_key}"} 

51 ) 

52 response.raise_for_status() 

53 data = response.json() 

54 

55 # Filter to chat models only 

56 models = [] 

57 for model in data.get("data", []): 

58 model_id = model["id"] 

59 if any(x in model_id for x in ["gpt", "o1"]): 

60 models.append({ 

61 "id": model_id, 

62 "name": model_id, 

63 "created": model.get("created", 0), 

64 }) 

65 

66 return sorted(models, key=lambda x: x["created"], reverse=True) 

67 except Exception as e: 

68 console.print(f"[yellow]Warning: Could not fetch OpenAI models: {e}[/yellow]") 

69 return self._get_fallback_openai_models() 

70 

71 async def fetch_gemini_models(self, api_key: str) -> List[Dict]: 

72 """Fetch all models from Gemini API""" 

73 try: 

74 async with httpx.AsyncClient(timeout=10.0) as client: 

75 response = await client.get( 

76 f"https://generativelanguage.googleapis.com/v1beta/models?key={api_key}" 

77 ) 

78 response.raise_for_status() 

79 data = response.json() 

80 

81 models = [] 

82 for model in data.get("models", []): 

83 model_name = model["name"].replace("models/", "") 

84 if "generateContent" in model.get("supportedGenerationMethods", []): 

85 models.append({ 

86 "id": model_name, 

87 "name": model.get("displayName", model_name), 

88 "description": model.get("description", ""), 

89 }) 

90 

91 return models 

92 except Exception as e: 

93 console.print(f"[yellow]Warning: Could not fetch Gemini models: {e}[/yellow]") 

94 return self._get_fallback_gemini_models() 

95 

96 async def fetch_together_models(self, api_key: str) -> List[Dict]: 

97 """Fetch all models from Together AI API""" 

98 try: 

99 async with httpx.AsyncClient(timeout=10.0) as client: 

100 response = await client.get( 

101 "https://api.together.xyz/v1/models", 

102 headers={"Authorization": f"Bearer {api_key}"} 

103 ) 

104 response.raise_for_status() 

105 data = response.json() 

106 

107 models = [] 

108 for model in data: 

109 if model.get("type") == "chat": 

110 models.append({ 

111 "id": model["id"], 

112 "name": model.get("display_name", model["id"]), 

113 "context_length": model.get("context_length", 0), 

114 }) 

115 

116 return models 

117 except Exception as e: 

118 console.print(f"[yellow]Warning: Could not fetch Together models: {e}[/yellow]") 

119 return self._get_fallback_together_models() 

120 

121 async def fetch_groq_models(self, api_key: str) -> List[Dict]: 

122 """Fetch all models from Groq API""" 

123 try: 

124 async with httpx.AsyncClient(timeout=10.0) as client: 

125 response = await client.get( 

126 "https://api.groq.com/openai/v1/models", 

127 headers={"Authorization": f"Bearer {api_key}"} 

128 ) 

129 response.raise_for_status() 

130 data = response.json() 

131 

132 models = [] 

133 for model in data.get("data", []): 

134 models.append({ 

135 "id": model["id"], 

136 "name": model["id"], 

137 "context_window": model.get("context_window", 0), 

138 }) 

139 

140 return models 

141 except Exception as e: 

142 console.print(f"[yellow]Warning: Could not fetch Groq models: {e}[/yellow]") 

143 return self._get_fallback_groq_models() 

144 

145 async def fetch_perplexity_models(self, api_key: str) -> List[Dict]: 

146 """Fetch models from Perplexity API""" 

147 try: 

148 async with httpx.AsyncClient(timeout=10.0) as client: 

149 response = await client.get( 

150 "https://api.perplexity.ai/models", 

151 headers={"Authorization": f"Bearer {api_key}"} 

152 ) 

153 if response.status_code != 200: 

154 return self._get_fallback_perplexity_models() 

155 

156 data = response.json() 

157 models = [] 

158 # Perplexity often returns a list directly or in 'data' 

159 items = data.get("data", data) if isinstance(data, dict) else data 

160 

161 if isinstance(items, list): 

162 for model in items: 

163 model_id = model.get("id", str(model)) 

164 models.append({ 

165 "id": model_id, 

166 "name": model_id, 

167 }) 

168 

169 return models if models else self._get_fallback_perplexity_models() 

170 except Exception as e: 

171 console.print(f"[yellow]Warning: Could not fetch Perplexity models: {e}[/yellow]") 

172 return self._get_fallback_perplexity_models() 

173 

174 async def fetch_generic_models(self, base_url: str, api_key: Optional[str] = None) -> List[Dict]: 

175 """Fetch models from any OpenAI-compatible API""" 

176 try: 

177 headers = {} 

178 if api_key: 

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

180 

181 async with httpx.AsyncClient(timeout=10.0) as client: 

182 url = base_url.rstrip("/") + "/models" 

183 response = await client.get(url, headers=headers) 

184 response.raise_for_status() 

185 data = response.json() 

186 

187 models = [] 

188 for model in data.get("data", []): 

189 models.append({ 

190 "id": model["id"], 

191 "name": model["id"], 

192 }) 

193 return models 

194 except Exception: 

195 return [] 

196 

197 async def fetch_ollama_models(self) -> List[Dict]: 

198 """Fetch all models from local Ollama instance""" 

199 try: 

200 async with httpx.AsyncClient(timeout=5.0) as client: 

201 response = await client.get("http://localhost:11434/api/tags") 

202 response.raise_for_status() 

203 data = response.json() 

204 

205 models = [] 

206 for model in data.get("models", []): 

207 models.append({ 

208 "id": model["name"], 

209 "name": model["name"], 

210 "size": model.get("size", 0), 

211 "modified": model.get("modified_at", ""), 

212 }) 

213 

214 return models 

215 except Exception as e: 

216 console.print(f"[yellow]Warning: Could not fetch Ollama models: {e}[/yellow]") 

217 return self._get_fallback_ollama_models() 

218 

219 # Fallback models (in case API fetch fails) 

220 

221 def _get_fallback_openrouter_models(self) -> List[Dict]: 

222 return [ 

223 {"id": "meta-llama/llama-3.3-70b-instruct", "name": "Llama 3.3 70B (free)"}, 

224 {"id": "meta-llama/llama-3.1-8b-instruct:free", "name": "Llama 3.1 8B (free)"}, 

225 {"id": "google/gemini-2.0-flash-exp:free", "name": "Gemini 2.0 Flash (free)"}, 

226 {"id": "openai/gpt-4o", "name": "GPT-4o"}, 

227 {"id": "anthropic/claude-3.5-sonnet", "name": "Claude 3.5 Sonnet"}, 

228 {"id": "x-ai/grok-2-1212", "name": "Grok 2"}, 

229 {"id": "deepseek/deepseek-chat", "name": "DeepSeek Chat"}, 

230 {"id": "qwen/qwen-2.5-72b-instruct", "name": "Qwen 2.5 72B"}, 

231 ] 

232 

233 def _get_fallback_openai_models(self) -> List[Dict]: 

234 return [ 

235 {"id": "gpt-4o", "name": "GPT-4o (latest)"}, 

236 {"id": "gpt-4o-mini", "name": "GPT-4o Mini"}, 

237 {"id": "gpt-4-turbo", "name": "GPT-4 Turbo"}, 

238 {"id": "gpt-4", "name": "GPT-4"}, 

239 {"id": "gpt-3.5-turbo", "name": "GPT-3.5 Turbo"}, 

240 {"id": "o1", "name": "O1 (reasoning)"}, 

241 {"id": "o1-mini", "name": "O1 Mini"}, 

242 {"id": "o1-preview", "name": "O1 Preview"}, 

243 ] 

244 

245 def _get_fallback_gemini_models(self) -> List[Dict]: 

246 return [ 

247 {"id": "gemini-2.0-flash-exp", "name": "Gemini 2.0 Flash (experimental)"}, 

248 {"id": "gemini-1.5-pro", "name": "Gemini 1.5 Pro"}, 

249 {"id": "gemini-1.5-flash", "name": "Gemini 1.5 Flash"}, 

250 ] 

251 

252 def _get_fallback_together_models(self) -> List[Dict]: 

253 return [ 

254 {"id": "meta-llama/Meta-Llama-3.1-70B-Instruct-Turbo", "name": "Llama 3.1 70B"}, 

255 {"id": "meta-llama/Meta-Llama-3.1-8B-Instruct-Turbo", "name": "Llama 3.1 8B"}, 

256 {"id": "mistralai/Mixtral-8x7B-Instruct-v0.1", "name": "Mixtral 8x7B"}, 

257 {"id": "Qwen/Qwen2.5-72B-Instruct-Turbo", "name": "Qwen 2.5 72B"}, 

258 ] 

259 

260 def _get_fallback_groq_models(self) -> List[Dict]: 

261 return [ 

262 {"id": "llama-3.3-70b-versatile", "name": "Llama 3.3 70B"}, 

263 {"id": "llama-3.1-70b-versatile", "name": "Llama 3.1 70B"}, 

264 {"id": "mixtral-8x7b-32768", "name": "Mixtral 8x7B"}, 

265 {"id": "gemma2-9b-it", "name": "Gemma 2 9B"}, 

266 ] 

267 

268 def _get_fallback_perplexity_models(self) -> List[Dict]: 

269 return [ 

270 {"id": "sonar", "name": "Sonar (Online)"}, 

271 {"id": "sonar-pro", "name": "Sonar Pro (Online)"}, 

272 {"id": "sonar-reasoning", "name": "Sonar Reasoning"}, 

273 {"id": "sonar-reasoning-pro", "name": "Sonar Reasoning Pro"}, 

274 {"id": "sonar-deep-research", "name": "Sonar Deep Research"}, 

275 ] 

276 

277 def _get_fallback_ollama_models(self) -> List[Dict]: 

278 return [ 

279 {"id": "llama3.2", "name": "Llama 3.2"}, 

280 {"id": "llama3.1", "name": "Llama 3.1"}, 

281 {"id": "qwen2.5", "name": "Qwen 2.5"}, 

282 {"id": "deepseek-r1", "name": "DeepSeek R1"}, 

283 {"id": "phi3", "name": "Phi-3"}, 

284 {"id": "mistral", "name": "Mistral"}, 

285 {"id": "codellama", "name": "Code Llama"}, 

286 ] 

287 

288 

289# Global fetcher instance 

290model_fetcher = ModelFetcher()