Update app.py

app.py

@@ -61,6 +61,7 @@ def gradio_ui():
             with gr.Column(scale=2):
                 compute_time = gr.Textbox(label="Compute Time", interactive=False)
                 chain_of_thought_display = gr.Textbox(label="Chain of Thought", interactive=False, lines=10)
+                token_usage_display = gr.Textbox(label="Token Usage", interactive=False)
 
         user_input = gr.Textbox(label="Type your message", placeholder="Ask me anything...", lines=2)
 
@@ -70,20 +71,22 @@ def gradio_ui():
 
         def handle_chat(chat_history, user_input):
             if not user_input.strip():
-                return chat_history, "", "Please enter a valid message."
+                return chat_history, "", "", "Please enter a valid message."
             ai_response, chain_of_thought, compute_info = chat_with_cerebras(user_input)
+            # Placeholder for token usage (to be replaced with actual logic if token count available)
+            token_usage = "Tokens used: Not available in this version"
             chat_history.append((user_input, ai_response))
-            return chat_history, chain_of_thought, compute_info
+            return chat_history, chain_of_thought, compute_info, token_usage
 
         def clear_chat():
-            return [], "", ""
+            return [], "", "", ""
 
-        send_button.click(handle_chat, inputs=[chat_history, user_input], outputs=[chat_history, chain_of_thought_display, compute_time])
-        clear_button.click(clear_chat, outputs=[chat_history, chain_of_thought_display, compute_time])
+        send_button.click(handle_chat, inputs=[chat_history, user_input], outputs=[chat_history, chain_of_thought_display, compute_time, token_usage_display])
+        clear_button.click(clear_chat, outputs=[chat_history, chain_of_thought_display, compute_time, token_usage_display])
 
-        user_input.submit(handle_chat, inputs=[chat_history, user_input], outputs=[chat_history, chain_of_thought_display, compute_time])
+        user_input.submit(handle_chat, inputs=[chat_history, user_input], outputs=[chat_history, chain_of_thought_display, compute_time, token_usage_display])
 
-        gr.Markdown("""---\n### 🌟 Features:\n- **Advanced Reasoning**: Chain-of-thought explanations for complex queries.\n- **Real-Time Performance Metrics**: Measure response compute time instantly.\n- **Insightful Chain of Thought**: See the reasoning process behind AI decisions.\n- **User-Friendly Design**: Intuitive chatbot interface with powerful features.\n- **Powered by IntellijMind**: Setting new standards for AI interaction.\n""")
+        gr.Markdown("""---\n### 🌟 Features:\n- **Advanced Reasoning**: Chain-of-thought explanations for complex queries.\n- **Real-Time Performance Metrics**: Measure response compute time instantly.\n- **Token Usage Tracking**: (Future-ready) Display token usage per response.\n- **User-Friendly Design**: Intuitive chatbot interface with powerful features.\n- **Insightful Chain of Thought**: See the reasoning process behind AI decisions.\n- **Powered by IntellijMind**: Setting new standards for AI interaction.\n- **Submit on Enter**: Seamless interaction with keyboard support.\n""")
 
     return demo