Compound AI Systems：复合式 AI 系统架构详解

from dataclasses import dataclass, field
from typing import List, Dict, Any, Optional, Callable
import json
import time
from enum import Enum

# ===== 核心组件定义 =====

class ComponentType(Enum):
    RETRIEVAL = "retrieval"
    LLM = "llm"
    TOOL = "tool"
    EXECUTOR = "executor"
    ORCHESTRATOR = "orchestrator"


@dataclass
class ComponentResult:
    """组件执行结果"""
    success: bool
    data: Any = None
    error: Optional[str] = None
    latency_ms: float = 0.0
    metadata: Dict[str, Any] = field(default_factory=dict)


class BaseComponent:
    """所有组件的基类——实现标准化接口"""

    def __init__(self, name: str, timeout_ms: float = 30000):
        self.name = name
        self.timeout_ms = timeout_ms

    def execute(self, input_data: Dict[str, Any]) -> ComponentResult:
        """执行组件——子类必须实现"""
        start = time.time()
        try:
            result = self._process(input_data)
            return ComponentResult(
                success=True, data=result,
                latency_ms=(time.time() - start) * 1000
            )
        except Exception as e:
            return ComponentResult(
                success=False, error=str(e),
                latency_ms=(time.time() - start) * 1000
            )

    def _process(self, input_data: Dict[str, Any]) -> Any:
        raise NotImplementedError


# ===== 检索组件 =====

class RetrievalComponent(BaseComponent):
    """RAG 检索组件——从知识库检索相关文档"""

    def __init__(self, knowledge_base: List[Dict], embedding_fn=None):
        super().__init__("Retrieval")
        self.knowledge_base = knowledge_base
        self.embedding_fn = embedding_fn or self._default_embed

    def _default_embed(self, text: str) -> List[float]:
        """简化版嵌入——生产环境应使用真实的 embedding 模型"""
        return [hash(text + str(i)) % 100 / 100 for i in range(128)]

    def _process(self, input_data: Dict) -> Dict:
        query = input_data["query"]
        query_vec = self.embedding_fn(query)

        # 计算余弦相似度并排序
        scored_docs = []
        for doc in self.knowledge_base:
            doc_vec = self.embedding_fn(doc["content"])
            similarity = self._cosine_similarity(query_vec, doc_vec)
            scored_docs.append((similarity, doc))

        scored_docs.sort(key=lambda x: x[0], reverse=True)
        top_k = scored_docs[:input_data.get("top_k", 5)]

        return {
            "documents": [d[1] for d in top_k],
            "scores": [d[0] for d in top_k],
            "retrieval_strategy": "vector_similarity"
        }

    def _cosine_similarity(self, a: List[float], b: List[float]) -> float:
        dot = sum(x * y for x, y in zip(a, b))
        norm_a = sum(x  2 for x in a)  0.5
        norm_b = sum(x  2 for x in b)  0.5
        return dot / (norm_a * norm_b) if norm_a * norm_b > 0 else 0.0


# ===== LLM 组件 =====

class LLMComponent(BaseComponent):
    """大语言模型组件——负责生成和推理"""

    def __init__(self, model: str = "gpt-4", system_prompt: str = ""):
        super().__init__("LLM")
        self.model = model
        self.system_prompt = system_prompt

    def _process(self, input_data: Dict) -> str:
        """生产环境应调用真实的 LLM API"""
        messages = [{"role": "system", "content": self.system_prompt}]

        if "retrieved_docs" in input_data:
            context = "\n".join(
                [d["content"] for d in input_data["retrieved_docs"]]
            )
            messages.append({
                "role": "user",
                "content": f"基于以下知识回答问题：\n{context}\n\n{input_data['query']}"
            })
        else:
            messages.append({"role": "user", "content": input_data["query"]})

        # 这里模拟 LLM 调用
        return self._call_llm(messages)

    def _call_llm(self, messages: List[Dict]) -> str:
        # 生产环境：使用 openai、anthropic 等 SDK
        return "模拟LLM响应：基于检索结果生成的分析报告"


# ===== 工具执行组件 =====

class ToolExecutorComponent(BaseComponent):
    """代码执行组件——在沙箱中执行生成的代码"""

    def __init__(self, sandbox: str = "local"):
        super().__init__("ToolExecutor")
        self.sandbox = sandbox
        self.execution_count = 0
        self.max_executions = 10  # 安全限制

    def _process(self, input_data: Dict) -> Dict:
        if self.execution_count >= self.max_executions:
            raise RuntimeError("超过最大执行次数限制")

        code = input_data.get("code", "")
        if not code:
            raise ValueError("未提供可执行代码")

        # 安全检查：禁止危险操作
        dangerous_patterns = ["os.system", "subprocess", "__import__"]
        for pattern in dangerous_patterns:
            if pattern in code:
                raise ValueError(f"代码包含危险操作: {pattern}")

        self.execution_count += 1
        local_ns = {"result": None}

        try:
            exec(code, {"__builtins__": {}}, local_ns)
            return {
                "output": str(local_ns.get("result", "无输出")),
                "execution_time": "模拟执行时间",
                "status": "success"
            }
        except Exception as e:
            return {"error": str(e), "status": "failed"}


