import gradio as gr
from pangea.services import Redact
import openai

def redact_input(input: str, pangea_token: str):
    redact = Redact(token=pangea_token)
    return redact.redact(text=input).result.redacted_text

def gpt_completion(prompt: str, openai_api_key: str):
    openai.api_key = openai_api_key
    response = openai.Completion.create(model="text-davinci-003",
                                      prompt=prompt,
                                      temperature=0,
                                      max_tokens=100)
    
    return response.choices[0].text.strip()


input_text = gr.components.Textbox(label="Enter your GPT prompt here (User Input)", value='Rewrite and continue this email to let the specific user know the next steps for their lease, "hey user with email dummyuser@email.com (last logged in from 127.0.0.1), here\'s the lease agreement for your new property":')
output_text1 = gr.components.Textbox(label="Redacted User Input")
output_text2 = gr.components.Textbox(label="GPT-3 Prompt Completion (on Redacted Input)")


def main(oai_api_key, pangea_token, input_text):
    redacted_text = redact_input(input_text, pangea_token)
    completion = gpt_completion(redacted_text, oai_api_key)
    return [redacted_text, completion]

openai_api_key_input = gr.components.Textbox(label="OpenAI API Key (Not stored)", type="password", placeholder="sk-......")
pangea_token_input = gr.components.Textbox(label="Pangea Redact Token (Not stored)", type="password", placeholder="pts_.....")

iface = gr.Interface(
    fn=main,
    inputs=[openai_api_key_input, pangea_token_input, input_text],
    outputs=[output_text1, output_text2],
    title="GPT-3 Input Redaction Playground",
    description="<center>Enter your OpenAI key and your Pangea redact token, then put in your prompt and watch the magic happen 🪄. <br/>\
    Data redaction is powered by <a href=\"https://pangea.cloud/?utm_source=huggingface&utm_medium=redact-gpt-prompt-demo\">Pangea's APIs</a> in this demo. To learn more about how to get your API keys and tokens to run this, watch this short <a href=\"https://www.youtube.com/watch?v=LNN5s_6G3Cc\" target=\"_blank\">walkthrough demo</a></center>",
    allow_flagging="never"
)
iface.launch(
    favicon_path="./favicon-32x32.png"
)