Update app.py

app.py

@@ -1,87 +1,169 @@
 import os
 import time
+import uuid
+from typing import List, Tuple, Optional, Union
+
 import google.generativeai as genai
-from typing import List
-
-# Configurar la API con la clave de entorno
-genai.configure(api_key=os.environ["GEMINI_API_KEY"])
-
-def upload_to_gemini(path, mime_type=None):
-    """Sube el archivo dado a Gemini."""
-    file = genai.upload_file(path, mime_type=mime_type)
-    print(f"Uploaded file '{file.display_name}' as: {file.uri}")
-    return file
-
-def wait_for_files_active(files: List):
-    """Espera que los archivos cargados estén activos y procesados."""
-    print("Waiting for file processing...")
-    for name in (file.name for file in files):
-        file = genai.get_file(name)
-        while file.state.name == "PROCESSING":
-            print(".", end="", flush=True)
-            time.sleep(10)
-            file = genai.get_file(name)
-        if file.state.name != "ACTIVE":
-            raise Exception(f"File {file.name} failed to process")
-    print("...all files ready")
-    print()
-
-# Configuración del modelo y parámetros de generación
-generation_config = {
-    "temperature": 1,
-    "top_p": 0.95,
-    "top_k": 40,
-    "max_output_tokens": 8192,
-    "response_mime_type": "text/plain",
-}
-
-model = genai.GenerativeModel(
-    model_name="gemini-1.5-flash",
-    generation_config=generation_config,
-)
+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]]]
 
-def upload_files(files: List[str]) -> List:
-    """Sube múltiples archivos (imágenes, PDFs, etc.) a Gemini."""
-    uploaded_files = []
+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))
+
+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
+
+def upload(files: Optional[List[str]], chatbot: CHAT_HISTORY) -> CHAT_HISTORY:
     for file in files:
-        file_extension = os.path.splitext(file)[1].lower()
-        if file_extension in [".jpg", ".jpeg", ".png", ".gif"]:
-            # Para imágenes, usamos el mime_type correspondiente
-            mime_type = "image/jpeg" if file_extension in [".jpg", ".jpeg"] else "image/png"
-        elif file_extension == ".pdf":
-            mime_type = "application/pdf"
-        elif file_extension in [".txt", ".csv"]:
-            mime_type = "text/plain"
-        else:
-            mime_type = None  # Para otros tipos de archivo, determinamos el mime_type si es necesario
+        # Detectar el tipo de archivo y manejarlo adecuadamente
+        mime_type = file.type if hasattr(file, 'type') else None
         
-        uploaded_files.append(upload_to_gemini(file, mime_type))
-    
-    return uploaded_files
-
-# Archivos a cargar (Ejemplo: agregar los archivos a subir)
-files_to_upload = [
-    "Report 2025 - Jan 15th.pdf",  # Ejemplo de archivo PDF
-    "image_example.jpg",           # Ejemplo de archivo de imagen
-    "text_example.txt",            # Ejemplo de archivo de texto
-]
+        # Si es una imagen, la procesamos con PIL
+        if mime_type and mime_type.startswith('image'):
+            image = Image.open(file).convert('RGB')
+            image_preview = preprocess_image(image)
+            if image_preview:
+                # Mostrar una vista previa de la imagen cargada
+                gr.Image(image_preview).render()
+            image_path = cache_pil_image(image)
+            chatbot.append(((image_path,), None))
+        else:
+            # Si no es una imagen, se guarda el archivo tal cual
+            file_path = cache_file(file)
+            chatbot.append(((file_path,), None))
+    return chatbot
+
+def cache_file(file: str) -> str:
+    file_filename = f"{uuid.uuid4()}_{os.path.basename(file.name)}"
+    os.makedirs(IMAGE_CACHE_DIRECTORY, exist_ok=True)
+    file_path = os.path.join(IMAGE_CACHE_DIRECTORY, file_filename)
+    with open(file_path, 'wb') as f:
+        f.write(file.read())
+    return file_path
 
-# Subir los archivos
-uploaded_files = upload_files(files_to_upload)
+def user(text_prompt: str, chatbot: CHAT_HISTORY):
+    if text_prompt:
+        chatbot.append((text_prompt, None))
+    return "", chatbot
 
-# Esperar a que los archivos sean procesados y estén activos
-wait_for_files_active(uploaded_files)
+def bot(
+    files: Optional[List[str]],
+    model_choice: str,
+    chatbot: CHAT_HISTORY
+):
+    if not GOOGLE_API_KEY:
+        raise ValueError("GOOGLE_API_KEY is not set.")
 
-# Iniciar sesión de chat con el modelo
-chat_session = model.start_chat(
-    history=[{
-        "role": "user",
-        "parts": uploaded_files,  # Incluir archivos subidos en la conversación
-    }],
+    # Configurar la API con la clave
+    genai.configure(api_key=GOOGLE_API_KEY)
+    generation_config = genai.types.GenerationConfig(
+        temperature=0.7,  # Valor predeterminado
+        max_output_tokens=8192,  # Fijar el límite de tokens a 8,192
+        top_k=10,  # Valor predeterminado
+        top_p=0.9   # Valor predeterminado
+    )
+
+    text_prompt = [chatbot[-1][0]] if chatbot and chatbot[-1][0] and isinstance(chatbot[-1][0], str) else []
+    files_paths = [preprocess_image(Image.open(file).convert('RGB')) if file.type.startswith('image') else cache_file(file) for file in files]
+    model = genai.GenerativeModel(model_choice)
+    response = model.generate_content(text_prompt + files_paths, stream=True, generation_config=generation_config)
+
+    chatbot[-1][1] = ""
+    for chunk in response:
+        for i in range(0, len(chunk.text), 10):
+            section = chunk.text[i:i + 10]
+            chatbot[-1][1] += section
+            time.sleep(0.01)
+            yield chatbot
+
+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 Files", file_count="multiple", 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,
+    chatbot_component
+]
+
+with gr.Blocks() as demo:
+    gr.HTML(TITLE)
+    gr.HTML(SUBTITLE)
+    with gr.Column():
+        chatbot_component.render()
+        with gr.Row():
+            text_prompt_component.render()
+            upload_button_component.render()
+            run_button_component.render()
+        model_choice_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],
+    )
 
-# Enviar mensaje al modelo
-response = chat_session.send_message("INSERT_INPUT_HERE")
+    upload_button_component.upload(
+        fn=upload,
+        inputs=[upload_button_component, chatbot_component],
+        outputs=[chatbot_component],
+        queue=False
+    )
 
-# Imprimir la respuesta
-print(response.text)
+demo.queue(max_size=99).launch(debug=False, show_error=True)