Coverage for src\agents_sdk_models\flow.py: 80%

212 statements  

« prev     ^ index     » next       coverage.py v7.8.0, created at 2025-06-04 17:38 +0900

1from __future__ import annotations 

2 

3"""Flow — Workflow orchestration engine for Step-based workflows. 

4 

5Flowはステップベースワークフロー用のワークフローオーケストレーションエンジンです。 

6同期・非同期両方のインターフェースを提供し、CLI、GUI、チャットボット対応します。 

7""" 

8 

9import asyncio 

10import logging 

11from typing import Any, Dict, List, Optional, Callable, Union 

12from datetime import datetime 

13import traceback 

14 

15from .context import Context 

16from .step import Step 

17 

18 

19logger = logging.getLogger(__name__) 

20 

21 

22class FlowExecutionError(Exception): 

23 """ 

24 Exception raised during flow execution 

25 フロー実行中に発生する例外 

26 """ 

27 pass 

28 

29 

30class Flow: 

31 """ 

32 Flow orchestration engine for Step-based workflows 

33 ステップベースワークフロー用フローオーケストレーションエンジン 

34  

35 This class provides: 

36 このクラスは以下を提供します: 

37 - Declarative step-based workflow definition / 宣言的ステップベースワークフロー定義 

38 - Synchronous and asynchronous execution modes / 同期・非同期実行モード 

39 - User input coordination for interactive workflows / 対話的ワークフロー用ユーザー入力調整 

40 - Error handling and observability / エラーハンドリングとオブザーバビリティ 

41 """ 

42 

43 def __init__( 

44 self, 

45 start: str, 

46 steps: Dict[str, Step], 

47 context: Optional[Context] = None, 

48 max_steps: int = 1000, 

49 trace_id: Optional[str] = None 

50 ): 

51 """ 

52 Initialize Flow with start step and step definitions 

53 開始ステップとステップ定義でFlowを初期化 

54  

55 Args: 

56 start: Start step label / 開始ステップラベル 

57 steps: Dictionary of step definitions / ステップ定義の辞書 

58 context: Initial context (optional) / 初期コンテキスト(オプション) 

59 max_steps: Maximum number of steps to prevent infinite loops / 無限ループ防止のための最大ステップ数 

60 trace_id: Trace ID for observability / オブザーバビリティ用トレースID 

61 """ 

62 self.start = start 

63 self.steps = steps 

64 self.context = context or Context() 

65 self.max_steps = max_steps 

66 self.trace_id = trace_id or f"flow_{datetime.now().strftime('%Y%m%d_%H%M%S')}" 

67 

68 # Initialize context 

69 # コンテキストを初期化 

70 self.context.trace_id = self.trace_id 

71 self.context.next_label = start 

72 

73 # Execution state 

74 # 実行状態 

75 self._running = False 

76 self._run_loop_task: Optional[asyncio.Task] = None 

77 self._execution_lock = asyncio.Lock() 

78 

79 # Hooks for observability 

80 # オブザーバビリティ用フック 

81 self.before_step_hooks: List[Callable[[str, Context], None]] = [] 

82 self.after_step_hooks: List[Callable[[str, Context, Any], None]] = [] 

83 self.error_hooks: List[Callable[[str, Context, Exception], None]] = [] 

84 

85 @property 

86 def finished(self) -> bool: 

87 """ 

88 Check if flow is finished 

89 フローが完了しているかチェック 

90  

91 Returns: 

92 bool: True if finished / 完了している場合True 

93 """ 

94 return self.context.is_finished() 

95 

96 @property 

97 def current_step_name(self) -> Optional[str]: 

98 """ 

99 Get current step name 

100 現在のステップ名を取得 

101  

102 Returns: 

103 str | None: Current step name / 現在のステップ名 

104 """ 

105 return self.context.current_step 

106 

107 @property 

108 def next_step_name(self) -> Optional[str]: 

109 """ 

110 Get next step name 

111 次のステップ名を取得 

112  

113 Returns: 

114 str | None: Next step name / 次のステップ名 

115 """ 

116 return self.context.next_label 

117 

118 async def run(self, initial_input: Optional[str] = None) -> Context: 

119 """ 

120 Run flow to completion without user input coordination 

121 ユーザー入力調整なしでフローを完了まで実行 

122  

123 This is for non-interactive workflows that don't require user input. 

124 これはユーザー入力が不要な非対話的ワークフロー用です。 

125  

126 Args: 

127 initial_input: Initial input to the flow / フローへの初期入力 

128  

129 Returns: 

130 Context: Final context / 最終コンテキスト 

131  

132 Raises: 

133 FlowExecutionError: If execution fails / 実行失敗時 

134 """ 

135 async with self._execution_lock: 

136 try: 

137 self._running = True 

138 

139 # Reset context for new execution 

