from flask import Flask
from flask import request
from groq import Groq

import os

app = Flask(__name__)
client = Groq(
    api_key=os.environ.get("GROQ_API_KEY")
)


@app.route("/api/generate", methods=['POST'])
def completion():
    """
    {
        "model": "llama3-70b-8192",
        "prompt": "why is the sky blue?"
    }
    """

    message = request.get_json()
    
    model = message['model']
    prompt = message['prompt']

    chat_completion = client.chat.completions.create(
        messages=[
            {
                "role": "user",
                "content": prompt,
            }
        ],
        model=model,
    )

    return chat_completion.to_dict()

# curl -v -X POST 'https://robinroy03-fury-bot.hf.space/api/generate' --header 'Content-Type: application/json' --data '{"model": "llama3-70b-8192", "prompt": "why is sky blue?"}'