Coverage for src/dataknobs_fsm/execution/history.py: 16%

1"""Execution history tracking for FSM state machines."""

2

3 import time

4 from enum import Enum

5 from typing import Any, Dict, List

6

7 from dataknobs_structures import Tree

8

9 from dataknobs_fsm.core.data_modes import DataHandlingMode

10

11

12 class ExecutionStatus(Enum):

13 """Status of an execution step."""

14 PENDING = "pending"

15 IN_PROGRESS = "in_progress"

16 COMPLETED = "completed"

17 FAILED = "failed"

18 SKIPPED = "skipped"

19

20

21 class ExecutionStep:

22 """Represents a single step in the execution history."""

23

24 def __init__(

25 self,

26 step_id: str,

27 state_name: str,

28 network_name: str,

29 timestamp: float,

30 data_mode: DataHandlingMode = DataHandlingMode.COPY,

31 status: ExecutionStatus = ExecutionStatus.PENDING

32 ):

33 """Initialize execution step.

34

35 Args:

36 step_id: Unique identifier for this step.

37 state_name: Name of the state.

38 network_name: Name of the network.

39 timestamp: Unix timestamp when step was created.

40 data_mode: Data mode used for this step.

41 status: Current status of the step.

42 """

43 self.step_id = step_id

44 self.state_name = state_name

45 self.network_name = network_name

46 self.timestamp = timestamp

47 self.start_time: float | None = None

48 self.end_time: float | None = None

49 self.data_mode = data_mode

50 self.status = status

51

52 # Execution details

53 self.arc_taken: str | None = None

54 self.data_snapshot: Any | None = None

55 self.error: Exception | None = None

56 self.metrics: Dict[str, Any] = {}

57 self.resource_usage: Dict[str, Any] = {}

58

59 # Stream tracking

60 self.stream_progress: Dict[str, Any] | None = None

61 self.chunks_processed: int = 0

62 self.records_processed: int = 0

63

64 def start(self) -> None:

65 """Mark step as started."""

66 self.start_time = time.time()

67 self.status = ExecutionStatus.IN_PROGRESS

68

69 def complete(self, arc_taken: str | None = None) -> None:

70 """Mark step as completed.

71

72 Args:

73 arc_taken: The arc that was taken from this state.

74 """

75 self.end_time = time.time()

76 self.status = ExecutionStatus.COMPLETED

77 self.arc_taken = arc_taken

78

79 def fail(self, error: Exception) -> None:

80 """Mark step as failed.

81

82 Args:

83 error: The exception that caused the failure.

84 """

85 self.end_time = time.time()

86 self.status = ExecutionStatus.FAILED

87 self.error = error

88

89 def skip(self, reason: str) -> None:

90 """Mark step as skipped.

91

92 Args:

93 reason: Reason for skipping.

94 """

95 self.end_time = time.time()

96 self.status = ExecutionStatus.SKIPPED

97 self.metrics['skip_reason'] = reason

98

99 def add_metric(self, key: str, value: Any) -> None:

100 """Add a metric to this step.

101

102 Args:

103 key: Metric key.

104 value: Metric value.

105 """

106 self.metrics[key] = value

107

108 def add_resource_usage(self, resource_type: str, usage: Dict[str, Any]) -> None:

109 """Track resource usage for this step.

110

111 Args:

112 resource_type: Type of resource (e.g., "database", "llm").

113 usage: Usage metrics.

114 """

115 self.resource_usage[resource_type] = usage

116

117 def update_stream_progress(

118 self,

119 chunks: int,

120 records: int,

121 current_position: int | None = None

122 ) -> None:

123 """Update streaming progress.

124

125 Args:

126 chunks: Number of chunks processed.

127 records: Number of records processed.

128 current_position: Current position in stream.

129 """

130 self.chunks_processed = chunks

131 self.records_processed = records

132 if current_position is not None:

133 if self.stream_progress is None:

134 self.stream_progress = {}

135 self.stream_progress['position'] = current_position

136

137 @property

138 def duration(self) -> float | None:

139 """Get execution duration in seconds."""

140 if self.start_time and self.end_time:

141 return self.end_time - self.start_time

