Update app.py

app.py

@@ -1,85 +1,212 @@
 import gradio as gr
 from huggingface_hub import InferenceClient
 import re
-
-client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
-
+import os # Good practice to import os if needed, though not strictly used here yet
+
+# --- 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
+
+# --- Inference Client ---
+# Consider adding error handling for client initialization if needed
+try:
+    client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
+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 ---
 def parse_files(raw_response):
-    # Parse files and code from raw AI output
-    pattern = re.compile(r"(?:(?:\n|^)([\w\-.\/\\]+)\n)(.+?)(?=\n[\w\-.\/\\]+\n|\Z)", re.DOTALL)
-    return pattern.findall(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 = 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)
+        content_cleaned = re.sub(r"^\s*```[a-zA-Z]*\n?", "", content, flags=re.MULTILINE)
+        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
+    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
+            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."""
+    if not client:
+         # Return an error structure if client failed to initialize
+         return "Error: Inference Client not available.", []
+
     full_sys_msg = f"""
-You are a code-only AI. Given a prompt, generate a full website using the {backend} backend.
+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.
+Respond ONLY with filenames and the raw code for each file.
+Each file must start with its filename on a new line. Example:
 
-Respond ONLY with filenames and raw code. NO commentary, NO backticks, NO markdown.
-
-Example:
 index.html
-<html>...</html>
+<!DOCTYPE html>
+<html>
+<head><title>My Site</title></head>
+<body><h1>Hello</h1></body>
+</html>
 
 style.css
-body {{ ... }}
+body {{
+  font-family: sans-serif;
+}}
+
+script.js
+console.log('Hello World!');
 
-Start now.
+Ensure the code is complete and functional for each file. NO commentary, NO explanations, NO markdown formatting like backticks (```).
+Start generating the files now.
 """.strip()
 
     messages = [
-        {"role": "system", "content": full_sys_msg + "\n" + system_message},
+        {"role": "system", "content": full_sys_msg + ("\n" + system_message if system_message else "")},
         {"role": "user", "content": prompt}
     ]
 
-    response = client.chat_completion(
-        messages=messages,
-        max_tokens=max_tokens,
-        temperature=temperature,
-        top_p=top_p,
-    )
+    try:
+        response = client.chat_completion(
+            messages=messages,
+            max_tokens=int(max_tokens), # Ensure max_tokens is int
+            temperature=temperature,
+            top_p=top_p,
+            stream=False # Ensure streaming is off for this logic
+        )
+        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
+
+    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)
 
-    raw = response.choices[0].message.content
-    files = parse_files(raw)
+    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
 
-    return files
+    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
 
-def on_generate(prompt, backend, system_message, max_tokens, temperature, top_p):
-    files = generate_code(prompt, backend, system_message, max_tokens, temperature, top_p)
     tabs = []
     for name, content in files:
-        lang = "html" if name.endswith(".html") else \
-               "css" if name.endswith(".css") else \
-               "javascript" if name.endswith(".js") else \
-               "python" if name.endswith(".py") else "text"
-        tabs.append(gr.TabItem(label=name.strip(), elem_id=name.strip(), children=[
-            gr.Code(value=content.strip(), language=lang)
-        ]))
-    return tabs  # Return tabs directly
-
+        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
+        ])
+        tabs.append(tab_item)
+
+    # *** 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.Row():
-        prompt = gr.Textbox(label="Describe your website", placeholder="E.g. a portfolio site with dark mode")
-        backend = gr.Dropdown(["Static", "Flask", "Node.js"], value="Static", label="Backend")
+        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", value="")
-        max_tokens = gr.Slider(256, 2048, value=1024, label="Max Tokens")
-        temperature = gr.Slider(0.1, 2.0, value=0.7, label="Temperature")
-        top_p = gr.Slider(0.1, 1.0, value=0.95, label="Top-p")
+        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)")
 
-    out_tabs = gr.Tabs()  # Define the Tabs component
+    generate_button = gr.Button("✨ Generate Code ✨", variant="primary")
 
-    generate_button = gr.Button("Generate Code")
+    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
     generate_button.click(
         on_generate,
         inputs=[prompt, backend, system_message, max_tokens, temperature, top_p],
-        outputs=out_tabs,  # Output goes directly to the Tabs component
-        show_progress=True
+        outputs=[out_tabs], # Output the new Tabs component to replace the placeholder
+        show_progress="full" # Show progress during generation
     )
 
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(debug=True) # Use debug=True for more detailed error messages in console