app.py

import gradio as gr
from huggingface_hub import InferenceClient
import os
import re

API_TOKEN = os.getenv("HF_TOKEN", None)
MODEL = "Qwen/Qwen2.5-Coder-32B-Instruct"

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 for {MODEL}. Error: {e}. Check HF_TOKEN and model availability.")

def parse_code_into_files(raw_response: str) -> dict:
    files = {}
    default_first_filename = "index.html"
    separator_pattern = r'\.TAB\[NAME=([^\]]+)\]\n?'

    matches = list(re.finditer(separator_pattern, raw_response))

    start_index = 0
    first_separator_pos = matches[0].start() if matches else len(raw_response)
    first_block = raw_response[start_index:first_separator_pos].strip()
    if first_block:
         files[default_first_filename] = first_block

    if matches:
        backend_filename = matches[0].group(1).strip()
        start_index = matches[0].end()

        second_separator_pos = matches[1].start() if len(matches) > 1 else len(raw_response)
        backend_code = raw_response[start_index:second_separator_pos].strip()

        if backend_code:
            files['backend_file'] = backend_code
            files['backend_filename'] = backend_filename
            if backend_filename.endswith(".py"):
                files['backend_language'] = 'python'
            elif backend_filename.endswith(".js"):
                files['backend_language'] = 'javascript'
            elif backend_filename.endswith(".css"):
                 files['backend_language'] = 'css'
            else:
                 files['backend_language'] = None

    return files

def generate_code(
    prompt: str,
    backend_choice: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
    progress=gr.Progress(track_tqdm=True)
):
    print(f"Generating code for: {prompt[:100]}... | Backend: {backend_choice}")
    progress(0, desc="Initializing Request...")

    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. "
         "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, 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]'). "
         "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 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 = f"USER_PROMPT = {prompt}\nUSER_BACKEND = {backend_choice}"

    messages = [
        {"role": "system", "content": system_message},
        {"role": "user", "content": user_prompt}
    ]

    full_response = ""
    token_count = 0
    est_total_tokens = max_tokens

    try:
        progress(0.1, desc="Sending Request to Model...")
        stream = client.chat_completion(
            messages=messages,
            max_tokens=max_tokens,
            stream=True,
            temperature=temperature,
            top_p=top_p,
        )

        progress(0.2, desc="Receiving Stream...")
        for message in stream:
            token = message.choices[0].delta.content
            if isinstance(token, str):
                full_response += token
                token_count += 1
                prog = min(0.2 + (token_count / est_total_tokens) * 0.7, 0.9)
                progress(prog, desc="Generating Code...")


        progress(0.9, desc="Processing Response...")
        cleaned_response = full_response.strip()
        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|system|assistant)\s*\|?\s*>", "", cleaned_response, flags=re.IGNORECASE).strip()
        common_phrases = [
            "Here is the code:", "Okay, here is the code:", "Here's the code:",
            "Sure, here is the code you requested:",
        ]
        temp_response_lower = cleaned_response.lower()
        for phrase in common_phrases:
            if temp_response_lower.startswith(phrase.lower()):
                cleaned_response = cleaned_response[len(phrase):].lstrip()
                temp_response_lower = cleaned_response.lower()

        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:
             progress(1, desc="Refusal Message Generated")
             return gr.update(value=refusal_message, language=None, visible=True), gr.update(value="", visible=False, label="Backend")

        parsed_files = parse_code_into_files(cleaned_response)

        html_code = parsed_files.get("index.html", "")
        backend_code = parsed_files.get("backend_file", "")
        backend_filename = parsed_files.get("backend_filename", "Backend")
        backend_language = parsed_files.get("backend_language", None)

        html_update = gr.update(value=html_code, language='html', visible=True)

        if backend_code:
            backend_update = gr.update(value=backend_code, language=backend_language, label=backend_filename, visible=True)
        else:
            backend_update = gr.update(value="", visible=False, label="Backend")

        progress(1, desc="Done")
        return html_update, backend_update

    except Exception as e:
        print(f"ERROR during code generation: {e}")
        progress(1, desc="Error Occurred")
        error_message = f"## Error\n\nFailed to generate or process code.\n**Reason:** {e}"
        return gr.update(value=error_message, language=None, visible=True), gr.update(value="", visible=False, label="Backend")


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"
        "If multiple files are generated (e.g., for Flask/Node.js), they will appear in separate tabs below."
    )

    with gr.Row():
        with gr.Column(scale=2):
            prompt_input = gr.Textbox(
                label="Website Description",
                placeholder="e.g., A Flask app with a simple chat using Socket.IO",
                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."
            )
            generate_button = gr.Button("✨ Generate Website Code", variant="primary")

        with gr.Column(scale=3):
            with gr.Tabs(elem_id="code-tabs") as code_tabs:
                with gr.Tab("index.html", elem_id="html-tab") as html_tab:
                     html_code_output = gr.Code(
                         label="index.html",
                         language="html",
                         lines=25,
                         interactive=False,
                         elem_id="html_code",
                     )
                with gr.Tab("Backend", elem_id="backend-tab", visible=False) as backend_tab:
                     backend_code_output = gr.Code(
                         label="Backend",
                         language=None,
                         lines=25,
                         interactive=False,
                         elem_id="backend_code",
                         visible=False
                     )

    with gr.Accordion("Advanced Settings", open=False):
        max_tokens_slider = gr.Slider(
            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"
        )
        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=[html_code_output, backend_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()