Update app.py

app.py

@@ -5,7 +5,7 @@ import time # For potential brief pauses if needed
 
 # --- Hugging Face Token (Optional but Recommended) ---
 # from huggingface_hub import login
-# login("YOUR_HUGGINGFACE_TOKEN")
+# login("YOUR_HUGGINGFACE_TOKEN") # Replace with your token if needed
 
 # --- Inference Client ---
 try:
@@ -15,9 +15,9 @@ try:
     client.timeout = 120 # Increase timeout for potentially long generations
 except Exception as e:
     print(f"Error initializing InferenceClient: {e}")
-    client = None
+    client = None # Set client to None if initialization fails
 
-# --- Parsing Function (from previous good version) ---
+# --- Parsing Function ---
 def parse_files(raw_response):
     """
     Parses filenames and code blocks from the raw AI output.
@@ -60,16 +60,20 @@ def stream_and_parse_code(prompt, backend, system_message, max_tokens, temperatu
     """
     Streams raw output to one component and generates final tabs for another.
     This function acts as the main callback for the button click.
+    Yields dictionary updates for Gradio components.
     """
+    # Check if client initialized correctly
     if not client:
-        error_msg = "Error: Inference Client not available."
+        error_msg = "Error: Inference Client not available. Check API token or model name."
+        # Yield updates to both components indicating the error
         yield {
-            live_output: error_msg,
+            live_output: gr.update(value=error_msg),
             final_tabs: gr.Tabs(tabs=[gr.TabItem(label="Error", children=[gr.Textbox(value=error_msg)])])
         }
         return # Stop execution
 
     # --- Prepare for Streaming ---
+    # Construct the system prompt dynamically
     full_sys_msg = f"""
 You are a code generation AI. Given a prompt, generate the necessary files for a website using the {backend} backend.
 Always include an index.html file.
@@ -83,9 +87,14 @@ index.html
 style.css
 body {{}}
 
