Update Gradio_UI.py

Gradio_UI.py

@@ -18,14 +18,14 @@ import os
 import re
 import shutil
 from typing import Optional
-import tempfile # Added for PIL image saving
-from PIL import Image as PILImage # Added for PIL image handling
+import tempfile 
+from PIL import Image as PILImage 
 
 from smolagents.agent_types import AgentAudio, AgentImage, AgentText, handle_agent_output_types
 from smolagents.agents import ActionStep, MultiStepAgent 
 from smolagents.memory import MemoryStep
 from smolagents.utils import _is_package_available
-import gradio as gr # Ensure gradio is imported at the top level
+import gradio as gr 
 
 
 def pull_messages_from_step_dict(step_log: MemoryStep):
@@ -66,7 +66,7 @@ def pull_messages_from_step_dict(step_log: MemoryStep):
                 obs_content = step_log.observations.strip()
                 obs_content = re.sub(r"^Execution logs:\s*", "", obs_content).strip()
                 if obs_content: 
-                    tool_info_md += f"📝 **Tool Output/Logs:**\n```text\n{obs_content}\n```\n" # Use text for generic logs
+                    tool_info_md += f"📝 **Tool Output/Logs:**\n```text\n{obs_content}\n```\n"
             
             if hasattr(step_log, "error") and step_log.error:
                 tool_info_md += f"💥 **Error:** {str(step_log.error)}\n"