# ===== 编排器 =====

class CompoundOrchestrator:
    """复合系统编排器——管理组件生命周期和执行流程"""

    def __init__(self):
        self.components: Dict[str, BaseComponent] = {}
        self.pipeline: List[str] = []
        self.fallback_strategy: str = "skip"  # skip / retry / degrade

    def register(self, name: str, component: BaseComponent):
        """注册组件"""
        self.components[name] = component

    def set_pipeline(self, stages: List[str]):
        """设置执行管道"""
        self.pipeline = stages

    def execute_pipeline(self, input_data: Dict) -> Dict:
        """执行完整管道"""
        context = {"input": input_data}

        for stage_name in self.pipeline:
            if stage_name not in self.components:
                raise ValueError(f"未注册的组件: {stage_name}")

            component = self.components[stage_name]
            result = component.execute(context)

            context[stage_name] = result

            if not result.success:
                if self.fallback_strategy == "skip":
                    continue
                elif self.fallback_strategy == "degrade":
                    context[stage_name] = self._degrade(stage_name, context)

        return context

    def _degrade(self, stage_name: str, context: Dict) -> ComponentResult:
        """降级处理"""
        return ComponentResult(
            success=False,
            data={"degraded": True, "stage": stage_name},
            metadata={"fallback": "applied"}
        )


# ===== 使用示例 =====

if __name__ == "__main__":
    # 构建知识库
    kb = [
        {"content": "pandas groupby 用于数据分组聚合", "id": "doc1"},
        {"content": "matplotlib 用于数据可视化图表", "id": "doc2"},
        {"content": "scipy stats 用于统计检验", "id": "doc3"},
    ]

    # 创建组件
    retrieval = RetrievalComponent(knowledge_base=kb)
    llm = LLMComponent(model="gpt-4", system_prompt="你是数据分析专家")
    executor = ToolExecutorComponent()

    # 构建编排器
    orchestrator = CompoundOrchestrator()
    orchestrator.register("retrieval", retrieval)
    orchestrator.register("llm", llm)
    orchestrator.register("executor", executor)
    orchestrator.set_pipeline(["retrieval", "llm", "executor"])

    # 执行
    result = orchestrator.execute_pipeline({
        "query": "分析销售数据的趋势并生成可视化图表",
        "top_k": 3
    })

    for stage, res in result.items():
        if isinstance(res, ComponentResult):
            print(f"[{stage}] {'✅' if res.success else '❌'} "
                  f"({res.latency_ms:.0f}ms)")

// ===== 类型定义 =====

interface ComponentConfig {
  name: string;
  type: 'llm' | 'retrieval' | 'tool' | 'executor';
  timeoutMs: number;
  retryCount: number;
}

interface PipelineConfig {
  id: string;
  description: string;
  stages: string[]; // 组件名称列表
  router_condition: (input: RequestInput) => boolean;
}

interface RequestInput {
  query: string;
  context?: Record<string, unknown>;
  userId: string;
  priority: 'low' | 'normal' | 'high';
}

interface PipelineResult {
  pipelineId: string;
  success: boolean;
  output: string;
  stages: StageResult[];
  totalLatencyMs: number;
  cost: { tokens: number; dollars: number };
}

interface StageResult {
  stage: string;
  success: boolean;
  latencyMs: number;
}

// ===== 组件注册表 =====

class ComponentRegistry {
  private components = new Map<string, ComponentConfig>();

  register(config: ComponentConfig): void {
    this.components.set(config.name, config);
  }

  get(name: string): ComponentConfig | undefined {
    return this.components.get(name);
  }

  getAll(): ComponentConfig[] {
    return Array.from(this.components.values());
  }
}

// ===== 路由器 =====

class Router {
  private pipelines: PipelineConfig[] = [];
  private defaultPipeline: string | null = null;

  registerPipeline(config: PipelineConfig): void {
    this.pipelines.push(config);
  }

  setDefaultPipeline(id: string): void {
    this.defaultPipeline = id;
  }

  route(input: RequestInput): PipelineConfig {
    for (const pipeline of this.pipelines) {
      if (pipeline.router_condition(input)) {
        return pipeline;
      }
    }

    // 回退到默认管道
    const defaultPipe = this.pipelines.find(p => p.id === this.defaultPipeline);
    if (defaultPipe) return defaultPipe;

    throw new Error('没有匹配的处理管道');
  }

