Update app.py

app.py

@@ -3,30 +3,59 @@ from huggingface_hub import InferenceClient
 import os
 import re
 
+# --- Configuration ---
 API_TOKEN = os.getenv("HF_TOKEN", None)
-MODEL = "Qwen/Qwen2.5-Coder-32B-Instruct"
+# 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
 
+# --- 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:
-    raise gr.Error(f"Failed to initialize model client. Error: {e}")
-
-def generate_code(prompt: str, backend_choice: str, max_tokens: int, temperature: float, top_p: float):
+    # 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.")
+
+# --- Core Code Generation Function ---
+def generate_code(
+    prompt: str,
+    backend_choice: str,
+    max_tokens: int,
+    temperature: float,
+    top_p: float,
+):
+    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 assistant programmed to generate website codes only. "
-        "You must not use triple backticks (```html, ```python, etc.). "
-        "If multiple files are needed, separate them clearly using:\n"
-        "TAB.NAME={filename}\n"
-        "Only generate code. No explanations, no phrases like 'Here is the code'. "
-        "If user asks non-website code, reply:\n"
-        "'hey there! am here to create websites for you unfortunately am programmed to not create codes! otherwise I would go on the naughty list :-('."
+        "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). "
+        "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). "
+        "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
+        "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
+        "'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 = [
         {"role": "system", "content": system_message},
         {"role": "user", "content": user_prompt}
     ]
+
+    response_stream = ""
     full_response = ""
+
     try:
         stream = client.chat_completion(
             messages=messages,
@@ -38,205 +67,114 @@ def generate_code(prompt: str, backend_choice: str, max_tokens: int, temperature
         for message in stream:
             token = message.choices[0].delta.content
             if isinstance(token, str):
+                response_stream += token
                 full_response += token
+                # Yield intermediate stream for responsiveness
+                yield response_stream
 
+        # --- 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.
         cleaned_response = full_response.strip()
+
+        # Fallback fence removal (hopefully not needed often with the new prompt)
         cleaned_response = re.sub(r"^\s*```[a-z]*\s*\n?", "", cleaned_response)
         cleaned_response = re.sub(r"\n?\s*```\s*$", "", cleaned_response)
-        cleaned_response = re.sub(r"<\s*\|?\s*(user|assistant|system|endoftext)\s*\|?\s*>", "", cleaned_response, flags=re.IGNORECASE)
-        cleaned_response = cleaned_response.replace("<|im_end|>", "").replace("<|im_start|>", "").strip()
 
+        # Remove potential chat markers (like <|user|>, <|assistant|>)
+        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
         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.",
-            "Here is the website code you requested:", "Here are the files for your website:",
-            "Okay, here are the files:"
+            "```html", "```python", "```javascript", "```", # Adding fences here just in case they appear standalone
         ]
-        lower_response = cleaned_response.lower()
+        # Use lower() for case-insensitive matching of leading phrases
+        temp_response_lower = cleaned_response.lower()
         for phrase in common_phrases:
-            if lower_response.startswith(phrase.lower()):
+            if temp_response_lower.startswith(phrase.lower()):
+                # Use original case length for slicing
                 cleaned_response = cleaned_response[len(phrase):].lstrip()
-                lower_response = cleaned_response.lower()
+                temp_response_lower = cleaned_response.lower() # Update lower version after stripping
 
-        if not cleaned_response:
-             return "Error: Empty response from model after cleaning."
+        # Ensure the specific refusal message isn't accidentally cleaned
+        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
+             yield refusal_message # Yield the exact refusal message
+        else:
+             yield cleaned_response # Yield the cleaned code
 
-        return cleaned_response
     except Exception as e:
-        return f"## Error\n\nFailed to generate code.\n**Reason:** {e}"
-
-def split_files(full_code_text):
-    file_blocks = []
-    splits = re.split(r'(TAB\.NAME=\{.+?\})', full_code_text)
-    initial_content = splits[0].strip()
-
-    if len(splits) == 1:
-        if initial_content:
-            default_name = "index.html"
-            if "def " in initial_content or "import " in initial_content: default_name = "app.py"
-            elif "function " in initial_content or "const " in initial_content or "let " in initial_content: default_name = "script.js"
-            elif "<!DOCTYPE html>" in initial_content or "<html" in initial_content: default_name = "index.html"
-            elif "@app.route" in initial_content: default_name = "app.py"
-            elif "require(" in initial_content or "module.exports" in initial_content: default_name = "server.js"
-            elif "<?php" in initial_content: default_name = "index.php"
-            elif "package main" in initial_content: default_name = "main.go"
-
-            file_blocks.append((default_name, initial_content))
-    else:
-        for i in range(1, len(splits), 2):
-            marker = splits[i]
-            content = splits[i+1].strip() if (i+1) < len(splits) else ""
-            filename_match = re.search(r'TAB\.NAME=\{(.+?)\}', marker)
-            if filename_match:
-                filename = filename_match.group(1).strip()
-                if content:
-                    file_blocks.append((filename, content))
-            elif i == 1 and initial_content: # Handle content before the first explicit marker
-                 file_blocks.append(("file_0.txt", initial_content))
-                 if content: # Add the content after the first marker if it exists
-                     filename_match_fallback = re.search(r'TAB\.NAME=\{(.+?)\}', marker)
-                     if filename_match_fallback:
-                         filename = filename_match_fallback.group(1).strip()
-                         file_blocks.append((filename, content))
-
-    return file_blocks
+        # 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}")
 
+
+# --- Build Gradio Interface ---
 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 website code.\n\n"
+        "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"
+        "**Output Format:**\n"
+        "- No explanations, just code.\n"
+        "- Multiple files separated by `.TAB[NAME=filename.ext]` on its own line.\n"
+        "- Minimal necessary comments only.\n\n"
         "**Rules:**\n"
-        "- Provide a backend hint (Static / Flask / Node.js).\n"
-        "- Generated code should be functional and SFW.\n"
-        "- Only generates website-related code.\n"
-        "- If multiple files are generated, they will be separated below using the format `TAB.NAME={filename}`."
+        "- 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."
     )
