import gradio as gr
import cohere
import os

co = cohere.ClientV2(os.getenv("COHERE_API_KEY"))

def chat(message, history, system_prompt):
    messages = []
    if system_prompt:
        messages.append({"role": "system", "content": system_prompt})
    for item in history:
        messages.append({"role": item["role"], "content": item["content"]})
    messages.append({"role": "user", "content": message})

    stream = co.chat_stream(
        model="command-a-03-2025",
        messages=messages
    )

    response = ""
    for event in stream:
        if event.type == "content-delta":
            response += event.delta.message.content.text
            yield response


with gr.Blocks() as demo:
    gr.Markdown("## Chatbot Cohere (Gradio)")

    system_prompt = gr.Textbox(
        value="Tu es un assistant utile.",
        label="Prompt système"
    )

    chatbot = gr.Chatbot()
    msg = gr.Textbox(label="Message")

    def user_fn(message, history):
        return "", history + [{"role": "user", "content": message}]

    def bot_fn(history, system_prompt):
        message = history[-1]["content"]
        previous = history[:-1]
        history = history + [{"role": "assistant", "content": ""}]
        for chunk in chat(message, previous, system_prompt):
            history[-1]["content"] = chunk
            yield history

    msg.submit(user_fn, [msg, chatbot], [msg, chatbot], queue=False).then(
        bot_fn, [chatbot, system_prompt], chatbot
    )

demo.queue()
demo.launch()