Update Gradio_UI.py

Gradio_UI.py

@@ -13,10 +13,10 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-import mimetypes
+# import mimetypes # No longer needed if file upload is removed
 import os
 import re
-import shutil
+# import shutil # No longer needed if file upload is removed
 from typing import Optional
 import tempfile 
 from PIL import Image as PILImage 
@@ -123,14 +123,16 @@ def stream_to_gradio(
     final_answer_content = all_step_logs[-1] 
 
     actual_content_for_handling = final_answer_content
-    if hasattr(final_answer_content, 'final_answer') and not isinstance(final_answer_content, (str, PILImage.Image)):
+    # Check if final_answer_content is a wrapper like FinalAnswerStep and extract the core content
+    if hasattr(final_answer_content, 'final_answer') and not isinstance(final_answer_content, (str, PILImage.Image, tuple)): # Added tuple to avoid unwrapping already formatted image content
          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)}")
 
+    # Priority 1: Handle raw PIL Image object for direct display
     if isinstance(actual_content_for_handling, PILImage.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:
+            with tempfile.NamedTemporaryFile(delete=False, suffix=".png") as tmp_file: # delete=False is important
                 actual_content_for_handling.save(tmp_file, format="PNG")
                 image_path_for_gradio = tmp_file.name
             print(f"DEBUG Gradio: Saved PIL image to temp path: {image_path_for_gradio}")
@@ -141,6 +143,7 @@ def stream_to_gradio(
             yield {"role": "assistant", "content": f"**Final Answer (Error displaying image):** {e}"}
             return
 
+    # Priority 2: Use smolagents' type handling if not a raw PIL image
     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)}")
 
