import gradio as gr
import requests
import json

def http_bot(prompt, history, system_prompt, endpoint_url):
    print(prompt, history, system_prompt, endpoint_url)
    prompt = f"### Instruction:\n{system_prompt}\n\n### Input:\n{prompt}\n\n### Response:\n"

    headers = {"User-Agent": "vLLM Client"}
    pload = {
        "prompt": prompt,
        "stream": True,
        "max_tokens": 3000,
    }
    response = requests.post(endpoint_url,
                             headers=headers,
                             json=pload,
                             stream=True)

    for chunk in response.iter_lines(chunk_size=8192,
                                     decode_unicode=False,
                                     delimiter=b"\0"):
        if chunk:
            data = json.loads(chunk.decode("utf-8"))
            output = data["text"][0]
        
            yield output[len(prompt):]




with gr.Blocks(theme=gr.themes.Soft(), title="DanskGPT") as demo:
    gr.Markdown("# DanskGPT v0.3")
    gr.Markdown("Et dansk alternativ til ChatGPT der kører lokalt.")
    system_prompt = gr.Textbox(value="Du er en dansk AI-assistent. Du vil blive givet en opgave. Du skal hjælpe så meget du kan.",
                                label="System besked")
    endpoint_url = gr.Textbox(label="Endpoint url", value="https://h6lsu4k84eqhh4-8000.proxy.runpod.net/generate", render=False)


    gr.ChatInterface(
        http_bot,
        additional_inputs=[system_prompt, endpoint_url],
        clear_btn=None,
        undo_btn=None,
        retry_btn=None,
        submit_btn="Send",
        
    )

    gr.Markdown("Version 0.3 - instruktionsmodel trænet på data op til 31-06-2023.")
    gr.Markdown("Lavet af Mads Henrichsen - Kontakt: mads.gade.henrichsen@live.dk")
    gr.Markdown("Bemærk: DanskGPT producerer ikke nødvendigvis sandfærdig information - der er en ny model på vej der er meget bedre til dette.")

demo.queue(concurrency_count=100).launch()