File size: 5,853 Bytes
d184447 | 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 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | import asyncio
import random
import time
import aiohttp
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
failed_keys: set = set() # 记录 401 失败的 key，不再重复使用
for attempt in range(self.max_retries):
# 从尚未 401 失败的 key 中选择
available_keys = [k for k in self.api_keys if k not in failed_keys]
if not available_keys:
print(f"All API keys returned 401, stopping retries")
break
api_key = random.choice(available_keys)
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) — key excluded"
failed_keys.add(api_key)
time.sleep(self.retry_delay) # 短暂等待后换 key 重试
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 []
async def async_get_models(self) -> List[Dict[str, Any]]:
"""异步获取模型列表，使用 aiohttp"""
url = "https://openrouter.ai/api/v1/models"
api_key = self.get_random_key()
headers = {
"Authorization": f"Bearer {api_key}",
"Content-Type": "application/json"
}
for attempt in range(self.max_retries):
try:
async with aiohttp.ClientSession() as session:
async with session.get(url, headers=headers, timeout=aiohttp.ClientTimeout(total=self.timeout)) as response:
if response.status == 200:
data = await response.json()
return data.get("data", [])
elif response.status == 429:
print(f"[async_get_models] Rate limited (429), attempt {attempt + 1}")
await asyncio.sleep(self.retry_delay * (attempt + 1))
api_key = self.get_random_key()
headers["Authorization"] = f"Bearer {api_key}"
elif response.status == 401:
print(f"[async_get_models] Unauthorized (401), trying next key")
api_key = self.get_random_key()
headers["Authorization"] = f"Bearer {api_key}"
await asyncio.sleep(self.retry_delay)
else:
body = await response.text()
print(f"[async_get_models] HTTP {response.status}: {body[:100]}")
await asyncio.sleep(self.retry_delay)
except (aiohttp.ClientError, asyncio.TimeoutError) as e:
print(f"[async_get_models] Request error: {e}")
await asyncio.sleep(self.retry_delay)
print(f"[async_get_models] Failed after {self.max_retries} attempts")
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
|