FastAPI服务设计

# api/main.py - FastAPI应用主文件
from fastapi import FastAPI, HTTPException, Depends, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.middleware.gzip import GZipMiddleware
from fastapi.responses import JSONResponse, FileResponse
from pydantic import BaseModel, Field
from typing import List, Optional, Dict, Any
import logging
import asyncio
from pathlib import Path
import json
import yaml

# 导入核心服务
from inference.generate_ui import UIGenerator
from render.render_to_image import UIRenderer
from render.render_to_vue import VueRenderer

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

# 创建FastAPI应用
app = FastAPI(
    title="AI UI生成系统API",
    description="基于AI的UI设计和代码生成系统",
    version="1.0.0",
    docs_url="/docs",
    redoc_url="/redoc"
)

# 添加中间件
app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

app.add_middleware(GZipMiddleware, minimum_size=1000)

# 全局变量存储服务实例
ui_generator = None
ui_renderer = None
vue_renderer = None

class GenerateUIRequest(BaseModel):
    """UI生成请求模型"""
    prompt: str = Field(..., description="中文UI描述", min_length=1, max_length=500)
    theme: str = Field(default="obsidian-gold", description="主题名称")
    page_type: str = Field(default="home_page", description="页面类型")
    use_model: bool = Field(default=True, description="是否使用AI模型")
    max_length: int = Field(default=512, description="最大生成长度")
    temperature: float = Field(default=0.7, description="生成温度")
    
    class Config:
        schema_extra = {
            "example": {
                "prompt": "设计一个电商首页，包含轮播图、商品分类、热门商品推荐",
                "theme": "obsidian-gold",
                "page_type": "home_page",
                "use_model": True,
                "max_length": 512,
                "temperature": 0.7
            }
        }

class RenderRequest(BaseModel):
    """渲染请求模型"""
    dsl: Dict[str, Any] = Field(..., description="UI-DSL数据")
    output_format: str = Field(default="both", description="输出格式: png, vue, both")
    theme: str = Field(default="obsidian-gold", description="主题名称")
    optimize: bool = Field(default=True, description="是否优化输出")
    
    class Config:
        schema_extra = {
            "example": {
                "dsl": {
                    "page": {
                        "name": "home_page",
                        "theme": "obsidian-gold",
                        "sections": [
                            {
                                "type": "topbar",
                                "props": {
                                    "title": "首页",
                                    "logo": "Logo",
                                    "actions": ["search", "bell"]
                                }
                            }
                        ]
                    }
                },
                "output_format": "both",
                "theme": "obsidian-gold",
                "optimize": True
            }
        }

class ThemeRequest(BaseModel):
    """主题请求模型"""
    theme_name: str = Field(..., description="主题名称")
    action: str = Field(..., description="操作类型: get, list, preview, apply")
    dsl: Optional[Dict[str, Any]] = Field(default=None, description="UI-DSL数据（apply时必需）")
    
    class Config:
        schema_extra = {
            "example": {
                "theme_name": "obsidian-gold",
                "action": "get",
                "dsl": None
            }
        }

class GenerateUIResponse(BaseModel):
    """UI生成响应模型"""
    success: bool = Field(..., description="是否成功")
    dsl: Optional[Dict[str, Any]] = Field(default=None, description="生成的UI-DSL")
    prompt: str = Field(..., description="原始提示")
    theme: str = Field(..., description="使用的主题")
    page_type: str = Field(..., description="页面类型")
    generation_method: str = Field(..., description="生成方法: model, rules")
    processing_time: float = Field(..., description="处理时间（秒）")
    error_message: Optional[str] = Field(default=None, description="错误信息")
    
    class Config:
        schema_extra = {
            "example": {
                "success": True,
                "dsl": {
                    "page": {
                        "name": "home_page",
                        "theme": "obsidian-gold",
                        "sections": []
                    }
                },
                "prompt": "设计一个电商首页",
                "theme": "obsidian-gold",
                "page_type": "home_page",
                "generation_method": "model",
                "processing_time": 1.23,
                "error_message": None
            }
        }

