Spaces:
Running
Running
| import gradio as gr | |
| import time | |
| import logging | |
| import os | |
| import re | |
| from model_handler import ModelHandler | |
| from config import CODE_FRAMEWORK_SPECS, STATIC_PAGE | |
| from tab_code_prompts.framework_system_prompts import get_framework_system_prompt | |
| from code_kit.code_trimmer import trim_html_markdown_blocks # Import the trimmer | |
| # Configure logging | |
| logger = logging.getLogger(__name__) | |
| # Read the content of the JavaScript file for error catching | |
| try: | |
| with open("static/catch-error.js", "r", encoding="utf-8") as f: | |
| CATCH_ERROR_JS_SCRIPT = f.read() | |
| except FileNotFoundError: | |
| logger.error("Error: static/catch-error.js not found. The error catching overlay will not work.") | |
| CATCH_ERROR_JS_SCRIPT = "" | |
| def get_spinner_html(): | |
| """Return HTML with a CSS spinner animation""" | |
| return """ | |
| <div style="width: 100%; height: 600px; display: flex; justify-content: center; align-items: center; border: 1px solid #ddd; background-color: #f9f9f9;"> | |
| <div class="spinner"></div> | |
| </div> | |
| <style> | |
| .spinner { | |
| border: 4px solid rgba(0, 0, 0, 0.1); | |
| width: 36px; | |
| height: 36px; | |
| border-radius: 50%; | |
| border-left-color: #09f; | |
| animation: spin 1s ease infinite; | |
| } | |
| @keyframes spin { | |
| 0% { transform: rotate(0deg); } | |
| 100% { transform: rotate(360deg); } | |
| } | |
| </style> | |
| """ | |
| def inject_error_catcher(html_code): | |
| """ | |
| Injects the error catching script into the HTML code. | |
| It tries to insert it after the <head> tag. | |
| If <head> is not found, it prepends it to the code (fallback). | |
| """ | |
| if not CATCH_ERROR_JS_SCRIPT: | |
| return html_code | |
| script_tag = f"<script>{CATCH_ERROR_JS_SCRIPT}</script>" | |
| # Try to find <head> (case insensitive) | |
| head_match = re.search(r"<head>", html_code, re.IGNORECASE) | |
| if head_match: | |
| # Insert after <head> | |
| insert_pos = head_match.end() | |
| return html_code[:insert_pos] + script_tag + html_code[insert_pos:] | |
| # Try to find <body> (case insensitive) | |
| body_match = re.search(r"<body>", html_code, re.IGNORECASE) | |
| if body_match: | |
| # Insert before <body> | |
| insert_pos = body_match.start() | |
| return html_code[:insert_pos] + script_tag + html_code[insert_pos:] | |
| # Fallback: Just prepend if it looks somewhat like HTML or empty | |
| return script_tag + html_code | |
| def code_generation_agent(code_type_display_name, model_choice, user_prompt, color_palette, decoration_style, overall_style, chatbot_history): | |
| """Generate code and provide a preview, updating a log stream chatbot.""" | |
| logger.info(f"--- [Code Generation] Start ---") | |
| logger.info(f"Code Type (Display Name): {code_type_display_name}, Model: {model_choice}, Prompt: '{user_prompt}'") | |
| # Map display name back to constant key | |
| code_framework_key = next((k for k, v in CODE_FRAMEWORK_SPECS.items() if v["display_name"] == code_type_display_name), STATIC_PAGE) | |
| logger.info(f"Mapped Code Framework Key: {code_framework_key}") | |
| if not user_prompt: | |
| chatbot_history.append({"role": "assistant", "content": "🚨 **错误**: 请输入提示词。"}) | |
| yield "", gr.update(value="<p>预览将在此处显示。</p>"), chatbot_history, gr.update() | |
| return | |
| chatbot_history.append({"role": "assistant", "content": "⏳ 开始生成代码..."}) | |
| yield "", gr.HTML(get_spinner_html()), chatbot_history, gr.update() | |
| if user_prompt == "create an error" or user_prompt == "创建一个报错示例": | |
| error_code = f"""<h1>This will create an error</h1><script>{CATCH_ERROR_JS_SCRIPT}</script><script>nonExistentFunction();</script>""" | |
| escaped_code = error_code.replace("'", "'").replace('"', """) | |
| final_preview_html = f""" | |
| <div style="width: 100%; height: 600px; border: 1px solid #ddd; overflow: hidden; position: relative; background-color: #f9f9f9;"> | |
| <iframe srcdoc='{escaped_code}' | |
| style="position: absolute; top: 0; left: 0; width: 200%; height: 200%; transform: scale(0.5); transform-origin: 0 0; border: none;"> | |
| </iframe> | |
| </div> | |
| """ | |
| chatbot_history.append({"role": "assistant", "content": "✅ **成功**: 已生成一个用于测试的错误页面。"}) | |
| yield error_code, gr.update(value=final_preview_html), chatbot_history, gr.Tabs(selected=0) | |
| return | |
| # --- Append Style Prompt --- | |
| full_user_prompt = user_prompt | |
| if color_palette or decoration_style or overall_style: | |
| full_user_prompt += "\n\n--- Visual Style Requirements (Strictly Follow These) ---\n" | |
| if overall_style: | |
| full_user_prompt += f"Overall Theme/Vibe: {overall_style}\n" | |
| if decoration_style: | |
| full_user_prompt += f"UI Decoration Style: {decoration_style}\n" | |
| if color_palette: | |
| full_user_prompt += f"Color Palette Mapping (Use these colors for their assigned roles): {color_palette}\n" | |
| logger.info(f"Full Prompt with Style: {full_user_prompt}") | |
| # --------------------------- | |
| start_time = time.time() | |
| model_handler = ModelHandler() | |
| # Get the system prompt based on the selected framework AND specific color palette | |
| # Passing color_palette allows it to be baked into the system prompt instructions | |
| system_prompt = get_framework_system_prompt(code_framework_key, color_palette) | |
| full_code_with_think = "" | |
| full_code_for_preview = "" | |
| buffer = "" | |
| is_thinking = False | |
| for code_chunk in model_handler.generate_code(system_prompt, full_user_prompt, model_choice): | |
| full_code_with_think += code_chunk | |
| buffer += code_chunk | |
| while True: | |
| if is_thinking: | |
| end_index = buffer.find("</think>") | |
| if end_index != -1: | |
| is_thinking = False | |
| buffer = buffer[end_index + len("</think>"):] | |
| else: | |
| break | |
| else: | |
| start_index = buffer.find("<think>") | |
| if start_index != -1: | |
| part_to_add = buffer[:start_index] | |
| full_code_for_preview += part_to_add | |
| is_thinking = True | |
| buffer = buffer[start_index:] | |
| else: | |
| full_code_for_preview += buffer | |
| buffer = "" | |
| break | |
| elapsed_time = time.time() - start_time | |
| generated_length = len(full_code_with_think) | |
| speed = generated_length / elapsed_time if elapsed_time > 0 else 0 | |
| log_message = f""" | |
| **⏳ 正在生成中...** | |
| - **时间:** {elapsed_time:.2f}s | |
| - **长度:** {generated_length} chars | |
| - **速度:** {speed:.2f} char/s | |
| """ | |
| if len(chatbot_history) > 0 and "正在生成中" in chatbot_history[-1]["content"]: | |
| chatbot_history[-1] = {"role": "assistant", "content": log_message} | |
| else: | |
| chatbot_history.append({"role": "assistant", "content": log_message}) | |
| yield full_code_with_think, gr.update(), chatbot_history, gr.update() | |
| # Apply trimming to the final generated code | |
| trimmed_full_code_with_think = trim_html_markdown_blocks(full_code_with_think) | |
| trimmed_full_code_for_preview = trim_html_markdown_blocks(full_code_for_preview) | |
| # Inject error catching script before final render | |
| final_code_with_error_catcher = inject_error_catcher(trimmed_full_code_for_preview) | |
| escaped_code = final_code_with_error_catcher.replace("'", "'").replace('"', """) | |
| final_preview_html = f""" | |
| <div style="width: 100%; height: 600px; border: 1px solid #ddd; overflow: hidden; position: relative; background-color: #f9f9f9;"> | |
| <iframe srcdoc='{escaped_code}' | |
| style="position: absolute; top: 0; left: 0; width: 200%; height: 200%; transform: scale(0.5); transform-origin: 0 0; border: none;"> | |
| </iframe> | |
| </div> | |
| """ | |
| chatbot_history.append({"role": "assistant", "content": "✅ **成功**: 代码生成完成！"}) | |
| # Yield the trimmed code for display in the "Generated Source Code" tab | |
| yield trimmed_full_code_with_think, gr.HTML(final_preview_html), chatbot_history, gr.Tabs(selected=0) | |
| logger.info("Code generation streaming finished.") |