Spaces:

MINEOGO
/

pro-zephyr-coder

Running

File size: 15,236 Bytes

import gradio as gr
from huggingface_hub import InferenceClient
import os
import re # For post-processing fallback

# --- FastAPI & Pydantic Imports ---
from fastapi import FastAPI, Request, HTTPException, Depends, Body
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
from typing import Literal, Optional

# --- Configuration ---
API_TOKEN = os.getenv("HF_TOKEN", None)
MODEL = "HuggingFaceH4/zephyr-7b-beta"
# --- Define the Secret Key for the API ---
API_SECRET_KEY = "onlyfordearygt" # Keep this secure in a real application (e.g., env variable)

# --- Initialize Inference Client ---
try:
    print(f"Attempting to initialize Inference Client for model: {MODEL}")
    if API_TOKEN:
        print("Using HF Token found in environment.")
        client = InferenceClient(model=MODEL, token=API_TOKEN)
    else:
        print("HF Token not found. Running without token (may lead to rate limits).")
        client = InferenceClient(model=MODEL)
    print("Inference Client initialized successfully.")
except Exception as e:
    print(f"Error initializing Inference Client: {e}")
    # We still want the Gradio app to potentially load, but maybe show an error
    # raise gr.Error(f"Failed to initialize the AI model client for '{MODEL}'. Check model name, network, and HF_TOKEN secret if applicable. Error: {e}")
    client = None # Set client to None so we can check later
    print("WARNING: AI Client initialization failed. API/Generation will not work.")


# --- Pydantic Model for API Request Body ---
class GenerateRequest(BaseModel):
    prompt: str
    backend_choice: Literal["Static", "Flask", "Node.js"] = "Static"
    file_structure: Literal["Single File", "Multiple Files"] = "Multiple Files"
    max_tokens: Optional[int] = Field(default=3072, gt=128, le=4096) # Add validation
    temperature: Optional[float] = Field(default=0.7, gt=0.0, le=2.0)
    top_p: Optional[float] = Field(default=0.9, gt=0.0, le=1.0)
    secret_key: str # Required for authentication