class RenderResponse(BaseModel):
    """渲染响应模型"""
    success: bool = Field(..., description="是否成功")
    output_format: str = Field(..., description="输出格式")
    theme: str = Field(..., description="使用的主题")
    files: Dict[str, str] = Field(..., description="生成的文件路径")
    processing_time: float = Field(..., description="处理时间（秒）")
    error_message: Optional[str] = Field(default=None, description="错误信息")
    
    class Config:
        schema_extra = {
            "example": {
                "success": True,
                "output_format": "both",
                "theme": "obsidian-gold",
                "files": {
                    "png": "/output/images/ui_design.png",
                    "vue": "/output/vue/ui_design.vue"
                },
                "processing_time": 2.45,
                "error_message": None
            }
        }

class ThemeResponse(BaseModel):
    """主题响应模型"""
    success: bool = Field(..., description="是否成功")
    theme_name: str = Field(..., description="主题名称")
    theme_config: Optional[Dict[str, Any]] = Field(default=None, description="主题配置")
    available_themes: Optional[List[str]] = Field(default=None, description="可用主题列表")
    preview_url: Optional[str] = Field(default=None, description="预览URL")
    processing_time: float = Field(..., description="处理时间（秒）")
    error_message: Optional[str] = Field(default=None, description="错误信息")
    
    class Config:
        schema_extra = {
            "example": {
                "success": True,
                "theme_name": "obsidian-gold",
                "theme_config": {
                    "colors": {
                        "primary": "#FFD700",
                        "secondary": "#B8860B"
                    }
                },
                "available_themes": ["obsidian-gold", "silver-white", "minimal"],
                "preview_url": "/themes/obsidian-gold/preview",
                "processing_time": 0.12,
                "error_message": None
            }
        }

class HealthResponse(BaseModel):
    """健康检查响应模型"""
    status: str = Field(..., description="服务状态")
    version: str = Field(..., description="版本号")
    uptime: float = Field(..., description="运行时间（秒）")
    services: Dict[str, str] = Field(..., description="服务状态")
    timestamp: str = Field(..., description="时间戳")
    
    class Config:
        schema_extra = {
            "example": {
                "status": "healthy",
                "version": "1.0.0",
                "uptime": 3600.0,
                "services": {
                    "ui_generator": "healthy",
                    "ui_renderer": "healthy",
                    "vue_renderer": "healthy"
                },
                "timestamp": "2024-01-01T12:00:00Z"
            }
        }

async def initialize_services():
    """初始化所有服务"""
    global ui_generator, ui_renderer, vue_renderer
    
    try:
        logger.info("正在初始化UI生成器...")
        ui_generator = UIGenerator()
        await ui_generator.load_model()
        logger.info("UI生成器初始化完成")
        
        logger.info("正在初始化UI渲染器...")
        ui_renderer = UIRenderer()
        logger.info("UI渲染器初始化完成")
        
        logger.info("正在初始化Vue渲染器...")
        vue_renderer = VueRenderer()
        logger.info("Vue渲染器初始化完成")
        
        logger.info("所有服务初始化完成")
        
    except Exception as e:
        logger.error(f"服务初始化失败: {e}")
        raise e

@app.on_event("startup")
async def startup_event():
    """应用启动事件"""
    logger.info("应用启动中...")
    await initialize_services()
    logger.info("应用启动完成")

@app.on_event("shutdown")
async def shutdown_event():
    """应用关闭事件"""
    logger.info("应用关闭中...")
    # 清理资源
    logger.info("应用关闭完成")

def get_ui_generator() -> UIGenerator:
    """获取UI生成器实例"""
    if ui_generator is None:
        raise HTTPException(status_code=503, detail="UI生成器未初始化")
    return ui_generator

def get_ui_renderer() -> UIRenderer:
    """获取UI渲染器实例"""
    if ui_renderer is None:
        raise HTTPException(status_code=503, detail="UI渲染器未初始化")
    return ui_renderer

def get_vue_renderer() -> VueRenderer:
    """获取Vue渲染器实例"""
    if vue_renderer is None:
        raise HTTPException(status_code=503, detail="Vue渲染器未初始化")
    return vue_renderer

def get_config() -> Dict[str, Any]:
    """获取配置"""
    config_path = "config/model_config.yaml"
    with open(config_path, 'r', encoding='utf-8') as f:
        return yaml.safe_load(f)