142 return None

143

144 def to_dict(self) -> Dict[str, Any]:

145 """Convert to dictionary representation."""

146 return {

147 'step_id': self.step_id,

148 'state_name': self.state_name,

149 'network_name': self.network_name,

150 'timestamp': self.timestamp,

151 'start_time': self.start_time,

152 'end_time': self.end_time,

153 'duration': self.duration,

154 'data_mode': self.data_mode.value,

155 'status': self.status.value,

156 'arc_taken': self.arc_taken,

157 'error': str(self.error) if self.error else None,

158 'metrics': self.metrics,

159 'resource_usage': self.resource_usage,

160 'stream_progress': self.stream_progress,

161 'chunks_processed': self.chunks_processed,

162 'records_processed': self.records_processed

163 }

164

165 @classmethod

166 def from_dict(cls, data: Dict[str, Any]) -> 'ExecutionStep':

167 """Create ExecutionStep from dictionary representation.

168

169 Args:

170 data: Dictionary with step data.

171

172 Returns:

173 ExecutionStep instance.

174 """

175 step = cls(

176 step_id=data['step_id'],

177 state_name=data['state_name'],

178 network_name=data['network_name'],

179 timestamp=data['timestamp'],

180 data_mode=DataHandlingMode(data['data_mode']),

181 status=ExecutionStatus(data['status'])

182 )

183

184 # Restore timing info

185 step.start_time = data.get('start_time')

186 step.end_time = data.get('end_time')

187

188 # Restore execution details

189 step.arc_taken = data.get('arc_taken')

190

191 # Restore error (as string - can't fully reconstruct Exception)

192 error_str = data.get('error')

193 if error_str:

194 step.error = Exception(error_str)

195

196 # Restore metrics and usage

197 step.metrics = data.get('metrics', {})

198 step.resource_usage = data.get('resource_usage', {})

199

200 # Restore stream progress

201 step.stream_progress = data.get('stream_progress')

202 step.chunks_processed = data.get('chunks_processed', 0)

203 step.records_processed = data.get('records_processed', 0)

204

205 return step

206

207

208 class ExecutionHistory:

209 """Tracks execution history using a tree structure.

210

211 This class manages the execution history of an FSM, tracking:

212 - State transitions

213 - Data modifications based on mode

214 - Resource usage

215 - Stream progress

216 - Execution metrics

217 """

218

219 def __init__(

220 self,

221 fsm_name: str,

222 execution_id: str,

223 data_mode: DataHandlingMode = DataHandlingMode.COPY,

224 max_depth: int | None = None,

225 enable_data_snapshots: bool = False

226 ):

227 """Initialize execution history.

228

229 Args:

230 fsm_name: Name of the FSM.

231 execution_id: Unique execution identifier.

232 data_mode: Default data mode for execution.

233 max_depth: Maximum tree depth (for pruning).

234 enable_data_snapshots: Whether to store data snapshots.

235 """

236 self.fsm_name = fsm_name

237 self.execution_id = execution_id

238 self.data_mode = data_mode

239 self.max_depth = max_depth

240 self.enable_data_snapshots = enable_data_snapshots

241

242 # Create tree structure - Tree stores root nodes

243 self.tree_roots: List[Tree] = []

244 self.current_node: Tree | None = None

245

246 # Tracking

247 self.start_time = time.time()

248 self.end_time: float | None = None

249 self.total_steps = 0

250 self.failed_steps = 0

251 self.skipped_steps = 0

252

253 # Mode-specific storage

254 self._mode_storage: Dict[DataHandlingMode, List[ExecutionStep]] = {

255 DataHandlingMode.COPY: [],

256 DataHandlingMode.REFERENCE: [],

257 DataHandlingMode.DIRECT: []

258 }

259

260 # Resource tracking

261 self.resource_summary: Dict[str, Dict[str, Any]] = {}

262

263 # Stream tracking

264 self.stream_summary: Dict[str, Any] = {

265 'total_chunks': 0,

266 'total_records': 0,

267 'streams_processed': 0

268 }

269

270 def add_step(

271 self,

272 state_name: str,

273 network_name: str,

274 data: Any | None = None,

275 parent_step_id: str | None = None

276 ) -> ExecutionStep:

