Spaces:

MINEOGO
/

pro-zephyr-coder

Running

File size: 13,154 Bytes

import gradio as gr
from huggingface_hub import InferenceClient
import re
import time # For potential brief pauses if needed

# --- Hugging Face Token (Optional but Recommended) ---
# from huggingface_hub import login
# login("YOUR_HUGGINGFACE_TOKEN") # Replace with your token if needed

# --- Inference Client ---
try:
    # You might need to specify the model URL directly if the alias isn't working
    # client = InferenceClient(model="https://api-inference.huggingface.co/models/HuggingFaceH4/zephyr-7b-beta")
    client = InferenceClient("HuggingFaceH4/zephyr-7b-beta")
    client.timeout = 120 # Increase timeout for potentially long generations
except Exception as e:
    print(f"Error initializing InferenceClient: {e}")
    client = None # Set client to None if initialization fails

# --- Parsing Function ---
def parse_files(raw_response):
    """
    Parses filenames and code blocks from the raw AI output.
    """
    if not raw_response:
        return []

    # Pattern: Look for a filename line followed by content until the next filename line or end of string.
    pattern = re.compile(
        r"^\s*([\w\-.\/\\]+\.\w+)\s*\n"  # Filename line (must have an extension)
        r"(.*?)"                        # Capture content (non-greedy)
        r"(?=\n\s*[\w\-.\/\\]+\.\w+\s*\n|\Z)", # Lookahead for next filename or end of string
        re.DOTALL | re.MULTILINE
    )
    files = pattern.findall(raw_response)

    cleaned_files = []
    for name, content in files:
        # Remove common code block markers (``` optionally followed by lang)
        content_cleaned = re.sub(r"^\s*```[a-zA-Z]*\n?", "", content, flags=re.MULTILINE)
        content_cleaned = re.sub(r"\n?```\s*$", "", content_cleaned, flags=re.MULTILINE)
        cleaned_files.append((name.strip(), content_cleaned.strip()))

    # Fallback if no files parsed but content exists
    if not cleaned_files and raw_response.strip():
        if any(c in raw_response for c in ['<','>','{','}',';','(',')']):
            print("Warning: No filenames found, defaulting to index.html")
            lang = "html"
            if "{" in raw_response and "}" in raw_response and ":" in raw_response: lang = "css"
            elif "function" in raw_response or "const" in raw_response or "let" in raw_response: lang = "javascript"
            default_filename = "index.html"
            if lang == "css": default_filename = "style.css"
            elif lang == "javascript": default_filename = "script.js"
            cleaned_files.append((default_filename, raw_response.strip()))

    return cleaned_files

