import gradio as gr
import requests
import os
import json

api_key = os.getenv('API_KEY')

def call_llama_guard_api(content, assistant_response):
    invoke_url = "https://api.nvcf.nvidia.com/v2/nvcf/pexec/functions/b34280ac-24e4-4081-bfaa-501e9ee16b6f"
    headers = {
        "Authorization": f"Bearer {api_key}",
        "Accept": "application/json",
    }
    payload = {
        "messages": [
            {"content": content, "role": "user"},
            {"content": assistant_response, "role": "assistant"}
        ]
    }

    session = requests.Session()
    response = session.post(invoke_url, headers=headers, json=payload)

    while response.status_code == 202:
        request_id = response.headers.get("NVCF-REQID")
        fetch_url = f"https://api.nvcf.nvidia.com/v2/nvcf/pexec/status/{request_id}"
        response = session.get(fetch_url, headers=headers)

    response.raise_for_status()
    response_body = response.json()
    print(response_body)
    return response_body

content_input = gr.Textbox(lines=2, placeholder="Enter your content here...", label="User Content")
assistant_response_input = gr.Textbox(lines=2, placeholder="Enter assistant's response here...", label="Assistant Response")

iface = gr.Interface(fn=call_llama_guard_api,
                     inputs=[content_input, assistant_response_input],
                     outputs="text",
                     title="Llama Guard Safety Classifier",
                     description="Classify the safety of LLM prompts and responses using Llama Guard"
                    )

iface.launch()