277 """Add a new execution step.

278

279 Args:

280 state_name: Name of the state.

281 network_name: Name of the network.

282 data: Optional data snapshot.

283 parent_step_id: Parent step ID for branching.

284

285 Returns:

286 The created ExecutionStep.

287 """

288 import uuid

289

290 step_id = str(uuid.uuid4())

291 step = ExecutionStep(

292 step_id=step_id,

293 state_name=state_name,

294 network_name=network_name,

295 timestamp=time.time(),

296 data_mode=self.data_mode

297 )

298

299 # Store data snapshot if enabled

300 if self.enable_data_snapshots and data is not None:

301 step.data_snapshot = self._snapshot_data(data)

302

303 # Add to tree

304 if parent_step_id:

305 parent_node = self._find_node_by_step_id(parent_step_id)

306 if parent_node:

307 self.current_node = Tree(step, parent=parent_node)

308 elif self.current_node:

309 self.current_node = Tree(step, parent=self.current_node)

310 else:

311 # Create root node

312 self.current_node = Tree(step)

313 self.tree_roots.append(self.current_node)

314

315 # Track in mode-specific storage

316 self._mode_storage[self.data_mode].append(step)

317

318 self.total_steps += 1

319

320 # Prune if needed

321 if self.max_depth and self._get_max_depth() > self.max_depth:

322 self._prune_old_branches()

323

324 return step

325

326 def update_step(

327 self,

328 step_id: str,

329 status: ExecutionStatus | None = None,

330 arc_taken: str | None = None,

331 error: Exception | None = None,

332 metrics: Dict[str, Any] | None = None

333 ) -> bool:

334 """Update an existing step.

335

336 Args:

337 step_id: Step ID to update.

338 status: New status.

339 arc_taken: Arc taken from this step.

340 error: Error if failed.

341 metrics: Additional metrics.

342

343 Returns:

344 True if step was found and updated.

345 """

346 node = self._find_node_by_step_id(step_id)

347 if not node:

348 return False

349

350 step = node.data

351

352 if status == ExecutionStatus.COMPLETED:

353 step.complete(arc_taken)

354 elif status == ExecutionStatus.FAILED:

355 if error:

356 step.fail(error)

357 self.failed_steps += 1

358 elif status == ExecutionStatus.SKIPPED:

359 step.skip(metrics.get('reason', 'Unknown') if metrics else 'Unknown')

360 self.skipped_steps += 1

361 elif status:

362 step.status = status

363

364 if metrics:

365 for key, value in metrics.items():

366 step.add_metric(key, value)

367

368 return True

369

370 def get_path_to_current(self) -> List[ExecutionStep]:

371 """Get the path from root to current step.

372

373 Returns:

374 List of steps from root to current.

375 """

376 if not self.current_node:

377 return []

378

379 path = []

380 node = self.current_node

381 while node:

382 path.insert(0, node.data)

383 node = node.parent

384

385 return path

386

387 def get_all_paths(self) -> List[List[ExecutionStep]]:

388 """Get all execution paths in the tree.

389

390 Returns:

391 List of all paths from root to leaves.

392 """

393 paths = []

394

395 def collect_paths(node: Tree, current_path: List[ExecutionStep]):

396 current_path.append(node.data)

397

398 if not node.children:

399 # Leaf node - save path

400 paths.append(current_path.copy())

401 else:

402 # Continue down each branch

403 for child in node.children:

404 collect_paths(child, current_path.copy())

405

406 for root in self.tree_roots:

407 collect_paths(root, [])

408

409 return paths

410

411 @property

412 def steps(self) -> List[ExecutionStep]:

413 """Get all execution steps from the history tree.

414

415 Returns:

416 List of all execution steps in order.

417 """

418 all_steps = []

419

420 def collect_steps(node: Tree):

421 all_steps.append(node.data)

422 if node.children:

423 for child in node.children:

424 collect_steps(child)

425

426 for root in self.tree_roots:

427 collect_steps(root)

428

429 return all_steps

430

431 def get_steps_by_state(self, state_name: str) -> List[ExecutionStep]:

