Update app.py

app.py

@@ -8,21 +8,40 @@ TEMP_DIR = "generated_output"
 os.makedirs(TEMP_DIR, exist_ok=True)
 
 def save_files(html: str, log: str):
+    TEMP_DIR = "generated_output"
+    os.makedirs(TEMP_DIR, exist_ok=True)
+
+    # Extract CSS
+    import re
+    css_match = re.search(r"<style[^>]*>(.*?)</style>", html, flags=re.DOTALL)
+    css = css_match.group(1).strip() if css_match else ""
+
+    # Remove <style> tag and replace with <link>
+    html_clean = re.sub(r"<style[^>]*>.*?</style>", '<link rel="stylesheet" href="styles.css">', html, flags=re.DOTALL)
+
+    # Paths
     html_path = os.path.join(TEMP_DIR, "ui.html")
+    css_path = os.path.join(TEMP_DIR, "styles.css")
     log_path = os.path.join(TEMP_DIR, "agent_log.txt")
     zip_path = os.path.join(TEMP_DIR, "output_bundle.zip")
 
+    # Save files
     with open(html_path, "w", encoding="utf-8") as f:
-        f.write(html)
+        f.write(html_clean)
+
+    with open(css_path, "w", encoding="utf-8") as f:
+        f.write(css)
 
     with open(log_path, "w", encoding="utf-8") as f:
         f.write(log)
 
     with zipfile.ZipFile(zip_path, "w") as zipf:
         zipf.write(html_path, arcname="ui.html")
+        zipf.write(css_path, arcname="styles.css")
         zipf.write(log_path, arcname="agent_log.txt")
 
-    return html_path, log_path, zip_path
+    return html_path, css_path, log_path, zip_path
+
 
 def process(user_input):
     messages = [HumanMessage(content=user_input)]
@@ -31,13 +50,11 @@ def process(user_input):
     html = final_state["html_output"]
     log = "\n\n".join([f"{msg.type.upper()}: {msg.content}" for msg in final_state["messages"]])
 
-    html_file, log_file, zip_file = save_files(html, log)
+    html_file, css_file, log_file, zip_file = save_files(html, log)
+    preview = f"<iframe src='file/{html_file}' width='100%' height='500px' style='border:1px solid #ccc;'></iframe>"
 
-    # Escape quotes for safe iframe embedding
-    escaped_html = html.replace('"', "&quot;")
-    preview = f"<iframe srcdoc='{escaped_html}' width='100%' height='600px' style='border:1px solid #ccc;'></iframe>"
+    return preview, html, log, html_file, css_file, log_file, zip_file
 
-    return preview, html, log, html_file, log_file, zip_file
 
 with gr.Blocks() as demo:
     gr.Markdown("# 🤖 Multi-Agent UI Generator")
@@ -58,11 +75,24 @@ with gr.Blocks() as demo:
     log_download = gr.File(label="Download Agent Log")
     zip_download = gr.File(label="Download All (.zip)")
 
+    html_download = gr.File(label="Download HTML")
+    css_download = gr.File(label="Download CSS")
+    log_download = gr.File(label="Download Agent Log")
+    zip_download = gr.File(label="Download All (.zip)")
+
     run_btn.click(
         fn=process,
         inputs=[user_prompt],
-        outputs=[preview_html, html_output, log_output, html_download, log_download, zip_download]
-    )
+        outputs=[
+        preview_html,
+        html_output,
+        log_output,
+        html_download,
+        css_download,
+        log_download,
+        zip_download
+    ]
+)
 
 if __name__ == "__main__":
     demo.launch()