screenshot-to-code扩展开发教程：从零构建AI代码生成中间件

【免费下载链接】screenshot-to-code 上传一张屏幕截图并将其转换为整洁的代码（HTML/Tailwind/React/Vue） 项目地址: https://gitcode.com/GitHub_Trending/sc/screenshot-to-code

你是否曾经想过，如何为强大的screenshot-to-code项目添加自定义功能？本文将深入解析该项目的中间件架构，教你如何从零开始构建扩展功能，实现个性化的代码生成流程。

项目架构概览

screenshot-to-code采用基于中间件（Middleware）的管道（Pipeline）架构，这种设计模式使得功能扩展变得异常简单。整个系统由以下几个核心组件构成：

核心中间件详解

1. 基础中间件类

所有中间件都继承自抽象的Middleware基类，必须实现process方法：

class Middleware(ABC):
    """Base class for all pipeline middleware"""
    
    @abstractmethod
    async def process(
        self, context: PipelineContext, next_func: Callable[[], Awaitable[None]]
    ) -> None:
        """Process the context and call the next middleware"""
        pass

2. 管道执行器

Pipeline类负责管理中间件的执行顺序：

class Pipeline:
    """Pipeline for processing WebSocket code generation requests"""
    
    def __init__(self):
        self.middlewares: List[Middleware] = []
    
    def use(self, middleware: Middleware) -> "Pipeline":
        """Add a middleware to the pipeline"""
        self.middlewares.append(middleware)
        return self
    
    async def execute(self, websocket: WebSocket) -> None:
        """Execute the pipeline with the given WebSocket"""
        context = PipelineContext(websocket=websocket)
        # ... 构建中间件链并执行

3. 上下文对象

PipelineContext承载整个处理流程的状态：

@dataclass
class PipelineContext:
    """Context object that carries state through the pipeline"""
    
    websocket: WebSocket
    ws_comm: "WebSocketCommunicator | None" = None
    params: Dict[str, str] = field(default_factory=dict)
    extracted_params: "ExtractedParams | None" = None
    prompt_messages: List[ChatCompletionMessageParam] = field(default_factory=list)
    image_cache: Dict[str, str] = field(default_factory=dict)
    variant_models: List[Llm] = field(default_factory=list)
    completions: List[str] = field(default_factory=list)
    variant_completions: Dict[int, str] = field(default_factory=dict)
    metadata: Dict[str, Any] = field(default_factory=dict)

实战：创建自定义中间件

示例1：日志记录中间件

让我们创建一个简单的日志记录中间件，记录每个请求的处理时间：

import time
from datetime import datetime

class LoggingMiddleware(Middleware):
    """Middleware for request logging and timing"""
    
    async def process(
        self, context: PipelineContext, next_func: Callable[[], Awaitable[None]]
    ) -> None:
        start_time = time.time()
        request_id = datetime.now().strftime("%Y%m%d%H%M%S")
        
        print(f"[{request_id}] Request started at {datetime.now()}")
        
        try:
            await next_func()
        finally:
            end_time = time.time()
            duration = end_time - start_time
            print(f"[{request_id}] Request completed in {duration:.2f} seconds")

示例2：输入验证中间件

创建一个输入验证中间件，确保传入的参数符合预期：

class ValidationMiddleware(Middleware):
    """Middleware for input validation"""
    
    def __init__(self, allowed_stacks: List[str] = None):
        self.allowed_stacks = allowed_stacks or [
            "html_tailwind", "react_tailwind", "vue_tailwind", "bootstrap", "ionic_tailwind", "svg"
        ]
    
    async def process(
        self, context: PipelineContext, next_func: Callable[[], Awaitable[None]]
    ) -> None:
        assert context.extracted_params is not None
        
        # 验证技术栈是否被支持
        if context.extracted_params.stack not in self.allowed_stacks:
            await context.throw_error(
                f"Unsupported stack: {context.extracted_params.stack}. "
                f"Allowed stacks: {', '.join(self.allowed_stacks)}"
            )
            return
        
        # 验证输入模式
        if context.extracted_params.input_mode not in ["image", "text", "video"]:
            await context.throw_error("Invalid input mode")
            return
        
        await next_func()

