Spaces:
Running
Running
File size: 2,692 Bytes
49cf970 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 | import os
from groq import Groq
from dotenv import load_dotenv
# Load environment variables
load_dotenv()
class GroqClient:
def __init__(self, api_key=None, model=None):
self.api_key = api_key or os.getenv("GROQ_API_KEY")
self.model = model or os.getenv("GROQ_MODEL", "meta-llama/llama-4-scout-17b-16e-instruct")
if not self.api_key:
raise ValueError("Groq API Key not found. Please set GROQ_API_KEY in your .env file.")
self.client = Groq(api_key=self.api_key)
def get_completion(self, prompt: str, system_message: str = "You are a helpful AI assistant."):
"""
Calls the Groq API to get a completion for the given prompt.
"""
try:
chat_completion = self.client.chat.completions.create(
messages=[
{
"role": "system",
"content": system_message,
},
{
"role": "user",
"content": prompt,
}
],
model=self.model,
)
return chat_completion.choices[0].message.content
except Exception as e:
return f"Error calling Groq API: {e}"
def get_json_completion(self, prompt: str, system_message: str = "You are a helpful AI assistant."):
"""
Calls the Groq API with JSON mode enabled.
"""
try:
chat_completion = self.client.chat.completions.create(
messages=[
{
"role": "system",
"content": system_message,
},
{
"role": "user",
"content": prompt,
}
],
model=self.model,
response_format={"type": "json_object"},
)
return chat_completion.choices[0].message.content
except Exception as e:
return f"{{\"error\": \"{e}\"}}"
def list_models(self):
"""
Lists available models from Groq.
"""
try:
models = self.client.models.list()
return [model.id for model in models.data]
except Exception as e:
print(f"Error listing models: {e}")
return []
if __name__ == "__main__":
# Test LLM client (requires API key)
try:
client = GroqClient()
response = client.get_completion("Hello, how are you?")
print(f"Groq Response: {response}")
except ValueError as e:
print(e)
|