import gradio as gr
from huggingface_hub import InferenceClient
import os
# --- Configuration ---
# Use environment variable for token, fallback to default if not set
# For Spaces, set the HF_TOKEN secret
API_TOKEN = os.getenv("HF_TOKEN", None)
MODEL = "HuggingFaceH4/zephyr-7b-beta" # Or choose another suitable model
# --- 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)
# Optional: Add a quick health check if needed, though client init usually suffices
# client.get_model_status(MODEL) # Example check, might raise if model invalid
print("Inference Client initialized successfully.")
except Exception as e:
print(f"Error initializing Inference Client: {e}")
# Provide a more informative error in the Gradio interface
raise gr.Error(f"Failed to initialize the AI model client for '{MODEL}'. Please check the model name and ensure network connectivity. If running locally without a token, you might hit rate limits. If using HF Spaces, ensure the HF_TOKEN secret is set correctly. Original Error: {e}")
# --- Core Code Generation Function ---
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.
Yields the code token by token for live updates.
"""
print(f"--- Generating Code ---")
print(f"Prompt: {prompt[:100]}...") # Log truncated prompt
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}")
# --- System Message (Internal) ---
# Guides the AI's behavior. Not user-editable in the UI.
# Refined system prompt for clarity and stricter output formatting
system_message = (
"You are an expert frontend web developer AI. Your task is to generate HTML, CSS, and potentially JavaScript code "
"for a website based ONLY on the user's description. Adhere strictly to the following constraints:\n"
"1. **Output ONLY Code:** Generate only the raw code for the requested files (`index.html`, `style.css`, `script.js`). Do NOT include any introductory text, explanations, apologies, markdown formatting (like ```html), or closing remarks. Your response must start *immediately* with the code (e.g., `` or ``).\n"
"2. **index.html is Mandatory:** ALWAYS generate a complete `index.html` file.\n"
"3. **File Structure:**\n"
f" - If '{file_structure}' is 'Multiple Files':\n"
" - Use clear markers EXACTLY as follows:\n"
" ``\n"
" `/* style.css */`\n"
" `// script.js` (only include if JavaScript is necessary for the described functionality)\n"
" - Place the corresponding code directly after each marker.\n"
" - Link the CSS (``) in the `` of `index.html`.\n"
" - Include the JS (``) just before the closing `` tag in `index.html` if `script.js` is generated.\n"
f" - If '{file_structure}' is 'Single File':\n"
" - Embed ALL CSS within `