import gradio as gr
import requests
import os
import json

api_key = os.getenv('API_KEY')

def call_mistral_7b_api(content, system_prompt, temperature, top_p, max_tokens):
    invoke_url = "https://api.nvcf.nvidia.com/v2/nvcf/pexec/functions/8f4118ba-60a8-4e6b-8574-e38a4067a4a3"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "accept": "text/event-stream",
        "content-type": "application/json",
    }
    payload = {
        "messages": [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": content}
        ],
        "temperature": temperature,
        "top_p": top_p,
        "max_tokens": max_tokens,
        "stream": True
    }
    response = requests.post(invoke_url, headers=headers, json=payload, stream=True)
    if response.status_code != 200:
        print(f"Erro na requisição: {response.status_code}")
        try:
            error_details = response.json()
            print(error_details)
        except ValueError:
            print(response.text)
    else:
        response_text = ""
        for line in response.iter_lines():
            if line:
                decoded_line = line.decode('utf-8').strip()
                if decoded_line.startswith('data: {'):
                    json_str = decoded_line[6:]
                    try:
                        json_line = json.loads(json_str)
                        content_parts = json_line.get("choices", [{}])[0].get("delta", {}).get("content", "")
                        response_text += content_parts
                    except json.JSONDecodeError as e:
                        print(f"Erro ao decodificar JSON: {e}")
                        print(f"Linha problemática: {decoded_line}")
                elif decoded_line == 'data: [DONE]':
                    print("Recebido sinal de conclusão da API.")
                    break
                else:
                    print(f"Linha ignorada (não é JSON ou sinal de conclusão): {decoded_line}")
        return response_text

content_input = gr.Textbox(lines=2, placeholder="Enter your content here...", label="Content")
system_prompt_input = gr.Textbox(value="I carefully provide accurate, factual, thoughtful, nuanced answers and am brilliant at reasoning.", label="System Prompt")
temperature_input = gr.Slider(minimum=0, maximum=1, step=0.01, value=0.2, label="Temperature")
top_p_input = gr.Slider(minimum=0, maximum=1, step=0.01, value=0.7, label="Top P")
max_tokens_input = gr.Slider(minimum=1, maximum=1024, step=1, value=1024, label="Max Tokens")

iface = gr.Interface(fn=call_mistral_7b_api,
                     inputs=[content_input, system_prompt_input, temperature_input, top_p_input, max_tokens_input],
                     outputs="text",
                     title="Mixtral 8x7B Instruct Free Demo",
                     description="Explore the capabilities of Mixtral 8x7B Instruct Free Demo"
                    )

iface.launch()