432 """Get all steps for a specific state.

433

434 Args:

435 state_name: Name of the state.

436

437 Returns:

438 List of steps for that state.

439 """

440 steps = []

441

442 def collect_steps(node: Tree):

443 if node.data.state_name == state_name:

444 steps.append(node.data)

445 if node.children:

446 for child in node.children:

447 collect_steps(child)

448

449 for root in self.tree_roots:

450 collect_steps(root)

451

452 return steps

453

454 def get_steps_by_mode(self, mode: DataHandlingMode) -> List[ExecutionStep]:

455 """Get all steps executed in a specific data mode.

456

457 Args:

458 mode: Data mode to filter by.

459

460 Returns:

461 List of steps in that mode.

462 """

463 return self._mode_storage.get(mode, []).copy()

464

465 def get_resource_usage(self) -> Dict[str, Any]:

466 """Get aggregated resource usage.

467

468 Returns:

469 Resource usage summary.

470 """

471 usage = {}

472

473 def aggregate_usage(node: Tree):

474 step = node.data

475 for resource_type, metrics in step.resource_usage.items():

476 if resource_type not in usage:

477 usage[resource_type] = {

478 'total_calls': 0,

479 'total_duration': 0,

480 'steps': []

481 }

482

483 usage[resource_type]['total_calls'] += 1 # type: ignore

484 if 'duration' in metrics:

485 usage[resource_type]['total_duration'] += metrics['duration']

486 usage[resource_type]['steps'].append(step.step_id)

487

488 if node.children:

489 for child in node.children:

490 aggregate_usage(child)

491

492 for root in self.tree_roots:

493 aggregate_usage(root)

494

495 return usage

496

497 def get_stream_progress(self) -> Dict[str, Any]:

498 """Get streaming progress summary.

499

500 Returns:

501 Stream progress information.

502 """

503 total_chunks = 0

504 total_records = 0

505

506 def aggregate_stream(node: Tree):

507 nonlocal total_chunks, total_records

508 step = node.data

509 total_chunks += step.chunks_processed

510 total_records += step.records_processed

511

512 if node.children:

513 for child in node.children:

514 aggregate_stream(child)

515

516 for root in self.tree_roots:

517 aggregate_stream(root)

518

519 return {

520 'total_chunks': total_chunks,

521 'total_records': total_records,

522 'streams_processed': self.stream_summary['streams_processed']

523 }

524

525 def finalize(self) -> None:

526 """Mark execution as complete."""

527 self.end_time = time.time()

528

529 def get_summary(self) -> Dict[str, Any]:

530 """Get execution summary.

531

532 Returns:

533 Summary of the execution.

534 """

535 return {

536 'fsm_name': self.fsm_name,

537 'execution_id': self.execution_id,

538 'start_time': self.start_time,

539 'end_time': self.end_time,

540 'duration': self.end_time - self.start_time if self.end_time else None,

541 'total_steps': self.total_steps,

542 'failed_steps': self.failed_steps,

543 'skipped_steps': self.skipped_steps,

544 'completed_steps': self.total_steps - self.failed_steps - self.skipped_steps,

545 'data_mode': self.data_mode.value,

546 'tree_depth': self._get_max_depth(),

547 'total_paths': len(self.get_all_paths()),

548 'resource_usage': self.get_resource_usage(),

549 'stream_progress': self.get_stream_progress()

550 }

551

552 def to_dict(self) -> Dict[str, Any]:

553 """Convert to dictionary representation.

554

555 Returns:

556 Dictionary representation of history.

557 """

558 paths = []

559 for path in self.get_all_paths():

560 paths.append([step.to_dict() for step in path])

561

562 return {

563 'summary': self.get_summary(),

564 'paths': paths,

565 'mode_storage': {

566 mode.value: [s.to_dict() for s in steps]

567 for mode, steps in self._mode_storage.items()

568 }

569 }

570

571 @classmethod

572 def from_dict(cls, data: Dict[str, Any]) -> 'ExecutionHistory':

573 """Create ExecutionHistory from dictionary representation.

574

575 Args:

576 data: Dictionary with history data.

577

578 Returns:

579 ExecutionHistory instance.

580 """

