tab_code.py

import gradio as gr
import subprocess
import threading
import queue
import uuid
import os
import tempfile
import sys
import logging
import time
from models import generate_code_for_tab

# 配置日志
logger = logging.getLogger(__name__)

# 用于存储当前运行的 Gradio 子进程
running_processes = {}

def stop_process(session_id):
    """停止与特定会话关联的子进程"""
    process = running_processes.get(session_id)
    if process and process.poll() is None:
        process.terminate()
        process.wait()
        print(f"Terminated process for session {session_id}")
    if session_id in running_processes:
        del running_processes[session_id]

def get_gradio_sys_prompt():
    """获取用于生成 Gradio 应用的 System Prompt"""
    return """
You are an expert Gradio developer. Create a complete, runnable, single-file Gradio application based on the user's request.
The code must be self-contained in a single Python script.
The script must end with the app launch command, like `demo.launch()`.
Do not include any explanations, just the raw Python code.
"""

def get_html_sys_prompt():
    """获取用于生成静态页面的 System Prompt"""
    return """
You are an expert front-end developer. Create a complete, modern, and responsive single HTML file based on the user's request.
The file must be self-contained, including all necessary HTML, CSS, and JavaScript.
Do not include any explanations, just the raw HTML code.
"""

def run_gradio_in_thread(code, url_queue, session_id):
    """在单独的线程中运行Gradio应用，以避免阻塞主应用"""
    temp_dir = tempfile.mkdtemp()
    file_path = os.path.join(temp_dir, "app.py")
    with open(file_path, "w") as f:
        f.write(code)

    python_executable = sys.executable
    process = subprocess.Popen(
        [python_executable, file_path],
        stdout=subprocess.PIPE,
        stderr=subprocess.PIPE,
        text=True,
        bufsize=1,
        universal_newlines=True
    )
    running_processes[session_id] = process

    for line in process.stdout:
        print(f"Gradio App stdout: {line.strip()}")
        if "Running on local URL:" in line:
            url = line.split("Running on local URL:")[1].strip()
            url_queue.put(url)
            break
    
    process.wait()
    try:
        os.remove(file_path)
        os.rmdir(temp_dir)
    except OSError as e:
        print(f"Error cleaning up temp files: {e}")

def get_spinner_html():
    """返回带 CSS 旋转动画的 HTML"""
    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 generate_code(code_type, model_choice, user_prompt, session_id: gr.State):
    """生成代码并根据类型决定如何展示"""
    logger.info(f"--- [Code Generation] Start ---")
    logger.info(f"Code Type: {code_type}, Model: {model_choice}, Prompt: '{user_prompt}'")
    
    stop_process(session_id)

    if not user_prompt:
        yield "Please enter a prompt.", gr.HTML("Preview will appear here.")
        return

    if code_type == "静态页面":
        system_prompt = get_html_sys_prompt()
        full_code_with_think = ""
        full_code_for_preview = ""
        buffer = ""
        is_thinking = False
        last_update_time = 0
        
        yield "", gr.HTML(get_spinner_html())

        # The model's raw output is streamed here
        for code_chunk in generate_code_for_tab(system_prompt, user_prompt, code_type, model_choice):
            full_code_with_think += code_chunk
            buffer += code_chunk
            
            # Process the buffer to filter out think tags for the preview
            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
            
            current_time = time.time()
            if current_time - last_update_time >= 5:
                escaped_code = full_code_for_preview.replace("'", "&apos;").replace('"', '&quot;')
                preview_html = f"""
                <div style="width: 100%; height: 600px; border: 1px solid #ddd; overflow: hidden; position: relative; background-color: #f9f9f9;">
                    <div style="position: absolute; top: 10px; right: 10px; z-index: 10;">
                        <div class="spinner-small"></div>
                    </div>
                    <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>
                <style>.spinner-small {{ border: 2px solid rgba(0,0,0,0.1); width: 18px; height: 18px; border-radius: 50%; border-left-color: #09f; animation: spin 1s ease infinite; }} @keyframes spin {{ 0% {{ transform: rotate(0deg); }} 100% {{ transform: rotate(360deg); }} }}</style>
                """
                yield full_code_with_think, gr.HTML(preview_html)
                last_update_time = current_time
            else:
                yield full_code_with_think, gr.update()
        
        # Final update for the preview without the spinner
        escaped_code = full_code_for_preview.replace("'", "&apos;").replace('"', '&quot;')
        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>
        """
        yield full_code_with_think, gr.HTML(final_preview_html)
        logger.info("Static page streaming finished.")

def toggle_fullscreen(is_fullscreen):
    """切换全屏模式的可见性"""
    is_fullscreen = not is_fullscreen
    new_button_text = "退出全屏" if is_fullscreen else "全屏预览"
    panel_visibility = not is_fullscreen
    return is_fullscreen, gr.update(value=new_button_text), gr.update(visible=panel_visibility)

def create_code_tab():
    """创建代码生成功能的UI Tab"""
    session_id = str(uuid.uuid4())
    session_state = gr.State(session_id)
    fullscreen_state = gr.State(False)

    with gr.Blocks() as demo:
        with gr.Row():
            with gr.Column(scale=1) as left_panel:
                gr.Markdown("### 1. 选择代码类型")
                code_type_radio = gr.Radio(["静态页面", "Gradio 应用"], value="静态页面", label="Code Type")
                
                gr.Markdown("### 2. 选择模型")
                model_choice_radio = gr.Radio(
                    ["效果更好 (使用 Ling-1T)", "更快速 (使用 Ring-flash-2.0)"],
                    value="效果更好 (使用 Ling-1T)",
                    label="Model Selection"
                )

                gr.Markdown("### 3. 输入你的需求")
                prompt_input = gr.Textbox(lines=5, placeholder="例如：创建一个带有标题和按钮的简单页面", label="Prompt")
                gr.Examples(
                    examples=[
                        "创建一个在黑色背景上不断绽放五彩烟花的 Canvas 动画。",
                        "生成一个具有流光溢彩效果的 Canvas 特效。",
                        "设计一个能与鼠标交互的粒子系统 Canvas 动画。",
                        "用 HTML Canvas 实现一个经典的贪吃蛇游戏。",
                        "创建一个模拟'大象牙膏'化学实验的 Canvas 动画：一个容器中，彩色泡沫不断快速涌出、膨胀、溢出，充满整个屏幕。",
                        "创建一个梦幻的低多边形漂浮岛屿场景，带有动态光照和柔和的动画，在一个单一的HTML文件中。使用 d3.js 。"
                    ],
                    inputs=prompt_input,
                    label="✨ 不妨试试这些酷炫的例子"
                )
                generate_button = gr.Button("生成代码", variant="primary")
            
            with gr.Column(scale=2):
                with gr.Tabs():
                    with gr.TabItem("实时预览"):
                        with gr.Row():
                            gr.Markdown("### 3. 实时预览")
                            fullscreen_button = gr.Button("全屏预览", scale=0)
                        preview_output = gr.HTML(value="<p>Preview will appear here.</p>")
                    with gr.TabItem("生成的源代码"):
                        gr.Markdown("### 4. 生成的源代码")
                        code_output = gr.Code(language="html", label="Generated Code")

        generate_button.click(
            fn=generate_code,
            inputs=[code_type_radio, model_choice_radio, prompt_input, session_state],
            outputs=[code_output, preview_output]
        )
        
        fullscreen_button.click(
            fn=toggle_fullscreen,
            inputs=[fullscreen_state],
            outputs=[fullscreen_state, fullscreen_button, left_panel]
        )

        demo.unload(fn=lambda: stop_process(session_id))

    return demo