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 uuid
from typing import List, Tuple, Optional, Union

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")

# Verificar que la clave de la API esté configurada
if not GOOGLE_API_KEY:
    raise ValueError("GOOGLE_API_KEY is not set in environment variables.")

IMAGE_CACHE_DIRECTORY = "/tmp"
IMAGE_WIDTH = 512
CHAT_HISTORY = List[Tuple[Optional[Union[Tuple[str], str]], Optional[str]]]

# Configurar la API de Gemini
genai.configure(api_key=GOOGLE_API_KEY)

# Función para transformar el historial del chat
def transform_history(history: CHAT_HISTORY):
    """
    Transforma el historial del chat en el formato necesario para el modelo.
    """
    transformed = []
    for user_input, bot_response in history:
        if user_input:
            transformed.append({"role": "user", "content": user_input})
        if bot_response:
            transformed.append({"role": "assistant", "content": bot_response})
    return transformed

# Función de generación de respuesta
def response(
    message: str, history: CHAT_HISTORY, model_choice: str, system_instruction: str
) -> str:
    """
    Genera una respuesta basada en el historial del chat y el mensaje del usuario.
    """
    generation_config = genai.types.GenerationConfig(
        temperature=0.7,
        max_output_tokens=8192,
        top_k=10,
        top_p=0.9
    )

    model = genai.GenerativeModel(
        model_name=model_choice,
        generation_config=generation_config,
        system_instruction=system_instruction,
    )

    transformed_history = transform_history(history)
    model_response = model.chat(messages=transformed_history + [{"role": "user", "content": message}])
    return model_response.text

# Preprocesamiento de imágenes
def preprocess_image(image: Image.Image) -> Optional[Image.Image]:
    if image:
        image_height = int(image.height * IMAGE_WIDTH / image.width)
        return image.resize((IMAGE_WIDTH, image_height))

# Guardar imágenes en caché
def cache_pil_image(image: Image.Image) -> str:
    image_filename = f"{uuid.uuid4()}.jpeg"
    os.makedirs(IMAGE_CACHE_DIRECTORY, exist_ok=True)
    image_path = os.path.join(IMAGE_CACHE_DIRECTORY, image_filename)
    image.save(image_path, "JPEG")
    return image_path

# Subir imágenes
def upload(files: Optional[List[str]], chatbot: CHAT_HISTORY) -> CHAT_HISTORY:
    for file in files:
        image = Image.open(file).convert("RGB")
        image_preview = preprocess_image(image)
        if image_preview:
            gr.Image(image_preview).render()
        image_path = cache_pil_image(image)
        chatbot.append(((image_path,), None))
    return chatbot

# Manejo del usuario
def user(text_prompt: str, chatbot: CHAT_HISTORY):
    if text_prompt:
        chatbot.append((text_prompt, None))
    return "", chatbot

# Manejo del bot con historial
def bot(
    files: Optional[List[str]],
    model_choice: str,
    system_instruction: str,
    chatbot: CHAT_HISTORY,
):
    text_prompt = chatbot[-1][0] if chatbot and chatbot[-1][0] and isinstance(chatbot[-1][0], str) else ""
    bot_reply = response(text_prompt, chatbot, model_choice, system_instruction)
    chatbot[-1] = (text_prompt, bot_reply)
    return chatbot

# Componentes de la interfaz
system_instruction_component = gr.Textbox(
    placeholder="Enter system instruction...", show_label=True, scale=8
)
chatbot_component = gr.Chatbot(
    label="Gemini",
    bubble_full_width=False,
    scale=2,
    height=300
)
text_prompt_component = gr.Textbox(
    placeholder="Message...", show_label=False, autofocus=True, scale=8
)
upload_button_component = gr.UploadButton(
    label="Upload Images", file_count="multiple", file_types=["image"], scale=1
)
run_button_component = gr.Button(value="Run", variant="primary", scale=1)
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",
    scale=2
)

user_inputs = [
    text_prompt_component,
    chatbot_component,
]

bot_inputs = [
    upload_button_component,
    model_choice_component,
    system_instruction_component,
    chatbot_component,
]

# Interfaz de usuario
with gr.Blocks() as demo:
    gr.HTML(TITLE)
    gr.HTML(SUBTITLE)
    with gr.Column():
        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=user,
        inputs=user_inputs,
        outputs=[text_prompt_component, chatbot_component],
        queue=False
    ).then(
        fn=bot, inputs=bot_inputs, outputs=[chatbot_component],
    )

    text_prompt_component.submit(
        fn=user,
        inputs=user_inputs,
        outputs=[text_prompt_component, chatbot_component],
        queue=False
    ).then(
        fn=bot, inputs=bot_inputs, outputs=[chatbot_component],
    )

    upload_button_component.upload(
        fn=upload,
        inputs=[upload_button_component, chatbot_component],
        outputs=[chatbot_component],
        queue=False
    )

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