示例3：缓存中间件

创建一个简单的缓存中间件，避免重复处理相同的请求：

import hashlib
import json

class CacheMiddleware(Middleware):
    """Middleware for response caching"""
    
    def __init__(self, cache_size: int = 100):
        self.cache = {}
        self.cache_size = cache_size
        self.cache_keys = []  # LRU缓存管理
    
    def _generate_cache_key(self, params: Dict[str, Any]) -> str:
        """Generate a unique cache key from request parameters"""
        param_str = json.dumps(params, sort_keys=True)
        return hashlib.md5(param_str.encode()).hexdigest()
    
    async def process(
        self, context: PipelineContext, next_func: Callable[[], Awaitable[None]]
    ) -> None:
        cache_key = self._generate_cache_key(context.params)
        
        # 检查缓存
        if cache_key in self.cache:
            print(f"Cache hit for key: {cache_key}")
            cached_response = self.cache[cache_key]
            
            # 发送缓存的响应
            for i, completion in enumerate(cached_response["completions"]):
                await context.send_message("setCode", completion, i)
                await context.send_message("variantComplete", "From cache", i)
            
            return
        
        # 缓存未命中，继续处理
        original_send_message = context.send_message
        
        # 拦截发送的消息以收集响应
        responses = {i: "" for i in range(4)}  # 假设最多4个变体
        
        async def intercepted_send_message(type: str, value: str, variant_index: int):
            if type == "setCode":
                responses[variant_index] = value
            await original_send_message(type, value, variant_index)
        
        # 临时替换发送方法
        context.send_message = intercepted_send_message
        
        try:
            await next_func()
            
            # 缓存成功的响应
            if all(responses.values()):  # 所有变体都有响应
                if len(self.cache) >= self.cache_size:
                    # LRU淘汰
                    oldest_key = self.cache_keys.pop(0)
                    del self.cache[oldest_key]
                
                self.cache[cache_key] = {
                    "completions": list(responses.values()),
                    "timestamp": time.time()
                }
                self.cache_keys.append(cache_key)
                
        finally:
            # 恢复原始发送方法
            context.send_message = original_send_message

中间件注册与使用

注册自定义中间件

在backend/routes/generate_code.py中找到stream_code函数，添加你的中间件：

@router.websocket("/generate-code")
async def stream_code(websocket: WebSocket):
    """WebSocket endpoint for streaming code generation"""
    pipeline = (
        Pipeline()
        .use(WebSocketSetupMiddleware())
        .use(ParameterExtractionMiddleware())
        .use(StatusBroadcastMiddleware())
        .use(LoggingMiddleware())  # 添加日志中间件
        .use(ValidationMiddleware())  # 添加验证中间件
        .use(PromptCreationMiddleware())
        .use(CodeGenerationMiddleware())
        .use(CacheMiddleware(cache_size=50))  # 添加缓存中间件
        .use(PostProcessingMiddleware())
    )
    
    await pipeline.execute(websocket)

中间件执行顺序表

顺序	中间件名称	功能描述	是否必需
1	WebSocketSetupMiddleware	WebSocket连接建立	是
2	ParameterExtractionMiddleware	参数提取与验证	是
3	StatusBroadcastMiddleware	状态广播	是
4	LoggingMiddleware	请求日志记录	否
5	ValidationMiddleware	输入验证	否
6	PromptCreationMiddleware	提示词创建	是
7	CodeGenerationMiddleware	代码生成	是
8	CacheMiddleware	响应缓存	否
9	PostProcessingMiddleware	后处理	是

高级扩展技巧

1. 错误处理中间件

创建一个统一的错误处理中间件：

