import gradio as gr
from blindbox.requests import SecureSession

DEMO_SERVER = "4.208.9.167:80"
text = "<h1 style='text-align: center; color: white; font-size: 30px;'>🎅 SantaCoder with <span style='color: #f0ba2d;'>BlindBox:</span> Confidential Coding Assistant</h1><p style='text-align: left; color: white; font-size: 18px;'></br>This is a demo to show what Zero Trust LLM usage looks like for the use case of Confidential Coding Assistant.</p><p style='text-align: left; color: white; font-size: 18px;'>Here we can leverage a remotely hosted <a style='color: #f0ba2d'; href='https://huggingface.co/bigcode/santacoder#training'>SantaCoder</a>, a state-of-the-art code completion LLM, inside a secure enclave, which ensures code sent for completion is not exposed to anyone else, including us, thanks to end-to-end protection! Therefore LLMs can be leveraged easily to help boost productivity without worrying about IP exposure.</p><p style='text-align: left; color: white; font-size: 18px;'>To learn more about how data is secured, you can find out more in our <a style='color: #f0ba2d'; href='https://blindbox.mithrilsecurity.io/en/latest/'>docs</a>.</p><p style='text-align: left; color: white; font-size: 18px;'>You can see how we deployed SantaCoder with an Azure Confidential VM by checking out the relevant <a style='color: #f0ba2d;', href='https://blindbox.mithrilsecurity.io/en/latest/docs/how-to-guides/santacoder/'>integration guide</a> in our docs.</p><p style='text-align: left; color: white; font-size: 18px;'>⚠️ <a style='color: #f0ba2d;', href='https://github.com/mithril-security/blindbox'>BlindBox</a> is still under development. We have implemented attestation and deployment on Confidential VMs, but we recommend not to send production data on this demo yet.</p><p style='text-align: left; color: white; font-size: 18px;'>If you are interested in pentesting, improving security or knowing more about Confidential LLMs, <a style='color: #f0ba2d' href='https://www.mithrilsecurity.io/contact'>reach out to us!</a></p>"
def run_query(prompt):
    POLICY = "./cce_policy.txt"
    if prompt == None:
        return ("⛔ Error: please provide input code")
    message = "\n\n✅ Secure query succesful"
    try:
        with SecureSession(f"http://{DEMO_SERVER}", POLICY) as secure_session:
            res = secure_session.post(endpoint="/generate", json={"input_text": prompt})
            cleaned = res.text.replace('\\n', '\n').split('\n\n')[0].split(':"')[1]
            cleaned = cleaned.replace('\\', '')
            return(cleaned + message)
    except Exception as err:
        return(f"⛔ Query failed!\n{err}")

demo = gr.Blocks(css=".gradio-container { background-color: #20233fff;} .app.svelte-1mya07g.svelte-1mya07g {max-width: 900px !important;}")
with demo:
    gr.Markdown(value=text)
    _, colum_2, _ = gr.Column(scale=1), gr.Column(scale=6), gr.Column(scale=1)
    with colum_2:
        prompt = gr.Code(lines=3, language="python", label="Input code", value="def hello_name(name):")
        trigger = gr.Button(label="Run query")
    with gr.Column():
        output = gr.Textbox(placeholder="Output", label="Output")
    trigger.click(fn=run_query, inputs=[prompt], outputs=[output])
    gr.HTML(label="Contact", value="<img src='https://github.com/mithril-security/blindbox/blob/laura-images/docs/assets/contact-us.png?raw=true.png' alt='contact' style='display: block; margin: auto; max-width: 600px;'>")
if __name__ == "__main__":
    demo.launch()