import gradio as gr
import json
from openai import OpenAI


def stream_chat(
    message,
    history: list[tuple[str, str]],
    system_message,
    max_tokens,
    temperature,
    top_p,
    base_url,
    api_key,
    model_name,
    ):

    messages = [{"role": "system", "content": system_message}]

    for val in history:
        if val[0]:
            messages.append({"role": "user", "content": val[0]})
        if val[1]:
            messages.append({"role": "assistant", "content": val[1]})

    messages.append({"role": "user", "content": message})

    reply = ""

    client = OpenAI(
        base_url=base_url,
        api_key=api_key
    )

    # 发送带有流式输出的请求
    for chunk in client.chat.completions.create(
        model=model_name,
        messages=messages,
        max_tokens=max_tokens,
        temperature=temperature,
        top_p=top_p,
        stream=True  # 启用流式输出
    ):
        chunk_message = chunk.choices[0].delta.content
        if chunk_message is not None:
            reply += chunk_message
        else:
            reply += ""
        yield reply


chatapp = gr.ChatInterface(
        stream_chat,
        additional_inputs=[
            gr.Textbox(value="你是一个乐于助人的AI助手.", label="System message"),
            gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
            gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
            gr.Slider(
                minimum=0.1,
                maximum=1.0,
                value=0.95,
                step=0.05,
                label="Top-p (nucleus sampling)",
            ),
            gr.Textbox(value="", label="基础URL", type="text"),
            gr.Textbox(value="", label="API Key", type="password"),
            gr.Textbox(value="", label="模型名称", type="text"),
        ]
    )

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