Update Gradio_UI.py

Gradio_UI.py

@@ -111,9 +111,11 @@ def pull_messages_from_step(
         # Calculate duration and token information
         step_footnote = f"{step_number}"
         if hasattr(step_log, "input_token_count") and hasattr(step_log, "output_token_count"):
-            token_str = (
-                f" | Input-tokens:{step_log.input_token_count:,} | Output-tokens:{step_log.output_token_count:,}"
-            )
+            token_str = ""
+                if hasattr(step_log, "input_token_count") and step_log.input_token_count is not None:
+                    token_str += f" | Input-tokens:{step_log.input_token_count:,}"
+                if hasattr(step_log, "output_token_count") and step_log.output_token_count is not None:
+                    token_str += f" | Output-tokens:{step_log.output_token_count:,}"
             step_footnote += token_str
         if hasattr(step_log, "duration"):
             step_duration = f" | Duration: {round(float(step_log.duration), 2)}" if step_log.duration else None