@@ -105,10 +105,9 @@ def stream_to_gradio(
         agent.interaction_logs.clear()
         print("DEBUG Gradio: Cleared agent interaction_logs for new run.")
 
-    # This will collect all step_log objects from the agent run
     all_step_logs = []
     for step_log in agent.run(task, stream=True, reset=reset_agent_memory, additional_args=additional_args):
-        all_step_logs.append(step_log) # Store the log
+        all_step_logs.append(step_log) 
         if hasattr(agent.model, "last_input_token_count") and agent.model.last_input_token_count is not None:
             if isinstance(step_log, ActionStep):
                 step_log.input_token_count = agent.model.last_input_token_count
@@ -117,79 +116,58 @@ def stream_to_gradio(
         for msg_dict in pull_messages_from_step_dict(step_log): 
             yield msg_dict
     
-    # After the loop, the last item in all_step_logs is the final output/state from agent.run
-    if not all_step_logs: # Should not happen if agent.run yields at least one thing
+    if not all_step_logs:
         yield {"role": "assistant", "content": "Agent did not produce any output."}
         return
 
-    final_answer_content = all_step_logs[-1] # This is what final_answer tool returns or the last ActionStep.final_answer
+    final_answer_content = all_step_logs[-1] 
 
-    # --- Handle final answer for type="messages" ---
-    if isinstance(final_answer_content, PILImage.Image):
-        print("DEBUG Gradio (stream_to_gradio): Final answer content IS a raw PIL Image.")
-        try:
-            # delete=False is crucial for Gradio to access the file before it's cleaned up
-            with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp_file:
-                final_answer_content.save(tmp_file, format="PNG")
-                image_path_for_gradio = tmp_file.name
-            print(f"DEBUG Gradio: Saved PIL image to temp path for display: {image_path_for_gradio}")
-            yield {"role": "assistant", "content": image_path_for_gradio} 
-            return 
-        except Exception as e:
-            print(f"DEBUG Gradio: Error saving PIL image from final_answer_content: {e}")
-            yield {"role": "assistant", "content": f"**Final Answer (Error displaying image):** {e}"}
-            return
-
-    # If not a raw PIL Image, then try smolagents processing from handle_agent_output_types
-    # The 'final_answer_content' here could be a FinalAnswerStep object or similar
-    # We need to extract the actual content from it if it's a wrapper.
     actual_content_for_handling = final_answer_content
     if hasattr(final_answer_content, 'final_answer') and not isinstance(final_answer_content, (str, PILImage.Image)):
          actual_content_for_handling = final_answer_content.final_answer
          print(f"DEBUG Gradio: Extracted actual_content_for_handling from FinalAnswerStep: {type(actual_content_for_handling)}")
 
-
-    # Re-check if the extracted content is a PIL Image
     if isinstance(actual_content_for_handling, PILImage.Image):
-        print("DEBUG Gradio (stream_to_gradio): Extracted content IS a raw PIL Image.")
+        print("DEBUG Gradio (stream_to_gradio): Actual content IS a raw PIL Image.")
         try:
             with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp_file:
                 actual_content_for_handling.save(tmp_file, format="PNG")
                 image_path_for_gradio = tmp_file.name
-            print(f"DEBUG Gradio: Saved extracted PIL image to temp path: {image_path_for_gradio}")
-            yield {"role": "assistant", "content": image_path_for_gradio} 
+            print(f"DEBUG Gradio: Saved PIL image to temp path: {image_path_for_gradio}")
+            # MODIFIED: Yield tuple (filepath, alt_text)
+            yield {"role": "assistant", "content": (image_path_for_gradio, "Generated Image")} 
             return
         except Exception as e:
             print(f"DEBUG Gradio: Error saving extracted PIL image: {e}")
-            yield {"role": "assistant", "content": f"**Final Answer (Error displaying image from extracted content):** {e}"}
+            yield {"role": "assistant", "content": f"**Final Answer (Error displaying image):** {e}"}
             return
 
     final_answer_processed = handle_agent_output_types(actual_content_for_handling)
     print(f"DEBUG Gradio: final_answer_processed type after handle_agent_output_types: {type(final_answer_processed)}")
 
-
     if isinstance(final_answer_processed, AgentText):
         yield {"role": "assistant", "content": f"**Final Answer:**\n{final_answer_processed.to_string()}"}
     elif isinstance(final_answer_processed, AgentImage):
         image_path = final_answer_processed.to_string()
         print(f"DEBUG Gradio (stream_to_gradio): final_answer_processed is AgentImage. Path: {image_path}")
         if image_path and os.path.exists(image_path):
-             yield {"role": "assistant", "content": image_path} 
+             # MODIFIED: Yield tuple (filepath, alt_text)
+             yield {"role": "assistant", "content": (image_path, "Generated Image (from AgentImage)")} 
         else:
-            err_msg = f"Error: Image path from AgentImage ('{image_path}') not found or invalid after smolagents processing."
+            err_msg = f"Error: Image path from AgentImage ('{image_path}') not found or invalid."
             print(f"DEBUG Gradio: {err_msg}")
             yield {"role": "assistant", "content": f"**Final Answer ({err_msg})**"}
     elif isinstance(final_answer_processed, AgentAudio):
         audio_path = final_answer_processed.to_string()
         print(f"DEBUG Gradio (stream_to_gradio): AgentAudio path: {audio_path}")
         if audio_path and os.path.exists(audio_path):
-            yield {"role": "assistant", "content": audio_path}
+            # MODIFIED: Yield tuple (filepath, alt_text) for consistency, though Gradio might just use path for audio
+            yield {"role": "assistant", "content": (audio_path, "Generated Audio")}
         else:
             err_msg = f"Error: Audio path from AgentAudio ('{audio_path}') not found"
             print(f"DEBUG Gradio: {err_msg}")
             yield {"role": "assistant", "content": f"**Final Answer ({err_msg})**"}
     else: 
-        # This will display the string representation of FinalAnswerStep if not handled above
         yield {"role": "assistant", "content": f"**Final Answer:**\n{str(final_answer_processed)}"}
 
 
@@ -211,12 +189,8 @@ class GradioUI:
             for log_entry in reversed(self.agent.interaction_logs): 
                 if isinstance(log_entry, ActionStep):
                     observations = getattr(log_entry, 'observations', None)
-                    tool_calls = getattr(log_entry, 'tool_calls', [])
-
-                    # Check if python_interpreter was used AND its code involved create_document
-                    # For simplicity, we'll primarily rely on parsing observations for the path pattern
                     if observations and isinstance(observations, str):
-                        # This regex should match paths printed by your create_document tool
+                        print(f"DEBUG Gradio UI: Checking observations: {observations[:200]}") 
                         path_match = re.search(r"(/tmp/[a-zA-Z0-9_]+/generated_document\.(?:docx|pdf|txt))", observations)
                         if path_match:
                             extracted_path = path_match.group(1)
@@ -234,7 +208,6 @@ class GradioUI:
         print(f"DEBUG Gradio: interact_with_agent called with prompt: '{prompt_text}'")
         print(f"DEBUG Gradio: Current chat history (input type {type(current_chat_history)}): {current_chat_history}")
 
-        # current_chat_history from gr.Chatbot(type="messages") is already a list of dicts
         updated_chat_history = current_chat_history + [{"role": "user", "content": prompt_text}]
         
         yield updated_chat_history, gr.update(visible=False), gr.update(value=None, visible=False)
@@ -275,7 +248,6 @@ class GradioUI:
         sanitized_name = re.sub(r"[^\w\-.]", "_", original_name)
         base_name, current_ext = os.path.splitext(sanitized_name)
         
-        # Updated mimetypes to extension mapping
         common_mime_to_ext = {
             "application/pdf": ".pdf",
             "application/vnd.openxmlformats-officedocument.wordprocessingml.document": ".docx",
@@ -302,7 +274,7 @@ class GradioUI:
         if current_file_uploads:
             files_str = ", ".join([os.path.basename(f) for f in current_file_uploads])
             full_prompt += f"\n\n[Uploaded files for context: {files_str}]"
-        print(f"DEBUG Gradio: Prepared prompt for agent: {full_prompt[:300]}...") # Log snippet
+        print(f"DEBUG Gradio: Prepared prompt for agent: {full_prompt[:300]}...") 
         return full_prompt, "" 
 
     def prepare_and_show_download_file(self):
@@ -313,46 +285,44 @@ class GradioUI:
                                   visible=True)
         else:
             print("DEBUG Gradio UI: No valid file path to prepare for download component.")
-            # gr.Warning("No file available for download or path is invalid.") # Causes JS error if used as return
-            return gr.File.update(visible=False, value=None) # Ensure value is None if not visible
+            return gr.File.update(visible=False, value=None) 
 
     def launch(self, **kwargs):
         with gr.Blocks(fill_height=True, theme=gr.themes.Soft(primary_hue=gr.themes.colors.blue)) as demo:
             file_uploads_log_state = gr.State([]) 
             prepared_prompt_for_agent = gr.State("")
 
-            gr.Markdown("## Smol Talk with your Agent") # Changed title slightly
+            gr.Markdown("## Smol Talk with your Agent") 
 
-            with gr.Row(equal_height=False): # Allow columns to size independently
+            with gr.Row(equal_height=False): 
                 with gr.Column(scale=3):
                     chatbot_display = gr.Chatbot(
                         label="Agent Interaction",
                         type="messages", 
                         avatar_images=(None, "https://huggingface.co/datasets/huggingface/brand-assets/resolve/main/hf-logo-round.png"),
-                        height=700, # Increased height
+                        height=700, 
                         show_copy_button=True,
                         bubble_full_width=False,
-                        show_label=False # Hide the "Agent Interaction" label above chatbot
+                        show_label=False 
                     )
                     text_message_input = gr.Textbox(
                         lines=1, 
                         label="Your Message to the Agent", 
                         placeholder="Type your message and press Enter, or Shift+Enter for new line...",
-                        show_label=False # Hide label for text input
+                        show_label=False 
                     )
                 
                 with gr.Column(scale=1):
-                    if self.file_upload_folder is not None:
-                        with gr.Accordion("File Upload", open=False): # Collapsible section
-                            file_uploader = gr.File(label="Upload a supporting file (PDF, DOCX, TXT, JPG, PNG)")
-                            upload_status_text = gr.Textbox(label="Upload Status", interactive=False, lines=1) # single line
-                            file_uploader.upload( # Changed from .change to .upload for gr.File
-                                self.upload_file,
-                                [file_uploader, file_uploads_log_state],
-                                [upload_status_text, file_uploads_log_state],
-                            )
+                    with gr.Accordion("File Upload", open=False): 
+                        file_uploader = gr.File(label="Upload a supporting file (PDF, DOCX, TXT, JPG, PNG)")
+                        upload_status_text = gr.Textbox(label="Upload Status", interactive=False, lines=1) 
+                        file_uploader.upload( 
+                            self.upload_file,
+                            [file_uploader, file_uploads_log_state],
+                            [upload_status_text, file_uploads_log_state],
+                        )
                     
-                    with gr.Accordion("Generated File", open=True): # Collapsible, open by default
+                    with gr.Accordion("Generated File", open=True): 
                         download_action_button = gr.Button("Download Generated File", visible=False) 
                         file_download_display_component = gr.File(label="Downloadable Document", visible=False, interactive=False) 
 
@@ -362,8 +332,8 @@ class GradioUI:
                 [prepared_prompt_for_agent, text_message_input] 
             ).then(
                 self.interact_with_agent, 
-                [prepared_prompt_for_agent, chatbot_display], # chatbot_display is input here
-                [chatbot_display, download_action_button, file_download_display_component] # chatbot_display is output here
+                [prepared_prompt_for_agent, chatbot_display], 
+                [chatbot_display, download_action_button, file_download_display_component] 
             )
 
             download_action_button.click(
@@ -371,7 +341,6 @@ class GradioUI:
                 [], 
                 [file_download_display_component] 
             )
-        # Default share=False, can be overridden by kwargs
         demo.launch(debug=True, share=kwargs.get("share", False), **kwargs)
 
 __all__ = ["stream_to_gradio", "GradioUI"]