Update app.py

app.py

@@ -14,6 +14,17 @@ st.set_page_config(
     initial_sidebar_state="expanded",
 )
 
+# Add the missing translation function
+def translate_role_to_streamlit(role):
+    """
+    Translate Gemini message roles to Streamlit chat roles
+    """
+    role_mapping = {
+        'model': 'assistant',
+        'user': 'user'
+    }
+    return role_mapping.get(role, 'assistant')
+
 # Archivo para almacenamiento permanente
 PROMPTS_FILE = "training_prompts.json"
 
@@ -120,27 +131,81 @@ elif selected == "Chatbot":
             st.info(current_prompt)
     
     # Display chat history
-    for message in st.session_state.chat_session.history:
-        with st.chat_message(translate_role_to_streamlit(message.role)):
-            st.markdown(message.parts[0].text)
+    if hasattr(st.session_state.chat_session, 'history'):
+        for message in st.session_state.chat_session.history:
+            role = translate_role_to_streamlit(message.role)
+            with st.chat_message(role):
+                st.markdown(message.parts[0].text)
 
     # Chat input
     user_prompt = st.chat_input("Preguntame algo...")
     if user_prompt:
         st.chat_message("user").markdown(user_prompt)
-        gemini_response = st.session_state.chat_session.send_message(user_prompt)
-        with st.chat_message("assistant"):
-            st.markdown(gemini_response.text)
+        try:
+            gemini_response = st.session_state.chat_session.send_message(user_prompt)
+            with st.chat_message("assistant"):
+                st.markdown(gemini_response.text)
+        except Exception as e:
+            st.error(f"Error en la respuesta: {str(e)}")
 
 elif selected == "Image Captioning":
     st.title("Image Caption Generation📸")
+    
+    # Add custom prompt option
+    use_custom_prompt = st.checkbox("Usar prompt personalizado")
+    if use_custom_prompt:
+        custom_prompt = st.text_area(
+            "Escribe tu prompt personalizado",
+            value="Write a caption for this image",
+            help="Puedes personalizar las instrucciones para el análisis de la imagen"
+        )
+    
     upload_image = st.file_uploader("Upload an image...", type=["jpg", "jpeg", "png"])
-    if upload_image and st.button("Generate"):
-        image = Image.open(upload_image)
-        col1, col2 = st.columns(2)
-        with col1:
-            st.image(image, caption="Uploaded Image", use_column_width=True)
-        default_prompt = "Write a caption for this image"
-        caption = gemini_pro_vision_responce(default_prompt, image)
-        with col2:
-            st.info(caption)
+    
+    if upload_image:
+        try:
+            image = Image.open(upload_image)
+            col1, col2 = st.columns(2)
+            
+            with col1:
+                st.image(image, caption="Uploaded Image", use_column_width=True)
+                image_info = f"Dimensiones: {image.size[0]}x{image.size[1]}"
+                st.text(image_info)
+            
+            if st.button("Generate", key="generate_caption"):
+                with st.spinner("Generando descripción..."):
+                    prompt = custom_prompt if use_custom_prompt else "Write a caption for this image"
+                    try:
+                        caption = gemini_pro_vision_responce(prompt, image)
+                        with col2:
+                            st.success("¡Descripción generada!")
+                            st.info(caption)
+                            
+                            # Opción para guardar la descripción
+                            if st.button("Guardar descripción"):
+                                timestamp = time.strftime("%Y%m%d-%H%M%S")
+                                filename = f"caption_{timestamp}.txt"
+                                try:
+                                    with open(filename, "w", encoding="utf-8") as f:
+                                        f.write(caption)
+                                    st.success(f"Descripción guardada en {filename}")
+                                except Exception as e:
+                                    st.error(f"Error al guardar la descripción: {str(e)}")
+                    
+                    except Exception as e:
+                        st.error(f"Error al generar la descripción: {str(e)}")
+                        st.error("Por favor, intenta con otra imagen o revisa tu conexión")
+        
+        except Exception as e:
+            st.error(f"Error al procesar la imagen: {str(e)}")
+            st.error("Por favor, asegúrate de que el archivo es una imagen válida")
+
+# Añadir footer
+st.markdown("---")
+col1, col2, col3 = st.columns(3)
+with col1:
+    st.markdown("**GnosticDev AI**")
+with col2:
+    st.markdown("Versión 1.0.0")
+with col3:
+    st.markdown("Made with ❤️ by GnosticDev")