Coverage for src/chat_limiter/models.py: 92%
130 statements
« prev ^ index » next coverage.py v7.9.2, created at 2025-07-07 15:28 +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 datetime import datetime, timedelta
11from typing import Any
13import httpx
15logger = logging.getLogger(__name__)
17# Cache for model lists to avoid hitting APIs too frequently
18_model_cache: dict[str, dict[str, Any]] = {}
19_cache_duration = timedelta(hours=1) # Cache models for 1 hour
22class ModelDiscovery:
23 """Dynamic model discovery from provider APIs."""
25 @staticmethod
26 async def get_openai_models(api_key: str) -> set[str]:
27 """Get available OpenAI models from the API."""
28 cache_key = f"openai_models_{hash(api_key)}"
30 # Check cache first
31 if _model_cache.get(cache_key):
32 cache_entry = _model_cache[cache_key]
33 if datetime.now() - cache_entry["timestamp"] < _cache_duration:
34 return cache_entry["models"] # type: ignore[no-any-return]
36 try:
37 async with httpx.AsyncClient() as client:
38 response = await client.get(
39 "https://api.openai.com/v1/models",
40 headers={"Authorization": f"Bearer {api_key}"},
41 timeout=10.0
42 )
43 response.raise_for_status()
45 data = response.json()
46 models = set()
48 for model in data.get("data", []):
49 model_id = model.get("id", "")
50 # Filter for chat completion models
51 if any(keyword in model_id.lower() for keyword in ["gpt", "chat"]):
52 models.add(model_id)
54 # Cache the result
55 _model_cache[cache_key] = {
56 "models": models,
57 "timestamp": datetime.now()
58 }
60 logger.info(f"Retrieved {len(models)} OpenAI models from API")
61 return models
63 except Exception as e:
64 logger.warning(f"Failed to fetch OpenAI models: {e}")
65 raise
67 @staticmethod
68 async def get_anthropic_models(api_key: str) -> set[str]:
69 """Get available Anthropic models from the API."""
70 cache_key = f"anthropic_models_{hash(api_key)}"
72 # Check cache first
73 if _model_cache.get(cache_key):
74 cache_entry = _model_cache[cache_key]
75 if datetime.now() - cache_entry["timestamp"] < _cache_duration:
76 return cache_entry["models"] # type: ignore[no-any-return]
78 try:
79 async with httpx.AsyncClient() as client:
80 response = await client.get(
81 "https://api.anthropic.com/v1/models",
82 headers={
83 "x-api-key": api_key,
84 "anthropic-version": "2023-06-01"
85 },
86 timeout=10.0
87 )
88 response.raise_for_status()
90 data = response.json()
91 models = set()
93 for model in data.get("data", []):
94 model_id = model.get("id", "")
95 # Filter for Claude models
96 if "claude" in model_id.lower():
97 models.add(model_id)
99 # Cache the result
100 _model_cache[cache_key] = {
101 "models": models,
102 "timestamp": datetime.now()
103 }
105 logger.info(f"Retrieved {len(models)} Anthropic models from API")
106 return models
108 except Exception as e:
109 logger.warning(f"Failed to fetch Anthropic models: {e}")
110 raise
112 @staticmethod
113 async def get_openrouter_models(api_key: str | None = None) -> set[str]:
114 """Get available OpenRouter models from the API."""
115 cache_key = "openrouter_models"
117 # Check cache first
118 if _model_cache.get(cache_key):
119 cache_entry = _model_cache[cache_key]
120 if datetime.now() - cache_entry["timestamp"] < _cache_duration:
121 return cache_entry["models"] # type: ignore[no-any-return]
123 try:
124 headers = {}
125 if api_key:
126 headers["Authorization"] = f"Bearer {api_key}"
128 async with httpx.AsyncClient() as client:
129 response = await client.get(
130 "https://openrouter.ai/api/v1/models",
131 headers=headers,
132 timeout=10.0
133 )
134 response.raise_for_status()
136 data = response.json()
137 models = set()
139 for model in data.get("data", []):
140 model_id = model.get("id", "")
141 if model_id:
142 models.add(model_id)
144 # Cache the result
145 _model_cache[cache_key] = {
146 "models": models,
147 "timestamp": datetime.now()
148 }
150 logger.info(f"Retrieved {len(models)} OpenRouter models from API")
151 return models
153 except Exception as e:
154 logger.warning(f"Failed to fetch OpenRouter models: {e}")
155 raise
157 @staticmethod
158 def get_openai_models_sync(api_key: str) -> set[str]:
159 """Synchronous version of get_openai_models."""
160 return asyncio.run(ModelDiscovery.get_openai_models(api_key))
162 @staticmethod
163 def get_anthropic_models_sync(api_key: str) -> set[str]:
164 """Synchronous version of get_anthropic_models."""
165 return asyncio.run(ModelDiscovery.get_anthropic_models(api_key))
167 @staticmethod
168 def get_openrouter_models_sync(api_key: str | None = None) -> set[str]:
169 """Synchronous version of get_openrouter_models."""
170 return asyncio.run(ModelDiscovery.get_openrouter_models(api_key))
173async def detect_provider_from_model_async(
174 model: str,
175 api_keys: dict[str, str] | None = None
176) -> str | None:
177 """
178 Detect provider from model name using live API queries.
180 Args:
181 model: The model name to check
182 api_keys: Dictionary of API keys {"openai": "sk-...", "anthropic": "sk-ant-..."}
184 Returns:
185 Provider name or None if not found
186 """
187 if not api_keys:
188 api_keys = {}
190 # First try simple pattern matching for known formats
191 if "/" in model: # OpenRouter format
192 return "openrouter"
194 # Create all tasks
195 tasks = []
197 if api_keys.get("openai"):
198 tasks.append(("openai", ModelDiscovery.get_openai_models(api_keys["openai"])))
200 if api_keys.get("anthropic"):
201 tasks.append(("anthropic", ModelDiscovery.get_anthropic_models(api_keys["anthropic"])))
203 if api_keys.get("openrouter"):
204 tasks.append(("openrouter", ModelDiscovery.get_openrouter_models(api_keys["openrouter"])))
205 else:
206 # OpenRouter doesn't require API key for model listing
207 tasks.append(("openrouter", ModelDiscovery.get_openrouter_models()))
209 # Use asyncio.gather to run all tasks concurrently and properly handle them
210 try:
211 # Extract just the coroutines for gather
212 coroutines = [task[1] for task in tasks]
213 provider_names = [task[0] for task in tasks]
215 # Wait for all results
216 results = await asyncio.gather(*coroutines, return_exceptions=True)
218 # Check results in order
219 for provider_name, result in zip(provider_names, results):
220 if isinstance(result, Exception):
221 logger.debug(f"Failed to check {provider_name} for model {model}: {result}")
222 continue
223 if model in result:
224 return provider_name
226 except Exception as e:
227 logger.debug(f"Failed to run dynamic discovery for model {model}: {e}")
229 return None
232def detect_provider_from_model_sync(
233 model: str,
234 api_keys: dict[str, str] | None = None
235) -> str | None:
236 """Synchronous version of detect_provider_from_model_async."""
237 # Check if we're already in an async context
238 try:
239 loop = asyncio.get_running_loop()
240 # We're in an async context, but need to run in sync mode
241 # Create a new event loop in a thread
242 import concurrent.futures
244 def run_in_thread():
245 return asyncio.run(detect_provider_from_model_async(model, api_keys))
247 with concurrent.futures.ThreadPoolExecutor() as executor:
248 future = executor.submit(run_in_thread)
249 return future.result(timeout=30) # 30 second timeout
251 except RuntimeError:
252 # No running loop, safe to use asyncio.run
253 return asyncio.run(detect_provider_from_model_async(model, api_keys))
256def clear_model_cache() -> None:
257 """Clear the model cache to force fresh API queries."""
258 global _model_cache
259 _model_cache.clear()
260 logger.info("Model cache cleared")