546 lines
18 KiB
Python
546 lines
18 KiB
Python
"""Pipeline执行引擎 - 编排多阶段Agent任务链的执行"""
|
||
|
||
import asyncio
|
||
import logging
|
||
import time
|
||
import uuid
|
||
from datetime import datetime, timezone
|
||
from typing import Any, Optional
|
||
|
||
from .loader import PipelineLoader
|
||
from .schema import (
|
||
Pipeline,
|
||
PipelineResult,
|
||
PipelineStage,
|
||
StageResult,
|
||
StageStatus,
|
||
)
|
||
|
||
logger = logging.getLogger(__name__)
|
||
|
||
|
||
class PipelineExecutionError(Exception):
|
||
"""Pipeline执行错误"""
|
||
|
||
def __init__(self, pipeline_name: str, reason: str = ""):
|
||
self.pipeline_name = pipeline_name
|
||
self.reason = reason
|
||
super().__init__(f"Pipeline '{pipeline_name}' execution error: {reason}")
|
||
|
||
|
||
class PipelineEngine:
|
||
"""
|
||
Pipeline执行引擎。
|
||
|
||
负责按照YAML定义的依赖关系（DAG）编排多阶段Agent任务的执行，
|
||
支持变量传递、超时控制、重试、条件执行等。
|
||
|
||
如果dispatcher为None，则使用模拟执行（dry-run模式），
|
||
用于测试和调试Pipeline定义。
|
||
"""
|
||
|
||
def __init__(self, dispatcher=None):
|
||
"""
|
||
Args:
|
||
dispatcher: TaskDispatcher实例（可选）。
|
||
如果为None，使用模拟执行（dry-run模式）。
|
||
"""
|
||
self.dispatcher = dispatcher
|
||
self.loader = PipelineLoader()
|
||
|
||
async def execute(
|
||
self,
|
||
pipeline: Pipeline,
|
||
context: dict | None = None,
|
||
) -> PipelineResult:
|
||
"""
|
||
执行整个Pipeline。
|
||
|
||
执行流程：
|
||
1. 合并全局变量和传入context
|
||
2. 拓扑排序确定执行顺序
|
||
3. 按依赖顺序逐stage执行
|
||
4. 变量解析：将${...}替换为实际值
|
||
5. 调用dispatcher分发任务给对应agent
|
||
6. 收集结果，传递给下游stage
|
||
|
||
Args:
|
||
pipeline: Pipeline定义对象
|
||
context: 外部传入的上下文变量（覆盖pipeline.variables中的默认值）
|
||
|
||
Returns:
|
||
PipelineResult包含所有阶段的执行结果
|
||
"""
|
||
start_time = time.monotonic()
|
||
pipeline_name = pipeline.name
|
||
|
||
# 构建执行上下文：合并全局变量和外部context
|
||
exec_context: dict[str, Any] = dict(pipeline.variables)
|
||
if context:
|
||
exec_context.update(context)
|
||
|
||
# 初始化stages结果存储
|
||
stages_results: dict[str, StageResult] = {}
|
||
|
||
# 构建stages上下文（用于变量解析）
|
||
stages_context: dict[str, Any] = {}
|
||
|
||
# 拓扑排序
|
||
try:
|
||
sorted_stages = self._topological_sort(pipeline.stages)
|
||
except ValueError as e:
|
||
duration_ms = int((time.monotonic() - start_time) * 1000)
|
||
return PipelineResult(
|
||
pipeline_name=pipeline_name,
|
||
status=StageStatus.FAILED,
|
||
stages_results=stages_results,
|
||
duration_ms=duration_ms,
|
||
error=str(e),
|
||
)
|
||
|
||
logger.info(
|
||
f"Pipeline '{pipeline_name}' starting execution "
|
||
f"({len(sorted_stages)} stages, order: {[s.name for s in sorted_stages]})"
|
||
)
|
||
|
||
# 逐阶段执行
|
||
overall_status = StageStatus.COMPLETED
|
||
failed_stages: set[str] = set()
|
||
skipped_stages: set[str] = set()
|
||
|
||
for stage in sorted_stages:
|
||
# 检查是否应该跳过此阶段
|
||
if self._should_skip(stage, failed_stages, skipped_stages):
|
||
skipped_stages.add(stage.name)
|
||
stages_results[stage.name] = StageResult(
|
||
stage_name=stage.name,
|
||
status=StageStatus.SKIPPED,
|
||
)
|
||
# 将跳过阶段的信息写入stages_context
|
||
stages_context[stage.name] = {"outputs": {}, "status": StageStatus.SKIPPED}
|
||
logger.info(f"Stage '{stage.name}' skipped (dependency failed)")
|
||
continue
|
||
|
||
# 检查条件表达式
|
||
if stage.condition and not self._evaluate_condition(stage.condition, exec_context, stages_context):
|
||
skipped_stages.add(stage.name)
|
||
stages_results[stage.name] = StageResult(
|
||
stage_name=stage.name,
|
||
status=StageStatus.SKIPPED,
|
||
)
|
||
stages_context[stage.name] = {"outputs": {}, "status": StageStatus.SKIPPED}
|
||
logger.info(f"Stage '{stage.name}' skipped (condition not met: {stage.condition})")
|
||
continue
|
||
|
||
# 执行阶段
|
||
stage_result = await self._execute_stage(stage, exec_context, stages_context)
|
||
stages_results[stage.name] = stage_result
|
||
|
||
# 将结果写入stages_context供下游引用
|
||
stages_context[stage.name] = {
|
||
"outputs": stage_result.outputs,
|
||
"status": stage_result.status,
|
||
}
|
||
|
||
if stage_result.status == StageStatus.FAILED:
|
||
if stage.continue_on_failure:
|
||
logger.warning(
|
||
f"Stage '{stage.name}' failed but continue_on_failure=True"
|
||
)
|
||
else:
|
||
failed_stages.add(stage.name)
|
||
overall_status = StageStatus.FAILED
|
||
logger.error(
|
||
f"Stage '{stage.name}' failed: {stage_result.error}"
|
||
)
|
||
|
||
logger.info(
|
||
f"Stage '{stage.name}' {stage_result.status.value} "
|
||
f"({stage_result.duration_ms}ms)"
|
||
)
|
||
|
||
duration_ms = int((time.monotonic() - start_time) * 1000)
|
||
|
||
result = PipelineResult(
|
||
pipeline_name=pipeline_name,
|
||
status=overall_status,
|
||
stages_results=stages_results,
|
||
duration_ms=duration_ms,
|
||
)
|
||
|
||
logger.info(
|
||
f"Pipeline '{pipeline_name}' finished with status={overall_status.value} "
|
||
f"({duration_ms}ms)"
|
||
)
|
||
|
||
return result
|
||
|
||
async def _execute_stage(
|
||
self,
|
||
stage: PipelineStage,
|
||
exec_context: dict[str, Any],
|
||
stages_context: dict[str, Any],
|
||
) -> StageResult:
|
||
"""
|
||
执行单个Stage。
|
||
|
||
1. 解析inputs中的变量引用
|
||
2. 通过dispatcher向agent发送任务
|
||
3. 等待结果（或超时）
|
||
4. 如果dispatcher为None，返回模拟结果（dry-run模式）
|
||
|
||
Args:
|
||
stage: 阶段定义
|
||
exec_context: 全局执行上下文
|
||
stages_context: 已完成阶段的结果上下文
|
||
|
||
Returns:
|
||
StageResult
|
||
"""
|
||
start_time = time.monotonic()
|
||
|
||
# 构建变量解析的完整上下文
|
||
resolve_ctx = dict(exec_context)
|
||
resolve_ctx["stages"] = stages_context
|
||
|
||
# 解析inputs中的变量引用
|
||
resolved_inputs = self._resolve_stage_inputs(stage.inputs, resolve_ctx)
|
||
|
||
logger.info(
|
||
f"Stage '{stage.name}' starting "
|
||
f"(agent={stage.agent}, action={stage.action})"
|
||
)
|
||
|
||
# 执行（带重试）
|
||
last_error: str | None = None
|
||
for attempt in range(stage.retry_count + 1):
|
||
if attempt > 0:
|
||
logger.info(
|
||
f"Stage '{stage.name}' retry attempt {attempt}/{stage.retry_count}"
|
||
)
|
||
|
||
try:
|
||
if self.dispatcher is not None:
|
||
# 实际执行：通过dispatcher分发任务
|
||
result = await self._dispatch_and_wait(stage, resolved_inputs)
|
||
else:
|
||
# Dry-run模式：返回模拟结果
|
||
result = self._dry_run_stage(stage, resolved_inputs)
|
||
|
||
# 成功
|
||
duration_ms = int((time.monotonic() - start_time) * 1000)
|
||
result.duration_ms = duration_ms
|
||
return result
|
||
|
||
except asyncio.TimeoutError:
|
||
last_error = f"Stage timed out after {stage.timeout_seconds}s"
|
||
logger.warning(
|
||
f"Stage '{stage.name}' timeout (attempt {attempt + 1}/{stage.retry_count + 1})"
|
||
)
|
||
|
||
except Exception as e:
|
||
last_error = str(e)
|
||
logger.warning(
|
||
f"Stage '{stage.name}' error (attempt {attempt + 1}/{stage.retry_count + 1}): {e}"
|
||
)
|
||
|
||
# 所有重试都失败
|
||
duration_ms = int((time.monotonic() - start_time) * 1000)
|
||
return StageResult(
|
||
stage_name=stage.name,
|
||
status=StageStatus.FAILED,
|
||
error=last_error,
|
||
duration_ms=duration_ms,
|
||
)
|
||
|
||
async def _dispatch_and_wait(
|
||
self,
|
||
stage: PipelineStage,
|
||
resolved_inputs: dict[str, Any],
|
||
) -> StageResult:
|
||
"""
|
||
通过TaskDispatcher分发任务并等待结果。
|
||
|
||
Args:
|
||
stage: 阶段定义
|
||
resolved_inputs: 已解析的输入参数
|
||
|
||
Returns:
|
||
StageResult
|
||
|
||
Raises:
|
||
asyncio.TimeoutError: 超时
|
||
Exception: 分发或执行错误
|
||
"""
|
||
from app.agent_framework.protocol import TaskMessage, TaskResult, TaskStatus
|
||
|
||
task_id = str(uuid.uuid4())
|
||
task_message = TaskMessage(
|
||
task_id=task_id,
|
||
agent_name=stage.agent,
|
||
task_type=stage.action,
|
||
priority=0,
|
||
input_data=resolved_inputs,
|
||
callback_url=None,
|
||
created_at=datetime.now(timezone.utc),
|
||
timeout_seconds=stage.timeout_seconds,
|
||
)
|
||
|
||
# 分发任务
|
||
dispatched_id = await self.dispatcher.dispatch(task_message)
|
||
|
||
# 等待结果（轮询方式）
|
||
elapsed = 0.0
|
||
poll_interval = 1.0 # 1秒轮询一次
|
||
while elapsed < stage.timeout_seconds:
|
||
await asyncio.sleep(poll_interval)
|
||
elapsed += poll_interval
|
||
|
||
task_status = await self.dispatcher.get_task_status(dispatched_id)
|
||
|
||
if task_status.get("status") == TaskStatus.COMPLETED:
|
||
output_data = task_status.get("output_data", {})
|
||
# 构建输出：只提取stage.outputs中声明的变量
|
||
outputs = self._extract_outputs(stage, output_data)
|
||
return StageResult(
|
||
stage_name=stage.name,
|
||
status=StageStatus.COMPLETED,
|
||
outputs=outputs,
|
||
)
|
||
|
||
elif task_status.get("status") == TaskStatus.FAILED:
|
||
error_msg = task_status.get("error_message", "Unknown error")
|
||
return StageResult(
|
||
stage_name=stage.name,
|
||
status=StageStatus.FAILED,
|
||
error=error_msg,
|
||
)
|
||
|
||
elif task_status.get("status") == TaskStatus.CANCELLED:
|
||
return StageResult(
|
||
stage_name=stage.name,
|
||
status=StageStatus.FAILED,
|
||
error="Task was cancelled",
|
||
)
|
||
|
||
# 超时
|
||
raise asyncio.TimeoutError()
|
||
|
||
def _dry_run_stage(
|
||
self,
|
||
stage: PipelineStage,
|
||
resolved_inputs: dict[str, Any],
|
||
) -> StageResult:
|
||
"""
|
||
Dry-run模式执行：模拟Agent返回结果。仅用于测试/开发环境。
|
||
|
||
在没有dispatcher的环境下使用，用于测试和调试Pipeline定义。
|
||
如果在生产环境中触发，则记录 ERROR 级别告警。
|
||
|
||
Args:
|
||
stage: 阶段定义
|
||
resolved_inputs: 已解析的输入参数
|
||
|
||
Returns:
|
||
模拟的StageResult
|
||
"""
|
||
import os
|
||
if os.environ.get("ENV", "development") == "production":
|
||
logger.error(
|
||
f"Pipeline 进入 dry-run 模式（stage={stage.name}）！"
|
||
"生产环境中 TaskDispatcher 未正确初始化，请检查系统配置。"
|
||
)
|
||
else:
|
||
logger.warning(f"[DRY-RUN] stage={stage.name} 返回模拟输出")
|
||
|
||
# 为声明的输出变量生成模拟值
|
||
mock_outputs: dict[str, Any] = {}
|
||
for output_name in stage.outputs:
|
||
mock_outputs[output_name] = f"[dry-run] {output_name} from {stage.agent}.{stage.action}"
|
||
|
||
logger.info(
|
||
f"Stage '{stage.name}' dry-run completed "
|
||
f"(agent={stage.agent}, action={stage.action}, "
|
||
f"inputs={list(resolved_inputs.keys())}, "
|
||
f"outputs={stage.outputs})"
|
||
)
|
||
|
||
return StageResult(
|
||
stage_name=stage.name,
|
||
status=StageStatus.COMPLETED,
|
||
outputs=mock_outputs,
|
||
)
|
||
|
||
def _topological_sort(self, stages: list[PipelineStage]) -> list[PipelineStage]:
|
||
"""
|
||
拓扑排序（Kahn算法）。
|
||
|
||
根据depends_on依赖关系确定执行顺序。
|
||
无依赖的阶段先执行，有依赖的后执行。
|
||
支持并行执行无依赖关系的阶段（本实现为串行，但顺序正确）。
|
||
|
||
Args:
|
||
stages: PipelineStage列表
|
||
|
||
Returns:
|
||
排序后的PipelineStage列表
|
||
|
||
Raises:
|
||
ValueError: 存在循环依赖
|
||
"""
|
||
name_to_stage = {s.name: s for s in stages}
|
||
stage_names = set(name_to_stage.keys())
|
||
|
||
# 构建邻接表和入度表
|
||
in_degree: dict[str, int] = {name: 0 for name in stage_names}
|
||
adj: dict[str, list[str]] = {name: [] for name in stage_names}
|
||
|
||
for stage in stages:
|
||
for dep in stage.depends_on:
|
||
if dep in stage_names:
|
||
adj[dep].append(stage.name)
|
||
in_degree[stage.name] += 1
|
||
|
||
# 入度为0的节点入队
|
||
queue: list[str] = sorted(
|
||
[name for name, deg in in_degree.items() if deg == 0]
|
||
)
|
||
sorted_names: list[str] = []
|
||
|
||
while queue:
|
||
# 按名称排序保证确定性
|
||
node = queue.pop(0)
|
||
sorted_names.append(node)
|
||
|
||
for neighbor in sorted(adj[node]):
|
||
in_degree[neighbor] -= 1
|
||
if in_degree[neighbor] == 0:
|
||
queue.append(neighbor)
|
||
queue.sort() # 保证确定性顺序
|
||
|
||
if len(sorted_names) != len(stage_names):
|
||
raise ValueError(
|
||
f"Cyclic dependency detected in pipeline stages: "
|
||
f"only {len(sorted_names)}/{len(stage_names)} stages sorted"
|
||
)
|
||
|
||
return [name_to_stage[name] for name in sorted_names]
|
||
|
||
def _resolve_stage_inputs(
|
||
self,
|
||
inputs: dict[str, Any],
|
||
context: dict[str, Any],
|
||
) -> dict[str, Any]:
|
||
"""
|
||
解析stage的输入变量。
|
||
|
||
将inputs中的${...}变量引用替换为context中的实际值。
|
||
|
||
Args:
|
||
inputs: 原始输入参数（可能包含变量引用）
|
||
context: 完整上下文（全局变量 + stages上下文）
|
||
|
||
Returns:
|
||
解析后的输入参数
|
||
"""
|
||
return PipelineLoader.resolve_variables(inputs, context)
|
||
|
||
def _should_skip(
|
||
self,
|
||
stage: PipelineStage,
|
||
failed_stages: set[str],
|
||
skipped_stages: set[str],
|
||
) -> bool:
|
||
"""
|
||
判断是否应该跳过此阶段。
|
||
|
||
如果某个依赖的阶段失败了且不是continue_on_failure，
|
||
则此阶段应跳过。
|
||
|
||
Args:
|
||
stage: 阶段定义
|
||
failed_stages: 已失败的阶段集合
|
||
skipped_stages: 已跳过的阶段集合
|
||
|
||
Returns:
|
||
True如果应该跳过
|
||
"""
|
||
for dep in stage.depends_on:
|
||
if dep in failed_stages or dep in skipped_stages:
|
||
return True
|
||
return False
|
||
|
||
def _evaluate_condition(
|
||
self,
|
||
condition: str,
|
||
exec_context: dict[str, Any],
|
||
stages_context: dict[str, Any],
|
||
) -> bool:
|
||
"""
|
||
评估条件表达式。
|
||
|
||
当前支持简单的变量存在性检查和等值比较：
|
||
- "${var}" → 变量存在且非空
|
||
- "${var} == 'value'" → 变量等于指定值
|
||
- "${var} != 'value'" → 变量不等于指定值
|
||
|
||
Args:
|
||
condition: 条件表达式字符串
|
||
exec_context: 全局执行上下文
|
||
stages_context: 阶段结果上下文
|
||
|
||
Returns:
|
||
True如果条件满足
|
||
"""
|
||
resolve_ctx = dict(exec_context)
|
||
resolve_ctx["stages"] = stages_context
|
||
|
||
# 先解析变量
|
||
resolved = PipelineLoader.resolve_variables(condition, resolve_ctx)
|
||
|
||
if isinstance(resolved, bool):
|
||
return resolved
|
||
|
||
resolved_str = str(resolved).strip()
|
||
|
||
# 等值比较
|
||
if "==" in resolved_str:
|
||
left, right = resolved_str.split("==", 1)
|
||
return left.strip().strip("'\"") == right.strip().strip("'\"")
|
||
|
||
# 不等比较
|
||
if "!=" in resolved_str:
|
||
left, right = resolved_str.split("!=", 1)
|
||
return left.strip().strip("'\"") != right.strip().strip("'\"")
|
||
|
||
# 非空判断
|
||
return bool(resolved_str) and resolved_str.lower() not in ("false", "0", "none", "")
|
||
|
||
def _extract_outputs(
|
||
self,
|
||
stage: PipelineStage,
|
||
output_data: dict[str, Any] | None,
|
||
) -> dict[str, Any]:
|
||
"""
|
||
从Agent的output_data中提取stage声明的输出变量。
|
||
|
||
Args:
|
||
stage: 阶段定义
|
||
output_data: Agent返回的完整输出数据
|
||
|
||
Returns:
|
||
仅包含声明输出的字典
|
||
"""
|
||
if not output_data:
|
||
return {}
|
||
|
||
if not stage.outputs:
|
||
# 如果未声明outputs，返回全部output_data
|
||
return output_data
|
||
|
||
return {
|
||
key: output_data.get(key)
|
||
for key in stage.outputs
|
||
if key in output_data
|
||
}
|