Coverage for pydantic_graph/pydantic_graph/persistence/in_mem.py: 98.21%

1"""In memory state persistence. 

2 

3This module provides simple in memory state persistence for graphs. 

4""" 

5 

6from __future__ import annotations as _annotations (empty)

7 

8import copy (empty)

9from collections.abc import AsyncIterator (empty)

10from contextlib import asynccontextmanager (empty)

11from dataclasses import dataclass, field (empty)

12from time import perf_counter (empty)

13from typing import Any (empty)

14 

15import pydantic (empty)

16 

17from .. import exceptions (empty)

18from ..nodes import BaseNode, End (empty)

19from . import ( (empty)

20 BaseStatePersistence, 

21 EndSnapshot, 

22 NodeSnapshot, 

23 RunEndT, 

24 Snapshot, 

25 StateT, 

26 _utils, 

27 build_snapshot_list_type_adapter, 

28) 

29 

30 

31@dataclass (empty)

32class SimpleStatePersistence(BaseStatePersistence[StateT, RunEndT]): (empty)

33 """Simple in memory state persistence that just hold the latest snapshot. 

34 

35 If no state persistence implementation is provided when running a graph, this is used by default. 

36 """ 

37 

38 last_snapshot: Snapshot[StateT, RunEndT] | None = None (empty)

39 """The last snapshot.""" (empty)

40 

41 async def snapshot_node(self, state: StateT, next_node: BaseNode[StateT, Any, RunEndT]) -> None: (empty)

42 self.last_snapshot = NodeSnapshot(state=state, node=next_node) (empty)

43 

44 async def snapshot_node_if_new( (empty)

45 self, snapshot_id: str, state: StateT, next_node: BaseNode[StateT, Any, RunEndT] 

46 ) -> None: 

47 if self.last_snapshot and self.last_snapshot.id == snapshot_id: 47 ↛ 48line 47 didn't jump to line 48 because the condition on line 47 was never true(empty)

48 return 

49 else: 

50 await self.snapshot_node(state, next_node) (empty)

51 

52 async def snapshot_end(self, state: StateT, end: End[RunEndT]) -> None: (empty)

53 self.last_snapshot = EndSnapshot(state=state, result=end) (empty)

54 

55 @asynccontextmanager (empty)

56 async def record_run(self, snapshot_id: str) -> AsyncIterator[None]: (empty)

57 if self.last_snapshot is None or snapshot_id != self.last_snapshot.id: (empty)

58 raise LookupError(f'No snapshot found with id={snapshot_id!r}') (empty)

59 

60 assert isinstance(self.last_snapshot, NodeSnapshot), 'Only NodeSnapshot can be recorded' (empty)

61 exceptions.GraphNodeStatusError.check(self.last_snapshot.status) (empty)

62 self.last_snapshot.status = 'running' (empty)

63 self.last_snapshot.start_ts = _utils.now_utc() (empty)

64 

65 start = perf_counter() (empty)

66 try: (empty)

67 yield (empty)

68 except Exception: (empty)

69 self.last_snapshot.duration = perf_counter() - start (empty)

70 self.last_snapshot.status = 'error' (empty)

71 raise (empty)

72 else: 

73 self.last_snapshot.duration = perf_counter() - start (empty)

74 self.last_snapshot.status = 'success' (empty)

75 

76 async def load_next(self) -> NodeSnapshot[StateT, RunEndT] | None: (empty)

77 if isinstance(self.last_snapshot, NodeSnapshot) and self.last_snapshot.status == 'created': (empty)

78 self.last_snapshot.status = 'pending' (empty)

79 return self.last_snapshot (empty)

80 

81 async def load_all(self) -> list[Snapshot[StateT, RunEndT]]: (empty)

82 raise NotImplementedError('load is not supported for SimpleStatePersistence') 

83 

84 

85@dataclass (empty)

86class FullStatePersistence(BaseStatePersistence[StateT, RunEndT]): (empty)

87 """In memory state persistence that hold a list of snapshots.""" 

88 

89 deep_copy: bool = True (empty)

90 """Whether to deep copy the state and nodes when storing them. (empty)

91 

92 Defaults to `True` so even if nodes or state are modified after the snapshot is taken, 

93 the persistence history will record the value at the time of the snapshot. 

94 """ 