class ErrorHandlingMiddleware(Middleware):
    """Global error handling middleware"""
    
    async def process(
        self, context: PipelineContext, next_func: Callable[[], Awaitable[None]]
    ) -> None:
        try:
            await next_func()
        except Exception as e:
            error_message = f"Internal server error: {str(e)}"
            print(f"Error: {error_message}")
            print(traceback.format_exc())
            
            if context.ws_comm and not context.ws_comm.is_closed:
                await context.throw_error("An unexpected error occurred. Please try again.")

2. 性能监控中间件

class PerformanceMonitoringMiddleware(Middleware):
    """Middleware for performance monitoring"""
    
    async def process(
        self, context: PipelineContext, next_func: Callable[[], Awaitable[None]]
    ) -> None:
        timings = {}
        stages = [
            "parameter_extraction",
            "prompt_creation", 
            "code_generation",
            "post_processing"
        ]
        
        for stage in stages:
            stage_start = time.time()
            # 这里需要根据具体阶段进行更精细的计时
            await next_func()
            timings[stage] = time.time() - stage_start
        
        # 记录性能数据
        print(f"Performance metrics: {json.dumps(timings, indent=2)}")

3. 自定义元数据中间件

class MetadataMiddleware(Middleware):
    """Middleware for adding custom metadata"""
    
    async def process(
        self, context: PipelineContext, next_func: Callable[[], Awaitable[None]]
    ) -> None:
        # 添加请求元数据
        context.metadata.update({
            "request_timestamp": datetime.now().isoformat(),
            "user_agent": context.websocket.headers.get("user-agent", "unknown"),
            "client_ip": context.websocket.client.host if context.websocket.client else "unknown"
        })
        
        await next_func()
        
        # 添加响应元数据
        context.metadata.update({
            "completion_count": len(context.completions),
            "processing_time": time.time() - float(context.metadata["request_timestamp"])
        })

测试与调试

单元测试示例

为你的中间件编写单元测试：

import pytest
from unittest.mock import AsyncMock, MagicMock

@pytest.mark.asyncio
async def test_logging_middleware():
    """Test that logging middleware works correctly"""
    middleware = LoggingMiddleware()
    context = MagicMock()
    next_func = AsyncMock()
    
    await middleware.process(context, next_func)
    
    # 验证next_func被调用
    next_func.assert_awaited_once()
    # 验证日志输出（可能需要捕获stdout）

集成测试

@pytest.mark.asyncio 
async def test_full_pipeline_with_custom_middleware():
    """Test the full pipeline with custom middleware"""
    pipeline = (
        Pipeline()
        .use(WebSocketSetupMiddleware())
        .use(LoggingMiddleware())
        .use(ParameterExtractionMiddleware())
        .use(ValidationMiddleware())
    )
    
    # 使用测试WebSocket
    test_websocket = create_test_websocket()
    
    await pipeline.execute(test_websocket)
    
    # 验证中间件执行顺序和结果

最佳实践总结

1. 保持中间件单一职责：每个中间件只负责一个明确的功能
2. 错误处理要全面：确保中间件能够妥善处理各种异常情况
3. 性能考虑：避免在中间件中执行耗时的同步操作
4. 可测试性：设计中间件时要考虑如何编写单元测试
5. 文档完善：为每个自定义中间件提供清晰的文档说明

通过本文的教程，你应该已经掌握了screenshot-to-code项目的扩展开发技巧。这种中间件架构不仅提供了极大的灵活性，还使得功能扩展变得简单而优雅。现在就开始创建你自己的中间件，为这个强大的AI代码生成工具添加个性化功能吧！

记住，良好的扩展设计应该遵循开闭原则（Open-Closed Principle）——对扩展开放，对修改关闭。通过中间件模式，你可以在不修改核心代码的情况下，为系统添加无限可能的新功能。

【免费下载链接】screenshot-to-code 上传一张屏幕截图并将其转换为整洁的代码（HTML/Tailwind/React/Vue） 项目地址: https://gitcode.com/GitHub_Trending/sc/screenshot-to-code