@app.post("/generate-ui", response_model=GenerateUIResponse)
async def generate_ui(
    request: GenerateUIRequest,
    generator: UIGenerator = Depends(get_ui_generator),
    config: Dict[str, Any] = Depends(get_config)
):
    """生成UI设计"""
    start_time = time.time()
    
    try:
        logger.info(f"收到UI生成请求: {request.prompt}")
        
        # 验证主题
        available_themes = config.get("ui", {}).get("themes", [])
        if request.theme not in available_themes:
            raise HTTPException(
                status_code=400, 
                detail=f"不支持的主题: {request.theme}"
            )
        
        # 生成UI-DSL
        if request.use_model:
            dsl = await generator.generate_with_model(
                prompt=request.prompt,
                theme=request.theme,
                page_type=request.page_type,
                max_length=request.max_length,
                temperature=request.temperature
            )
            generation_method = "model"
        else:
            dsl = await generator.generate_with_rules(
                prompt=request.prompt,
                theme=request.theme,
                page_type=request.page_type
            )
            generation_method = "rules"
        
        processing_time = time.time() - start_time
        
        logger.info(f"UI生成完成，耗时: {processing_time:.2f}秒")
        
        return GenerateUIResponse(
            success=True,
            dsl=dsl,
            prompt=request.prompt,
            theme=request.theme,
            page_type=request.page_type,
            generation_method=generation_method,
            processing_time=processing_time,
            error_message=None
        )
        
    except Exception as e:
        processing_time = time.time() - start_time
        error_message = str(e)
        logger.error(f"UI生成失败: {error_message}")
        
        return GenerateUIResponse(
            success=False,
            dsl=None,
            prompt=request.prompt,
            theme=request.theme,
            page_type=request.page_type,
            generation_method="error",
            processing_time=processing_time,
            error_message=error_message
        )

@app.post("/render", response_model=RenderResponse)
async def render_ui(
    request: RenderRequest,
    ui_renderer: UIRenderer = Depends(get_ui_renderer),
    vue_renderer: VueRenderer = Depends(get_vue_renderer),
    config: Dict[str, Any] = Depends(get_config)
):
    """渲染UI设计"""
    start_time = time.time()
    
    try:
        logger.info(f"收到渲染请求: {request.output_format}")
        
        # 验证输出格式
        valid_formats = ["png", "vue", "both"]
        if request.output_format not in valid_formats:
            raise HTTPException(
                status_code=400,
                detail=f"不支持的输出格式: {request.output_format}"
            )
        
        # 验证主题
        available_themes = config.get("ui", {}).get("themes", [])
        if request.theme not in available_themes:
            raise HTTPException(
                status_code=400,
                detail=f"不支持的主题: {request.theme}"
            )
        
        files = {}
        
        # 渲染PNG图片
        if request.output_format in ["png", "both"]:
            png_path = await render_png(request.dsl, request.theme, ui_renderer)
            files["png"] = png_path
        
        # 渲染Vue页面
        if request.output_format in ["vue", "both"]:
            vue_path = await render_vue(request.dsl, request.theme, vue_renderer, request.optimize)
            files["vue"] = vue_path
        
        processing_time = time.time() - start_time
        
        logger.info(f"渲染完成，耗时: {processing_time:.2f}秒")
        
        return RenderResponse(
            success=True,
            output_format=request.output_format,
            theme=request.theme,
            files=files,
            processing_time=processing_time,
            error_message=None
        )
        
    except Exception as e:
        processing_time = time.time() - start_time
        error_message = str(e)
        logger.error(f"渲染失败: {error_message}")
        
        return RenderResponse(
            success=False,
            output_format=request.output_format,
            theme=request.theme,
            files={},
            processing_time=processing_time,
            error_message=error_message
        )

async def render_png(dsl: Dict[str, Any], theme: str, renderer: UIRenderer) -> str:
    """渲染PNG图片"""
    # 设置主题
    renderer.set_theme(theme)
    
    # 渲染图片
    image_path = renderer.render(dsl)
    
    return image_path

async def render_vue(dsl: Dict[str, Any], theme: str, renderer: VueRenderer, optimize: bool) -> str:
    """渲染Vue页面"""
    # 设置主题
    renderer.set_theme(theme)
    
    # 渲染Vue页面
    vue_path = renderer.render(dsl, optimize=optimize)
    
    return vue_path

