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 """
"""
def inject_error_catcher(html_code):
"""
Injects the error catching script into the HTML code.
It tries to insert it after the tag.
If is not found, it prepends it to the code (fallback).
"""
if not CATCH_ERROR_JS_SCRIPT:
return html_code
script_tag = f""
# Try to find (case insensitive)
head_match = re.search(r"", html_code, re.IGNORECASE)
if head_match:
# Insert after
insert_pos = head_match.end()
return html_code[:insert_pos] + script_tag + html_code[insert_pos:]
# Try to find (case insensitive)
body_match = re.search(r"", html_code, re.IGNORECASE)
if body_match:
# Insert before
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="预览将在此处显示。
"), 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"""This will create an error
"""
escaped_code = error_code.replace("'", "'").replace('"', """)
final_preview_html = f"""
"""
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("")
if end_index != -1:
is_thinking = False
buffer = buffer[end_index + len(""):]
else:
break
else:
start_index = buffer.find("")
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"""
"""
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.")