Spaces:
Running
Running
File size: 3,176 Bytes
4c40bac | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 | import random
import time
import requests
from typing import Optional, Dict, Any, List
from . import config
class OpenRouterClient:
def __init__(self):
self.api_keys = config.get_api_keys()
self.max_retries = config.get_max_retries()
self.retry_delay = config.get_retry_delay()
self.timeout = config.get_request_timeout()
self.current_key_index = 0
def get_random_key(self) -> str:
return random.choice(self.api_keys)
def get_next_key(self) -> str:
key = self.api_keys[self.current_key_index]
self.current_key_index = (self.current_key_index + 1) % len(self.api_keys)
return key
def _make_request(
self,
method: str,
url: str,
api_key: str,
**kwargs
) -> requests.Response:
headers = kwargs.pop("headers", {})
headers["Authorization"] = f"Bearer {api_key}"
headers["Content-Type"] = "application/json"
return requests.request(
method=method,
url=url,
headers=headers,
timeout=self.timeout,
**kwargs
)
def request_with_retry(
self,
method: str,
url: str,
**kwargs
) -> Optional[requests.Response]:
last_error = None
for attempt in range(self.max_retries):
api_key = self.get_random_key()
try:
response = self._make_request(method, url, api_key, **kwargs)
if response.status_code == 200:
return response
elif response.status_code == 429:
last_error = f"Rate limited (429)"
time.sleep(self.retry_delay * (attempt + 1))
elif response.status_code == 401:
last_error = f"Unauthorized (401)"
else:
last_error = f"HTTP {response.status_code}: {response.text[:100]}"
time.sleep(self.retry_delay)
except requests.exceptions.Timeout:
last_error = "Request timeout"
time.sleep(self.retry_delay)
except requests.exceptions.RequestException as e:
last_error = f"Request error: {str(e)}"
time.sleep(self.retry_delay)
print(f"Request failed after {self.max_retries} attempts: {last_error}")
return None
def get_models(self) -> List[Dict[str, Any]]:
url = "https://openrouter.ai/api/v1/models"
response = self.request_with_retry("GET", url)
if response and response.status_code == 200:
data = response.json()
return data.get("data", [])
return []
def test_model(self, model_id: str, prompt: str = "hi") -> bool:
url = "https://openrouter.ai/api/v1/chat/completions"
payload = {
"model": model_id,
"messages": [{"role": "user", "content": prompt}],
"max_tokens": 10
}
response = self.request_with_retry("POST", url, json=payload)
if response and response.status_code == 200:
return True
return False
|