# --- Streaming and Parsing Orchestrator ---
def stream_and_parse_code(prompt, backend, system_message, max_tokens, temperature, top_p):
    """
    Streams raw output to one component and generates final tabs for another.
    This function acts as the main callback for the button click.
    Yields dictionary updates for Gradio components.
    """
    # Check if client initialized correctly
    if not client:
        error_msg = "Error: Inference Client not available. Check API token or model name."
        # Yield updates to both components indicating the error
        yield {
            live_output: gr.update(value=error_msg),
            final_tabs: gr.Tabs(tabs=[gr.TabItem(label="Error", children=[gr.Textbox(value=error_msg)])])
        }
        return # Stop execution

    # --- Prepare for Streaming ---
    # Construct the system prompt dynamically
    full_sys_msg = f"""
You are a code generation AI. Given a prompt, generate the necessary files for a website using the {backend} backend.
Always include an index.html file.
Respond ONLY with filenames and the raw code for each file.
Each file must start with its filename on a new line. Example:

index.html
<!DOCTYPE html>
<html></html>

style.css
body {{}}

script.js
console.log("Hello");

Ensure the code is complete. NO commentary, NO explanations, NO markdown formatting like backticks (```).
Start generating the files now.
""".strip()
    if system_message: # Append user's system message if provided
        full_sys_msg += "\n\n" + system_message

    messages = [
        {"role": "system", "content": full_sys_msg},
        {"role": "user", "content": prompt}
    ]

    full_raw_response = ""
    error_occurred = False
    error_message = ""

    # Initial state update: Clear previous output and show generating status
    yield {
        live_output: gr.update(value="Generating stream..."),
        final_tabs: gr.Tabs(tabs=[gr.TabItem(label="Generating...")]) # Indicate loading in tabs
    }

    # --- Streaming Loop ---
    try:
        # Start the streaming call
        stream = client.chat_completion(
            messages,
            max_tokens=int(max_tokens), # Ensure max_tokens is an integer
            stream=True,
            temperature=temperature,
            top_p=top_p
        )
        # Process each chunk received from the stream
        for chunk in stream:
            token = chunk.choices[0].delta.content
            if token:
                full_raw_response += token
                # Yield updates for the live raw output component only
                # Keep tabs in a 'streaming' state during the stream
                yield {
                    live_output: gr.update(value=full_raw_response),
                    # No update needed for final_tabs here, or keep showing streaming state
                    # final_tabs: gr.Tabs(tabs=[gr.TabItem(label="Streaming...")]) # Optional: update tabs state
                }
                # time.sleep(0.01) # Optional: small delay if updates are too fast and causing UI lag

    except Exception as e:
        # Handle errors during the API call or streaming process
        print(f"Error during AI streaming: {e}")
        error_message = f"Error during AI generation: {e}\n\nPartial Response (if any):\n{full_raw_response}"
        error_occurred = True
        # Update live output with error, prepare error tab
        yield {
            live_output: gr.update(value=error_message),
            final_tabs: gr.Tabs(tabs=[gr.TabItem(label="Error")]) # Indicate error state in tabs
        }

    # --- Post-Streaming: Parsing and Final Tab Generation ---
    if error_occurred:
        # If an error happened during stream, create a final error tab
         final_tabs_update = gr.Tabs(tabs=[
            gr.TabItem(label="Error", children=[gr.Textbox(value=error_message, label="Generation Error", lines=10)])
        ])
    else:
        # If streaming succeeded, parse the complete raw response
        print("\n--- Final Raw AI Response ---")
        print(full_raw_response)
        print("--------------------------\n")
        files = parse_files(full_raw_response)

        if not files:
            # Handle case where parsing failed or AI gave empty/invalid response
            no_files_msg = "AI finished, but did not return recognizable file content or the response was empty. See raw output above."
            final_tabs_update = gr.Tabs(tabs=[
                gr.TabItem(label="Output", children=[gr.Textbox(value=no_files_msg, label="Result")])
            ])
            # Update live output as well to make the message clear
            yield { live_output: gr.update(value=full_raw_response + "\n\n" + no_files_msg), final_tabs: final_tabs_update }
            return # Exit if no files

        # --- Create Tabs (if files were parsed successfully) ---
        tabs_content = []
        for name, content in files:
            name = name.strip()
            content = content.strip()
            # Skip if filename or content is empty after stripping
            if not name or not content:
                print(f"Skipping file with empty name or content: Name='{name}'")
                continue

            # Determine language for syntax highlighting
            lang = "plaintext" # Default
            if name.endswith((".html", ".htm")): lang = "html"
            elif name.endswith(".css"): lang = "css"
            elif name.endswith(".js"): lang = "javascript"
            elif name.endswith(".py"): lang = "python"
            elif name.endswith(".json"): lang = "json"
            elif name.endswith(".md"): lang = "markdown"
            elif name.endswith((".sh", ".bash")): lang = "bash"
            elif name.endswith((".xml", ".xaml", ".svg")): lang = "xml"
            elif name.endswith(".yaml") or name.endswith(".yml"): lang = "yaml"

            # Ensure elem_id is unique and valid (replace problematic characters)
            elem_id = f"tab_{re.sub(r'[^a-zA-Z0-9_-]', '_', name)}"

            tab_item = gr.TabItem(label=name, elem_id=elem_id, children=[
                gr.Code(value=content, language=lang, label=name, interactive=False) # Show code in Code block
            ])
            tabs_content.append(tab_item)

        # Handle case where parsing found files, but they were all filtered out (empty name/content)
        if not tabs_content:
             final_tabs_update = gr.Tabs(tabs=[gr.TabItem(label="Output", children=[gr.Textbox(value="No valid files generated after filtering.", label="Result")])])
        else:
            final_tabs_update = gr.Tabs(tabs=tabs_content) # Create the final Tabs component with content

    # --- Final Update ---
    # Yield the final state for both components
    # Use gr.update for live_output if you only want to set its value without recreating it
    # Directly return the new final_tabs component
    yield {
        live_output: gr.update(value=full_raw_response if not error_occurred else error_message), # Show final raw response or error
        final_tabs: final_tabs_update # Update the final_tabs component completely
    }


