Spaces:
Running
Running
| 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) |