@app.post("/themes", response_model=ThemeResponse)
async def manage_themes(
    request: ThemeRequest,
    config: Dict[str, Any] = Depends(get_config)
):
    """主题管理"""
    start_time = time.time()
    
    try:
        logger.info(f"收到主题请求: {request.action}")
        
        available_themes = config.get("ui", {}).get("themes", [])
        
        if request.action == "list":
            return ThemeResponse(
                success=True,
                theme_name="",
                theme_config=None,
                available_themes=available_themes,
                preview_url=None,
                processing_time=time.time() - start_time,
                error_message=None
            )
        
        elif request.action == "get":
            if request.theme_name not in available_themes:
                raise HTTPException(
                    status_code=404,
                    detail=f"主题不存在: {request.theme_name}"
                )
            
            # 加载主题配置
            theme_config = load_theme_config(request.theme_name)
            
            return ThemeResponse(
                success=True,
                theme_name=request.theme_name,
                theme_config=theme_config,
                available_themes=None,
                preview_url=f"/themes/{request.theme_name}/preview",
                processing_time=time.time() - start_time,
                error_message=None
            )
        
        elif request.action == "preview":
            if request.theme_name not in available_themes:
                raise HTTPException(
                    status_code=404,
                    detail=f"主题不存在: {request.theme_name}"
                )
            
            # 生成主题预览
            preview_url = await generate_theme_preview(request.theme_name)
            
            return ThemeResponse(
                success=True,
                theme_name=request.theme_name,
                theme_config=None,
                available_themes=None,
                preview_url=preview_url,
                processing_time=time.time() - start_time,
                error_message=None
            )
        
        elif request.action == "apply":
            if request.dsl is None:
                raise HTTPException(
                    status_code=400,
                    detail="应用主题时需要提供DSL数据"
                )
            
            # 应用主题到DSL
            themed_dsl = apply_theme_to_dsl(request.dsl, request.theme_name)
            
            return ThemeResponse(
                success=True,
                theme_name=request.theme_name,
                theme_config=themed_dsl,
                available_themes=None,
                preview_url=None,
                processing_time=time.time() - start_time,
                error_message=None
            )
        
        else:
            raise HTTPException(
                status_code=400,
                detail=f"不支持的操作: {request.action}"
            )
        
    except Exception as e:
        processing_time = time.time() - start_time
        error_message = str(e)
        logger.error(f"主题管理失败: {error_message}")
        
        return ThemeResponse(
            success=False,
            theme_name=request.theme_name,
            theme_config=None,
            available_themes=None,
            preview_url=None,
            processing_time=processing_time,
            error_message=error_message
        )

def load_theme_config(theme_name: str) -> Dict[str, Any]:
    """加载主题配置"""
    tokens_path = "config/ui_tokens.json"
    with open(tokens_path, 'r', encoding='utf-8') as f:
        tokens = json.load(f)
    
    return tokens.get("themes", {}).get(theme_name, {})

async def generate_theme_preview(theme_name: str) -> str:
    """生成主题预览"""
    # 生成主题预览HTML
    preview_html = generate_theme_preview_html(theme_name)
    
    # 保存预览文件
    preview_path = f"output/themes/{theme_name}_preview.html"
    Path(preview_path).parent.mkdir(parents=True, exist_ok=True)
    
    with open(preview_path, 'w', encoding='utf-8') as f:
        f.write(preview_html)
    
    return f"/themes/{theme_name}/preview"

def apply_theme_to_dsl(dsl: Dict[str, Any], theme_name: str) -> Dict[str, Any]:
    """应用主题到DSL"""
    themed_dsl = dsl.copy()
    
    # 更新页面主题
    if "page" in themed_dsl:
        themed_dsl["page"]["theme"] = theme_name
    
    return themed_dsl

