Coverage for agentos/tools/pipeline.py: 39%
145 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-06 11:37 +0800
1"""
2Pipeline — composable data processing pipeline for AgentOS.
4Features:
5- Linear and branching pipelines
6- Fan-out / fan-in patterns
7- Backpressure and quota control
8- Stage-level error handling and retry
9- Pipeline serialization for checkpoint/resume
10"""
12import threading
13import time
14from abc import ABC, abstractmethod
15from dataclasses import dataclass, field
16from enum import Enum, auto
17from typing import Any, Callable, Dict, Generic, List, Optional, TypeVar
19T = TypeVar("T")
20U = TypeVar("U")
23# ============================================================================
24# Core Types
25# ============================================================================
27class StageStatus(Enum):
28 IDLE = auto()
29 RUNNING = auto()
30 PAUSED = auto()
31 STOPPED = auto()
32 ERROR = auto()
35@dataclass
36class PipelineContext:
37 """Shared context flowing through the pipeline."""
38 data: Dict[str, Any] = field(default_factory=dict)
39 metadata: Dict[str, Any] = field(default_factory=dict)
41 def get(self, key: str, default: Any = None) -> Any:
42 return self.data.get(key, default)
44 def set(self, key: str, value: Any) -> None:
45 self.data[key] = value
48# ============================================================================
49# Stage
50# ============================================================================
52_NO_FALLBACK = object()
55class Stage(Generic[T, U], ABC):
56 """Abstract pipeline stage. Transforms T → U."""
58 def __init__(self, name: str = "", max_retries: int = 0):
59 self.name = name or self.__class__.__name__
60 self.max_retries = max_retries
61 self.status: StageStatus = StageStatus.IDLE
62 self._error: Optional[Exception] = None
63 self._items_processed: int = 0
64 self._items_errored: int = 0
66 @abstractmethod
67 def process(self, item: T, ctx: PipelineContext) -> U: ...
69 def on_error(self, item: T, error: Exception, ctx: PipelineContext) -> Any:
70 """Override to provide fallback on error. Return _NO_FALLBACK to propagate."""
71 return _NO_FALLBACK
73 def execute(self, item: T, ctx: PipelineContext) -> U:
74 self.status = StageStatus.RUNNING
75 for attempt in range(self.max_retries + 1):
76 try:
77 result = self.process(item, ctx)
78 self._items_processed += 1
79 self.status = StageStatus.IDLE
80 return result
81 except Exception as e:
82 self._items_errored += 1
83 if attempt < self.max_retries:
84 time.sleep(0.01 * (attempt + 1))
85 continue
86 fallback = self.on_error(item, e, ctx)
87 if fallback is not _NO_FALLBACK:
88 self.status = StageStatus.IDLE
89 return fallback
90 self._error = e
91 self.status = StageStatus.ERROR
92 raise
94 @property
95 def stats(self) -> Dict[str, Any]:
96 return {
97 "name": self.name,
98 "status": self.status.name,
99 "items_processed": self._items_processed,
100 "items_errored": self._items_errored,
101 }
104class LambdaStage(Stage[T, U]):
105 """Convenience stage from a callable."""
106 def __init__(self, fn: Callable[[T, PipelineContext], U], name: str = "", max_retries: int = 0):
107 super().__init__(name=name, max_retries=max_retries)
108 self._fn = fn
110 def process(self, item: T, ctx: PipelineContext) -> U:
111 return self._fn(item, ctx)
114# ============================================================================
115# Pipeline
116# ============================================================================
118class Pipeline(Generic[T, U]):
119 """Linear pipeline: a sequence of stages T → ? → ... → U."""
121 def __init__(self, name: str = "pipeline"):
122 self.name = name
123 self._stages: List[Stage] = []
124 self._lock = threading.Lock()
125 self._ctx = PipelineContext()
126 self.status: StageStatus = StageStatus.IDLE
128 def add_stage(self, stage: Stage) -> "Pipeline":
129 with self._lock:
130 self._stages.append(stage)
131 return self
133 def then(self, fn: Callable[[Any, PipelineContext], Any], name: str = "", max_retries: int = 0) -> "Pipeline":
134 """Fluent API: add a lambda stage."""
135 return self.add_stage(LambdaStage(fn, name=name, max_retries=max_retries))
137 def run(self, input_item: T) -> U:
138 """Run pipeline on a single item."""
139 current = input_item
140 self.status = StageStatus.RUNNING
141 try:
142 for stage in self._stages:
143 current = stage.execute(current, self._ctx)
144 return current
145 finally:
146 all_idle = all(s.status == StageStatus.IDLE for s in self._stages)
147 self.status = StageStatus.IDLE if all_idle else StageStatus.ERROR
149 def run_batch(self, items: List[T]) -> List[U]:
150 """Run pipeline on a batch."""
151 results = []
152 for item in items:
153 results.append(self.run(item))
154 return results
156 @property
157 def context(self) -> PipelineContext:
158 return self._ctx
160 @property
161 def stats(self) -> Dict[str, Any]:
162 return {
163 "name": self.name,
164 "status": self.status.name,
165 "stages": [s.stats for s in self._stages],
166 }
169# ============================================================================
170# ParallelPipeline — Fan-out / Fan-in
171# ============================================================================
173class ParallelPipeline(Generic[T, U]):
174 """Branches: split input across parallel stages, then merge results.
176 Fan-out: single input → all branches simultaneously.
177 Fan-in: all branch outputs → merge function → single output.
178 """
180 def __init__(self, name: str = "parallel_pipeline"):
181 self.name = name
182 self._branches: List[Pipeline] = []
183 self._merge: Optional[Callable[[List[Any], PipelineContext], U]] = None
184 self._lock = threading.Lock()
185 self._ctx = PipelineContext()
187 def branch(self, pipeline: Pipeline) -> "ParallelPipeline":
188 with self._lock:
189 self._branches.append(pipeline)
190 return self
192 def merge(self, fn: Callable[[List[Any], PipelineContext], U]) -> "ParallelPipeline":
193 self._merge = fn
194 return self
196 def run(self, input_item: T) -> U:
197 import concurrent.futures
199 with concurrent.futures.ThreadPoolExecutor(max_workers=len(self._branches)) as pool:
200 futures = {
201 pool.submit(branch.run, input_item): i
202 for i, branch in enumerate(self._branches)
203 }
204 results = [None] * len(self._branches)
205 for future in concurrent.futures.as_completed(futures):
206 idx = futures[future]
207 results[idx] = future.result()
209 if self._merge:
210 return self._merge(results, self._ctx)
211 return results # type: ignore
213 @property
214 def context(self) -> PipelineContext:
215 return self._ctx
218# ============================================================================
219# Stage helpers
220# ============================================================================
222class FilterStage(Stage[T, T]):
223 """Pass-through stage that filters items."""
224 def __init__(self, predicate: Callable[[T, PipelineContext], bool], name: str = "filter", max_retries: int = 0):
225 super().__init__(name=name, max_retries=max_retries)
226 self._predicate = predicate
228 def process(self, item: T, ctx: PipelineContext) -> T:
229 if not self._predicate(item, ctx):
230 raise FilterDrop()
231 return item
233 def on_error(self, item: T, error: Exception, ctx: PipelineContext) -> Optional[T]:
234 if isinstance(error, FilterDrop):
235 return None
236 return super().on_error(item, error, ctx)
239class FilterDrop(Exception):
240 """Signal that an item should be filtered out."""
241 pass
244class BatchStage(Stage[List[T], List[U]]):
245 """Accumulates items into batches before processing."""
246 def __init__(self, batch_size: int, fn: Callable[[List[T], PipelineContext], List[U]], name: str = "batch", max_retries: int = 0):
247 super().__init__(name=name, max_retries=max_retries)
248 self.batch_size = batch_size
249 self._buffer: List[T] = []
250 self._fn = fn
252 def process(self, item: List[T], ctx: PipelineContext) -> List[U]:
253 return self._fn(item, ctx)