Update Gradio_UI.py

Gradio_UI.py

@@ -109,12 +109,23 @@ def pull_messages_from_step(
             yield gr.ChatMessage(role="assistant", content=str(step_log.error), metadata={"title": "💥 Error"})
 
         # 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:,}"
+        #    )
+        #    step_footnote += token_str
+
         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:,}"
-            )
-            step_footnote += token_str
+        input_tokens = getattr(step_log, "input_token_count", None)
+        output_tokens = getattr(step_log, "output_token_count", None)
+        
+        if input_tokens is not None:
+            step_footnote += f" | Input-tokens:{input_tokens:,}"
+        if output_tokens is not None:
+            step_footnote += f" | Output-tokens:{output_tokens:,}"
+
+    
         if hasattr(step_log, "duration"):
             step_duration = f" | Duration: {round(float(step_log.duration), 2)}" if step_log.duration else None
             step_footnote += step_duration