95 history: list[Snapshot[StateT, RunEndT]] = field(default_factory=list) (empty)

96 """List of snapshots taken during the graph run.""" (empty)

97 _snapshots_type_adapter: pydantic.TypeAdapter[list[Snapshot[StateT, RunEndT]]] | None = field( (empty)

98 default=None, init=False, repr=False 

99 ) 

100 

101 async def snapshot_node(self, state: StateT, next_node: BaseNode[StateT, Any, RunEndT]) -> None: (empty)

102 snapshot = NodeSnapshot( (empty)

103 state=self._prep_state(state), 

104 node=next_node.deep_copy() if self.deep_copy else next_node, 

105 ) 

106 self.history.append(snapshot) (empty)

107 

108 async def snapshot_node_if_new( (empty)

109 self, snapshot_id: str, state: StateT, next_node: BaseNode[StateT, Any, RunEndT] 

110 ) -> None: 

111 if not any(s.id == snapshot_id for s in self.history): (empty)

112 await self.snapshot_node(state, next_node) (empty)

113 

114 async def snapshot_end(self, state: StateT, end: End[RunEndT]) -> None: (empty)

115 snapshot = EndSnapshot( (empty)

116 state=self._prep_state(state), 

117 result=end.deep_copy_data() if self.deep_copy else end, 

118 ) 

119 self.history.append(snapshot) (empty)

120 

121 @asynccontextmanager (empty)

122 async def record_run(self, snapshot_id: str) -> AsyncIterator[None]: (empty)

123 try: (empty)

124 snapshot = next(s for s in self.history if s.id == snapshot_id) (empty)

125 except StopIteration as e: (empty)

126 raise LookupError(f'No snapshot found with id={snapshot_id!r}') from e (empty)

127 

128 assert isinstance(snapshot, NodeSnapshot), 'Only NodeSnapshot can be recorded' (empty)

129 exceptions.GraphNodeStatusError.check(snapshot.status) (empty)

130 snapshot.status = 'running' (empty)

131 snapshot.start_ts = _utils.now_utc() (empty)

132 start = perf_counter() (empty)

133 try: (empty)

134 yield (empty)

135 except Exception: (empty)

136 snapshot.duration = perf_counter() - start (empty)

137 snapshot.status = 'error' (empty)

138 raise (empty)

139 else: 

140 snapshot.duration = perf_counter() - start (empty)

141 snapshot.status = 'success' (empty)

142 

143 async def load_next(self) -> NodeSnapshot[StateT, RunEndT] | None: (empty)

144 if snapshot := next((s for s in self.history if isinstance(s, NodeSnapshot) and s.status == 'created'), None): (empty)

145 snapshot.status = 'pending' (empty)

146 return snapshot (empty)

147 

148 async def load_all(self) -> list[Snapshot[StateT, RunEndT]]: (empty)

149 return self.history (empty)

150 

151 def should_set_types(self) -> bool: (empty)

152 return self._snapshots_type_adapter is None (empty)

153 

154 def set_types(self, state_type: type[StateT], run_end_type: type[RunEndT]) -> None: (empty)

155 self._snapshots_type_adapter = build_snapshot_list_type_adapter(state_type, run_end_type) (empty)

156 

157 def dump_json(self, *, indent: int | None = None) -> bytes: (empty)

158 """Dump the history to JSON bytes.""" 

159 assert self._snapshots_type_adapter is not None, 'type adapter must be set to use `dump_json`' (empty)

160 return self._snapshots_type_adapter.dump_json(self.history, indent=indent) (empty)

161 

162 def load_json(self, json_data: str | bytes | bytearray) -> None: (empty)

163 """Load the history from JSON.""" 

164 assert self._snapshots_type_adapter is not None, 'type adapter must be set to use `load_json`' (empty)

165 self.history = self._snapshots_type_adapter.validate_json(json_data) (empty)

166 

167 def _prep_state(self, state: StateT) -> StateT: (empty)

168 """Prepare state for snapshot, uses [`copy.deepcopy`][copy.deepcopy] by default.""" 

169 if not self.deep_copy or state is None: (empty)

170 return state (empty)

171 else: 

172 return copy.deepcopy(state) (empty)