DocETL项目中LLM代理任务的策略模式重构实践

背景与问题分析

在现代ETL（Extract-Transform-Load）系统中，LLM（大语言模型）正被越来越多地用于优化查询执行计划。在DocETL项目中，我们观察到多个优化器类（如MapOptimizer、ReduceOptimizer和JoinOptimizer）都嵌入了大量与LLM交互的重复代码模式。这种实现方式虽然功能完整，但存在几个显著问题：

1. 代码重复：每个优化任务都重复实现相似的LLM交互流程
2. 紧耦合：优化器业务逻辑与LLM交互细节深度耦合
3. 维护困难：修改LLM交互方式需要多处同步变更
4. 测试复杂：难以对LLM交互逻辑进行独立测试

现有实现的问题模式

当前实现中，每个LLM代理任务都遵循相似的7步模式：

1. 任务特定上下文准备
2. 代理角色定义
3. 详细用户提示构建
4. 输出模式定义
5. LLM客户端调用
6. 响应解析与提取
7. 任务特定验证/重试

以JoinOptimizer中的_blocking_rules生成为例，这些步骤都直接嵌入在优化器方法中，导致方法臃肿且难以维护。

策略模式解决方案

我们采用策略模式对这部分代码进行重构，核心思想是将每个LLM代理任务封装为独立的策略类。具体实现包括：

1. 抽象基类设计

首先定义LLMAgentTask抽象基类，提供统一的接口：

class LLMAgentTask(ABC):
    def __init__(self, llm_client: LLMClient, console: Console):
        self.llm_client = llm_client
        self.console = console

    @abstractmethod
    def execute(self, *args, **kwargs) -> Any:
        """执行特定的LLM代理任务"""

2. 具体策略实现

为每个LLM代理任务创建具体策略类。以生成阻塞规则为例：

class GenerateBlockingRuleTask(LLMAgentTask):
    SYSTEM_PROMPT = "您是实体解析规则生成专家..."
    OUTPUT_SCHEMA = {...}

    def _build_user_prompt(self, sample_datas, blocking_keys, comparison_prompt):
        """封装提示构建逻辑"""
        return f"分析这些样本比较..."

    def execute(self, sample_datas, blocking_keys, comparison_prompt, max_retries):
        """完整封装7步交互流程"""
        # 实现细节...

3. 优化器重构

优化器类只需实例化并使用相应策略：

def _generate_blocking_rules(self, blocking_keys, input_data, comparisons):
    sample_datas = self._sample_comparison_data(input_data, comparisons)
    rule_generator = GenerateBlockingRuleTask(self.llm_client, self.console)
    blocking_rule = rule_generator.execute(...)
    return [blocking_rule] if blocking_rule else []

架构优势与实施效果

这种重构带来了多方面改进：

1. 关注点分离：优化器专注于业务逻辑，策略类处理LLM交互细节
2. 代码复用：公共LLM交互模式被统一封装
3. 易于测试：每个策略类可独立测试
4. 扩展性强：新增LLM任务只需添加新策略类
5. 维护简单：修改特定任务的LLM交互不影响其他部分

实践经验与建议

在实际重构过程中，我们总结了以下经验：

1. 上下文传递：策略类可能需要访问优化器的部分上下文，需设计合理的参数传递机制
2. 错误处理：将重试逻辑封装在策略类中，保持优化器代码简洁
3. 提示工程：将系统提示和用户提示模板集中管理，便于统一调整
4. 输出验证：每个策略类应包含针对其任务的特定输出验证逻辑

未来展望

虽然策略模式解决了当前问题，但随着系统演进，我们还可以考虑：

1. 引入LLM交互框架，进一步标准化常见模式
2. 开发提示模板管理系统
3. 实现LLM交互的监控和性能分析
4. 探索自动生成策略类的可能性

这种基于策略模式的重构不仅提升了DocETL项目的代码质量，也为其他类似系统中LLM集成提供了可借鉴的架构模式。通过清晰的职责划分和模块化设计，系统在面对不断变化的LLM技术和业务需求时，将展现出更强的适应性和可维护性。