+
     with gr.Row():
         with gr.Column(scale=2):
             prompt_input = gr.Textbox(
                 label="Website Description",
-                placeholder="e.g., A simple Flask app with one route that displays 'Hello World'.",
+                placeholder="e.g., A Flask app with a form that stores data in a variable.",
                 lines=6,
             )
             backend_radio = gr.Radio(
                 ["Static", "Flask", "Node.js"],
-                label="Backend Context / Hint",
+                label="Backend Context",
                 value="Static",
+                info="Guides AI if backend code (like Python/JS) is needed alongside HTML." # Updated info text
             )
             generate_button = gr.Button("✨ Generate Website Code", variant="primary")
+
         with gr.Column(scale=3):
-            main_output_label = gr.Markdown("### Full Generated Code / Main File")
-            main_output_code = gr.Code(
-                label="Generated Code", # Label is less prominent now
-                language="html",
-                lines=15,
+            code_output = gr.Code(
+                label="Generated Code", # Changed label slightly
+                language=None, # Set language to None for plain text display, better for mixed content
+                lines=30,
                 interactive=False,
             )
-            extra_outputs_column = gr.Column(visible=False) # Initially hidden
 
     with gr.Accordion("Advanced Settings", open=False):
         max_tokens_slider = gr.Slider(
             minimum=512,
-            maximum=8192,
-            value=4096,
+            maximum=4096, # Adjust max based on model limits if necessary
+            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"
+            minimum=0.1, maximum=1.2, value=0.7, step=0.1, label="Temperature"
         )
         top_p_slider = gr.Slider(
-            minimum=0.1,
-            maximum=1.0,
-            value=0.9,
-            step=0.05,
-            label="Top-P"
+            minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-P"
         )
 
-    def get_language(filename):
-        if filename.endswith(".html") or filename.endswith(".htm"): return "html"
-        if filename.endswith(".css"): return "css"
-        if filename.endswith(".js"): return "javascript"
-        if filename.endswith(".py"): return "python"
-        if filename.endswith(".json"): return "json"
-        if filename.endswith(".sql"): return "sql"
-        if filename.endswith(".php"): return "php"
-        if filename.endswith(".go"): return "go"
-        if filename.endswith(".java"): return "java"
-        if filename.endswith(".rb"): return "ruby"
-        if filename.endswith(".sh"): return "shell"
-        if filename.endswith(".yml") or filename.endswith(".yaml"): return "yaml"
-        if filename.endswith(".md"): return "markdown"
-        return "text"
-
-    def generate_and_display(prompt, backend, max_tokens, temperature, top_p):
-        full_code = generate_code(prompt, backend, max_tokens, temperature, top_p)
-
-        if full_code.startswith("## Error"):
-            return {
-                main_output_label: gr.Markdown.update(value="### Error Occurred"),
-                main_output_code: gr.Code.update(value=full_code, language="markdown"),
-                extra_outputs_column: gr.Column.update(visible=False, children=[])
-            }
-        if full_code.startswith("Error: Empty response"):
-             return {
-                main_output_label: gr.Markdown.update(value="### Error Occurred"),
-                main_output_code: gr.Code.update(value=full_code, language="text"),
-                extra_outputs_column: gr.Column.update(visible=False, children=[])
-             }
-
-        files = split_files(full_code)
-        dynamic_components = []
-        main_file_content = full_code
-        main_file_lang = "text"
-        main_file_label = "### Full Generated Code"
-
-
-        if not files:
-             main_file_content = full_code # Show full code if split failed but we got output
-             main_file_lang = get_language("output.txt") # Basic guess
-             main_file_label = "### Full Generated Output (No Files Detected)"
-             return {
-                 main_output_label: gr.Markdown.update(value=main_file_label),
-                 main_output_code: gr.Code.update(value=main_file_content, language=main_file_lang),
-                 extra_outputs_column: gr.Column.update(visible=False, children=[])
-             }
-
-        if len(files) == 1:
-             main_file_content = files[0][1]
-             main_file_lang = get_language(files[0][0])
-             main_file_label = f"### File: {files[0][0]}"
-             return {
-                 main_output_label: gr.Markdown.update(value=main_file_label),
-                 main_output_code: gr.Code.update(value=main_file_content, language=main_file_lang),
-                 extra_outputs_column: gr.Column.update(visible=False, children=[])
-             }
-        else:
-            main_file_content = files[0][1]
-            main_file_lang = get_language(files[0][0])
-            main_file_label = f"### File: {files[0][0]}"
-
-            for i, (filename, content) in enumerate(files[1:], start=1): # Start from the second file for extras
-                lang = get_language(filename)
-                dynamic_components.append(gr.Markdown(f"### File: {filename}"))
-                dynamic_components.append(gr.Code(value=content, language=lang, label=filename, interactive=False))
-
-            return {
-                main_output_label: gr.Markdown.update(value=main_file_label),
-                main_output_code: gr.Code.update(value=main_file_content, language=main_file_lang),
-                extra_outputs_column: gr.Column.update(visible=True, children=dynamic_components)
-            }
-
-
     generate_button.click(
-        fn=generate_and_display,
+        fn=generate_code,
         inputs=[prompt_input, backend_radio, max_tokens_slider, temperature_slider, top_p_slider],
-        outputs=[main_output_label, main_output_code, extra_outputs_column]
+        outputs=code_output,
     )
 
 if __name__ == "__main__":
-    demo.queue().launch()
+    if not API_TOKEN:
+        print("Warning: HF_TOKEN environment variable not set. Using anonymous access.")
+    demo.queue(max_size=10).launch()