UI update

app.py

@@ -259,8 +259,10 @@ Answer:"""
                 formatted_response = f"<div class='response-title'>Response for {state}</div><div class='response-content'>{answer}</div>"
                 return f"<div class='output-box'>{formatted_response}</div>"
 
-        def toggle_theme(mode: str):
-            return gr.update(visible=mode == "dark"), gr.update(visible=mode == "light")
+        def toggle_theme():
+            is_dark = "dark-mode" in document.body.className
+            document.body.className = "dark-mode" if not is_dark else ""
+            return gr.update(visible=not is_dark), gr.update(visible=is_dark)
 
         try:
             available_states = self.get_states()
@@ -560,9 +562,7 @@ Answer:"""
             with gr.Row():
                 with gr.Column(scale=1):
                     gr.Markdown("<div class='header'><h1>Landlord-Tenant Rights Assistant</h1></div>")
-                    theme_light = gr.HTML(visible=True, value="<div class='theme-toggle'><button onclick='document.body.classList.remove(\"dark-mode\");'>Light Mode</button></div>")
-                    theme_dark = gr.HTML(visible=False, value="<div class='theme-toggle'><button onclick='document.body.classList.add(\"dark-mode\");'>Dark Mode</button></div>")
-                    gr.Button("Toggle Dark Mode", lambda: [gr.update(visible=True), gr.update(visible=False)], outputs=[theme_dark, theme_light])
+                    theme_toggle_btn = gr.Button("Toggle Dark Mode", elem_classes="theme-toggle")
 
             with gr.Column(elem_classes="container"):
                 with gr.Column():
@@ -587,6 +587,11 @@ Answer:"""
                 with gr.Column():
                     gr.Markdown("<div class='footer'><p><strong>Disclaimer:</strong> This tool is for informational purposes only and does not constitute legal advice. Consult a licensed attorney for specific guidance. Developed by Nischal Subedi.</p></div>")
 
+            theme_toggle_btn.click(
+                fn=toggle_theme,
+                outputs=[]
+            )
+
             submit_button.click(
                 fn=query_interface_wrapper,
                 inputs=[api_key_input, query_input, state_input],