Add code

app.py

@@ -109,29 +109,32 @@ def response(audio: tuple[int, np.ndarray], lepton_conversation: list[dict],
         else:
             yield AdditionalOutputs(lepton_conversation, gradio_conversation)
 
+
 def set_api_key(api_key):
     if not api_key:
         raise gr.Error("Please enter a valid API key.")
     client = create_client(api_key)
-    return client
+    gr.Info("Set API Key Successfully")
+    return client, gr.skip()
 
 
 with gr.Blocks() as demo:
     with gr.Group():
         with gr.Row():
             chatbot = gr.Chatbot(label="Conversation", type="messages")
-        with gr.Row():
-            with gr.Column():
+        with gr.Row(equal_height=True):
+            with gr.Column(scale=1):
                 format_dropdown = gr.Dropdown(choices=["mp3", "opus"], value="mp3", label="Output Audio Format")
                 api_key_input = gr.Textbox(type="password", label="Enter your Lepton API Key")
-            with gr.Column():
+                set_key_button = gr.Button("Set API Key", variant="primary")
+            with gr.Column(scale=3):
                 audio = WebRTC(modality="audio", mode="send-receive",
-                               label="Audio Stream")
-    
+                                label="Audio Stream")
+
     client_state = gr.State(None)
     lepton_conversation = gr.State([])
 
-    api_key_input.submit(set_api_key, inputs=[api_key_input], outputs=[client_state])
+    set_key_button.click(set_api_key, inputs=[api_key_input], outputs=[client_state, set_key_button])
 
     audio.stream(
         ReplyOnPause(response),