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

1 from __future__ import annotations

2

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

4

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

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

7"""

8

9 import asyncio

10 import logging

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

12 from datetime import datetime

13 import traceback

14

15 from .context import Context

16 from .step import Step

17

18

19 logger = logging.getLogger(__name__)

20

21

22 class FlowExecutionError(Exception):

23 """

24 Exception raised during flow execution

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

26 """

27 pass

28

29

30 class 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

494 def 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

528 def 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 )