import gradio as gr
import ai  # your custom model backend

# --- Response function ---
def respond(message, max_tokens, temperature, top_p):
    start = f"{message}"
    output_so_far = ""
    for chunk in ai.stream_response(
        message=message,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
    ):
        output_so_far = start + chunk
        yield output_so_far

# --- UI ---
with gr.Blocks(theme=gr.themes.Soft()) as demo:
    gr.Markdown("""
    # 🤖 Streaming Chatbot
    Type a message below and watch the model respond in real time.
    """)

    with gr.Row():
        with gr.Column(scale=3):
            output_box = gr.Textbox(label="Generated text", placeholder="Model output will appear here...", lines=20)
            msg = gr.Textbox(label="Your message", placeholder="Ask me anything...", lines=2)
        
        with gr.Column(scale=1):
            with gr.Accordion("⚙️ Advanced Settings", open=False):
                max_tokens = gr.Slider(
                    minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"
                )
                temperature = gr.Slider(
                    minimum=0.1, maximum=4.0, value=0.8, step=0.1, label="Temperature"
                )
                top_p = gr.Slider(
                    minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p"
                )

    msg.submit(respond, [msg, max_tokens, temperature, top_p], output_box)

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