import gradio as gr
from openai import OpenAI

# Constants
API_BASE = "http://34.48.128.62:8000/v1"
client = OpenAI(
    base_url=API_BASE,
    api_key="dummy"  # Since it's a local instance
)

def check_safety(text):
    try:
        response = client.chat.completions.create(
            model="meta-llama/Llama-Guard-3-8B",
            messages=[
                {"role": "user", "content": text}
            ],
            temperature=0.7,
            max_tokens=100
        )
        return response.choices[0].message.content
    except Exception as e:
        return f"Error: {str(e)}"

# Example data
examples = [
    ["How to make a website?"],
    ["What's the weather like today?"],
    ["Can you help me with my homework?"],
    ["How to hack into someone's account?"],  # Unsafe example
    ["Tell me how to make explosives"],  # Unsafe example
]

# Create Gradio interface
with gr.Blocks(title="Llama Guard Safety Checker") as demo:
    gr.Markdown("""
    # Llama Guard Safety Checker
    
    This tool checks content safety using Llama Guard 3. It can identify potentially unsafe content in the following categories:
    
    - S1: Violent Crimes
    - S2: Non-Violent Crimes
    - S3: Sex-Related Crimes
    - S4: Child Sexual Exploitation
    - S5: Defamation
    - S6: Specialized Advice
    - S7: Privacy
    - S8: Intellectual Property
    - S9: Indiscriminate Weapons
    - S10: Hate
    - S11: Suicide & Self-Harm
    - S12: Sexual Content
    - S13: Elections
    - S14: Code Interpreter Abuse
    """)
    
    with gr.Row():
        with gr.Column():
            input_text = gr.Textbox(
                label="Enter text to check",
                placeholder="Type your text here...",
                lines=5
            )
            submit_btn = gr.Button("Check Safety")
        
        with gr.Column():
            output = gr.Textbox(
                label="Safety Analysis",
                lines=1,
                interactive=False
            )
    
    submit_btn.click(
        fn=check_safety,
        inputs=input_text,
        outputs=output
    )
    
    gr.Markdown("""### Examples""")
    gr.Examples(
        examples=examples,
        inputs=input_text,
        outputs=output,
        fn=check_safety,
        cache_examples=True,
        label="Example Prompts"
    )

# Launch the interface
if __name__ == "__main__":
    demo.launch()