@app.get("/health", response_model=HealthResponse)
async def health_check():
    """健康检查"""
    start_time = time.time()
    
    try:
        # 检查服务状态
        services = {
            "ui_generator": "healthy" if ui_generator is not None else "unhealthy",
            "ui_renderer": "healthy" if ui_renderer is not None else "unhealthy",
            "vue_renderer": "healthy" if vue_renderer is not None else "unhealthy"
        }
        
        # 检查整体状态
        overall_status = "healthy" if all(
            status == "healthy" for status in services.values()
        ) else "unhealthy"
        
        return HealthResponse(
            status=overall_status,
            version="1.0.0",
            uptime=time.time() - start_time,
            services=services,
            timestamp=time.strftime("%Y-%m-%dT%H:%M:%SZ")
        )
        
    except Exception as e:
        logger.error(f"健康检查失败: {e}")
        
        return HealthResponse(
            status="unhealthy",
            version="1.0.0",
            uptime=0.0,
            services={},
            timestamp=time.strftime("%Y-%m-%dT%H:%M:%SZ")
        )

@app.get("/files/{file_path:path}")
async def serve_file(file_path: str):
    """提供文件下载服务"""
    try:
        # 构建完整文件路径
        full_path = Path("output") / file_path
        
        # 检查文件是否存在
        if not full_path.exists():
            raise HTTPException(status_code=404, detail="文件不存在")
        
        # 检查文件类型
        if full_path.suffix == ".png":
            media_type = "image/png"
        elif full_path.suffix == ".vue":
            media_type = "text/plain"
        elif full_path.suffix == ".html":
            media_type = "text/html"
        else:
            media_type = "application/octet-stream"
        
        return FileResponse(
            path=str(full_path),
            media_type=media_type,
            filename=full_path.name
        )
        
    except Exception as e:
        logger.error(f"文件服务失败: {e}")
        raise HTTPException(status_code=500, detail="文件服务失败")

@app.get("/themes/{theme_name}/preview")
async def serve_theme_preview(theme_name: str):
    """提供主题预览服务"""
    try:
        preview_path = Path("output/themes") / f"{theme_name}_preview.html"
        
        if not preview_path.exists():
            # 生成预览
            preview_html = generate_theme_preview_html(theme_name)
            preview_path.parent.mkdir(parents=True, exist_ok=True)
            
            with open(preview_path, 'w', encoding='utf-8') as f:
                f.write(preview_html)
        
        return FileResponse(
            path=str(preview_path),
            media_type="text/html"
        )
        
    except Exception as e:
        logger.error(f"主题预览服务失败: {e}")
        raise HTTPException(status_code=500, detail="主题预览服务失败")

@app.post("/batch/generate")
async def batch_generate_ui(
    requests: List[GenerateUIRequest],
    generator: UIGenerator = Depends(get_ui_generator)
):
    """批量生成UI"""
    try:
        logger.info(f"收到批量生成请求: {len(requests)}个")
        
        results = []
        for i, request in enumerate(requests):
            try:
                # 生成UI
                dsl = await generator.generate_with_model(
                    prompt=request.prompt,
                    theme=request.theme,
                    page_type=request.page_type,
                    max_length=request.max_length,
                    temperature=request.temperature
                )
                
                results.append({
                    "index": i,
                    "success": True,
                    "dsl": dsl,
                    "error": None
                })
                
            except Exception as e:
                results.append({
                    "index": i,
                    "success": False,
                    "dsl": None,
                    "error": str(e)
                })
        
        return {
            "success": True,
            "total": len(requests),
            "results": results
        }
        
    except Exception as e:
        logger.error(f"批量生成失败: {e}")
        raise HTTPException(status_code=500, detail="批量生成失败")

@app.post("/batch/render")
async def batch_render_ui(
    requests: List[RenderRequest],
    ui_renderer: UIRenderer = Depends(get_ui_renderer),
    vue_renderer: VueRenderer = Depends(get_vue_renderer)
):
    """批量渲染UI"""
    try:
        logger.info(f"收到批量渲染请求: {len(requests)}个")
        
        results = []
        for i, request in enumerate(requests):
            try:
                files = {}
                
                # 渲染PNG
                if request.output_format in ["png", "both"]:
                    png_path = await render_png(request.dsl, request.theme, ui_renderer)
                    files["png"] = png_path
                
                # 渲染Vue
                if request.output_format in ["vue", "both"]:
                    vue_path = await render_vue(request.dsl, request.theme, vue_renderer, request.optimize)
                    files["vue"] = vue_path
                
                results.append({
                    "index": i,
                    "success": True,
                    "files": files,
                    "error": None
                })
                
            except Exception as e:
                results.append({
                    "index": i,
                    "success": False,
                    "files": {},
                    "error": str(e)
                })
        
        return {
            "success": True,
            "total": len(requests),
            "results": results
        }
        
    except Exception as e:
        logger.error(f"批量渲染失败: {e}")
        raise HTTPException(status_code=500, detail="批量渲染失败")

