Spaces:
Running
Running
File size: 8,431 Bytes
3f80c6a b453cca 3f80c6a b453cca 3f80c6a b453cca 3f80c6a b453cca 3f80c6a b453cca 3f80c6a b453cca 3f80c6a b453cca 3f80c6a b453cca 3f80c6a b453cca 3f80c6a b453cca 3f80c6a b453cca |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 |
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.") |