app.py

import gradio as gr
from huggingface_hub import InferenceClient
import os
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

# --- 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.")

# --- 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 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,
            max_tokens=max_tokens,
            stream=True,
            temperature=temperature,
            top_p=top_p,
        )
        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)

        # 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.",
            "```html", "```python", "```javascript", "```", # Adding fences here just in case they appear standalone
        ]
        # 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

        # 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

    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}")


# --- 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 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"
        "- 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 Flask app with a form that stores data in a variable.",
                lines=6,
            )
            backend_radio = gr.Radio(
                ["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
            )
            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
                lines=30,
                interactive=False,
            )

    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"
        )
        temperature_slider = gr.Slider(
            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"
        )

    generate_button.click(
        fn=generate_code,
        inputs=[prompt_input, backend_radio, max_tokens_slider, temperature_slider, top_p_slider],
        outputs=code_output,
    )

if __name__ == "__main__":
    if not API_TOKEN:
        print("Warning: HF_TOKEN environment variable not set. Using anonymous access.")
    demo.queue(max_size=10).launch()