140 # 新しい実行用にコンテキストをリセット 

141 if self.context.step_count > 0: 

142 self.context = Context(trace_id=self.trace_id) 

143 self.context.next_label = self.start 

144 

145 # Add initial input if provided 

146 # 初期入力が提供されている場合は追加 

147 if initial_input: 

148 self.context.add_user_message(initial_input) 

149 

150 current_input = initial_input 

151 step_count = 0 

152 

153 while not self.finished and step_count < self.max_steps: 

154 step_name = self.context.next_label 

155 if not step_name or step_name not in self.steps: 

156 break 

157 

158 step = self.steps[step_name] 

159 

160 # Execute step 

161 # ステップを実行 

162 try: 

163 await self._execute_step(step, current_input) 

164 current_input = None # Only use initial input for first step 

165 step_count += 1 

166 

167 # If step is waiting for user input, break 

168 # ステップがユーザー入力を待機している場合、中断 

169 if self.context.awaiting_user_input: 

170 break 

171 

172 except Exception as e: 

173 logger.error(f"Error executing step {step_name}: {e}") 

174 self._handle_step_error(step_name, e) 

175 break 

176 

177 # Check for infinite loop 

178 # 無限ループのチェック 

179 if step_count >= self.max_steps: 

180 raise FlowExecutionError(f"Flow exceeded maximum steps ({self.max_steps})") 

181 

182 return self.context 

183 

184 finally: 

185 self._running = False 

186 

187 async def run_loop(self) -> None: 

188 """ 

189 Run flow as background task with user input coordination 

190 ユーザー入力調整を含むバックグラウンドタスクとしてフローを実行 

191  

192 This method runs the flow continuously, pausing when user input is needed. 

193 このメソッドはフローを継続的に実行し、ユーザー入力が必要な時に一時停止します。 

194 Use feed() to provide user input when the flow is waiting. 

195 フローが待機している時はfeed()を使用してユーザー入力を提供してください。 

196 """ 

197 async with self._execution_lock: 

198 try: 

199 self._running = True 

200 

201 # Reset context for new execution 

202 # 新しい実行用にコンテキストをリセット 

203 if self.context.step_count > 0: 

204 self.context = Context(trace_id=self.trace_id) 

205 self.context.next_label = self.start 

206 

207 step_count = 0 

208 current_input = None 

209 

210 while not self.finished and step_count < self.max_steps: 

211 step_name = self.context.next_label 

212 if not step_name or step_name not in self.steps: 

213 break 

214 

215 step = self.steps[step_name] 

216 

217 # Execute step 

218 # ステップを実行 

219 try: 

220 await self._execute_step(step, current_input) 

221 current_input = None 

222 step_count += 1 

223 

224 # If step is waiting for user input, wait for feed() 

225 # ステップがユーザー入力を待機している場合、feed()を待つ 

226 if self.context.awaiting_user_input: 

227 await self.context.wait_for_user_input() 

228 # After receiving input, continue with the same step 

229 # 入力受信後、同じステップで継続 

230 current_input = self.context.last_user_input 

231 continue 

232 

233 except Exception as e: 

234 logger.error(f"Error executing step {step_name}: {e}") 

235 self._handle_step_error(step_name, e) 

236 break 

237 

238 # Check for infinite loop 

239 # 無限ループのチェック 

240 if step_count >= self.max_steps: 

241 raise FlowExecutionError(f"Flow exceeded maximum steps ({self.max_steps})") 

242 

243 finally: 

244 self._running = False 

245 

246 def next_prompt(self) -> Optional[str]: 

247 """ 

248 Get next prompt for synchronous CLI usage 

249 同期CLI使用用の次のプロンプトを取得 

250  

251 Returns: 

252 str | None: Prompt if waiting for user input / ユーザー入力待ちの場合のプロンプト 

253 """ 

254 return self.context.clear_prompt() 

255 

256 def feed(self, user_input: str) -> None: 

257 """ 

258 Provide user input to the flow 

259 フローにユーザー入力を提供 

260  

261 Args: 

262 user_input: User input text / ユーザー入力テキスト 

263 """ 

264 self.context.provide_user_input(user_input) 

265 

266 def step(self) -> None: 

267 """ 

268 Execute one step synchronously 

269 1ステップを同期的に実行 

270  

271 This method executes one step and returns immediately. 

272 このメソッドは1ステップを実行してすぐに返ります。 

273 Use for synchronous CLI applications. 

274 同期CLIアプリケーション用に使用してください。 

275 """ 

276 if self.finished: 

277 return 

278 

279 step_name = self.context.next_label 

280 if not step_name or step_name not in self.steps: 

281 self.context.finish() 

282 return 

283 

284 step = self.steps[step_name] 

285 

286 # Run step in event loop 

287 # イベントループでステップを実行 

288 try: 