@app.exception_handler(HTTPException)
async def http_exception_handler(request, exc):
    """HTTP异常处理"""
    return JSONResponse(
        status_code=exc.status_code,
        content={
            "success": False,
            "error": {
                "type": "HTTPException",
                "message": exc.detail,
                "status_code": exc.status_code
            }
        }
    )

@app.exception_handler(Exception)
async def general_exception_handler(request, exc):
    """通用异常处理"""
    logger.error(f"未处理的异常: {exc}")
    
    return JSONResponse(
        status_code=500,
        content={
            "success": False,
            "error": {
                "type": "InternalServerError",
                "message": "内部服务器错误",
                "status_code": 500
            }
        }
    )

@app.middleware("http")
async def log_requests(request, call_next):
    """请求日志中间件"""
    start_time = time.time()
    
    # 记录请求
    logger.info(f"收到请求: {request.method} {request.url}")
    
    # 处理请求
    response = await call_next(request)
    
    # 记录响应
    process_time = time.time() - start_time
    logger.info(f"请求完成: {request.method} {request.url} - {response.status_code} - {process_time:.2f}s")
    
    # 添加响应头
    response.headers["X-Process-Time"] = str(process_time)
    
    return response

from slowapi import Limiter, _rate_limit_exceeded_handler
from slowapi.util import get_remote_address
from slowapi.errors import RateLimitExceeded

# 创建限流器
limiter = Limiter(key_func=get_remote_address)
app.state.limiter = limiter
app.add_exception_handler(RateLimitExceeded, _rate_limit_exceeded_handler)

# 应用限流
@app.post("/generate-ui")
@limiter.limit("10/minute")
async def generate_ui_with_rate_limit(
    request: Request,
    generate_request: GenerateUIRequest,
    generator: UIGenerator = Depends(get_ui_generator)
):
    """带限流的UI生成"""
    return await generate_ui(generate_request, generator)

@app.post("/render")
@limiter.limit("20/minute")
async def render_ui_with_rate_limit(
    request: Request,
    render_request: RenderRequest,
    ui_renderer: UIRenderer = Depends(get_ui_renderer),
    vue_renderer: VueRenderer = Depends(get_vue_renderer)
):
    """带限流的UI渲染"""
    return await render_ui(render_request, ui_renderer, vue_renderer)

from fastapi import BackgroundTasks

@app.post("/generate-ui-async")
async def generate_ui_async(
    request: GenerateUIRequest,
    background_tasks: BackgroundTasks,
    generator: UIGenerator = Depends(get_ui_generator)
):
    """异步UI生成"""
    # 立即返回任务ID
    task_id = str(uuid.uuid4())
    
    # 后台处理
    background_tasks.add_task(
        process_ui_generation,
        task_id,
        request,
        generator
    )
    
    return {
        "success": True,
        "task_id": task_id,
        "status": "processing",
        "message": "UI生成任务已提交，请稍后查询结果"
    }

async def process_ui_generation(
    task_id: str,
    request: GenerateUIRequest,
    generator: UIGenerator
):
    """处理UI生成任务"""
    try:
        # 生成UI
        dsl = await generator.generate_with_model(
            prompt=request.prompt,
            theme=request.theme,
            page_type=request.page_type,
            max_length=request.max_length,
            temperature=request.temperature
        )
        
        # 保存结果
        result_path = f"output/tasks/{task_id}.json"
        Path(result_path).parent.mkdir(parents=True, exist_ok=True)
        
        with open(result_path, 'w', encoding='utf-8') as f:
            json.dump({
                "task_id": task_id,
                "success": True,
                "dsl": dsl,
                "completed_at": time.strftime("%Y-%m-%dT%H:%M:%SZ")
            }, f, ensure_ascii=False, indent=2)
        
        logger.info(f"UI生成任务完成: {task_id}")
        
    except Exception as e:
        # 保存错误结果
        error_path = f"output/tasks/{task_id}_error.json"
        Path(error_path).parent.mkdir(parents=True, exist_ok=True)
        
        with open(error_path, 'w', encoding='utf-8') as f:
            json.dump({
                "task_id": task_id,
                "success": False,
                "error": str(e),
                "completed_at": time.strftime("%Y-%m-%dT%H:%M:%SZ")
            }, f, ensure_ascii=False, indent=2)
        
        logger.error(f"UI生成任务失败: {task_id} - {e}")