# --- Core Code Generation Function (Mostly Unchanged) ---
# Note: This function is now also used by the API
def generate_code(
    prompt: str,
    backend_choice: str,
    file_structure: str,
    max_tokens: int,
    temperature: float,
    top_p: float,
):
    """
    Generates website code based on user prompt and choices.
    Aims for richer CSS, emphasizes completeness, and strictly outputs ONLY raw code.
    Yields the code token by token for live updates (for Gradio UI).
    The *final* yielded value is the complete, cleaned code (for API).
    """
    # Check if client initialized properly
    if client is None:
         final_error_message = "## Error\n\nAI Model Client not initialized. Generation is unavailable."
         print(final_error_message)
         # Yield the error for Gradio, return it for API callers later
         yield final_error_message
         return # Stop execution for this generator

    print(f"--- Generating Code ---")
    print(f"Prompt: {prompt[:100]}...")
    print(f"Backend Context: {backend_choice}")
    print(f"File Structure: {file_structure}")
    print(f"Settings: Max Tokens={max_tokens}, Temp={temperature}, Top-P={top_p}")

    if file_structure == "Single File":
        file_structure_instruction = (
            "- **File Structure is 'Single File':** Generate ONLY a single, complete `index.html` file. "
            "Embed ALL CSS directly within `<style>` tags inside the `<head>`. "
            "Embed ALL necessary JavaScript directly within `<script>` tags just before the closing `</body>` tag. "
            "Do NOT use file separation markers."
        )
    else: # Multiple Files
        file_structure_instruction = (
            "- **File Structure is 'Multiple Files':** Generate code for `index.html`, `style.css`, and `script.js` (if JS is needed). "
            "Use these EXACT markers: `<!-- index.html -->`, `/* style.css */`, `// script.js` (only if JS is needed).\n"
            "- Place the corresponding code directly after each marker.\n"
            "- Inside `index.html`, link `style.css` in the `<head>` and include `script.js` before `</body>` if generated."
         )

    system_message = (
        "You are an expert frontend web developer AI. Your primary goal is to generate **complete, visually appealing, modern, and well-styled** frontend code (HTML, CSS, client-side JS) based *only* on the user's description and selected options. "
        "Follow ALL these rules with EXTREME STRICTNESS:\n"
        "1.  **STYLE & DETAIL:** Generate rich, detailed code. Use **plenty of CSS** for layout, spacing, typography, colors, and effects. Aim for a high-quality visual result.\n"
        "2.  **COMPLETENESS:** Generate the *entire* requested code structure. Ensure all files/sections are fully generated and properly closed. **DO NOT STOP GENERATING PREMATURELY.** Finish the whole task.\n"
        "3.  **RAW CODE ONLY:** Your *entire* response MUST consist *only* of the requested source code. NO extra text, NO explanations, NO apologies, NO introductions, NO summaries, NO comments about the code (except standard code comments), NO MARKDOWN formatting (like ```html), and ***ABSOLUTELY NO CONVERSATIONAL TEXT OR TAGS*** like `<|user|>` or `<|assistant|>`.\n"
        "4.  **IMMEDIATE CODE START:** The response MUST begin *directly* with the first character of the code (`<!DOCTYPE html>` or `<!-- index.html -->`).\n"
        "5.  **IMMEDIATE CODE END:** The response MUST end *immediately* after the very last character of the generated code. DO NOT add *any* text, spaces, or newlines after the code concludes.\n"
        "6.  **MANDATORY `index.html`:** Always generate the content for `index.html`.\n"
        f"7.  **FILE STRUCTURE ({file_structure}):** Strictly follow ONLY the instructions for the *selected* file structure:\n"
        f"    {file_structure_instruction}\n"
        "8.  **BACKEND CONTEXT ({backend_choice}):** Use as a hint for frontend structure only. Generate ONLY frontend code.\n"
        "9.  **FRONTEND ONLY:** Do NOT generate server-side code.\n"
        "10. **ACCURACY:** Generate functional code addressing the user's prompt.\n\n"
        "REMEMBER: Create COMPLETE, visually appealing code. Output ONLY raw code. START immediately with code. FINISH the entire code generation. END immediately with code. NO extra text/tags."
    )

    messages = [
        {"role": "system", "content": system_message},
        {"role": "user", "content": f"Generate the complete website frontend code for: {prompt}"}
    ]

    response_stream = ""
    full_response_for_cleaning = ""
    token_count = 0
    last_yielded_value = "" # Store the last value for the API

    try:
        print("Sending request to Hugging Face Inference API...")
        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):
                token_count += 1
                response_stream += token
                full_response_for_cleaning += token
                last_yielded_value = response_stream # Keep updating last value during stream
                yield response_stream # Yield cumulative response for live UI update

        print(f"API stream finished. Received ~{token_count} tokens. Raw length: {len(full_response_for_cleaning)}")
        if token_count >= max_tokens - 10:
             print(f"WARNING: Generation might have been cut short due to reaching max_tokens limit ({max_tokens}).")

        # --- Post-Processing ---
        cleaned_response = full_response_for_cleaning.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|assistant)\s*\|?\s*>", "", cleaned_response, flags=re.IGNORECASE)
        # Remove common phrases only if they are clearly at the start/end and unlikely to be code
        common_phrases_start = ["Here is the code:", "Okay, here is the code:", "Here's the code:", "Sure, here is the code you requested:"]
        for phrase in common_phrases_start:
            if cleaned_response.lower().startswith(phrase.lower()):
                cleaned_response = cleaned_response[len(phrase):].lstrip()

        last_yielded_value = cleaned_response.strip() # Final cleaned value
        yield last_yielded_value # Yield final cleaned response for Gradio UI

    except Exception as e:
        error_message = f"An error occurred during the API call: {e}"
        print(error_message)
        final_error_message = f"## Error\n\nFailed to generate code.\n**Reason:** {e}"
        # Yield error for Gradio UI
        yield final_error_message
        # Ensure the generator stops, API will handle the exception based on this
        # For the API, we will raise an exception in the route handler if needed