  // 基于关键词的简单路由——生产环境应使用 LLM 意图分类
  static keywordRouter(
    keywords: string[],
    pipelineId: string
  ): (input: RequestInput) => boolean {
    return (input) => keywords.some(kw =>
      input.query.toLowerCase().includes(kw.toLowerCase())
    );
  }
}

// ===== 网关主类 =====

class CompoundAIGateway {
  private registry = new ComponentRegistry();
  private router = new Router();
  private rateLimiters = new Map<string, number>();

  // 注册组件
  addComponent(config: ComponentConfig): this {
    this.registry.register(config);
    return this;
  }

  // 注册处理管道
  addPipeline(config: PipelineConfig): this {
    this.router.registerPipeline(config);
    return this;
  }

  // 处理请求
  async handleRequest(input: RequestInput): Promise<PipelineResult> {
    // 限流检查
    if (!this.checkRateLimit(input.userId)) {
      throw new Error('请求频率超限，请稍后重试');
    }

    // 路由决策
    const pipeline = this.router.route(input);

    // 执行管道（简化版——生产环境应实现完整的组件执行逻辑）
    const startTime = Date.now();
    const stageResults: StageResult[] = [];

    for (const stageName of pipeline.stages) {
      const stageStart = Date.now();
      const component = this.registry.get(stageName);

      if (!component) {
        stageResults.push({ stage: stageName, success: false, latencyMs: 0 });
        continue;
      }

      // 模拟组件执行——生产环境应调用真实组件
      await this.executeComponent(component, input);

      stageResults.push({
        stage: stageName,
        success: true,
        latencyMs: Date.now() - stageStart
      });
    }

    return {
      pipelineId: pipeline.id,
      success: true,
      output: '处理完成',
      stages: stageResults,
      totalLatencyMs: Date.now() - startTime,
      cost: { tokens: 1500, dollars: 0.015 }
    };
  }

  private async executeComponent(
    config: ComponentConfig,
    input: RequestInput
  ): Promise<void> {
    // 组件执行逻辑——包含超时控制和重试机制
    const timeout = new Promise((_, reject) =>
      setTimeout(() => reject(new Error('组件执行超时')), config.timeoutMs)
    );

    try {
      await Promise.race([this.runComponent(config, input), timeout]);
    } catch (error) {
      // 重试逻辑
      for (let i = 0; i < config.retryCount; i++) {
        try {
          await Promise.race([this.runComponent(config, input), timeout]);
          return;
        } catch {
          await this.delay(1000 * Math.pow(2, i)); // 指数退避
        }
      }
      throw error;
    }
  }

  private async runComponent(
    config: ComponentConfig,
    input: RequestInput
  ): Promise<void> {
    // 生产环境：根据 component.type 调用对应的执行逻辑
    await new Promise(resolve => setTimeout(resolve, 50)); // 模拟执行
  }

  private checkRateLimit(userId: string): boolean {
    const now = Date.now();
    const lastRequest = this.rateLimiters.get(userId) || 0;
    if (now - lastRequest < 1000) return false; // 每秒最多1次
    this.rateLimiters.set(userId, now);
    return true;
  }

  private delay(ms: number): Promise<void> {
    return new Promise(resolve => setTimeout(resolve, ms));
  }
}

// ===== 使用示例 =====

async function main() {
  const gateway = new CompoundAIGateway();

  // 注册组件
  gateway.addComponent({ name: 'retrieval', type: 'retrieval', timeoutMs: 5000, retryCount: 2 });
  gateway.addComponent({ name: 'code-generator', type: 'llm', timeoutMs: 30000, retryCount: 1 });
  gateway.addComponent({ name: 'code-executor', type: 'executor', timeoutMs: 10000, retryCount: 0 });
  gateway.addComponent({ name: 'report-generator', type: 'llm', timeoutMs: 15000, retryCount: 1 });

  // 注册管道
  gateway.addPipeline({
    id: 'data-analysis',
    description: '数据分析管道',
    stages: ['retrieval', 'code-generator', 'code-executor', 'report-generator'],
    router_condition: Router.keywordRouter(['分析', '数据', '趋势', '统计'], 'data-analysis')
  });

  gateway.addPipeline({
    id: 'qa-pipeline',
    description: '问答管道',
    stages: ['retrieval', 'report-generator'],
    router_condition: Router.keywordRouter(['是什么', '为什么', '如何', '查询'], 'qa-pipeline')
  });

  gateway.router.setDefaultPipeline('qa-pipeline');

  // 处理请求
  const result = await gateway.handleRequest({
    query: '分析上季度销售数据的趋势',
    userId: 'user-001',
    priority: 'normal'
  });

  console.log('处理结果:', JSON.stringify(result, null, 2));
}

main();