@app.get("/tasks/{task_id}")
async def get_task_result(task_id: str):
    """获取任务结果"""
    try:
        # 检查结果文件
        result_path = f"output/tasks/{task_id}.json"
        error_path = f"output/tasks/{task_id}_error.json"
        
        if Path(result_path).exists():
            with open(result_path, 'r', encoding='utf-8') as f:
                result = json.load(f)
            return result
        
        elif Path(error_path).exists():
            with open(error_path, 'r', encoding='utf-8') as f:
                result = json.load(f)
            return result
        
        else:
            return {
                "task_id": task_id,
                "status": "processing",
                "message": "任务仍在处理中"
            }
        
    except Exception as e:
        logger.error(f"获取任务结果失败: {e}")
        raise HTTPException(status_code=500, detail="获取任务结果失败")

from functools import lru_cache
import hashlib

@lru_cache(maxsize=1000)
def get_cached_ui_generation(prompt_hash: str, theme: str, page_type: str) -> Dict[str, Any]:
    """获取缓存的UI生成结果"""
    cache_path = f"output/cache/{prompt_hash}_{theme}_{page_type}.json"
    
    if Path(cache_path).exists():
        with open(cache_path, 'r', encoding='utf-8') as f:
            return json.load(f)
    
    return None

def cache_ui_generation(
    prompt: str,
    theme: str,
    page_type: str,
    dsl: Dict[str, Any]
) -> None:
    """缓存UI生成结果"""
    # 生成缓存键
    prompt_hash = hashlib.md5(prompt.encode()).hexdigest()
    
    # 保存缓存
    cache_path = f"output/cache/{prompt_hash}_{theme}_{page_type}.json"
    Path(cache_path).parent.mkdir(parents=True, exist_ok=True)
    
    with open(cache_path, 'w', encoding='utf-8') as f:
        json.dump({
            "prompt": prompt,
            "theme": theme,
            "page_type": page_type,
            "dsl": dsl,
            "cached_at": time.strftime("%Y-%m-%dT%H:%M:%SZ")
        }, f, ensure_ascii=False, indent=2)

@app.post("/generate-ui-cached")
async def generate_ui_cached(
    request: GenerateUIRequest,
    generator: UIGenerator = Depends(get_ui_generator)
):
    """带缓存的UI生成"""
    start_time = time.time()
    
    try:
        # 生成缓存键
        prompt_hash = hashlib.md5(request.prompt.encode()).hexdigest()
        
        # 检查缓存
        cached_result = get_cached_ui_generation(
            prompt_hash,
            request.theme,
            request.page_type
        )
        
        if cached_result:
            logger.info("使用缓存结果")
            return GenerateUIResponse(
                success=True,
                dsl=cached_result["dsl"],
                prompt=request.prompt,
                theme=request.theme,
                page_type=request.page_type,
                generation_method="cache",
                processing_time=time.time() - start_time,
                error_message=None
            )
        
        # 生成新结果
        dsl = await generator.generate_with_model(
            prompt=request.prompt,
            theme=request.theme,
            page_type=request.page_type,
            max_length=request.max_length,
            temperature=request.temperature
        )
        
        # 缓存结果
        cache_ui_generation(request.prompt, request.theme, request.page_type, dsl)
        
        processing_time = time.time() - start_time
        
        return GenerateUIResponse(
            success=True,
            dsl=dsl,
            prompt=request.prompt,
            theme=request.theme,
            page_type=request.page_type,
            generation_method="model",
            processing_time=processing_time,
            error_message=None
        )
        
    except Exception as e:
        processing_time = time.time() - start_time
        error_message = str(e)
        logger.error(f"UI生成失败: {error_message}")
        
        return GenerateUIResponse(
            success=False,
            dsl=None,
            prompt=request.prompt,
            theme=request.theme,
            page_type=request.page_type,
            generation_method="error",
            processing_time=processing_time,
            error_message=error_message
        )