# --- Build Gradio Interface using Blocks ---
# Define this *before* creating the FastAPI app that might mount it
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 **visually styled** frontend code (HTML, CSS, JS) using **plenty of CSS**. "
        "The code appears live below.\n"
        "**Important:**\n"
        "1. This generator creates code based *only* on your initial description. To refine, modify your description and generate again.\n"
        "2. **If the code output stops abruptly**, it likely hit the 'Max New Tokens' limit. **Increase the slider value below** and try again!\n"
        "3. An API endpoint is available at `/api/generate` (POST request, requires secret key)." # Notify about API
    )

    with gr.Row():
        with gr.Column(scale=2):
            prompt_input = gr.Textbox(label="Website Description", placeholder="e.g., A modern portfolio...", lines=6,)
            backend_radio = gr.Radio(["Static", "Flask", "Node.js"], label="Backend Context Hint", value="Static", info="Hint for AI - generates ONLY frontend code.")
            file_structure_radio = gr.Radio(["Multiple Files", "Single File"], label="Output File Structure", value="Multiple Files", info="Choose 'Single File' or 'Multiple Files'.")
            generate_button = gr.Button("🎨 Generate Stylish Website Code", variant="primary")

        with gr.Column(scale=3):
            code_output = gr.Code(label="Generated Code (Raw Output - Aiming for Style!)", language="html", lines=30, interactive=False,)

    with gr.Accordion("Advanced Generation Settings", open=False):
        max_tokens_slider = gr.Slider(minimum=512, maximum=4096, value=3072, step=256, label="Max New Tokens", info="Max length. Increase if output is cut off!")
        temperature_slider = gr.Slider(minimum=0.1, maximum=1.2, value=0.7, step=0.1, label="Temperature", info="Controls randomness.")
        top_p_slider = gr.Slider( minimum=0.1, maximum=1.0, value=0.9, step=0.05, label="Top-P", info="Alternative randomness control.")

    # --- Connect Gradio Inputs/Outputs ---
    generate_button.click(
        fn=generate_code,
        inputs=[prompt_input, backend_radio, file_structure_radio, max_tokens_slider, temperature_slider, top_p_slider,],
        outputs=code_output,
    )

    # --- Gradio Examples ---
    gr.Examples(
        examples=[
            ["A simple counter page...", "Static", "Single File"],
            ["A responsive product grid...", "Static", "Multiple Files"],
            ["A personal blog homepage...", "Flask", "Multiple Files"],
            ["A 'Coming Soon' page...", "Static", "Multiple Files"]
        ],
        inputs=[prompt_input, backend_radio, file_structure_radio],
        label="Example Prompts (Aiming for Style)"
    )


# --- Create FastAPI App and Mount Gradio ---
# This approach uses FastAPI as the main server and mounts Gradio onto it
app = FastAPI()

# --- Define the API Endpoint ---
@app.post("/api/generate")
async def api_generate_code(payload: GenerateRequest):
    """
    API endpoint to generate website code.
    Requires a secret key for authentication.
    """
    print(f"API Request received for prompt: {payload.prompt[:50]}...")

    # --- Authentication ---
    if payload.secret_key != API_SECRET_KEY:
        print("API Authentication failed: Invalid secret key.")
        raise HTTPException(status_code=403, detail="Invalid secret key")

    # Check if HF client is available
    if client is None:
        print("API Error: AI Client not initialized.")
        raise HTTPException(status_code=503, detail="AI Model Client not initialized. Service unavailable.")

    print("API Authentication successful.")

    # --- Call the generator function and consume it ---
    final_code = ""
    try:
        # Consume the generator to get the last yielded value (the complete code)
        code_generator = generate_code(
            prompt=payload.prompt,
            backend_choice=payload.backend_choice,
            file_structure=payload.file_structure,
            max_tokens=payload.max_tokens,
            temperature=payload.temperature,
            top_p=payload.top_p,
        )
        for code_chunk in code_generator:
            final_code = code_chunk # Keep overwriting until the last one

        # Check if the final result indicates an error from within generate_code
        if final_code.strip().startswith("## Error"):
             print(f"API Error during generation: {final_code}")
             # Extract reason if possible, otherwise return generic error
             reason = final_code.split("Reason:**")[-1].strip() if "Reason:**" in final_code else "Generation failed internally."
             raise HTTPException(status_code=500, detail=f"Code generation failed: {reason}")

        print(f"API generated code length: {len(final_code)}")
        # --- Return the final code ---
        return JSONResponse(content={"generated_code": final_code})

    except HTTPException as http_exc:
        # Re-raise HTTPException if it's already one (like auth failure or internal error)
        raise http_exc
    except Exception as e:
        # Catch any other unexpected errors during generation/consumption
        print(f"API - Unexpected Error during generation: {e}")
        raise HTTPException(status_code=500, detail=f"An unexpected error occurred: {str(e)}")


# --- Mount the Gradio app onto the FastAPI app ---
# The Gradio UI will be available at the root path "/"
app = gr.mount_gradio_app(app, demo, path="/")

# --- Launch ---
# Use Uvicorn to run the FastAPI app (which now includes Gradio)
if __name__ == "__main__":
    import uvicorn
    print("Starting FastAPI server with Gradio mounted...")
    # Recommended settings for Hugging Face Spaces: host="0.0.0.0", port=7860
    # You might need to adjust port if running locally and 7860 is taken.
    uvicorn.run(app, host="0.0.0.0", port=7860)
    # Note: demo.launch() is no longer used directly here, as uvicorn runs the combined app.