Update app.py

app.py

@@ -5,17 +5,57 @@ import re
 
 # --- Configuration ---
 API_TOKEN = os.getenv("HF_TOKEN", None)
-# Using a model known for better instruction following might be beneficial
-MODEL = "Qwen/Qwen2.5-Coder-32B-Instruct" # Kept your original choice, but consider testing others if needed
+MODEL = "Qwen/Qwen2.5-Coder-32B-Instruct"
 
 # --- Initialize Inference Client ---
 try:
     print(f"Initializing Inference Client for model: {MODEL}")
     client = InferenceClient(model=MODEL, token=API_TOKEN) if API_TOKEN else InferenceClient(model=MODEL)
 except Exception as e:
-    # Provide a more specific error message if possible
     raise gr.Error(f"Failed to initialize model client for {MODEL}. Error: {e}. Check HF_TOKEN and model availability.")
 
+# --- Helper Function to Parse and Format Code ---
+def parse_and_format_code(raw_response: str) -> str:
+    """
+    Parses raw AI output containing .TAB separators
+    and formats it for display in a single code block.
+    """
+    # Default filename for the first block if no TAB is present or before the first TAB
+    default_filename = "index.html"
+    separator_pattern = r'\.TAB\[NAME=([^\]]+)\]\n?'
+
+    filenames = re.findall(r'\.TAB\[NAME=([^\]]+)\]', raw_response)
+    code_blocks = re.split(separator_pattern, raw_response)
+
+    formatted_output = []
+
+    # Handle the first block (before any potential separator)
+    first_block = code_blocks[0].strip()
+    if first_block:
+        formatted_output.append(f"--- START FILE: {default_filename} ---\n\n{first_block}\n\n--- END FILE: {default_filename} ---")
+
+    # Handle subsequent blocks associated with filenames found by the separator
+    # re.split with capturing groups includes the separator AND the filename capture group
+    # in the results. We need to iterate carefully.
+    # Example: ['code1', 'app.py', 'code2', 'style.css', 'code3']
+    idx = 1
+    while idx < len(code_blocks) -1 :
+        filename = code_blocks[idx] # This should be the filename captured by the pattern
+        code = code_blocks[idx + 1].strip() # This should be the code after the separator
+        if code : # Only add if there's actual code content
+             formatted_output.append(f"--- START FILE: {filename} ---\n\n{code}\n\n--- END FILE: {filename} ---")
+        idx += 2 # Move past the filename and the code block
+
+    # If no separators were found, filenames list will be empty, and only the first block runs.
+    # If separators were found, the loop processes the rest.
+
+    if not formatted_output: # Handle case where input was empty or only whitespace
+        return raw_response # Return the original if parsing yields nothing
+
+    # Join the formatted blocks with clear separation
+    return "\n\n\n".join(formatted_output)
+
+
 # --- Core Code Generation Function ---
 def generate_code(
     prompt: str,
@@ -26,26 +66,22 @@ def generate_code(
 ):
     print(f"Generating code for: {prompt[:100]}... | Backend: {backend_choice}")
 
-    # --- Dynamically Build System Message ---
-    # Modified to include the specific formatting rules
     system_message = (
         "You are an AI that generates website code. You MUST ONLY output the raw code, without any conversational text like 'Here is the code' or explanations before or after the code blocks. "
-        "You MUST NOT wrap the code in markdown fences like ```html, ```python, or ```js. " # Explicit instruction to omit fences
-        "The user can select a backend hint (Static, Flask, Node.js). "
+        "You MUST NOT wrap the code in markdown fences like ```html, ```python, or ```js. "
         "If the user requests 'Static' or the prompt clearly implies only frontend code, generate ONLY the content for the `index.html` file. "
-        "If the user requests 'Flask' or 'Node.js' and the prompt requires backend logic (like handling forms, APIs, databases), you MUST generate both the `index.html` content AND the corresponding main backend file content (e.g., `app.py` for Flask, `server.js` or `app.js` for Node.js). "
+        "If the user requests 'Flask' or 'Node.js' and the prompt requires backend logic, you MUST generate both the `index.html` content AND the corresponding main backend file content (e.g., `app.py` for Flask, `server.js` or `app.js` for Node.js). "
         "When generating multiple files, you MUST separate them EXACTLY as follows: "
         "1. Output the complete code for the first file (e.g., `index.html`). "
-        "2. On a new line immediately after the first file's code, add the separator '.TAB[NAME=filename.ext]' (e.g., '.TAB[NAME=app.py]' or '.TAB[NAME=server.js]'). " # Specific separator format
+        "2. On a new line immediately after the first file's code, add the separator '.TAB[NAME=filename.ext]' (e.g., '.TAB[NAME=app.py]' or '.TAB[NAME=server.js]'). "
         "3. On the next line, immediately start the code for the second file. "
         "Generate only the necessary files (usually index.html and potentially one backend file). "
-        "The generated website code must be SFW (safe for work) and have minimal errors. "
-        "Only include comments where user modification is strictly required (e.g., API keys, database paths). Avoid explanatory comments. "
-        "If the user asks you to create code that is NOT for a website, you MUST respond ONLY with the exact phrase: " # Specific refusal phrase
+        "The generated website code must be SFW and have minimal errors. "
+        "Only include comments where user modification is strictly required. Avoid explanatory comments. "
+        "If the user asks you to create code that is NOT for a website, you MUST respond ONLY with the exact phrase: "
         "'hey there! am here to create websites for you unfortunately am programmed to not create codes! otherwise I would go on the naughty list :-('"
     )
 
-    # User prompt remains the same, passing the raw request and backend choice
     user_prompt = f"USER_PROMPT = {prompt}\nUSER_BACKEND = {backend_choice}"
 
     messages = [
@@ -69,49 +105,45 @@ def generate_code(
             if isinstance(token, str):
                 response_stream += token
                 full_response += token
-                # Yield intermediate stream for responsiveness
-                yield response_stream
+                # Yield intermediate stream for responsiveness during generation
+                yield response_stream # Show raw output as it comes
 
-        # --- Post-processing (Refined) ---
-        # Primarily focus on stripping whitespace and potential leftover model markers.
-        # The fence removal is kept as a fallback in case the model doesn't fully comply.
+        # --- Post-processing (After Stream Ends) ---
         cleaned_response = full_response.strip()
 
-        # Fallback fence removal (hopefully not needed often with the new prompt)
+        # Fallback fence removal
         cleaned_response = re.sub(r"^\s*```[a-z]*\s*\n?", "", cleaned_response)
         cleaned_response = re.sub(r"\n?\s*```\s*$", "", cleaned_response)
-
-        # Remove potential chat markers (like <|user|>, <|assistant|>)
+        # Remove potential chat markers
         cleaned_response = re.sub(r"<\s*\|?\s*(user|system|assistant)\s*\|?\s*>", "", cleaned_response, flags=re.IGNORECASE).strip()
-
-        # Remove common conversational phrases if they somehow slip through despite the prompt
+        # Remove common conversational phrases (if they slip through)
         common_phrases = [
             "Here is the code:", "Okay, here is the code:", "Here's the code:",
             "Sure, here is the code you requested:", "Let me know if you need anything else.",
-            "```html", "```python", "```javascript", "```", # Adding fences here just in case they appear standalone
+            "```html", "```python", "```javascript", "```",
         ]
-        # Use lower() for case-insensitive matching of leading phrases
         temp_response_lower = cleaned_response.lower()
         for phrase in common_phrases:
             if temp_response_lower.startswith(phrase.lower()):
-                # Use original case length for slicing
                 cleaned_response = cleaned_response[len(phrase):].lstrip()
-                temp_response_lower = cleaned_response.lower() # Update lower version after stripping
+                temp_response_lower = cleaned_response.lower()
 
-        # Ensure the specific refusal message isn't accidentally cleaned
+        # Check for refusal message
         refusal_message = "hey there! am here to create websites for you unfortunately am programmed to not create codes! otherwise I would go on the naughty list :-("
-        if refusal_message in full_response: # Check if the refusal message was generated
+        if refusal_message in full_response:
              yield refusal_message # Yield the exact refusal message
-        else:
-             yield cleaned_response # Yield the cleaned code
+             return # Stop processing
+
+        # --- PARSE and FORMAT the final cleaned response ---
+        formatted_final_code = parse_and_format_code(cleaned_response)
+
+        # Yield the final, formatted version replacing the streamed content
+        yield formatted_final_code
 
     except Exception as e:
-        # Log the full error for debugging on the server side
-        print(f"ERROR during code generation: {e}")
-        # Provide a user-friendly error message
-        yield f"## Error\n\nFailed to generate code.\n**Reason:** An unexpected error occurred. Please check the console logs or try again later."
-        # Consider raising a gr.Error for critical failures if preferred
-        # raise gr.Error(f"Code generation failed: {e}")
+        print(f"ERROR during code generation: {e}") # Log detailed error
+        # traceback.print_exc() # Uncomment for full traceback in logs
+        yield f"## Error\n\nFailed to generate or process code.\n**Reason:** {e}"
 
 
 # --- Build Gradio Interface ---
@@ -119,15 +151,15 @@ with gr.Blocks(css=".gradio-container { max-width: 90% !important; }") as demo:
     gr.Markdown("# ✨ Website Code Generator ✨")
     gr.Markdown(
         "Describe the website you want. The AI will generate the necessary code.\n"
-        "It will aim for `index.html` for 'Static', and potentially `index.html` + a backend file (like `app.py` or `server.js`) for 'Flask'/'Node.js'.\n"
+        "If multiple files are generated (e.g., for Flask/Node.js), they will be shown below separated by `--- START FILE: filename ---` markers.\n" # Updated description
         "**Output Format:**\n"
         "- No explanations, just code.\n"
-        "- Multiple files separated by `.TAB[NAME=filename.ext]` on its own line.\n"
+        "- Multiple files separated by file markers.\n" # Updated description
         "- Minimal necessary comments only.\n\n"
         "**Rules:**\n"
         "- Backend choice guides the AI on whether to include server-side code.\n"
         "- Always SFW and aims for minimal errors.\n"
-        "- Only generates website-related code. No other types of code."
+        "- Only generates website-related code."
     )
 
     with gr.Row():
@@ -141,25 +173,21 @@ with gr.Blocks(css=".gradio-container { max-width: 90% !important; }") as demo:
                 ["Static", "Flask", "Node.js"],
                 label="Backend Context",
                 value="Static",
-                info="Guides AI if backend code (like Python/JS) is needed alongside HTML." # Updated info text
+                info="Guides AI if backend code (like Python/JS) is needed alongside HTML."
             )
             generate_button = gr.Button("✨ Generate Website Code", variant="primary")
 
         with gr.Column(scale=3):
             code_output = gr.Code(
-                label="Generated Code", # Changed label slightly
-                language=None, # Set language to None for plain text display, better for mixed content
+                label="Generated Code (Scroll for multiple files)", # Updated label
+                language=None, # Keep as None for mixed/plain text display
                 lines=30,
-                interactive=False,
+                interactive=False, # Keep non-interactive
             )
 
     with gr.Accordion("Advanced Settings", open=False):
         max_tokens_slider = gr.Slider(
-            minimum=512,
-            maximum=4096, # Adjust max based on model limits if necessary
-            value=3072,
-            step=256,
-            label="Max New Tokens"
+            minimum=512, maximum=4096, value=3072, step=256, label="Max New Tokens"
         )
         temperature_slider = gr.Slider(
             minimum=0.1, maximum=1.2, value=0.7, step=0.1, label="Temperature"
@@ -168,6 +196,7 @@ with gr.Blocks(css=".gradio-container { max-width: 90% !important; }") as demo:
             minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-P"
         )
 
+    # The click function now yields the final formatted string AFTER streaming
     generate_button.click(
         fn=generate_code,
         inputs=[prompt_input, backend_radio, max_tokens_slider, temperature_slider, top_p_slider],
@@ -177,4 +206,5 @@ with gr.Blocks(css=".gradio-container { max-width: 90% !important; }") as demo:
 if __name__ == "__main__":
     if not API_TOKEN:
         print("Warning: HF_TOKEN environment variable not set. Using anonymous access.")
-    demo.queue(max_size=10).launch()
+    # Increased max_size slightly for potentially larger combined outputs
+    demo.queue(max_size=15).launch()