Update app.py

app.py

@@ -1,179 +1,87 @@
-TITLE = """<h1 align="center">Gemini Playground ✨</h1>"""
-SUBTITLE = """<h2 align="center">Play with Gemini Pro and Gemini Pro Vision</h2>"""
-
 import os
 import time
-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
-from PyPDF2 import PdfReader
-
-# 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 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 process_text_file(file_path: str):
-    with open(file_path, 'r') as file:
-        return file.read()
-
-def process_pdf_file(file_path: str):
-    reader = PdfReader(file_path)
-    text = ""
-    for page in reader.pages:
-        text += page.extract_text()
-    return text
+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,
+)
 
-def upload(files: Optional[List[str]], chatbot: CHAT_HISTORY) -> CHAT_HISTORY:
+def upload_files(files: List[str]) -> List:
+    """Sube múltiples archivos (imágenes, PDFs, etc.) a Gemini."""
+    uploaded_files = []
     for file in files:
-        file_extension = file.split('.')[-1].lower()
-        
-        if file_extension in ['jpg', 'jpeg', 'png', 'gif', 'bmp']:  # Imágenes
-            image = Image.open(file).convert('RGB')
-            image_preview = preprocess_image(image)
-            if image_preview:
-                # Display a preview of the uploaded image
-                gr.Image(image_preview).render()
-            image_path = cache_pil_image(image)
-            chatbot.append(((image_path,), None))
-
-        elif file_extension in ['txt']:  # Archivos de texto
-            text_content = process_text_file(file)
-            chatbot.append(((text_content,), None))
-
-        elif file_extension in ['pdf']:  # Archivos PDF
-            pdf_content = process_pdf_file(file)
-            chatbot.append(((pdf_content,), None))
+        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
         
-    return chatbot
-
-def user(text_prompt: str, chatbot: CHAT_HISTORY):
-    if text_prompt:
-        chatbot.append((text_prompt, None))
-    return "", chatbot
-
-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.")
-
-    # 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 []
-    image_prompt = [preprocess_image(Image.open(file).convert('RGB')) for file in files] if files else []
-    model = genai.GenerativeModel(model_choice)
-    response = model.generate_content(text_prompt + image_prompt, 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", file_types=["image", "text", "pdf"], 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
+        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
 ]
 
-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()
+# Subir los archivos
+uploaded_files = upload_files(files_to_upload)
 
-    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],
-    )
+# Esperar a que los archivos sean procesados y estén activos
+wait_for_files_active(uploaded_files)
 
-    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],
-    )
+# 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
+    }],
+)
 
-    upload_button_component.upload(
-        fn=upload,
-        inputs=[upload_button_component, chatbot_component],
-        outputs=[chatbot_component],
-        queue=False
-    )
+# Enviar mensaje al modelo
+response = chat_session.send_message("INSERT_INPUT_HERE")
 
-demo.queue(max_size=99).launch(debug=False, show_error=True)
+# Imprimir la respuesta
+print(response.text)