import gradio as gr

# Предполагается, что модель уже загружена и доступна через Gradio interface
interface = gr.load("models/mistralai/Mixtral-8x7B-Instruct-v0.1")

DEFAULT_SYSTEM_PROMPT = """
You are a helpful assistant in normal conversation.
When given a problem to solve, you are an expert problem-solving assistant.
Your task is to provide a detailed, step-by-step solution to a given question.
"""

# Создание интерфейса чата
with gr.Blocks() as demo:
    gr.Markdown("# Custom Chat Interface with AI")
    
    system_prompt = gr.Textbox(value=DEFAULT_SYSTEM_PROMPT, lines=5, label="System Prompt")
    chatbot = gr.Chatbot(label="Chat")

    # Функция для обработки сообщений пользователя и получения ответов от модели
    def chat(message):
        # Формируем входные данные для модели включая системный промпт и сообщение пользователя
        inputs = [system_prompt.value.strip(), message.strip()]
        # Вызов модели для генерации ответа с увеличением лимита токенов до 100
        response = interface(inputs, parameters={"max_tokens": 100})
        # Ожидаем, что ответ будет в формате списка строк, каждая строка - это сообщение
        chat_history = []
        for part in response:
            chat_history.append((part, "user" if part == message else "assistant"))
        return chat_history

    msg = gr.Textbox(label="Type your message here...", placeholder="Enter your message...")
    msg.submit(chat, inputs=msg, outputs=chatbot)

    demo.launch()