289 loop = asyncio.get_event_loop() 

290 if loop.is_running(): 

291 # If loop is running, create a task 

292 # ループが実行中の場合、タスクを作成 

293 task = asyncio.create_task(self._execute_step(step, None)) 

294 # This is a synchronous method, so we can't await 

295 # これは同期メソッドなので、awaitできない 

296 # The task will run in the background 

297 # タスクはバックグラウンドで実行される 

298 else: 

299 # If no loop is running, run until complete 

300 # ループが実行されていない場合、完了まで実行 

301 loop.run_until_complete(self._execute_step(step, None)) 

302 except Exception as e: 

303 logger.error(f"Error executing step {step_name}: {e}") 

304 self._handle_step_error(step_name, e) 

305 

306 async def _execute_step(self, step: Step, user_input: Optional[str]) -> None: 

307 """ 

308 Execute a single step with hooks and error handling 

309 フックとエラーハンドリングで単一ステップを実行 

310  

311 Args: 

312 step: Step to execute / 実行するステップ 

313 user_input: User input if any / ユーザー入力(あれば) 

314 """ 

315 step_name = step.name 

316 

317 # Before step hooks 

318 # ステップ前フック 

319 for hook in self.before_step_hooks: 

320 try: 

321 hook(step_name, self.context) 

322 except Exception as e: 

323 logger.warning(f"Before step hook error: {e}") 

324 

325 start_time = datetime.now() 

326 result = None 

327 error = None 

328 

329 try: 

330 # Execute step 

331 # ステップを実行 

332 result = await step.run(user_input, self.context) 

333 if result != self.context: 

334 # Step returned a new context, use it 

335 # ステップが新しいコンテキストを返した場合、それを使用 

336 self.context = result 

337 

338 logger.debug(f"Step {step_name} completed in {datetime.now() - start_time}") 

339 

340 except Exception as e: 

341 error = e 

342 logger.error(f"Step {step_name} failed: {e}") 

343 logger.debug(traceback.format_exc()) 

344 

345 # Add error to context 

346 # エラーをコンテキストに追加 

347 self.context.add_system_message(f"Step {step_name} failed: {str(e)}") 

348 

349 # Call error hooks 

350 # エラーフックを呼び出し 

351 for hook in self.error_hooks: 

352 try: 

353 hook(step_name, self.context, e) 

354 except Exception as hook_error: 

355 logger.warning(f"Error hook failed: {hook_error}") 

356 

357 raise e 

358 

359 finally: 

360 # After step hooks 

361 # ステップ後フック 

362 for hook in self.after_step_hooks: 

363 try: 

364 hook(step_name, self.context, result) 

365 except Exception as e: 

366 logger.warning(f"After step hook error: {e}") 

367 

368 def _handle_step_error(self, step_name: str, error: Exception) -> None: 

369 """ 

370 Handle step execution error 

371 ステップ実行エラーを処理 

372  

373 Args: 

374 step_name: Name of the failed step / 失敗したステップの名前 

375 error: The error that occurred / 発生したエラー 

376 """ 

377 # Mark flow as finished on error 

378 # エラー時はフローを完了としてマーク 

379 self.context.finish() 

380 self.context.set_artifact("error", { 

381 "step": step_name, 

382 "error": str(error), 

383 "type": type(error).__name__ 

384 }) 

385 

386 def add_hook( 

387 self, 

388 hook_type: str, 

389 callback: Callable 

390 ) -> None: 

391 """ 

392 Add observability hook 

393 オブザーバビリティフックを追加 

394  

395 Args: 

396 hook_type: Type of hook ("before_step", "after_step", "error") / フックタイプ 

397 callback: Callback function / コールバック関数 

398 """ 

399 if hook_type == "before_step": 

400 self.before_step_hooks.append(callback) 

401 elif hook_type == "after_step": 

402 self.after_step_hooks.append(callback) 

403 elif hook_type == "error": 

404 self.error_hooks.append(callback) 

405 else: 

406 raise ValueError(f"Unknown hook type: {hook_type}") 

407 

408 def get_step_history(self) -> List[Dict[str, Any]]: 

409 """ 

410 Get execution history 

411 実行履歴を取得 

412  

413 Returns: 

414 List[Dict[str, Any]]: Step execution history / ステップ実行履歴 

415 """ 

416 history = [] 

417 for msg in self.context.messages: 

418 if msg.role == "system" and "Step" in msg.content: 

419 history.append({ 

420 "timestamp": msg.timestamp, 

421 "message": msg.content, 

422 "metadata": msg.metadata 

423 }) 

424 return history 

425 

426 def get_flow_summary(self) -> Dict[str, Any]: 

427 """ 

428 Get flow execution summary 

429 フロー実行サマリーを取得 

430  

431 Returns: 

432 Dict[str, Any]: Flow summary / フローサマリー 

433 """ 