@@ -165,18 +168,20 @@ def stream_to_gradio(
             print(f"DEBUG Gradio: {err_msg}")
             yield {"role": "assistant", "content": f"**Final Answer ({err_msg})**"}
     else: 
+        # This will display the string representation of the object if not specifically handled
         yield {"role": "assistant", "content": f"**Final Answer:**\n{str(final_answer_processed)}"}
 
 
 class GradioUI:
-    def __init__(self, agent: MultiStepAgent, file_upload_folder: str | None = None):
+    def __init__(self, agent: MultiStepAgent, file_upload_folder: str | None = None): # file_upload_folder kept for potential future use
         if not _is_package_available("gradio"):
             raise ModuleNotFoundError("Install 'gradio': `pip install 'smolagents[gradio]'`")
         self.agent = agent
-        self.file_upload_folder = file_upload_folder
-        if self.file_upload_folder is not None:
-            if not os.path.exists(self.file_upload_folder):
-                os.makedirs(self.file_upload_folder, exist_ok=True)
+        # self.file_upload_folder = file_upload_folder # Commented out as per request to simplify
+        # if self.file_upload_folder is not None:
+        #     if not os.path.exists(self.file_upload_folder):
+        #         os.makedirs(self.file_upload_folder, exist_ok=True)
+        self.file_upload_folder = None # Explicitly disable for now
         self._latest_file_path_for_download = None
 
     def _check_for_created_file(self):
@@ -186,112 +191,69 @@ 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', []) # Get tool calls if any
                     
-                    if observations and isinstance(observations, str):
-                        # Log the full observations for the relevant step to help debug the regex
-                        # We are looking for the output of python_interpreter which contains create_document's print
-                        tool_calls = getattr(log_entry, 'tool_calls', [])
-                        is_python_interpreter_step = any(tc.name == "python_interpreter" for tc in tool_calls)
-
-                        if is_python_interpreter_step: # Only parse observations from python_interpreter
-                            print(f"DEBUG Gradio UI (_check_for_file): Python Interpreter Observations: '''{observations}'''") 
-                            
-                            # Regex to find paths specifically printed by our create_document tool
-                            # This pattern should match:
-                            # "Document created (docx): /tmp/random/generated_document.docx"
-                            # "Document created (txt): /tmp/random/generated_document.txt"
-                            # "Document converted to PDF: /tmp/random/generated_document.pdf"
-                            # It allows for optional "Execution logs:" prefix and surrounding newlines/whitespace.
-                            match = re.search(
-                                r"(?:Document created \((?:docx|pdf|txt)\):|Document converted to PDF:)\s*(/tmp/[a-zA-Z0-9_]+/generated_document\.(?:docx|pdf|txt))",
-                                observations,
-                                re.MULTILINE # Important if observations string has multiple lines
-                            )
-                            
-                            if match:
-                                extracted_path = match.group(1) 
-                                normalized_path = os.path.normpath(extracted_path)
-                                if os.path.exists(normalized_path):
-                                    self._latest_file_path_for_download = normalized_path
-                                    print(f"DEBUG Gradio UI: File path for download SET: {self._latest_file_path_for_download}")
-                                    return True # File found and path set
-                                else:
-                                    print(f"DEBUG Gradio UI: Path from create_document output ('{normalized_path}') does not exist.")
-                            # else: # Optional: log if the specific pattern isn't found in this observation
-                                # print(f"DEBUG Gradio UI: 'create_document' output pattern not found in this observation block.")
-        print("DEBUG Gradio UI: No valid generated file path (from create_document) found in agent logs for download.")
+                    # Check if this step involved python_interpreter
+                    is_python_interpreter_step = any(tc.name == "python_interpreter" for tc in tool_calls)
+
+                    if is_python_interpreter_step and observations and isinstance(observations, str):
+                        print(f"DEBUG Gradio UI (_check_for_file): Python Interpreter Observations: '''{observations[:500]}...'''") # Log snippet
+                        
+                        # Regex to find paths specifically printed by our create_document tool
+                        # This pattern expects: "Document created (docx): /tmp/random/generated_document.docx"
+                        match = re.search(
+                            # Capture group 1 is the prefix, group 2 is the path
+                            r"(Document created \((?:docx|pdf|txt)\):|Document converted to PDF:)\s*(/tmp/[a-zA-Z0-9_]+/generated_document\.(?:docx|pdf|txt))",
+                            observations,
+                            re.MULTILINE 
+                        )
+                        
+                        if match:
+                            # extracted_prefix = match.group(1) # e.g., "Document created (docx):"
+                            extracted_path = match.group(2) # The actual path
+                            print(f"DEBUG Gradio UI: Regex matched. Extracted path: '{extracted_path}'")
+                            normalized_path = os.path.normpath(extracted_path)
+                            if os.path.exists(normalized_path):
+                                self._latest_file_path_for_download = normalized_path
+                                print(f"DEBUG Gradio UI: File path for download SET: {self._latest_file_path_for_download}")
+                                return True 
+                            else:
+                                print(f"DEBUG Gradio UI: Path from create_document output ('{normalized_path}') does not exist.")
+                        # else:
+                            # print(f"DEBUG Gradio UI: 'create_document' output pattern not found in this observation block.")
+        print("DEBUG Gradio UI: No valid generated file path (from create_document) found for download button.")
         return False
 
     def interact_with_agent(self, prompt_text: str, current_chat_history: list):
         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}")
 
         updated_chat_history = current_chat_history + [{"role": "user", "content": prompt_text}]
         
+        # Initial yield: show user message, hide download components until agent run is complete
         yield updated_chat_history, gr.update(visible=False), gr.update(value=None, visible=False)
 
         agent_responses_for_history = []
         for msg_dict in stream_to_gradio(self.agent, task=prompt_text, reset_agent_memory=False):
             agent_responses_for_history.append(msg_dict)
+            # Yield progressively to update chat, keep download components hidden during streaming
             yield updated_chat_history + agent_responses_for_history, gr.update(visible=False), gr.update(value=None, visible=False) 
 
-        file_found = self._check_for_created_file() # This call is crucial
-        
-        final_chat_display = updated_chat_history + agent_responses_for_history
-        print(f"DEBUG Gradio: Final chat history for display: {len(final_chat_display)} messages. File found for download: {file_found}")
-        yield final_chat_display, gr.update(visible=file_found), gr.update(value=None, visible=False) # Update download button visibility
-
-    def upload_file(self, file, file_uploads_log_state):
-        if file is None: 
-            return gr.update(value="No file uploaded.", visible=True), file_uploads_log_state
-
-        if not self.file_upload_folder or not os.path.exists(self.file_upload_folder):
-            os.makedirs(self.file_upload_folder, exist_ok=True)
-
-        allowed_file_types = [
-            "application/pdf",
-            "application/vnd.openxmlformats-officedocument.wordprocessingml.document",
-            "text/plain", "image/jpeg", "image/png",
-        ]
+        # After all agent messages are processed and added to history
+        final_chat_display_content = updated_chat_history + agent_responses_for_history
         
