| | import os |
| | import logging |
| | from typing import List, Optional, Dict |
| | from dotenv import load_dotenv |
| |
|
| | load_dotenv() |
| |
|
| | logging.basicConfig(level=logging.INFO) |
| | logger = logging.getLogger(__name__) |
| |
|
| | class InstanceProvider: |
| | """Manages multiple Cerebras API keys with simple rotation""" |
| | |
| | def __init__(self): |
| | self.api_keys: List[str] = [] |
| | self.current_index = 0 |
| | self.base_url = os.getenv("CEREBRAS_BASE_URL") |
| | self.model_name = os.getenv("CEREBRAS_MODEL", "llama3.1-70b") |
| | self._initialize_instances() |
| | |
| | def _initialize_instances(self): |
| | """Load all API keys into a list""" |
| | keys_str = os.getenv("CEREBRAS_API_KEYS", "") |
| | self.api_keys = [k.strip() for k in keys_str.split(",") if k.strip()] |
| | |
| | if not self.api_keys: |
| | logger.error("No API keys found in CEREBRAS_API_KEYS") |
| |
|
| | def get_next_instance(self) -> Optional[Dict[str, str]]: |
| | """ |
| | Returns a dictionary with the credentials for the next instance. |
| | Returns: {'api_key': str, 'base_url': str, 'model': str} |
| | """ |
| | if not self.api_keys: |
| | return None |
| | |
| | |
| | key = self.api_keys[self.current_index] |
| | |
| | |
| | self.current_index = (self.current_index + 1) % len(self.api_keys) |
| | |
| | return { |
| | "api_key": key, |
| | "base_url": self.base_url, |
| | "model": self.model_name |
| | } |
| | |
| | def get_total_instances(self) -> int: |
| | return len(self.api_keys) |
