|
|
|
|
|
import os |
|
|
from typing import List, Optional |
|
|
from pydantic_ai.models.openai import OpenAIModel |
|
|
from pydantic_ai.providers.openai import OpenAIProvider |
|
|
from dotenv import load_dotenv |
|
|
|
|
|
load_dotenv() |
|
|
|
|
|
class InstanceProvider: |
|
|
"""Manages multiple Groq API instances with simple rotation""" |
|
|
|
|
|
def __init__(self): |
|
|
self.instances: List[OpenAIModel] = [] |
|
|
self.current_index = 0 |
|
|
self._initialize_instances() |
|
|
|
|
|
def _initialize_instances(self): |
|
|
"""Load all API keys and create instances""" |
|
|
api_keys = os.getenv("GROQ_API_KEYS", "").split(",") |
|
|
base_url = os.getenv("GROQ_OPENAI_URL") |
|
|
model_name = os.getenv("GROQ_LLM_MODEL") |
|
|
|
|
|
for key in api_keys: |
|
|
key = key.strip() |
|
|
if key: |
|
|
self.instances.append( |
|
|
OpenAIModel( |
|
|
model_name, |
|
|
provider=OpenAIProvider( |
|
|
base_url=base_url, |
|
|
api_key=key |
|
|
) |
|
|
) |
|
|
) |
|
|
|
|
|
def get_next_instance(self) -> Optional[OpenAIModel]: |
|
|
"""Get next instance in rotation""" |
|
|
if not self.instances: |
|
|
return None |
|
|
|
|
|
instance = self.instances[self.current_index] |
|
|
self.current_index = (self.current_index + 1) % len(self.instances) |
|
|
return instance |
|
|
|
|
|
def get_total_instances(self) -> int: |
|
|
"""Return total number of instances available""" |
|
|
return len(self.instances) |
