Coverage for agentos/tools/pipeline.py: 0%
145 statements
« prev ^ index » next coverage.py v7.14.3, created at 2026-07-08 12:22 +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 collections.abc import Callable
16from dataclasses import dataclass, field
17from enum import Enum, auto
18from typing import Any, Generic, TypeVar
20T = TypeVar("T")
21U = TypeVar("U")
24# ============================================================================
25# Core Types
26# ============================================================================
29class StageStatus(Enum):
30 IDLE = auto()
31 RUNNING = auto()
32 PAUSED = auto()
33 STOPPED = auto()
34 ERROR = auto()
37@dataclass
38class PipelineContext:
39 """Shared context flowing through the pipeline."""
41 data: dict[str, Any] = field(default_factory=dict)
42 metadata: dict[str, Any] = field(default_factory=dict)
44 def get(self, key: str, default: Any = None) -> Any:
45 return self.data.get(key, default)
47 def set(self, key: str, value: Any) -> None:
48 self.data[key] = value
51# ============================================================================
52# Stage
53# ============================================================================
55_NO_FALLBACK = object()
58class Stage(Generic[T, U], ABC):
59 """Abstract pipeline stage. Transforms T → U."""
61 def __init__(self, name: str = "", max_retries: int = 0):
62 self.name = name or self.__class__.__name__
63 self.max_retries = max_retries
64 self.status: StageStatus = StageStatus.IDLE
65 self._error: Exception | None = None
66 self._items_processed: int = 0
67 self._items_errored: int = 0
69 @abstractmethod
70 def process(self, item: T, ctx: PipelineContext) -> U: ...
72 def on_error(self, item: T, error: Exception, ctx: PipelineContext) -> Any:
73 """Override to provide fallback on error. Return _NO_FALLBACK to propagate."""
74 return _NO_FALLBACK
76 def execute(self, item: T, ctx: PipelineContext) -> U:
77 self.status = StageStatus.RUNNING
78 for attempt in range(self.max_retries + 1):
79 try:
80 result = self.process(item, ctx)
81 self._items_processed += 1
82 self.status = StageStatus.IDLE
83 return result
84 except Exception as e:
85 self._items_errored += 1
86 if attempt < self.max_retries:
87 time.sleep(0.01 * (attempt + 1))
88 continue
89 fallback = self.on_error(item, e, ctx)
90 if fallback is not _NO_FALLBACK:
91 self.status = StageStatus.IDLE
92 return fallback
93 self._error = e
94 self.status = StageStatus.ERROR
95 raise
97 @property
98 def stats(self) -> dict[str, Any]:
99 return {
100 "name": self.name,
101 "status": self.status.name,
102 "items_processed": self._items_processed,
103 "items_errored": self._items_errored,
104 }
107class LambdaStage(Stage[T, U]):
108 """Convenience stage from a callable."""
110 def __init__(self, fn: Callable[[T, PipelineContext], U], name: str = "", max_retries: int = 0):
111 super().__init__(name=name, max_retries=max_retries)
112 self._fn = fn
114 def process(self, item: T, ctx: PipelineContext) -> U:
115 return self._fn(item, ctx)
118# ============================================================================
119# Pipeline
120# ============================================================================
123class Pipeline(Generic[T, U]):
124 """Linear pipeline: a sequence of stages T → ? → ... → U."""
126 def __init__(self, name: str = "pipeline"):
127 self.name = name
128 self._stages: list[Stage] = []
129 self._lock = threading.Lock()
130 self._ctx = PipelineContext()
131 self.status: StageStatus = StageStatus.IDLE
133 def add_stage(self, stage: Stage) -> "Pipeline":
134 with self._lock:
135 self._stages.append(stage)
136 return self
138 def then(
139 self, fn: Callable[[Any, PipelineContext], Any], name: str = "", max_retries: int = 0
140 ) -> "Pipeline":
141 """Fluent API: add a lambda stage."""
142 return self.add_stage(LambdaStage(fn, name=name, max_retries=max_retries))
144 def run(self, input_item: T) -> U:
145 """Run pipeline on a single item."""
146 current = input_item
147 self.status = StageStatus.RUNNING
148 try:
149 for stage in self._stages:
150 current = stage.execute(current, self._ctx)
151 return current
152 finally:
153 all_idle = all(s.status == StageStatus.IDLE for s in self._stages)
154 self.status = StageStatus.IDLE if all_idle else StageStatus.ERROR
156 def run_batch(self, items: list[T]) -> list[U]:
157 """Run pipeline on a batch."""
158 results = []
159 for item in items:
160 results.append(self.run(item))
161 return results
163 @property
164 def context(self) -> PipelineContext:
165 return self._ctx
167 @property
168 def stats(self) -> dict[str, Any]:
169 return {
170 "name": self.name,
171 "status": self.status.name,
172 "stages": [s.stats for s in self._stages],
173 }
176# ============================================================================
177# ParallelPipeline — Fan-out / Fan-in
178# ============================================================================
181class ParallelPipeline(Generic[T, U]):
182 """Branches: split input across parallel stages, then merge results.
184 Fan-out: single input → all branches simultaneously.
185 Fan-in: all branch outputs → merge function → single output.
186 """
188 def __init__(self, name: str = "parallel_pipeline"):
189 self.name = name
190 self._branches: list[Pipeline] = []
191 self._merge: Callable[[list[Any], PipelineContext], U] | None = None
192 self._lock = threading.Lock()
193 self._ctx = PipelineContext()
195 def branch(self, pipeline: Pipeline) -> "ParallelPipeline":
196 with self._lock:
197 self._branches.append(pipeline)
198 return self
200 def merge(self, fn: Callable[[list[Any], PipelineContext], U]) -> "ParallelPipeline":
201 self._merge = fn
202 return self
204 def run(self, input_item: T) -> U:
205 import concurrent.futures
207 with concurrent.futures.ThreadPoolExecutor(max_workers=len(self._branches)) as pool:
208 futures = {
209 pool.submit(branch.run, input_item): i for i, branch in enumerate(self._branches)
210 }
211 results = [None] * len(self._branches)
212 for future in concurrent.futures.as_completed(futures):
213 idx = futures[future]
214 results[idx] = future.result()
216 if self._merge:
217 return self._merge(results, self._ctx)
218 return results # type: ignore
220 @property
221 def context(self) -> PipelineContext:
222 return self._ctx
225# ============================================================================
226# Stage helpers
227# ============================================================================
230class FilterStage(Stage[T, T]):
231 """Pass-through stage that filters items."""
233 def __init__(
234 self,
235 predicate: Callable[[T, PipelineContext], bool],
236 name: str = "filter",
237 max_retries: int = 0,
238 ):
239 super().__init__(name=name, max_retries=max_retries)
240 self._predicate = predicate
242 def process(self, item: T, ctx: PipelineContext) -> T:
243 if not self._predicate(item, ctx):
244 raise FilterDrop()
245 return item
247 def on_error(self, item: T, error: Exception, ctx: PipelineContext) -> T | None:
248 if isinstance(error, FilterDrop):
249 return None
250 return super().on_error(item, error, ctx)
253class FilterDrop(Exception): # noqa: N818
254 """Signal that an item should be filtered out."""
258class BatchStage(Stage[list[T], list[U]]):
259 """Accumulates items into batches before processing."""
261 def __init__(
262 self,
263 batch_size: int,
264 fn: Callable[[list[T], PipelineContext], list[U]],
265 name: str = "batch",
266 max_retries: int = 0,
267 ):
268 super().__init__(name=name, max_retries=max_retries)
269 self.batch_size = batch_size
270 self._buffer: list[T] = []
271 self._fn = fn
273 def process(self, item: list[T], ctx: PipelineContext) -> list[U]:
274 return self._fn(item, ctx)