# 自定义OpenAPI配置
def custom_openapi():
    if app.openapi_schema:
        return app.openapi_schema
    
    openapi_schema = get_openapi(
        title="AI UI生成系统API",
        version="1.0.0",
        description="""
        ## AI UI生成系统API文档
        
        这是一个基于AI的UI设计和代码生成系统，提供以下功能：
        
        * **UI生成**: 根据中文描述生成UI设计
        * **多格式渲染**: 支持PNG图片和Vue页面渲染
        * **主题管理**: 支持多种主题切换
        * **批量处理**: 支持批量生成和渲染
        
        ## 使用说明
        
        1. 首先调用 `/generate-ui` 生成UI设计
        2. 然后调用 `/render` 渲染为图片或代码
        3. 使用 `/themes` 管理主题
        4. 通过 `/health` 检查服务状态
        
        ## 认证
        
        当前版本无需认证，后续版本将支持API密钥认证。
        """,
        routes=app.routes,
    )
    
    # 添加自定义标签
    openapi_schema["tags"] = [
        {
            "name": "UI生成",
            "description": "UI设计和生成相关接口"
        },
        {
            "name": "渲染",
            "description": "UI渲染相关接口"
        },
        {
            "name": "主题",
            "description": "主题管理相关接口"
        },
        {
            "name": "系统",
            "description": "系统状态和健康检查"
        }
    ]
    
    app.openapi_schema = openapi_schema
    return app.openapi_schema

app.openapi = custom_openapi

# 测试用例
import pytest
from fastapi.testclient import TestClient

client = TestClient(app)

def test_generate_ui():
    """测试UI生成"""
    response = client.post("/generate-ui", json={
        "prompt": "设计一个简单的首页",
        "theme": "obsidian-gold",
        "page_type": "home_page"
    })
    
    assert response.status_code == 200
    data = response.json()
    assert data["success"] == True
    assert "dsl" in data

def test_render_ui():
    """测试UI渲染"""
    # 先生成UI
    generate_response = client.post("/generate-ui", json={
        "prompt": "设计一个简单的首页",
        "theme": "obsidian-gold",
        "page_type": "home_page"
    })
    
    dsl = generate_response.json()["dsl"]
    
    # 然后渲染
    render_response = client.post("/render", json={
        "dsl": dsl,
        "output_format": "both",
        "theme": "obsidian-gold"
    })
    
    assert render_response.status_code == 200
    data = render_response.json()
    assert data["success"] == True
    assert "files" in data

def test_health_check():
    """测试健康检查"""
    response = client.get("/health")
    
    assert response.status_code == 200
    data = response.json()
    assert data["status"] in ["healthy", "unhealthy"]
    assert "services" in data

def test_theme_management():
    """测试主题管理"""
    response = client.post("/themes", json={
        "theme_name": "obsidian-gold",
        "action": "get"
    })
    
    assert response.status_code == 200
    data = response.json()
    assert data["success"] == True
    assert "theme_config" in data

# 生产环境配置
import uvicorn
from multiprocessing import cpu_count

if __name__ == "__main__":
    # 生产环境配置
    uvicorn.run(
        "main:app",
        host="0.0.0.0",
        port=8000,
        workers=cpu_count(),
        reload=False,
        access_log=True,
        log_level="info"
    )

import os
from pydantic import BaseSettings

class Settings(BaseSettings):
    """应用配置"""
    app_name: str = "AI UI生成系统"
    app_version: str = "1.0.0"
    debug: bool = False
    
    # 服务配置
    host: str = "0.0.0.0"
    port: int = 8000
    workers: int = 1
    
    # 模型配置
    model_path: str = "models/ui-dsl-lora"
    config_path: str = "config/model_config.yaml"
    tokens_path: str = "config/ui_tokens.json"
    
    # 输出配置
    output_dir: str = "output"
    cache_dir: str = "output/cache"
    
    # 限流配置
    rate_limit: str = "100/minute"
    
    # 日志配置
    log_level: str = "INFO"
    log_file: str = "logs/api.log"
    
    class Config:
        env_file = ".env"

settings = Settings()

# 使用配置
app = FastAPI(
    title=settings.app_name,
    version=settings.app_version,
    debug=settings.debug
)