+script.js
+console.log("Hello");
+
 Ensure the code is complete. NO commentary, NO explanations, NO markdown formatting like backticks (```).
 Start generating the files now.
-""".strip() + ("\n" + system_message if system_message else "")
+""".strip()
+    if system_message: # Append user's system message if provided
+        full_sys_msg += "\n\n" + system_message
 
     messages = [
         {"role": "system", "content": full_sys_msg},
@@ -96,53 +105,55 @@ Start generating the files now.
     error_occurred = False
     error_message = ""
 
-    # Initial state update
+    # Initial state update: Clear previous output and show generating status
     yield {
-        live_output: "Generating stream...",
-        # Set initial tabs state to indicate loading
-        final_tabs: gr.Tabs(tabs=[gr.TabItem(label="Generating...")])
+        live_output: gr.update(value="Generating stream..."),
+        final_tabs: gr.Tabs(tabs=[gr.TabItem(label="Generating...")]) # Indicate loading in tabs
     }
 
     # --- Streaming Loop ---
     try:
+        # Start the streaming call
         stream = client.chat_completion(
             messages,
-            max_tokens=int(max_tokens),
+            max_tokens=int(max_tokens), # Ensure max_tokens is an integer
             stream=True,
             temperature=temperature,
             top_p=top_p
         )
+        # Process each chunk received from the stream
         for chunk in stream:
             token = chunk.choices[0].delta.content
             if token:
                 full_raw_response += token
-                # Yield updates for the live raw output component
-                # Keep tabs in a 'generating' state during the stream
+                # Yield updates for the live raw output component only
+                # Keep tabs in a 'streaming' state during the stream
                 yield {
-                    live_output: full_raw_response,
-                    final_tabs: gr.Tabs(tabs=[gr.TabItem(label="Streaming...")]) # Keep showing streaming
+                    live_output: gr.update(value=full_raw_response),
+                    # No update needed for final_tabs here, or keep showing streaming state
+                    # final_tabs: gr.Tabs(tabs=[gr.TabItem(label="Streaming...")]) # Optional: update tabs state
                 }
-                # time.sleep(0.01) # Optional small delay if updates are too fast
+                # time.sleep(0.01) # Optional: small delay if updates are too fast and causing UI lag
 
     except Exception as e:
+        # Handle errors during the API call or streaming process
         print(f"Error during AI streaming: {e}")
-        error_message = f"Error during AI generation: {e}\n\nPartial Response:\n{full_raw_response}"
+        error_message = f"Error during AI generation: {e}\n\nPartial Response (if any):\n{full_raw_response}"
         error_occurred = True
-        # Update live output with error, keep tabs showing error state
+        # Update live output with error, prepare error tab
         yield {
-            live_output: error_message,
-            final_tabs: gr.Tabs(tabs=[gr.TabItem(label="Error")])
+            live_output: gr.update(value=error_message),
+            final_tabs: gr.Tabs(tabs=[gr.TabItem(label="Error")]) # Indicate error state in tabs
         }
 
-
     # --- Post-Streaming: Parsing and Final Tab Generation ---
     if error_occurred:
-        # If an error happened during stream, create an error tab
+        # If an error happened during stream, create a final error tab
          final_tabs_update = gr.Tabs(tabs=[
-            gr.TabItem(label="Error", children=[gr.Textbox(value=error_message, label="Generation Error")])
+            gr.TabItem(label="Error", children=[gr.Textbox(value=error_message, label="Generation Error", lines=10)])
         ])
     else:
-        # Parse the complete raw response
+        # If streaming succeeded, parse the complete raw response
         print("\n--- Final Raw AI Response ---")
         print(full_raw_response)
         print("--------------------------\n")
@@ -150,12 +161,12 @@ Start generating the files now.
 
         if not files:
             # Handle case where parsing failed or AI gave empty/invalid response
-            no_files_msg = "AI finished, but did not return recognizable file content. See raw output above."
+            no_files_msg = "AI finished, but did not return recognizable file content or the response was empty. See raw output above."
             final_tabs_update = gr.Tabs(tabs=[
                 gr.TabItem(label="Output", children=[gr.Textbox(value=no_files_msg, label="Result")])
             ])
-            # Update live output as well if needed
-            yield { live_output: full_raw_response + "\n\n" + no_files_msg, final_tabs: final_tabs_update }
+            # Update live output as well to make the message clear
+            yield { live_output: gr.update(value=full_raw_response + "\n\n" + no_files_msg), final_tabs: final_tabs_update }
             return # Exit if no files
 
         # --- Create Tabs (if files were parsed successfully) ---
@@ -163,11 +174,13 @@ Start generating the files now.
         for name, content in files:
             name = name.strip()
             content = content.strip()
+            # Skip if filename or content is empty after stripping
             if not name or not content:
                 print(f"Skipping file with empty name or content: Name='{name}'")
                 continue
 
-            lang = "text" # Default
+            # Determine language for syntax highlighting
+            lang = "plaintext" # Default
             if name.endswith((".html", ".htm")): lang = "html"
             elif name.endswith(".css"): lang = "css"
             elif name.endswith(".js"): lang = "javascript"
@@ -175,50 +188,95 @@ Start generating the files now.
             elif name.endswith(".json"): lang = "json"
             elif name.endswith(".md"): lang = "markdown"
             elif name.endswith((".sh", ".bash")): lang = "bash"
+            elif name.endswith((".xml", ".xaml", ".svg")): lang = "xml"
+            elif name.endswith(".yaml") or name.endswith(".yml"): lang = "yaml"
+
+            # Ensure elem_id is unique and valid (replace problematic characters)
+            elem_id = f"tab_{re.sub(r'[^a-zA-Z0-9_-]', '_', name)}"
 
-            tab_item = gr.TabItem(label=name, elem_id=f"tab_{name.replace('.', '_').replace('/', '_')}", children=[
-                gr.Code(value=content, language=lang, label=name)
+            tab_item = gr.TabItem(label=name, elem_id=elem_id, children=[
+                gr.Code(value=content, language=lang, label=name, interactive=False) # Show code in Code block
             ])
             tabs_content.append(tab_item)
 
-        if not tabs_content: # Handle case where parsing found files, but they were filtered out
-             final_tabs_update = gr.Tabs(tabs=[gr.TabItem(label="Output", children=[gr.Textbox(value="No valid files generated.", label="Result")])])
+        # Handle case where parsing found files, but they were all filtered out (empty name/content)
+        if not tabs_content:
+             final_tabs_update = gr.Tabs(tabs=[gr.TabItem(label="Output", children=[gr.Textbox(value="No valid files generated after filtering.", label="Result")])])
         else:
-            final_tabs_update = gr.Tabs(tabs=tabs_content) # Create the final Tabs component
-
+            final_tabs_update = gr.Tabs(tabs=tabs_content) # Create the final Tabs component with content
 
     # --- Final Update ---
     # Yield the final state for both components
+    # Use gr.update for live_output if you only want to set its value without recreating it
+    # Directly return the new final_tabs component
     yield {
-        live_output: full_raw_response if not error_occurred else error_message, # Show final raw response or error
-        final_tabs: final_tabs_update # Show the generated tabs or error tab
+        live_output: gr.update(value=full_raw_response if not error_occurred else error_message), # Show final raw response or error
+        final_tabs: final_tabs_update # Update the final_tabs component completely
     }
 
 
 # --- Gradio UI Definition ---
-with gr.Blocks(css=".gradio-container { max-width: 95% !important; }") as demo: # Wider interface
+with gr.Blocks(css=".gradio-container { max-width: 95% !important; }") as demo: # Use more screen width
     gr.Markdown("## WebGen AI — One Prompt → Full Website Generator")
     gr.Markdown("Generates website code based on your description. Raw output streams live, final files appear in tabs below.")
 
     with gr.Row():
+        # Column for inputs and controls
         with gr.Column(scale=2):
-            prompt = gr.Textbox(label="Describe your website", placeholder="E.g., a simple portfolio site with a dark mode toggle", lines=3)
-            backend = gr.Dropdown(["Static", "Flask", "Node.js"], value="Static", label="Backend Technology")
+            prompt = gr.Textbox(
+                label="Describe your website",
+                placeholder="E.g., a simple landing page for a coffee shop with sections for menu, about, and contact.",
+                lines=3 # Allow more lines for the prompt
+            )
+            backend = gr.Dropdown(
+                ["Static", "Flask", "Node.js"],
+                value="Static",
+                label="Backend Technology"
+            )
             with gr.Accordion("Advanced Options", open=False):
-                system_message = gr.Textbox(label="Extra instructions for the AI (System Message)", placeholder="Optional: e.g., 'Use Bootstrap 5', 'Prefer functional components in React'", value="")
-                max_tokens = gr.Slider(minimum=256, maximum=4096, value=1536, step=64, label="Max Tokens (Length)")
-                temperature = gr.Slider(minimum=0.1, maximum=1.5, value=0.7, step=0.1, label="Temperature (Creativity)")
-                top_p = gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, label="Top-p (Sampling)")
-            generate_button = gr.Button("✨ Generate Code ✨", variant="primary")
-
+                system_message = gr.Textbox(
+                    label="Extra instructions for the AI (System Message)",
+                    placeholder="Optional: e.g., 'Use Tailwind CSS for styling', 'Make it responsive'",
+                    value="",
+                    lines=2
+                 )
+                max_tokens = gr.Slider(
+                    minimum=256,
+                    maximum=4096, # Increased max tokens for complex sites
+                    value=2048, # Increased default
+                    step=64,
+                    label="Max Tokens (Output Length)"
+                )
+                temperature = gr.Slider(
+                    minimum=0.1,
+                    maximum=1.5, # Allow slightly higher temperature
+                    value=0.7,
+                    step=0.1,
+                    label="Temperature (Creativity)"
+                )
+                top_p = gr.Slider(
+                    minimum=0.1,
+                    maximum=1.0,
+                    value=0.95,
+                    step=0.05,
+                    label="Top-p (Sampling Focus)"
+                )
+            generate_button = gr.Button("✨ Generate Code ✨", variant="primary") # Make button primary
+
+        # Column for live output
         with gr.Column(scale=3):
             gr.Markdown("#### Live Raw Output Stream")
-            # Component to show the live, unparsed stream
-            live_output = gr.Code(label="Raw AI Stream", language="text", lines=15, interactive=False)
-
-    gr.Markdown("---")
+            # Component to show the live, unparsed stream - CORRECTED LANGUAGE
+            live_output = gr.Code(
+                label="Raw AI Stream",
+                language="plaintext", # Use "plaintext" for generic text
+                lines=20, # Increased lines for visibility
+                interactive=False # Output only
+            )
+
+    gr.Markdown("---") # Separator
     gr.Markdown("#### Final Generated Files (Tabs)")
-    # Placeholder for the final structured tabs
+    # Placeholder for the final structured tabs - will be replaced by the output yield
     final_tabs = gr.Tabs(elem_id="output_tabs")
 
 
@@ -226,10 +284,11 @@ with gr.Blocks(css=".gradio-container { max-width: 95% !important; }") as demo:
     generate_button.click(
         stream_and_parse_code, # Call the main function that handles streaming and parsing
         inputs=[prompt, backend, system_message, max_tokens, temperature, top_p],
-        # Outputs dictionary maps function yields to components
+        # Outputs dictionary maps function yields to components by variable name
         outputs=[live_output, final_tabs],
-        show_progress="hidden" # Hide default progress bar as we show live stream
+        show_progress="hidden" # Hide default Gradio progress bar as we show live stream
     )
 
 if __name__ == "__main__":
+    # Launch the Gradio app with debug=True for development
     demo.launch(debug=True)