Update app.py

app.py

@@ -1,51 +1,39 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
 import re
-import os # Good practice to import os if needed, though not strictly used here yet
+import time # For potential brief pauses if needed
 
 # --- Hugging Face Token (Optional but Recommended) ---
-# It's better to use a token, especially for private models or higher rate limits
 # from huggingface_hub import login
-# login("YOUR_HUGGINGFACE_TOKEN") # Replace with your actual token or set HF_TOKEN env var
+# login("YOUR_HUGGINGFACE_TOKEN")
 
 # --- Inference Client ---
-# Consider adding error handling for client initialization if needed
 try:
+    # You might need to specify the model URL directly if the alias isn't working
+    # client = InferenceClient(model="https://api-inference.huggingface.co/models/HuggingFaceH4/zephyr-7b-beta")
     client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+    client.timeout = 120 # Increase timeout for potentially long generations
 except Exception as e:
     print(f"Error initializing InferenceClient: {e}")
-    # Optionally, raise the exception or handle it gracefully in the UI
-    # For now, we'll let it proceed and potentially fail later if client is None
     client = None
 
-# --- Parsing Function ---
+# --- Parsing Function (from previous good version) ---
 def parse_files(raw_response):
     """
     Parses filenames and code blocks from the raw AI output.
-    Assumes format:
-    filename1.ext
-    ```lang # Optional code block marker
-    code for file1
-    ``` # Optional code block marker
-    filename2.ext
-    code for file2
-    ...
     """
     if not raw_response:
         return []
 
-    # Improved pattern to handle optional code blocks and leading/trailing whitespace
-    # It looks for a filename line followed by content until the next filename line or end of string.
+    # Pattern: Look for a filename line followed by content until the next filename line or end of string.
     pattern = re.compile(
         r"^\s*([\w\-.\/\\]+\.\w+)\s*\n"  # Filename line (must have an extension)
         r"(.*?)"                        # Capture content (non-greedy)
         r"(?=\n\s*[\w\-.\/\\]+\.\w+\s*\n|\Z)", # Lookahead for next filename or end of string
         re.DOTALL | re.MULTILINE
     )
-
     files = pattern.findall(raw_response)
 
-    # Clean up content: remove potential code block markers and extra whitespace
     cleaned_files = []
     for name, content in files:
         # Remove common code block markers (``` optionally followed by lang)
@@ -53,34 +41,35 @@ def parse_files(raw_response):
         content_cleaned = re.sub(r"\n?```\s*$", "", content_cleaned, flags=re.MULTILINE)
         cleaned_files.append((name.strip(), content_cleaned.strip()))
 
-    # Handle case where the AI might just output code without filenames
+    # Fallback if no files parsed but content exists
     if not cleaned_files and raw_response.strip():
-         # Basic check if it looks like code (e.g., contains common HTML/CSS/JS chars)
         if any(c in raw_response for c in ['<','>','{','}',';','(',')']):
-             # Default to index.html if no files parsed but content exists
             print("Warning: No filenames found, defaulting to index.html")
-            lang = "html" # Guess language, default to html
-            if "{" in raw_response and "}" in raw_response and ":" in raw_response:
-                lang = "css"
-            elif "function" in raw_response or "const" in raw_response or "let" in raw_response:
-                lang = "javascript"
-            # Determine a default filename based on guessed language
+            lang = "html"
+            if "{" in raw_response and "}" in raw_response and ":" in raw_response: lang = "css"
+            elif "function" in raw_response or "const" in raw_response or "let" in raw_response: lang = "javascript"
             default_filename = "index.html"
             if lang == "css": default_filename = "style.css"
             elif lang == "javascript": default_filename = "script.js"
-
             cleaned_files.append((default_filename, raw_response.strip()))
 
-
     return cleaned_files
 
-# --- Code Generation Function ---
-def generate_code(prompt, backend, system_message, max_tokens, temperature, top_p):
-    """Generates code using the InferenceClient."""
+# --- Streaming and Parsing Orchestrator ---
+def stream_and_parse_code(prompt, backend, system_message, max_tokens, temperature, top_p):
+    """
+    Streams raw output to one component and generates final tabs for another.
+    This function acts as the main callback for the button click.
+    """
     if not client:
-         # Return an error structure if client failed to initialize
-         return "Error: Inference Client not available.", []
-
+        error_msg = "Error: Inference Client not available."
+        yield {
+            live_output: error_msg,
+            final_tabs: gr.Tabs(tabs=[gr.TabItem(label="Error", children=[gr.Textbox(value=error_msg)])])
+        }
+        return # Stop execution
+
+    # --- Prepare for Streaming ---
     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.
@@ -89,124 +78,158 @@ Each file must start with its filename on a new line. Example:
 
 index.html
 <!DOCTYPE html>
-<html>
-<head><title>My Site</title></head>
-<body><h1>Hello</h1></body>
-</html>
+<html></html>
 
 style.css
-body {{
-  font-family: sans-serif;
-}}
-
-script.js
-console.log('Hello World!');
+body {{}}
 
-Ensure the code is complete and functional for each file. NO commentary, NO explanations, NO markdown formatting like backticks (```).
+Ensure the code is complete. NO commentary, NO explanations, NO markdown formatting like backticks (```).
 Start generating the files now.
-""".strip()
+""".strip() + ("\n" + system_message if system_message else "")
 
     messages = [
-        {"role": "system", "content": full_sys_msg + ("\n" + system_message if system_message else "")},
+        {"role": "system", "content": full_sys_msg},
         {"role": "user", "content": prompt}
     ]
 
+    full_raw_response = ""
+    error_occurred = False
+    error_message = ""
+
+    # Initial state update
+    yield {
+        live_output: "Generating stream...",
+        # Set initial tabs state to indicate loading
+        final_tabs: gr.Tabs(tabs=[gr.TabItem(label="Generating...")])
+    }
+
+    # --- Streaming Loop ---
     try:
