import gradio as gr
from openai import OpenAI


def complete_text(prompt, max_tokens, temperature, top_p, openai_api_key):
    """
    Get a plain text completion from OpenAI gpt-3.5-turbo-instruct.
    """
    if not openai_api_key:
        return "⚠️ Please enter a valid OpenAI API key."

    client = OpenAI(api_key=openai_api_key)

    response_text = ""
    stream = client.completions.create(
        model="gpt-3.5-turbo-instruct",
        prompt=prompt,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        stream=True,
    )

    for event in stream:
        if hasattr(event, "choices") and event.choices:
            token = event.choices[0].text or ""
            response_text += token
            yield response_text


with gr.Blocks() as demo:
    gr.Markdown("## ✍️ Text Completion Demo (OpenAI gpt-3.5-turbo-instruct)")
    gr.Markdown("Enter a prompt, adjust decoding parameters, and watch the model complete your text.")

    with gr.Row():
        with gr.Column(scale=2):
            prompt = gr.Textbox(
                label="Prompt",
                placeholder="Type the beginning of your text...",
                lines=4,
            )
            max_tokens = gr.Slider(
                minimum=1, maximum=1024, value=100, step=1, label="Max tokens"
            )
            temperature = gr.Slider(
                minimum=0.0, maximum=2.0, value=0.7, step=0.1, label="Temperature"
            )
            top_p = gr.Slider(
                minimum=0.1, maximum=1.0, value=1.0, step=0.05, label="Top-p"
            )
            api_key = gr.Textbox(
                placeholder="sk-... Paste your OpenAI API key here",
                label="🔑 OpenAI API Key",
                type="password",
            )
            submit = gr.Button("Generate Completion")
        with gr.Column(scale=3):
            output = gr.Textbox(
                label="Generated Completion",
                lines=15,
            )

    submit.click(
        fn=complete_text,
        inputs=[prompt, max_tokens, temperature, top_p, api_key],
        outputs=output,
    )

if __name__ == "__main__":
    demo.launch()