# --- Gradio UI Definition ---
with gr.Blocks(css=".gradio-container { max-width: 95% !important; }") as demo: # Use more screen width
    gr.Markdown("## WebGen AI — One Prompt → Full Website Generator")
    gr.Markdown("Generates website code based on your description. Raw output streams live, final files appear in tabs below.")

    with gr.Row():
        # Column for inputs and controls
        with gr.Column(scale=2):
            prompt = gr.Textbox(
                label="Describe your website",
                placeholder="E.g., a simple landing page for a coffee shop with sections for menu, about, and contact.",
                lines=3 # Allow more lines for the prompt
            )
            backend = gr.Dropdown(
                ["Static", "Flask", "Node.js"],
                value="Static",
                label="Backend Technology"
            )
            with gr.Accordion("Advanced Options", open=False):
                system_message = gr.Textbox(
                    label="Extra instructions for the AI (System Message)",
                    placeholder="Optional: e.g., 'Use Tailwind CSS for styling', 'Make it responsive'",
                    value="",
                    lines=2
                 )
                max_tokens = gr.Slider(
                    minimum=256,
                    maximum=4096, # Increased max tokens for complex sites
                    value=2048, # Increased default
                    step=64,
                    label="Max Tokens (Output Length)"
                )
                temperature = gr.Slider(
                    minimum=0.1,
                    maximum=1.5, # Allow slightly higher temperature
                    value=0.7,
                    step=0.1,
                    label="Temperature (Creativity)"
                )
                top_p = gr.Slider(
                    minimum=0.1,
                    maximum=1.0,
                    value=0.95,
                    step=0.05,
                    label="Top-p (Sampling Focus)"
                )
            generate_button = gr.Button("✨ Generate Code ✨", variant="primary") # Make button primary

        # Column for live output
        with gr.Column(scale=3):
            gr.Markdown("#### Live Raw Output Stream")
            # Component to show the live, unparsed stream - CORRECTED LANGUAGE
            live_output = gr.Code(
                label="Raw AI Stream",
                language="plaintext", # Use "plaintext" for generic text
                lines=20, # Increased lines for visibility
                interactive=False # Output only
            )

    gr.Markdown("---") # Separator
    gr.Markdown("#### Final Generated Files (Tabs)")
    # Placeholder for the final structured tabs - will be replaced by the output yield
    final_tabs = gr.Tabs(elem_id="output_tabs")


    # Button click action - uses the orchestrator function
    generate_button.click(
        stream_and_parse_code, # Call the main function that handles streaming and parsing
        inputs=[prompt, backend, system_message, max_tokens, temperature, top_p],
        # Outputs dictionary maps function yields to components by variable name
        outputs=[live_output, final_tabs],
        show_progress="hidden" # Hide default Gradio progress bar as we show live stream
    )

if __name__ == "__main__":
    # Launch the Gradio app with debug=True for development
    demo.launch(debug=True)