import os
import gradio as gr
from mistralai.client import MistralClient
from mistralai.models.chat_completion import ChatMessage

api_key = os.environ["MISTRAL_API_KEY"]
model = "mistral-large-latest"

client = MistralClient(api_key=api_key)

header = r"""
<center>
    <h1>Mistral Chat</h1>
    <h2>
        This demo is for code reference, feel free to clone it and set your own Api key
    </h2>
    Get an Api Key: <a href='https://console.mistral.ai/api-keys'>Mistral Ai</a>
</center>
"""

footer = r"""
<center>
    <b>
        Demo for <a href='https://mistral.ai'>Mistral AI</a>
    </b>
</center>
"""

with gr.Blocks(title="Mistral Chat") as app:

    gr.HTML(header)

    with gr.Row(equal_height=True):
        with gr.Column():
            chatbot = gr.Chatbot()
            msg = gr.Textbox()
            clear = gr.Button("Clear")

            def user(user_message, history):
                return "", history + [[user_message, None]]

            def bot(history):
                
                messages = [
                    ChatMessage(role="user", content=history[-1][0])
                ]
                
                response = client.chat_stream(model=model, messages=messages)

                history[-1][1] = ""
                for chunk in response:
                    history[-1][1] += chunk.choices[0].delta.content
                    yield history

            msg.submit(user, [msg, chatbot], [msg, chatbot], queue=False).then(
                bot, chatbot, chatbot
            )
            clear.click(lambda: None, None, chatbot, queue=False)
    
    with gr.Row():
        gr.HTML(footer)
    
app.queue()
app.launch(share=False, debug=True, show_error=True)