app.py

import os
import gradio as gr
import google.generativeai as genai
import asyncio

###############################################################################
# 1. Настройка окружения и инициализация моделей
###############################################################################

# Подставьте свой ключ или берите из окружения
GEMINI_API_KEY = os.environ.get("GEMINI_API_KEY")
if not GEMINI_API_KEY:
    print("Error: GEMINI_API_KEY is not set. Please set the environment variable.")
    exit()

genai.configure(api_key=GEMINI_API_KEY)

# Выберите доступные модели (пример)
AVAILABLE_MODELS = [
    "gemini-1.5-flash",
    # Add other available models if needed
]

MODELS = {}
for model_name in AVAILABLE_MODELS:
    try:
        MODELS[model_name] = genai.GenerativeModel(model_name=model_name)
    except Exception as e:
        print(f"[Предупреждение] Не удалось инициализировать модель {model_name}: {e}")

###############################################################################
# 2. Дефолтные промпты (developer role) для каждой модели
###############################################################################
# Когда пользователь переключается на модель, мы добавляем это сообщение в историю.

DEFAULT_DEVELOPER_PROMPTS = {
    "gemini-1.5-flash": (
        "You are a helpful assistant (developer role). "
        "Provide direct answers."
    ),
    # Add default prompts for other models if needed
}

###############################################################################
# 3. Функция для определения роли ассистента (assistant vs model)
###############################################################################

def _assistant_role(model_name: str) -> str:
    """
    Некоторые новые модели не принимают 'assistant', а требуют 'model'.
    """
    # Adjust based on model requirements
    return "assistant"

###############################################################################
# 4. Преобразование истории из Gradio в формат Generative AI
###############################################################################

def _history_to_genai(history, model_name):
    """
    Gradio хранит [(user_msg, bot_msg), ...].
    Google GenAI ждёт [{"role": "...", "parts": "..."}].
    Дополнительно учитываем developer-сообщения, которые мы вставляем.
    """
    genai_history = []
    asst_role = _assistant_role(model_name)

    for user_text, bot_text in history:
        if user_text:
            if user_text.startswith("<developer>: "):
                genai_history.append({"role": "system", "parts": user_text.replace("<developer>: ", "", 1)})
            else:
                # Пользовательское сообщение
                genai_history.append({"role": "user", "parts": user_text})
        if bot_text:
            # Ответ ассистента (или model)
            genai_history.append({"role": asst_role, "parts": bot_text})
    return genai_history

###############################################################################
# 5. Генераторы для стрима обычных моделей и "thinking" моделей
###############################################################################

async def _respond_stream(model_name, user_message, history):
    """
    Стриминговый ответ для обычных моделей:
    - Кусочек за кусочком (partial_text).
    """
    if model_name not in MODELS:
        yield "Ошибка: модель не найдена."
        return

    model = MODELS[model_name]
    genai_history = _history_to_genai(history, model_name)

    try:
        chat = model.start_chat(history=genai_history)
        response = chat.send_message(user_message, stream=True)

        partial_text = ""
        async for chunk in response:
            partial_text += (chunk.text or "")
            yield partial_text

        return
    except Exception as e:
        yield f"Ошибка при запросе к API: {e}"
        return

async def _respond_thinking(model_name, user_message, history):
    """
    Для thinking-моделей:
    1) Выводим "Думаю..."
    2) После завершения — финальный ответ в формате {output: ...} + размышления.
    """
    if model_name not in MODELS:
        yield "Ошибка: модель не найдена.", ""
        return

    model = MODELS[model_name]
    genai_history = _history_to_genai(history, model_name)

    # Сначала "Думаю..."
    yield "Думаю...", ""

    try:
        chat = model.start_chat(history=genai_history)
        response = chat.send_message(user_message, stream=False)

        thinking_process_text = ""
        final_text = ""

        if response.candidates:
            parts = response.candidates[0].content.parts
            for p in parts:
                # Assuming a specific structure for "thinking" models' output
                # This part might need adjustments based on actual model output
                if isinstance(p, dict) and "output" not in p:  # Heuristic for thought
                    thinking_process_text += p.text or ""
                elif isinstance(p, dict) and "output" in p:
                    final_text += p["output"]  # Assuming output is directly in the dict
                else:
                    final_text += p.text or ""  # Fallback

        # Для thinking-моделей просили итоговый ответ в {output: ...}
        final_text_formatted = f"{{output: {final_text}}}"

        yield final_text_formatted, thinking_process_text
        return

    except Exception as e:
        yield f"Ошибка при запросе к API: {e}", ""
        return

###############################################################################
# 6. Основная функция для ввода пользователя
###############################################################################