-        response = client.chat_completion(
-            messages=messages,
-            max_tokens=int(max_tokens), # Ensure max_tokens is int
+        stream = client.chat_completion(
+            messages,
+            max_tokens=int(max_tokens),
+            stream=True,
             temperature=temperature,
-            top_p=top_p,
-            stream=False # Ensure streaming is off for this logic
+            top_p=top_p
         )
-        raw = response.choices[0].message.content
-        print("\n--- Raw AI Response ---")
-        print(raw)
-        print("----------------------\n")
-        files = parse_files(raw)
-        return None, files # Return None for error, and the list of files
+        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 {
+                    live_output: full_raw_response,
+                    final_tabs: gr.Tabs(tabs=[gr.TabItem(label="Streaming...")]) # Keep showing streaming
+                }
+                # time.sleep(0.01) # Optional small delay if updates are too fast
 
     except Exception as e:
-        print(f"Error during AI generation: {e}")
-        return f"Error during AI generation: {e}", [] # Return error message
-
-# --- Gradio Event Handler ---
-def on_generate(prompt, backend, system_message, max_tokens, temperature, top_p):
-    """Callback function for the generate button."""
-    error_msg, files = generate_code(prompt, backend, system_message, max_tokens, temperature, top_p)
-
-    if error_msg:
-        # Display error in a single tab if generation failed
-        error_tab = gr.TabItem(label="Error", children=[gr.Textbox(value=error_msg, label="Generation Error")])
-        return gr.Tabs(tabs=[error_tab]) # Return a Tabs component with the error tab
-
-    if not files:
-        # Display message if no files were parsed
-        no_files_tab = gr.TabItem(label="Output", children=[gr.Textbox(value="AI did not return recognizable file content. Check raw output in console.", label="Result")])
-        return gr.Tabs(tabs=[no_files_tab]) # Return a Tabs component with this message
-
-    tabs = []
-    for name, content in files:
-        name = name.strip()
-        content = content.strip()
-        if not name or not content: # Skip empty names or content
-            print(f"Skipping file with empty name or content: Name='{name}'")
-            continue
-
-        # Determine language for syntax highlighting
-        lang = "text" # Default
-        if name.endswith(".html") or name.endswith(".htm"):
-            lang = "html"
-        elif name.endswith(".css"):
-            lang = "css"
-        elif name.endswith(".js"):
-            lang = "javascript"
-        elif name.endswith(".py"):
-            lang = "python"
-        elif name.endswith(".json"):
-            lang = "json"
-        elif name.endswith(".md"):
-            lang = "markdown"
-        elif name.endswith(".sh") or name.endswith(".bash"):
-            lang = "bash"
-
-        tab_item = gr.TabItem(label=name, elem_id=f"tab_{name.replace('.', '_')}", children=[ # Ensure unique elem_id
-            gr.Code(value=content, language=lang, label=name) # Add label to Code block
+        print(f"Error during AI streaming: {e}")
+        error_message = f"Error during AI generation: {e}\n\nPartial Response:\n{full_raw_response}"
+        error_occurred = True
+        # Update live output with error, keep tabs showing error state
+        yield {
+            live_output: error_message,
+            final_tabs: gr.Tabs(tabs=[gr.TabItem(label="Error")])
+        }
+
+
+    # --- Post-Streaming: Parsing and Final Tab Generation ---
+    if error_occurred:
+        # If an error happened during stream, create an error tab
+         final_tabs_update = gr.Tabs(tabs=[
+            gr.TabItem(label="Error", children=[gr.Textbox(value=error_message, label="Generation Error")])
         ])
-        tabs.append(tab_item)
+    else:
+        # Parse the complete raw response
+        print("\n--- Final Raw AI Response ---")
+        print(full_raw_response)
+        print("--------------------------\n")
+        files = parse_files(full_raw_response)
+
+        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."
+            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 }
+            return # Exit if no files
+
+        # --- Create Tabs (if files were parsed successfully) ---
+        tabs_content = []
+        for name, content in files:
+            name = name.strip()
+            content = content.strip()
+            if not name or not content:
+                print(f"Skipping file with empty name or content: Name='{name}'")
+                continue
+
+            lang = "text" # Default
+            if name.endswith((".html", ".htm")): lang = "html"
+            elif name.endswith(".css"): lang = "css"
+            elif name.endswith(".js"): lang = "javascript"
+            elif name.endswith(".py"): lang = "python"
+            elif name.endswith(".json"): lang = "json"
+            elif name.endswith(".md"): lang = "markdown"
+            elif name.endswith((".sh", ".bash")): lang = "bash"
+
+            tab_item = gr.TabItem(label=name, elem_id=f"tab_{name.replace('.', '_').replace('/', '_')}", children=[
+                gr.Code(value=content, language=lang, label=name)
+            ])
+            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")])])
+        else:
+            final_tabs_update = gr.Tabs(tabs=tabs_content) # Create the final Tabs component
+
+
+    # --- Final Update ---
+    # Yield the final state for both components
+    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
+    }
 
-    # *** The Key Fix ***
-    # Return a new gr.Tabs component instance containing the generated TabItems
-    return gr.Tabs(tabs=tabs)
 
 # --- Gradio UI Definition ---
-with gr.Blocks() as demo:
-    gr.Markdown("### Website Generator (Static / Flask / Node.js)")
-    gr.Markdown("Describe the website you want to create. The AI will generate the necessary files.")
+with gr.Blocks(css=".gradio-container { max-width: 95% !important; }") as demo: # Wider interface
+    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():
-        prompt = gr.Textbox(label="Describe your website", placeholder="E.g., a simple portfolio site with a contact form", scale=3)
-        backend = gr.Dropdown(["Static", "Flask", "Node.js"], value="Static", label="Backend Technology", scale=1)
-
-    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)") # Increased max
-        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")
-
-    gr.Markdown("#### Generated Files")
-    # Define the Tabs component placeholder. It will be replaced by the output of on_generate.
-    out_tabs = gr.Tabs(elem_id="output_tabs")
-
-    # Button click action
+        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")
+            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")
+
+        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("---")
+    gr.Markdown("#### Final Generated Files (Tabs)")
+    # Placeholder for the final structured tabs
+    final_tabs = gr.Tabs(elem_id="output_tabs")
+
+
+    # Button click action - uses the orchestrator function
     generate_button.click(
-        on_generate,
+        stream_and_parse_code, # Call the main function that handles streaming and parsing
         inputs=[prompt, backend, system_message, max_tokens, temperature, top_p],
-        outputs=[out_tabs], # Output the new Tabs component to replace the placeholder
-        show_progress="full" # Show progress during generation
+        # Outputs dictionary maps function yields to components
+        outputs=[live_output, final_tabs],
+        show_progress="hidden" # Hide default progress bar as we show live stream
     )
 
 if __name__ == "__main__":
-    demo.launch(debug=True) # Use debug=True for more detailed error messages in console
+    demo.launch(debug=True)