581 summary = data['summary']

582

583 # Create history instance

584 history = cls(

585 fsm_name=summary['fsm_name'],

586 execution_id=summary['execution_id'],

587 data_mode=DataHandlingMode(summary['data_mode']),

588 max_depth=None, # Will be inferred

589 enable_data_snapshots=False # Will be inferred from data

590 )

591

592 # Restore properties from summary

593 history.start_time = summary['start_time']

594 history.end_time = summary.get('end_time')

595 history.total_steps = summary['total_steps']

596 history.failed_steps = summary['failed_steps']

597 history.skipped_steps = summary['skipped_steps']

598

599 # Rebuild tree structure from paths

600 for path_data in data.get('paths', []):

601 parent_node = None

602 for step_dict in path_data:

603 step = ExecutionStep.from_dict(step_dict)

604

605 # Create tree node

606 if parent_node is None:

607 # This is a root step

608 node = Tree(step)

609 history.tree_roots.append(node)

610 else:

611 # Child step

612 node = Tree(step, parent=parent_node)

613

614 parent_node = node

615

616 # Track in mode-specific storage

617 history._mode_storage[step.data_mode].append(step)

618

619 # Set current node to the last node in this path

620 if parent_node:

621 history.current_node = parent_node

622

623 return history

624

625 def _find_node_by_step_id(self, step_id: str) -> Tree | None:

626 """Find a node by step ID.

627

628 Args:

629 step_id: Step ID to find.

630

631 Returns:

632 Tree node if found, None otherwise.

633 """

634 def search_node(node: Tree) -> Tree | None:

635 if node.data.step_id == step_id:

636 return node

637 if node.children:

638 for child in node.children:

639 result = search_node(child)

640 if result:

641 return result

642 return None

643

644 for root in self.tree_roots:

645 result = search_node(root)

646 if result:

647 return result

648

649 return None

650

651 def _snapshot_data(self, data: Any) -> Any:

652 """Create a snapshot of data based on mode.

653

654 Args:

655 data: Data to snapshot.

656

657 Returns:

658 Snapshot of the data.

659 """

660 import copy

661

662 if self.data_mode == DataHandlingMode.COPY:

663 # Deep copy for COPY mode

664 return copy.deepcopy(data)

665 elif self.data_mode == DataHandlingMode.REFERENCE:

666 # Store reference info only

667 return {

668 'type': type(data).__name__,

669 'id': id(data),

670 'size': len(data) if hasattr(data, '__len__') else None

671 }

672 else: # DIRECT mode

673 # Store minimal info

674 return {'type': type(data).__name__}

675

676 def _get_max_depth(self) -> int:

677 """Get the maximum depth of the tree.

678

679 Returns:

680 Maximum depth across all trees.

681 """

682 max_depth = 0

683

684 def get_depth(node: Tree, depth: int = 0):

685 nonlocal max_depth

686 depth += 1

687 max_depth = max(max_depth, depth)

688 if node.children:

689 for child in node.children:

690 get_depth(child, depth)

691

692 for root in self.tree_roots:

693 get_depth(root)

694

695 return max_depth

696

697 def _prune_old_branches(self) -> None:

698 """Prune old branches based on mode-specific strategy."""

699 if self.data_mode == DataHandlingMode.COPY:

700 # Keep all branches for COPY mode (full history)

701 pass

702 elif self.data_mode == DataHandlingMode.REFERENCE:

703 # Prune branches not on critical path

704 # Keep only paths that lead to current node

705 pass

706 else: # DIRECT mode

707 # Aggressive pruning - keep only recent history

708 # Remove all but the current path

709 current_path = self.get_path_to_current()

710 if current_path and len(current_path) > self.max_depth: # type: ignore

711 # Keep only last max_depth steps

712 steps_to_keep = current_path[-self.max_depth:] # type: ignore

713 # Rebuild tree with only these steps

714 self.tree_roots = []

715 self.current_node = None

716 for step in steps_to_keep:

717 if not self.current_node:

718 self.current_node = Tree(step)

719 self.tree_roots.append(self.current_node)

720 else:

721 self.current_node = Tree(step, parent=self.current_node)