Spaces:
Running
Running
# ai_providers.py | |
from typing import Dict, Any | |
import time | |
from groq import Groq | |
from config import GROQ_API_KEY, AI_CONFIG | |
class GroqProvider: | |
def __init__(self): | |
if not GROQ_API_KEY: | |
raise ValueError("GROQ_API_KEY must be set in environment variables") | |
self.client = Groq(api_key=GROQ_API_KEY) | |
self.model = AI_CONFIG["model"] | |
def get_completion(self, prompt: str, **kwargs) -> str: | |
"""Get completion from Groq API with retry logic""" | |
params = { | |
"temperature": kwargs.get("temperature", AI_CONFIG["temperature"]), | |
"max_tokens": kwargs.get("max_tokens", AI_CONFIG["max_tokens"]), | |
"top_p": kwargs.get("top_p", AI_CONFIG["top_p"]), | |
"stop": kwargs.get("stop", None) | |
} | |
max_retries = 3 | |
retry_delay = 1 | |
for attempt in range(max_retries): | |
try: | |
completion = self.client.chat.completions.create( | |
model=self.model, | |
messages=[{"role": "user", "content": prompt}], | |
**params | |
) | |
return completion.choices[0].message.content | |
except Exception as e: | |
if attempt == max_retries - 1: | |
raise Exception(f"Failed to get completion from Groq after {max_retries} attempts: {str(e)}") | |
time.sleep(retry_delay * (attempt + 1)) | |
def get_config(self) -> Dict[str, Any]: | |
"""Get Groq configuration""" | |
return { | |
"config_list": [{ | |
"model": self.model, | |
"api_key": GROQ_API_KEY, | |
"temperature": AI_CONFIG["temperature"], | |
"max_tokens": AI_CONFIG["max_tokens"], | |
"top_p": AI_CONFIG["top_p"] | |
}] | |
} | |
# Global instance | |
groq_provider = GroqProvider() |