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

title = "Gaia Mistral Chat Demo"
description = "Example of simple chatbot with Gradio and Mistral AI via its API"
placeholder = "Posez moi une question sur l'agriculture"
examples = ["Comment fait on pour produire du maïs ?", "Rédige moi une lettre pour faire un stage dans une exploitation agricole", "Comment reprendre une exploitation agricole ?"]
api_key = os.environ.get("MISTRAL_API_KEY")

client = MistralClient(api_key=api_key)
model = 'mistral-small'

def chat_with_mistral(user_input, history):
    messages = [ChatMessage(role="user", content=user_input)]

    chat_response = client.chat(model=model, messages=messages)
    return chat_response.choices[0].message.content

iface = gr.ChatInterface(
    fn=chat_with_mistral,
    chatbot=gr.Chatbot(height=300),
    textbox=gr.Textbox(placeholder=placeholder, container=False, scale=7),
    title=title,
    description=description,
    theme="soft",
    examples=examples,
    cache_examples=True,
    retry_btn=None,
    undo_btn="Annuler",
    clear_btn="Effacer",
)

iface.launch(share=True)