Update Gradio_UI.py

Gradio_UI.py

@@ -153,25 +153,20 @@ def stream_to_gradio(
         ):
             yield message
 
-    final_step = step_log  # Final ActionStep
-
-    # Debugging tool calls to inspect final_answer structure
-    print("DEBUG TOOL CALLS:", final_step.tool_calls)
-
-    # Attempt to get the direct tool_output
-    final_output = getattr(final_step, "tool_output", None)
-
-    # Fallback if tool_output is missing
-    if not final_output and getattr(final_step, "tool_calls", None):
-        for call in final_step.tool_calls:
-            if call.name == "final_answer":
-                print("DEBUG FINAL ANSWER TOOL CALL:", call)
-                # Try the `output` field if present (some agents set this explicitly)
-                final_output = getattr(call, "output", None)
-                if final_output is None:
-                    # Fallback to arguments dict
+    from smolagents.agents import FinalAnswerStep
+
+    # Safely extract the final output from the last step
+    if isinstance(step_log, FinalAnswerStep):
+        final_output = step_log.final_answer
+    else:
+        # If not a FinalAnswerStep, fallback to tool_output or tool_calls
+        final_output = getattr(step_log, "tool_output", None)
+        if not final_output and hasattr(step_log, "tool_calls"):
+            for call in step_log.tool_calls:
+                if call.name == "final_answer":
                     final_output = call.arguments.get("answer")
-                break
+                    break
+
 
     final_output = handle_agent_output_types(final_output)