Updated Gradio_UI.py

Streaming the image in Gradio UI

Gradio_UI.py

@@ -20,7 +20,7 @@ import shutil
 from typing import Optional
 
 from smolagents.agent_types import AgentAudio, AgentImage, AgentText, handle_agent_output_types
-from smolagents.agents import ActionStep, MultiStepAgent
+from smolagents.agents import ActionStep, FinalAnswerStep, MultiStepAgent
 from smolagents.memory import MemoryStep
 from smolagents.utils import _is_package_available
 
@@ -138,8 +138,13 @@ def stream_to_gradio(
 
     total_input_tokens = 0
     total_output_tokens = 0
+    final_answer_step = None
 
     for step_log in agent.run(task, stream=True, reset=reset_agent_memory, additional_args=additional_args):
+        if isinstance(step_log, FinalAnswerStep):
+            final_answer_step = step_log
+            continue # Don't display the final answer step itself, process it at the end
+
         # Track tokens if model provides them
         if hasattr(agent.model, "last_input_token_count"):
             total_input_tokens += agent.model.last_input_token_count
@@ -152,27 +157,33 @@ def stream_to_gradio(
             step_log,
         ):
             yield message
+    
+    # FIX: Process the final answer correctly after the loop
+    if final_answer_step:
+        # Extract the actual value from the FinalAnswerStep object
+        final_answer_value = getattr(final_answer_step, 'final_answer', final_answer_step)
+        
+        # Convert the value (e.g., a file path string) into the correct Agent type (e.g., AgentImage)
+        processed_answer = handle_agent_output_types(final_answer_value)
 
-    final_answer = step_log  # Last log is the run's final_answer
-    final_answer = handle_agent_output_types(final_answer)
-
-    if isinstance(final_answer, AgentText):
-        yield gr.ChatMessage(
-            role="assistant",
-            content=f"**Final answer:**\n{final_answer.to_string()}\n",
-        )
-    elif isinstance(final_answer, AgentImage):
-        yield gr.ChatMessage(
-            role="assistant",
-            content={"path": final_answer.to_string(), "mime_type": "image/png"},
-        )
-    elif isinstance(final_answer, AgentAudio):
-        yield gr.ChatMessage(
-            role="assistant",
-            content={"path": final_answer.to_string(), "mime_type": "audio/wav"},
-        )
-    else:
-        yield gr.ChatMessage(role="assistant", content=f"**Final answer:** {str(final_answer)}")
+        if isinstance(processed_answer, AgentText):
+            yield gr.ChatMessage(
+                role="assistant",
+                content=f"**Final answer:**\n{processed_answer.to_string()}\n",
+            )
+        elif isinstance(processed_answer, AgentImage):
+            yield gr.ChatMessage(
+                role="assistant",
+                content=processed_answer.to_string(), # Gradio automatically handles file paths
+            )
+        elif isinstance(processed_answer, AgentAudio):
+            yield gr.ChatMessage(
+                role="assistant",
+                content=processed_answer.to_string(), # Gradio automatically handles file paths
+            )
+        else:
+            # Fallback for any other type
+            yield gr.ChatMessage(role="assistant", content=f"**Final answer:** {str(processed_answer)}")
 
 
 class GradioUI:
@@ -266,12 +277,13 @@ class GradioUI:
             file_uploads_log = gr.State([])
             chatbot = gr.Chatbot(
                 label="Agent",
-                type="messages",
+                render_markdown=True,
                 avatar_images=(
                     None,
-                    "https://huggingface.co/datasets/agents-course/course-images/resolve/main/en/communication/Alfred.png",
+                    "[https://huggingface.co/datasets/agents-course/course-images/resolve/main/en/communication/Alfred.png](https://huggingface.co/datasets/agents-course/course-images/resolve/main/en/communication/Alfred.png)",
                 ),
-                resizeable=True,
+                bubble_full_width=False,
+                show_copy_button=True,
                 scale=1,
             )
             # If an upload folder is provided, enable the upload feature
@@ -283,13 +295,16 @@ class GradioUI:
                     [upload_file, file_uploads_log],
                     [upload_status, file_uploads_log],
                 )
-            text_input = gr.Textbox(lines=1, label="Chat Message")
+            text_input = gr.Textbox(lines=1, label="Chat Message", scale=4)
             text_input.submit(
                 self.log_user_message,
                 [text_input, file_uploads_log],
                 [stored_messages, text_input],
             ).then(self.interact_with_agent, [stored_messages, chatbot], [chatbot])
-        
+
+        # Add allowed_paths to the launch command
+        # This gives Gradio permission to access and display files from the current directory
+        # and the specified file upload folder.
         allowed_paths = kwargs.pop("allowed_paths", [])
         if self.file_upload_folder and self.file_upload_folder not in allowed_paths:
             allowed_paths.append(self.file_upload_folder)