src/openrouter_client.py

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