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 `