434 return { 

435 "trace_id": self.trace_id, 

436 "start_step": self.start, 

437 "current_step": self.current_step_name, 

438 "next_step": self.next_step_name, 

439 "step_count": self.context.step_count, 

440 "finished": self.finished, 

441 "start_time": self.context.start_time, 

442 "artifacts": self.context.artifacts, 

443 "message_count": len(self.context.messages) 

444 } 

445 

446 def reset(self) -> None: 

447 """ 

448 Reset flow to initial state 

449 フローを初期状態にリセット 

450 """ 

451 self.context = Context(trace_id=self.trace_id) 

452 self.context.next_label = self.start 

453 self._running = False 

454 if self._run_loop_task: 

455 self._run_loop_task.cancel() 

456 self._run_loop_task = None 

457 

458 def stop(self) -> None: 

459 """ 

460 Stop flow execution 

461 フロー実行を停止 

462 """ 

463 self._running = False 

464 self.context.finish() 

465 if self._run_loop_task: 

466 self._run_loop_task.cancel() 

467 self._run_loop_task = None 

468 

469 async def start_background_task(self) -> asyncio.Task: 

470 """ 

471 Start flow as background task 

472 フローをバックグラウンドタスクとして開始 

473  

474 Returns: 

475 asyncio.Task: Background task / バックグラウンドタスク 

476 """ 

477 if self._run_loop_task and not self._run_loop_task.done(): 

478 raise RuntimeError("Flow is already running as background task") 

479 

480 self._run_loop_task = asyncio.create_task(self.run_loop()) 

481 return self._run_loop_task 

482 

483 def __str__(self) -> str: 

484 """String representation of flow""" 

485 return f"Flow(start={self.start}, steps={len(self.steps)}, finished={self.finished})" 

486 

487 def __repr__(self) -> str: 

488 return self.__str__() 

489 

490 

491# Utility functions for flow creation 

492# フロー作成用ユーティリティ関数 

493 

494def create_simple_flow( 

495 steps: List[tuple[str, Step]], 

496 context: Optional[Context] = None 

497) -> Flow: 

498 """ 

499 Create a simple linear flow from a list of steps 

500 ステップのリストから簡単な線形フローを作成 

501  

502 Args: 

503 steps: List of (name, step) tuples / (名前, ステップ)タプルのリスト 

504 context: Initial context / 初期コンテキスト 

505  

506 Returns: 

507 Flow: Created flow / 作成されたフロー 

508 """ 

509 if not steps: 

510 raise ValueError("At least one step is required") 

511 

512 step_dict = {} 

513 for i, (name, step) in enumerate(steps): 

514 # Set next step for each step 

515 # 各ステップの次ステップを設定 

516 if hasattr(step, 'next_step') and step.next_step is None: 

517 if i < len(steps) - 1: 

518 step.next_step = steps[i + 1][0] 

519 step_dict[name] = step 

520 

521 return Flow( 

522 start=steps[0][0], 

523 steps=step_dict, 

524 context=context 

525 ) 

526 

527 

528def create_conditional_flow( 

529 initial_step: Step, 

530 condition_step: Step, 

531 true_branch: List[Step], 

532 false_branch: List[Step], 

533 context: Optional[Context] = None 

534) -> Flow: 

535 """ 

536 Create a conditional flow with true/false branches 

537 true/falseブランチを持つ条件付きフローを作成 

538  

539 Args: 

540 initial_step: Initial step / 初期ステップ 

541 condition_step: Condition step / 条件ステップ 

542 true_branch: Steps for true branch / trueブランチのステップ 

543 false_branch: Steps for false branch / falseブランチのステップ 

544 context: Initial context / 初期コンテキスト 

545  

546 Returns: 

547 Flow: Created flow / 作成されたフロー 

548 """ 

549 steps = { 

550 "start": initial_step, 

551 "condition": condition_step 

552 } 

553 

554 # Add true branch steps 

555 # trueブランチステップを追加 

556 for i, step in enumerate(true_branch): 

557 step_name = f"true_{i}" 

558 steps[step_name] = step 

559 if i == 0 and hasattr(condition_step, 'if_true'): 

560 condition_step.if_true = step_name 

561 

562 # Add false branch steps 

563 # falseブランチステップを追加 

564 for i, step in enumerate(false_branch): 

565 step_name = f"false_{i}" 

566 steps[step_name] = step 

567 if i == 0 and hasattr(condition_step, 'if_false'): 

568 condition_step.if_false = step_name 

569 

570 # Connect initial step to condition 

571 # 初期ステップを条件に接続 

572 if hasattr(initial_step, 'next_step'): 

573 initial_step.next_step = "condition" 

574 

575 return Flow( 

576 start="start", 

577 steps=steps, 

578 context=context 

579 )