-        original_name = file.orig_name if hasattr(file, 'orig_name') and file.orig_name else os.path.basename(file.name)
+        # Now check for created files to decide visibility of download button
+        file_found_for_download = self._check_for_created_file()
         
-        mime_type, _ = mimetypes.guess_type(file.name) 
-        if mime_type is None: 
-            mime_type, _ = mimetypes.guess_type(original_name)
+        print(f"DEBUG Gradio: Final chat history for display: {len(final_chat_display_content)} messages. File found for download button: {file_found_for_download}")
+        # Final yield: update chat, set visibility of download button, keep file display component hidden (it's shown on button click)
+        yield final_chat_display_content, gr.update(visible=file_found_for_download), gr.update(value=None, visible=False)
 
-        if mime_type not in allowed_file_types:
-            return gr.update(value=f"File type '{mime_type or 'unknown'}' for '{original_name}' is disallowed.", visible=True), file_uploads_log_state
 
-        sanitized_name = re.sub(r"[^\w\-.]", "_", original_name)
-        base_name, current_ext = os.path.splitext(sanitized_name)
-        
-        common_mime_to_ext = {
-            "application/pdf": ".pdf",
-            "application/vnd.openxmlformats-officedocument.wordprocessingml.document": ".docx",
-            "text/plain": ".txt", "image/jpeg": ".jpg", "image/png": ".png"
-        }
-        expected_ext = common_mime_to_ext.get(mime_type)
-
-        if expected_ext and current_ext.lower() != expected_ext.lower():
-            sanitized_name = base_name + expected_ext
-        
-        destination_path = os.path.join(self.file_upload_folder, sanitized_name)
-        
-        try:
-            shutil.copy(file.name, destination_path) 
-            print(f"DEBUG Gradio: File '{original_name}' copied to '{destination_path}'")
-            updated_log = file_uploads_log_state + [destination_path]
-            return gr.update(value=f"Uploaded: {original_name}", visible=True), updated_log
-        except Exception as e:
-            print(f"DEBUG Gradio: Error copying uploaded file: {e}")
-            return gr.update(value=f"Error uploading {original_name}: {e}", visible=True), file_uploads_log_state
-
-    def log_user_message(self, text_input_value: str, current_file_uploads: list):
+    def log_user_message(self, text_input_value: str): # Removed current_file_uploads as upload is disabled
         full_prompt = text_input_value
-        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}]"
+        # if current_file_uploads: # This part is now disabled
+        #     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]}...") 
         return full_prompt, "" 
 
@@ -307,15 +269,14 @@ class GradioUI:
 
     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([]) 
+            # file_uploads_log_state = gr.State([]) # No longer needed as upload is disabled
             prepared_prompt_for_agent = gr.State("")
 
             gr.Markdown("## Smol Talk with your Agent") 
 
             with gr.Row(equal_height=False): 
-                with gr.Column(scale=3):
+                with gr.Column(scale=3): # Main chat column
                     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, 
@@ -325,28 +286,30 @@ class GradioUI:
                     )
                     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 
                     )
                 
-                with gr.Column(scale=1):
-                    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.Column(scale=1): # Sidebar column
+                    # --- File Upload Section Removed/Commented Out ---
+                    # if self.file_upload_folder is not None: 
+                    #     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): 
+                    with gr.Accordion("Generated File", open=True): # Keep this section
                         download_action_button = gr.Button("Download Generated File", visible=False) 
                         file_download_display_component = gr.File(label="Downloadable Document", visible=False, interactive=False) 
 
+            # Event Handling Chain for Text Submission
             text_message_input.submit(
                 self.log_user_message, 
-                [text_message_input, file_uploads_log_state],
+                [text_message_input], # Removed file_uploads_log_state from inputs
                 [prepared_prompt_for_agent, text_message_input] 
             ).then(
                 self.interact_with_agent,