src/openrouter_client.py

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