Coverage for src/chat_limiter/models.py: 92%
192 statements
« prev ^ index » next coverage.py v7.9.2, created at 2025-09-01 14:16 +0100
1"""
2Dynamic model discovery from provider APIs.
4This module provides functionality to query provider APIs for available models
5instead of relying on hardcoded lists.
6"""
8import asyncio
9import logging
10from dataclasses import dataclass
11from datetime import datetime, timedelta
12from typing import Any
14import httpx
16logger = logging.getLogger(__name__)
18# Cache for model lists to avoid hitting APIs too frequently
19_model_cache: dict[str, dict[str, Any]] = {}
20_cache_duration = timedelta(hours=1) # Cache models for 1 hour
23@dataclass
24class ModelDiscoveryResult:
25 """Result of model discovery process."""
27 # Discovery result
28 found_provider: str | None = None
29 model_found: bool = False
31 # All models found for each provider
32 openai_models: set[str] | None = None
33 anthropic_models: set[str] | None = None
34 openrouter_models: set[str] | None = None
36 # Errors encountered during discovery
37 errors: dict[str, str] | None = None
39 def get_all_models(self) -> dict[str, set[str]]:
40 """Get all models organized by provider."""
41 result = {}
42 if self.openai_models is not None:
43 result["openai"] = self.openai_models
44 if self.anthropic_models is not None:
45 result["anthropic"] = self.anthropic_models
46 if self.openrouter_models is not None:
47 result["openrouter"] = self.openrouter_models
48 return result
50 def get_total_models_found(self) -> int:
51 """Get total number of models found across all providers."""
52 total = 0
53 if self.openai_models:
54 total += len(self.openai_models)
55 if self.anthropic_models:
56 total += len(self.anthropic_models)
57 if self.openrouter_models:
58 total += len(self.openrouter_models)
59 return total
62class ModelDiscovery:
63 """Dynamic model discovery from provider APIs."""
65 @staticmethod
66 async def get_openai_models(api_key: str) -> set[str]:
67 """Get available OpenAI models from the API."""
68 cache_key = f"openai_models_{hash(api_key)}"
70 # Check cache first
71 if _model_cache.get(cache_key):
72 cache_entry = _model_cache[cache_key]
73 if datetime.now() - cache_entry["timestamp"] < _cache_duration:
74 return cache_entry["models"] # type: ignore[no-any-return]
76 try:
77 async with httpx.AsyncClient() as client:
78 response = await client.get(
79 "https://api.openai.com/v1/models",
80 headers={"Authorization": f"Bearer {api_key}"},
81 timeout=10.0
82 )
83 response.raise_for_status()
85 data = response.json()
86 models = set()
88 for model in data.get("data", []):
89 model_id = model.get("id", "")
90 models.add(model_id)
92 # Cache the result
93 _model_cache[cache_key] = {
94 "models": models,
95 "timestamp": datetime.now()
96 }
98 logger.info(f"Retrieved {len(models)} OpenAI models from API")
99 return models
101 except Exception as e:
102 logger.warning(f"Failed to fetch OpenAI models: {e}")
103 raise
105 @staticmethod
106 async def get_anthropic_models(api_key: str) -> set[str]:
107 """Get available Anthropic models from the API."""
108 cache_key = f"anthropic_models_{hash(api_key)}"
110 # Check cache first
111 if _model_cache.get(cache_key):
112 cache_entry = _model_cache[cache_key]
113 if datetime.now() - cache_entry["timestamp"] < _cache_duration:
114 return cache_entry["models"] # type: ignore[no-any-return]
116 try:
117 async with httpx.AsyncClient() as client:
118 response = await client.get(
119 "https://api.anthropic.com/v1/models",
120 headers={
121 "x-api-key": api_key,
122 "anthropic-version": "2023-06-01"
123 },
124 timeout=10.0
125 )
126 response.raise_for_status()
128 data = response.json()
129 models = set()
131 for model in data.get("data", []):
132 model_id = model.get("id", "")
133 models.add(model_id)
135 # Cache the result
136 _model_cache[cache_key] = {
137 "models": models,
138 "timestamp": datetime.now()
139 }
141 logger.info(f"Retrieved {len(models)} Anthropic models from API")
142 return models
144 except Exception as e:
145 logger.warning(f"Failed to fetch Anthropic models: {e}")
146 raise
148 @staticmethod
149 async def get_openrouter_models(api_key: str | None = None) -> set[str]:
150 """Get available OpenRouter models from the API."""
151 cache_key = "openrouter_models"
153 # Check cache first
154 if _model_cache.get(cache_key):
155 cache_entry = _model_cache[cache_key]
156 if datetime.now() - cache_entry["timestamp"] < _cache_duration:
157 return cache_entry["models"] # type: ignore[no-any-return]
159 try:
160 headers = {}
161 if api_key:
162 headers["Authorization"] = f"Bearer {api_key}"
164 async with httpx.AsyncClient() as client:
165 response = await client.get(
166 "https://openrouter.ai/api/v1/models",
167 headers=headers,
168 timeout=10.0
169 )
170 response.raise_for_status()
172 data = response.json()
173 models = set()
175 for model in data.get("data", []):
176 model_id = model.get("id", "")
177 if model_id:
178 models.add(model_id)
180 # Cache the result
181 _model_cache[cache_key] = {
182 "models": models,
183 "timestamp": datetime.now()
184 }
186 logger.info(f"Retrieved {len(models)} OpenRouter models from API")
187 return models
189 except Exception as e:
190 logger.warning(f"Failed to fetch OpenRouter models: {e}")
191 raise
193 @staticmethod
194 def get_openai_models_sync(api_key: str) -> set[str]:
195 """Synchronous version of get_openai_models."""
196 return asyncio.run(ModelDiscovery.get_openai_models(api_key))
198 @staticmethod
199 def get_anthropic_models_sync(api_key: str) -> set[str]:
200 """Synchronous version of get_anthropic_models."""
201 return asyncio.run(ModelDiscovery.get_anthropic_models(api_key))
203 @staticmethod
204 def get_openrouter_models_sync(api_key: str | None = None) -> set[str]:
205 """Synchronous version of get_openrouter_models."""
206 return asyncio.run(ModelDiscovery.get_openrouter_models(api_key))
209async def detect_provider_from_model_async(
210 model: str,
211 api_keys: dict[str, str] | None = None
212) -> ModelDiscoveryResult:
213 """
214 Detect provider from model name using live API queries.
216 Args:
217 model: The model name to check
218 api_keys: Dictionary of API keys {"openai": "sk-...", "anthropic": "sk-ant-..."}
220 Returns:
221 ModelDiscoveryResult with discovery information
222 """
223 if not api_keys:
224 api_keys = {}
226 result = ModelDiscoveryResult(errors={})
228 # Handle provider-prefixed models (e.g., "openai/o3", "anthropic/claude-3-sonnet")
229 preferred_provider = None
230 base_model = model
232 if "/" in model:
233 parts = model.split("/", 1)
234 if len(parts) == 2:
235 provider_prefix, base_model = parts
236 if provider_prefix == "openai":
237 preferred_provider = "openai"
238 elif provider_prefix == "anthropic":
239 preferred_provider = "anthropic"
241 # Create all tasks
242 tasks = []
244 if api_keys.get("openai"):
245 tasks.append(("openai", ModelDiscovery.get_openai_models(api_keys["openai"])))
247 if api_keys.get("anthropic"):
248 tasks.append(("anthropic", ModelDiscovery.get_anthropic_models(api_keys["anthropic"])))
250 if api_keys.get("openrouter"):
251 tasks.append(("openrouter", ModelDiscovery.get_openrouter_models(api_keys["openrouter"])))
252 else:
253 # OpenRouter doesn't require API key for model listing
254 tasks.append(("openrouter", ModelDiscovery.get_openrouter_models()))
256 # Use asyncio.gather to run all tasks concurrently and properly handle them
257 try:
258 # Extract just the coroutines for gather
259 coroutines = [task[1] for task in tasks]
260 provider_names = [task[0] for task in tasks]
262 # Wait for all results
263 results = await asyncio.gather(*coroutines, return_exceptions=True)
265 # Process results and store all model information
266 for provider_name, models_result in zip(provider_names, results, strict=False):
267 if isinstance(models_result, Exception):
268 logger.debug(f"Failed to check {provider_name} for model {model}: {models_result}")
269 if result.errors is not None:
270 result.errors[provider_name] = str(models_result)
271 continue
273 # Store models in result
274 if provider_name == "openai" and isinstance(models_result, set):
275 result.openai_models = models_result
276 elif provider_name == "anthropic" and isinstance(models_result, set):
277 result.anthropic_models = models_result
278 elif provider_name == "openrouter" and isinstance(models_result, set):
279 result.openrouter_models = models_result
281 # Determine the best provider to use
282 if preferred_provider and not result.model_found:
283 # Check if base model exists in preferred provider
284 provider_models = None
285 if preferred_provider == "openai" and result.openai_models:
286 provider_models = result.openai_models
287 elif preferred_provider == "anthropic" and result.anthropic_models:
288 provider_models = result.anthropic_models
290 if provider_models and base_model in provider_models:
291 result.found_provider = preferred_provider
292 result.model_found = True
293 elif result.openrouter_models and model in result.openrouter_models:
294 # Fallback to OpenRouter if base model not found in preferred provider
295 result.found_provider = "openrouter"
296 result.model_found = True
298 # For models without provider prefix, use original logic
299 if not result.model_found:
300 for provider_name, models_result in zip(provider_names, results, strict=False):
301 if isinstance(models_result, Exception):
302 continue
304 # Check if our target model was found
305 if isinstance(models_result, set) and model in models_result:
306 result.found_provider = provider_name
307 result.model_found = True
308 break
310 except Exception as e:
311 logger.debug(f"Failed to run dynamic discovery for model {model}: {e}")
312 if result.errors is not None:
313 result.errors["general"] = str(e)
315 return result
318def detect_provider_from_model_sync(
319 model: str,
320 api_keys: dict[str, str] | None = None
321) -> ModelDiscoveryResult:
322 """Synchronous version of detect_provider_from_model_async."""
323 # Check if we're already in an async context
324 try:
325 asyncio.get_running_loop()
326 # We're in an async context, but need to run in sync mode
327 # Create a new event loop in a thread
328 import concurrent.futures
330 def run_in_thread() -> ModelDiscoveryResult:
331 return asyncio.run(detect_provider_from_model_async(model, api_keys))
333 with concurrent.futures.ThreadPoolExecutor() as executor:
334 future = executor.submit(run_in_thread)
335 return future.result(timeout=30) # 30 second timeout
337 except RuntimeError:
338 # No running loop, safe to use asyncio.run
339 return asyncio.run(detect_provider_from_model_async(model, api_keys))
342def clear_model_cache() -> None:
343 """Clear the model cache to force fresh API queries."""
344 global _model_cache
345 _model_cache.clear()
346 logger.info("Model cache cleared")