Update Gradio_UI.py

Gradio_UI.py

@@ -153,29 +153,28 @@ def stream_to_gradio(
         ):
             yield message
 
-    final_answer = step_log  # Last log is the run's final_answer
-    final_answer = handle_agent_output_types(final_answer)
+    final_step = step_log  # this is the final ActionStep
+    final_output = getattr(final_step, "tool_output", None)
+    final_output = handle_agent_output_types(final_output)
     
-    # Strip any wrapper or formatting for exact match evaluation
-    if isinstance(final_answer, AgentText):
-        content = final_answer.to_string().strip()
+    if isinstance(final_output, AgentText):
         yield gr.ChatMessage(
             role="assistant",
-            content=content,
+            content=final_output.to_string().strip(),
         )
-    elif isinstance(final_answer, AgentImage):
+    elif isinstance(final_output, AgentImage):
         yield gr.ChatMessage(
             role="assistant",
-            content={"path": final_answer.to_string(), "mime_type": "image/png"},
+            content={"path": final_output.to_string(), "mime_type": "image/png"},
         )
-    elif isinstance(final_answer, AgentAudio):
+    elif isinstance(final_output, AgentAudio):
         yield gr.ChatMessage(
             role="assistant",
-            content={"path": final_answer.to_string(), "mime_type": "audio/wav"},
+            content={"path": final_output.to_string(), "mime_type": "audio/wav"},
         )
     else:
-        # Fall back to raw string with no prefix
-        yield gr.ChatMessage(role="assistant", content=str(final_answer).strip())
+        yield gr.ChatMessage(role="assistant", content=str(final_output).strip())
+