async def user_send_message(
    user_message: str,
    history: list[tuple[str, str]],
    model_name: str,
    thinking_text: str
):
    """
    Колбэк, когда пользователь отправляет запрос.
    Добавляем в history новый (user_msg, None), затем генерируем ответ.
    """
    # Пустой ввод
    if not user_message.strip():
        yield history, thinking_text
        return

    # Добавляем (user_message, None)
    history.append((user_message, None))

    # Если модель — thinking (adjust logic based on actual model names)
    if "thinking" in model_name.lower():
        async for assistant_text, thought_text in _respond_thinking(model_name, user_message, history):
            history[-1] = (user_message, assistant_text)
            yield history, thought_text
        return
    else:
        # Обычная модель
        partial_answer = ""
        async for chunk in _respond_stream(model_name, user_message, history):
            partial_answer = chunk
            history[-1] = (user_message, partial_answer)
            yield history, ""
        return

###############################################################################
# 7. Очистка диалога
###############################################################################

def clear_all():
    """Сброс истории и размышлений."""
    return [], ""

###############################################################################
# 8. Когда меняем модель в Dropdown
###############################################################################

def on_model_change(selected_model, history, thinking_text):
    """
    При переключении модели добавляем в историю developer-сообщение,
    + добавляем дефолтный промпт этой модели (тоже developer).
    """
    new_history = history.copy()

    # Cообщаем модели, что переключились (developer role)
    new_history.append((
        f"<developer>: Switched to model: {selected_model}",
        None
    ))

    # Добавляем дефолтный промпт (developer role)
    default_prompt = DEFAULT_DEVELOPER_PROMPTS.get(
        selected_model,
        "No default prompt for this model."
    )
    new_history.append((
        f"<developer>: {default_prompt}",
        None
    ))

    return new_history, thinking_text

###############################################################################
# 9. Функция конвертации истории с учётом developer role
###############################################################################

# This part is already implemented in _history_to_genai

###############################################################################
# 10. Построение интерфейса Gradio
###############################################################################

with gr.Blocks() as demo:
    gr.Markdown("## Chat с Gemini. Поддержка developer role, переключения моделей, JSON-ответа для thinking")

    with gr.Row():
        model_dropdown = gr.Dropdown(
            choices=AVAILABLE_MODELS,
            value=AVAILABLE_MODELS[0] if AVAILABLE_MODELS else None,
            label="Выберите модель"
        )
        clear_button = gr.Button("Очистить чат")

    history_state = gr.State([])
    thinking_store = gr.State("")

    chatbot = gr.Chatbot(label="Диалог с Gemini")
    user_input = gr.Textbox(label="Ваш вопрос", placeholder="Введите текст...")
    thinking_output = gr.Textbox(label="Размышления", interactive=False)
    send_btn = gr.Button("Отправить")

    ################################################
    # (A) Обработка переключения модели
    ################################################
    def handle_model_change(selected_model, history, thinking):
        new_history, new_thinking = on_model_change(selected_model, history, thinking)
        return new_history, new_thinking

    # Когда пользователь меняет модель:
    model_change = model_dropdown.change(
        fn=handle_model_change,
        inputs=[model_dropdown, history_state, thinking_store],
        outputs=[history_state, thinking_store],
        queue=False
    ).then(
        # После добавления developer-сообщений в историю → обновляем чат
        fn=lambda h: h,
        inputs=[history_state],
        outputs=[chatbot],
        queue=False
    )

    ################################################
    # (B) При нажатии «Отправить»
    ################################################
    send_chain = send_btn.click(
        fn=user_send_message,
        inputs=[user_input, history_state, model_dropdown, thinking_store],
        outputs=[history_state, thinking_store],
        queue=True
    )
    send_chain.then(
        fn=lambda h: h,
        inputs=[history_state],
        outputs=[chatbot],
        queue=True
    )
    send_chain.then(
        fn=lambda t: t,
        inputs=[thinking_store],
        outputs=[thinking_output],
        queue=True
    )
    # Очистка поля ввода
    send_chain.then(
        fn=lambda: "",
        inputs=[],
        outputs=[user_input],
        queue=False
    )

    ################################################
    # (C) При нажатии Enter в textbox
    ################################################
    submit_chain = user_input.submit(
        fn=user_send_message,
        inputs=[user_input, history_state, model_dropdown, thinking_store],
        outputs=[history_state, thinking_store],
        queue=True
    )
    submit_chain.then(
        fn=lambda h: h,
        inputs=[history_state],
        outputs=[chatbot],
        queue=True
    )
    submit_chain.then(
        fn=lambda t: t,
        inputs=[thinking_store],
        outputs=[thinking_output],
        queue=True
    )
    # Очистка поля ввода
    submit_chain.then(
        fn=lambda: "",
        inputs=[],
        outputs=[user_input],
        queue=False
    )

    ################################################
    # (D) Кнопка «Очистить»
    ################################################
    clear_chain = clear_button.click(
        fn=clear_all,
        inputs=[],
        outputs=[history_state, thinking_store],
        queue=False
    )
    clear_chain.then(
        fn=lambda h: h,
        inputs=[history_state],
        outputs=[chatbot]
    )
    clear_chain.then(
        fn=lambda _: "",
        inputs=[],
        outputs=[thinking_output]
    )
    clear_chain.then(
        fn=lambda: "",
        inputs=[],
        outputs=[user_input]
    )

# Запуск
if __name__ == "__main__":
    demo.launch()