File size: 615 Bytes
af2a153
 
 
 
 
63c4985
 
af2a153
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
import os
from groq import Groq

def init_groq(api_key: str=None):

    api_key = api_key or os.getenv("GROQ_API_KEY")
    return Groq(api_key=api_key)

def groq_chat(prompt: str, model_id: str, temperature: float = 0.7, max_tokens: int = 200):
    
    client = init_groq()
    response=client.chat.completions.create(
        model=model_id,
        messages=[
            {"role" : "system", "content" : "You are a helpful assistant."},
            {"role": "user", "content": prompt}
        ],
        temperature=temperature,
        max_tokens=max_tokens
    )

    return response.choices[0].message.content