app.py

TITLE = """<h1 align="center">Gemini Playground ✨</h1>"""
SUBTITLE = """<h2 align="center">Play with Gemini Pro and Gemini Pro Vision</h2>"""

import os
import time
from typing import List, Tuple, Optional

import google.generativeai as genai
import gradio as gr
from PIL import Image
from dotenv import load_dotenv

# Cargar las variables de entorno desde el archivo .env
load_dotenv()

print("google-generativeai:", genai.__version__)

# Obtener la clave de la API de las variables de entorno
GOOGLE_API_KEY = os.getenv("GOOGLE_API_KEY")
if not GOOGLE_API_KEY:
    raise ValueError("GOOGLE_API_KEY is not set in environment variables.")

# Configurar la API de Google Generative AI
genai.configure(api_key=GOOGLE_API_KEY)

# Variables globales
chat = None  # Sesión de chat
IMAGE_WIDTH = 512
CHAT_HISTORY = List[Tuple[Optional[str], Optional[str]]]


def preprocess_image(image: Image.Image) -> str:
    """Preprocesar la imagen y convertirla a texto descriptivo."""
    return "Image processed successfully."


def transform_history(history: CHAT_HISTORY) -> List[dict]:
    """Transformar el historial de Gradio al formato requerido por Gemini."""
    transformed = []
    for user_msg, model_msg in history:
        if user_msg:
            transformed.append({"role": "user", "content": user_msg})
        if model_msg:
            transformed.append({"role": "model", "content": model_msg})
    return transformed


def initialize_chat(model_name: str):
    """Inicializar una sesión de chat con el modelo seleccionado."""
    global chat
    model = genai.GenerativeModel(model_name=model_name)
    chat = model.start_chat(history=[])


def bot_with_logic(
    text_prompt: str,
    files: Optional[List[str]],
    model_choice: str,
    system_instruction: str,
    chatbot: CHAT_HISTORY,
):
    """Lógica del chatbot para manejar texto, imágenes o ambos."""
    global chat

    # Inicializar la sesión de chat si no existe
    if chat is None:
        initialize_chat(model_choice)

    # Configurar la instrucción del sistema
    chat.system_instruction = system_instruction or "You are a helpful assistant."

    # Caso 1: Solo texto
    if text_prompt and not files:
        response = chat.send_message(text_prompt)
        response.resolve()

        chatbot.append((text_prompt, ""))
        for i in range(len(response.text)):
            chatbot[-1] = (text_prompt, response.text[: i + 1])
            time.sleep(0.01)
            yield chatbot

    # Caso 2: Solo imágenes o texto + imágenes
    elif files:
        image_descriptions = [preprocess_image(Image.open(file)) for file in files]
        combined_prompt = f"{text_prompt}\n" + "\n".join(image_descriptions) if text_prompt else "\n".join(
            image_descriptions
        )

        response = chat.send_message(combined_prompt)
        response.resolve()

        chatbot.append((text_prompt or "[Images Uploaded]", ""))
        for i in range(len(response.text)):
            chatbot[-1] = (text_prompt or "[Images Uploaded]", response.text[: i + 1])
            time.sleep(0.01)
            yield chatbot


# Componentes de Gradio
chatbot_component = gr.Chatbot(label="Gemini Chat", height=400)
text_prompt_component = gr.Textbox(placeholder="Enter your message here...", show_label=False)
upload_button_component = gr.UploadButton(label="Upload Images", file_count="multiple", file_types=["image"])
run_button_component = gr.Button(value="Run", variant="primary")
model_choice_component = gr.Dropdown(
    choices=["gemini-1.5-flash", "gemini-2.0-flash-exp", "gemini-1.5-pro"],
    value="gemini-1.5-flash",
    label="Select Model",
)
system_instruction_component = gr.Textbox(placeholder="Enter system instruction...", label="System Instruction")


# Crear la interfaz
with gr.Blocks() as demo:
    gr.HTML(TITLE)
    gr.HTML(SUBTITLE)
    with gr.Row():
        model_choice_component.render()
    chatbot_component.render()
    with gr.Row():
        text_prompt_component.render()
        upload_button_component.render()
        run_button_component.render()
    system_instruction_component.render()

    run_button_component.click(
        fn=bot_with_logic,
        inputs=[
            text_prompt_component,
            upload_button_component,
            model_choice_component,
            system_instruction_component,
            chatbot_component,
        ],
        outputs=[chatbot_component],
    )

    text_prompt_component.submit(
        fn=bot_with_logic,
        inputs=[
            text_prompt_component,
            upload_button_component,
            model_choice_component,
            system_instruction_component,
            chatbot_component,
        ],
        outputs=[chatbot_component],
    )

# Lanzar la aplicación
demo.queue(max_size=99).launch(debug=True, show_error=True)