Agents-MCP-Hackathon/Website_Generator
19
1import gradio as gr2import base643import io4import json5from PIL import Image6import httpx7import html8from lzstring import LZString9import os10from openai import OpenAI11 12# --- 默认API配置 ---13DEFAULT_NEBIUS_API_KEY = "eyJhbGciOiJIUzI1NiIsImtpZCI6IlV6SXJWd1h0dnprLVRvdzlLZWstc0M1akptWXBvX1VaVkxUZlpnMDRlOFUiLCJ0eXAiOiJKV1QifQ.eyJzdWIiOiJnb29nbGUtb2F1dGgyfDEwNTA1MTQzMDg2MDMwMzIxNDEwMiIsInNjb3BlIjoib3BlbmlkIG9mZmxpbmVfYWNjZXNzIiwiaXNzIjoiYXBpX2tleV9pc3N1ZXIiLCJhdWQiOlsiaHR0cHM6Ly9uZWJpdXMtaW5mZXJlbmNlLmV1LmF1dGgwLmNvbS9hcGkvdjIvIl0sImV4cCI6MTkwNjU5ODA0NCwidXVpZCI6ImNkOGFiMWZlLTIxN2QtNDJlMy04OWUwLWM1YTg4MjcwMGVhNyIsIm5hbWUiOiJodW5nZ2luZyIsImV4cGlyZXNfYXQiOiIyMDMwLTA2LTAyVDAyOjM0OjA0KzAwMDAifQ.MA52QuIiNruK7_lX688RXAEI2TkcCOjcf_02XrpnhI8"14NEBIUS_BASE_URL = "https://api.studio.nebius.com/v1/"15 16# --- 核心工具函数 ---17def analyze_image(image: Image.Image, nebius_api_key: str = "") -> str:18 """19 Analyze an uploaded image and provide a detailed description of its content and layout.20 21 Args:22 image: The PIL Image object to analyze23 nebius_api_key: Nebius API key for image analysis24 25 Returns:26 A detailed description of the image content, layout, and website type27 """28 if image is None:29 return "Error: No image provided"30 31 # 使用提供的API密钥或默认密钥32 api_key = nebius_api_key.strip() if nebius_api_key.strip() else DEFAULT_NEBIUS_API_KEY33 34 if not api_key:35 return "Error: Nebius API key not provided"36 37 try:38 # 配置Nebius OpenAI客户端39 client = OpenAI(40 base_url=NEBIUS_BASE_URL,41 api_key=api_key,42 )43 44 # 转换图片为base6445 buffered = io.BytesIO()46 image.save(buffered, format="PNG")47 img_str = base64.b64encode(buffered.getvalue()).decode("utf-8")48 49 # 创建提示50 prompt = """51 Analyze this image and provide a concise description.52 Describe the main elements, colors, layout, and UI components.53 Identify what type of website or application this resembles.54 Focus on structural and visual elements that would be important for recreating the design.55 """56 57 # 使用Qwen2.5-VL-72B-Instruct模型进行图像分析58 response = client.chat.completions.create(59 model="Qwen/Qwen2.5-VL-72B-Instruct",60 messages=[61 {62 "role": "user",63 "content": [64 {"type": "text", "text": prompt},65 {66 "type": "image_url",67 "image_url": {68 "url": f"data:image/png;base64,{img_str}"69 }70 }71 ]72 }73 ],74 max_tokens=1000,75 temperature=0.776 )77 78 return response.choices[0].message.content79 80 except Exception as e:81 return f"Error analyzing image: {str(e)}"82 83def generate_html_code(description: str, nebius_api_key: str = "") -> str:84 """85 Generate HTML/CSS/JavaScript code based on a website description.86 87 Args:88 description: Detailed description of the website to generate89 nebius_api_key: Nebius API key for code generation90 91 Returns:92 Complete HTML code with embedded CSS and JavaScript93 """94 if not description or description.startswith("Error"):95 return "Error: Invalid or missing description"96 97 # 使用提供的API密钥或默认密钥98 api_key = nebius_api_key.strip() if nebius_api_key.strip() else DEFAULT_NEBIUS_API_KEY99 100 if not api_key:101 return "Error: Nebius API key not provided"102 103 prompt = f"""104 Generate a complete, responsive webpage based on this description:105 106 {description}107 108 Requirements:109 - Use modern HTML5, CSS3, and vanilla JavaScript only110 - Include TailwindCSS via CDN for styling111 - Make it responsive and visually appealing112 - Use placeholder images from https://unsplash.com/ if needed113 - Include proper semantic HTML structure114 - Add interactive elements where appropriate115 - Ensure the design matches the described layout and style116 117 Return only the complete HTML code starting with <!DOCTYPE html> and ending with </html>.118 """119 120 try:121 # 配置Nebius OpenAI客户端122 client = OpenAI(123 base_url=NEBIUS_BASE_URL,124 api_key=api_key,125 )126 127 # 使用DeepSeek-V3-0324模型进行代码生成128 response = client.chat.completions.create(129 model="deepseek-ai/DeepSeek-V3-0324",130 messages=[131 {"role": "user", "content": prompt}132 ],133 max_tokens=8000,134 temperature=0.7135 )136 137 html_code = response.choices[0].message.content138 139 # 清理代码格式140 if html_code.strip().startswith("```html"):141 html_code = html_code.split("```html", 1)[1].strip()142 if html_code.strip().endswith("```"):143 html_code = html_code.rsplit("```", 1)[0].strip()144 145 # 确保代码完整性146 if "<!DOCTYPE html>" in html_code and "</html>" in html_code:147 start = html_code.find("<!DOCTYPE html>")148 end = html_code.rfind("</html>") + 7149 return html_code[start:end]150 else:151 return html_code152 153 except Exception as e:154 return f"Error generating HTML code: {str(e)}"155 156def create_codesandbox(html_code: str) -> str:157 """158 Create a CodeSandbox project from HTML code.159 160 Args:161 html_code: Complete HTML code to upload to CodeSandbox162 163 Returns:164 CodeSandbox URL or error message165 """166 if not html_code or html_code.startswith("Error"):167 return "Error: No valid HTML code provided"168 169 try:170 # 准备文件结构171 files = {172 "index.html": {173 "content": html_code,174 "isBinary": False175 },176 "package.json": {177 "content": json.dumps({178 "name": "ai-generated-website",179 "version": "1.0.0",180 "description": "Website generated from image analysis",181 "main": "index.html",182 "scripts": {183 "start": "serve .",184 "build": "echo 'No build required'"185 },186 "devDependencies": {187 "serve": "^14.0.0"188 }189 }, indent=2),190 "isBinary": False191 }192 }193 194 # 准备参数195 parameters = {196 "files": files,197 "template": "static"198 }199 200 # 压缩数据201 json_str = json.dumps(parameters, separators=(',', ':'))202 lz = LZString()203 compressed = lz.compressToBase64(json_str)204 compressed = compressed.replace('+', '-').replace('/', '_').rstrip('=')205 206 # 生成URL207 codesandbox_url = f"https://codesandbox.io/api/v1/sandboxes/define?parameters={compressed}"208 209 # 尝试创建sandbox210 import requests211 212 # 尝试POST请求创建sandbox213 try:214 response = requests.post(215 "https://codesandbox.io/api/v1/sandboxes/define",216 json=parameters,217 timeout=10218 )219 except:220 # 如果POST失败，返回GET URL221 return codesandbox_url222 223 if response.status_code == 200:224 result = response.json()225 sandbox_id = result.get("sandbox_id")226 if sandbox_id:227 return f"https://codesandbox.io/s/{sandbox_id}"228 229 # 如果POST失败，返回GET URL230 return codesandbox_url231 232 except Exception as e:233 return f"Error creating CodeSandbox: {str(e)}"234 235def screenshot_to_code(image: Image.Image, nebius_api_key: str = "") -> tuple:236 """237 Complete pipeline: analyze image and generate corresponding HTML code.238 239 Args:240 image: Screenshot image to analyze241 nebius_api_key: Nebius API key for both image analysis and code generation242 243 Returns:244 Tuple of (description, html_code)245 """246 # 分析图片247 description = analyze_image(image, nebius_api_key)248 249 if description.startswith("Error"):250 return description, "Error: Cannot generate code due to image analysis failure"251 252 # 生成代码253 html_code = generate_html_code(description, nebius_api_key)254 255 return description, html_code256 257# --- Gradio界面 ---258with gr.Blocks(259 theme=gr.themes.Soft(),260 title="AI Website Generator - MCP Compatible (Nebius)",261 css="""262 .api-section { background: #f8f9fa; padding: 15px; border-radius: 8px; margin: 10px 0; }263 .tool-section { border: 1px solid #e0e0e0; padding: 15px; border-radius: 8px; margin: 10px 0; }264 """265) as app:266 267 gr.Markdown("""268 # 🚀 AI Website Generator (MCP Compatible - Nebius)269 270 Transform website screenshots into functional HTML code using Nebius AI.271 272 **Features:**273 - 📸 Image analysis with Qwen2.5-VL-72B-Instruct274 - 💻 HTML/CSS/JS code generation with DeepSeek-V3-0324275 - 🌐 Direct CodeSandbox deployment276 - 🔧 MCP (Model Context Protocol) compatible277 278 **Tools Available:**279 - `analyze_image`: Analyze website screenshots280 - `generate_html_code`: Generate HTML from descriptions281 - `create_codesandbox`: Deploy to CodeSandbox282 - `screenshot_to_code`: Complete pipeline283 """)284 285 with gr.Tab("🎯 Quick Generate"):286 with gr.Row():287 with gr.Column(scale=1):288 gr.Markdown("### 📤 Input", elem_classes=["tool-section"])289 290 # API配置291 with gr.Group():292 gr.Markdown("**Nebius API Key (Required)**")293 nebius_key = gr.Textbox(294 label="Nebius API Key",295 type="password",296 placeholder="Enter your Nebius API key",297 value=DEFAULT_NEBIUS_API_KEY298 )299 300 # 图片上传301 image_input = gr.Image(302 type="pil",303 label="Upload Website Screenshot",304 sources=["upload", "clipboard"]305 )306 307 generate_btn = gr.Button(308 "🎨 Generate Website",309 variant="primary",310 size="lg"311 )312 313 with gr.Column(scale=2):314 gr.Markdown("### 📋 Results", elem_classes=["tool-section"])315 316 description_output = gr.Textbox(317 label="📝 Image Analysis",318 lines=6,319 interactive=False320 )321 322 html_output = gr.Code(323 label="💻 Generated HTML Code",324 language="html",325 lines=15326 )327 328 with gr.Row():329 codesandbox_btn = gr.Button("🚀 Deploy to CodeSandbox")330 codesandbox_output = gr.Textbox(331 label="CodeSandbox URL",332 interactive=False333 )334 335 with gr.Tab("🔧 Individual Tools"):336 gr.Markdown("### Use individual MCP tools")337 338 with gr.Row():339 with gr.Column():340 gr.Markdown("#### 📸 Image Analysis Tool")341 img_tool = gr.Image(type="pil", label="Image")342 nebius_key_tool = gr.Textbox(label="Nebius API Key", type="password", value=DEFAULT_NEBIUS_API_KEY)343 analyze_btn = gr.Button("Analyze Image")344 analysis_result = gr.Textbox(label="Analysis Result", lines=5)345 346 with gr.Column():347 gr.Markdown("#### 💻 Code Generation Tool")348 desc_input = gr.Textbox(label="Description", lines=3)349 nebius_key_tool2 = gr.Textbox(label="Nebius API Key", type="password", value=DEFAULT_NEBIUS_API_KEY)350 code_btn = gr.Button("Generate Code")351 code_result = gr.Code(label="Generated Code", language="html")352 353 # 事件绑定354 generate_btn.click(355 fn=screenshot_to_code,356 inputs=[image_input, nebius_key],357 outputs=[description_output, html_output]358 )359 360 codesandbox_btn.click(361 fn=create_codesandbox,362 inputs=[html_output],363 outputs=[codesandbox_output]364 )365 366 analyze_btn.click(367 fn=analyze_image,368 inputs=[img_tool, nebius_key_tool],369 outputs=[analysis_result]370 )371 372 code_btn.click(373 fn=generate_html_code,374 inputs=[desc_input, nebius_key_tool2],375 outputs=[code_result]376 )377 378 # 示例379 gr.Examples(380 examples=[["1.jpg"]],381 inputs=[image_input],382 label="📷 Example Screenshots"383 )384 385if __name__ == "__main__":386 app.launch(mcp_server=True, share=False)