Coverage for pydantic_graph/pydantic_graph/nodes.py: 97.69%

1from __future__ import annotations as _annotations (empty)

2 

3import copy (empty)

4from abc import ABC, abstractmethod (empty)

5from dataclasses import dataclass, is_dataclass (empty)

6from functools import cache (empty)

7from typing import Any, ClassVar, Generic, get_type_hints (empty)

8from uuid import uuid4 (empty)

9 

10from typing_extensions import Never, Self, TypeVar, get_origin (empty)

11 

12from . import _utils, exceptions (empty)

13 

14__all__ = 'GraphRunContext', 'BaseNode', 'End', 'Edge', 'NodeDef', 'DepsT', 'StateT', 'RunEndT' (empty)

15 

16 

17StateT = TypeVar('StateT', default=None) (empty)

18"""Type variable for the state in a graph.""" (empty)

19RunEndT = TypeVar('RunEndT', covariant=True, default=None) (empty)

20"""Covariant type variable for the return type of a graph [`run`][pydantic_graph.graph.Graph.run].""" (empty)

21NodeRunEndT = TypeVar('NodeRunEndT', covariant=True, default=Never) (empty)

22"""Covariant type variable for the return type of a node [`run`][pydantic_graph.nodes.BaseNode.run].""" (empty)

23DepsT = TypeVar('DepsT', default=None, contravariant=True) (empty)

24"""Type variable for the dependencies of a graph and node.""" (empty)

25 

26 

27@dataclass (empty)

28class GraphRunContext(Generic[StateT, DepsT]): (empty)

29 """Context for a graph.""" 

30 

31 # TODO: Can we get rid of this struct and just pass both these things around..? 

32 

33 state: StateT (empty)

34 """The state of the graph.""" (empty)

35 deps: DepsT (empty)

36 """Dependencies for the graph.""" (empty)

37 

38 

39class BaseNode(ABC, Generic[StateT, DepsT, NodeRunEndT]): (empty)

40 """Base class for a node.""" 

41 

42 docstring_notes: ClassVar[bool] = False (empty)

43 """Set to `True` to generate mermaid diagram notes from the class's docstring. (empty)

44 

45 While this can add valuable information to the diagram, it can make diagrams harder to view, hence 

46 it is disabled by default. You can also customise notes overriding the 

47 [`get_note`][pydantic_graph.nodes.BaseNode.get_note] method. 

48 """ 

49 

50 @abstractmethod (empty)

51 async def run(self, ctx: GraphRunContext[StateT, DepsT]) -> BaseNode[StateT, DepsT, Any] | End[NodeRunEndT]: (empty)

52 """Run the node. 

53 

54 This is an abstract method that must be implemented by subclasses. 

55 

56 !!! note "Return types used at runtime" 

57 The return type of this method are read by `pydantic_graph` at runtime and used to define which 

58 nodes can be called next in the graph. This is displayed in [mermaid diagrams](mermaid.md) 

59 and enforced when running the graph. 

60 

61 Args: 

62 ctx: The graph context. 

63 

64 Returns: 

65 The next node to run or [`End`][pydantic_graph.nodes.End] to signal the end of the graph. 

66 """ 

67 ... 

68 

69 def get_snapshot_id(self) -> str: (empty)

70 if snapshot_id := getattr(self, '__snapshot_id', None): (empty)

71 return snapshot_id (empty)

72 else: 

73 self.__dict__['__snapshot_id'] = snapshot_id = generate_snapshot_id(self.get_node_id()) (empty)

74 return snapshot_id (empty)

75 

76 def set_snapshot_id(self, snapshot_id: str) -> None: (empty)

77 self.__dict__['__snapshot_id'] = snapshot_id (empty)

78 

79 @classmethod (empty)

80 @cache (empty)

81 def get_node_id(cls) -> str: (empty)

82 """Get the ID of the node.""" 

83 return cls.__name__ (empty)

84 

85 @classmethod (empty)

86 def get_note(cls) -> str | None: (empty)

87 """Get a note about the node to render on mermaid charts. 

88 

89 By default, this returns a note only if [`docstring_notes`][pydantic_graph.nodes.BaseNode.docstring_notes] 

90 is `True`. You can override this method to customise the node notes. 

91 """ 

92 if not cls.docstring_notes: (empty)

93 return None (empty)

94 docstring = cls.__doc__ (empty)

95 # dataclasses get an automatic docstring which is just their signature, we don't want that 

96 if docstring and is_dataclass(cls) and docstring.startswith(f'{cls.__name__}('): 96 ↛ 97line 96 didn't jump to line 97 because the condition on line 96 was never true(empty)

97 docstring = None 

98 if docstring: 98 ↛ 103line 98 didn't jump to line 103 because the condition on line 98 was always true(empty)

99 # remove indentation from docstring 

100 import inspect (empty)

101 

102 docstring = inspect.cleandoc(docstring) (empty)

103 return docstring (empty)

104 

105 @classmethod (empty)

106 def get_node_def(cls, local_ns: dict[str, Any] | None) -> NodeDef[StateT, DepsT, NodeRunEndT]: (empty)

107 """Get the node definition.""" 

108 type_hints = get_type_hints(cls.run, localns=local_ns, include_extras=True) (empty)

109 try: (empty)

110 return_hint = type_hints['return'] (empty)

111 except KeyError as e: (empty)

112 raise exceptions.GraphSetupError(f'Node {cls} is missing a return type hint on its `run` method') from e (empty)

113 

114 next_node_edges: dict[str, Edge] = {} (empty)

115 end_edge: Edge | None = None (empty)

116 returns_base_node: bool = False (empty)

117 for return_type in _utils.get_union_args(return_hint): (empty)

118 return_type, annotations = _utils.unpack_annotated(return_type) (empty)

119 edge = next((a for a in annotations if isinstance(a, Edge)), Edge(None)) (empty)

120 return_type_origin = get_origin(return_type) or return_type (empty)

121 if return_type_origin is End: (empty)

122 end_edge = edge (empty)

123 elif return_type_origin is BaseNode: (empty)

124 # TODO: Should we disallow this? 

125 returns_base_node = True (empty)

126 elif issubclass(return_type_origin, BaseNode): (empty)

127 next_node_edges[return_type.get_node_id()] = edge (empty)

128 else: 

129 raise exceptions.GraphSetupError(f'Invalid return type: {return_type}') (empty)

130 

131 return NodeDef( (empty)

132 cls, 

133 cls.get_node_id(), 

134 cls.get_note(), 

135 next_node_edges, 

136 end_edge, 

137 returns_base_node, 

138 ) 

139 

140 def deep_copy(self) -> Self: (empty)

141 """Returns a deep copy of the node.""" 

142 return copy.deepcopy(self) (empty)

143 

144 

145@dataclass (empty)

146class End(Generic[RunEndT]): (empty)

147 """Type to return from a node to signal the end of the graph.""" 

148 

149 data: RunEndT (empty)

150 """Data to return from the graph.""" (empty)

151 

152 def deep_copy_data(self) -> End[RunEndT]: (empty)

153 """Returns a deep copy of the end of the run.""" 

154 if self.data is None: (empty)

155 return self (empty)

156 else: 

157 end = End(copy.deepcopy(self.data)) (empty)

158 end.set_snapshot_id(self.get_snapshot_id()) (empty)

159 return end (empty)

160 

161 def get_snapshot_id(self) -> str: (empty)

162 if snapshot_id := getattr(self, '__snapshot_id', None): (empty)

163 return snapshot_id (empty)

164 else: 

165 self.__dict__['__snapshot_id'] = snapshot_id = generate_snapshot_id('end') (empty)

166 return snapshot_id (empty)

167 

168 def set_snapshot_id(self, set_id: str) -> None: (empty)

169 self.__dict__['__snapshot_id'] = set_id (empty)

170 

171 

172def generate_snapshot_id(node_id: str) -> str: (empty)

173 # module method to allow mocking 

174 return f'{node_id}:{uuid4().hex}' (empty)

175 

176 

177@dataclass (empty)

178class Edge: (empty)

179 """Annotation to apply a label to an edge in a graph.""" 

180 

181 label: str | None (empty)

182 """Label for the edge.""" (empty)

183 

184 

185@dataclass (empty)

186class NodeDef(Generic[StateT, DepsT, NodeRunEndT]): (empty)

187 """Definition of a node. 

188 

189 This is a primarily internal representation of a node; in general, it shouldn't be necessary to use it directly. 

190 

191 Used by [`Graph`][pydantic_graph.graph.Graph] to store information about a node, and when generating 

192 mermaid graphs. 

193 """ 

194 

195 node: type[BaseNode[StateT, DepsT, NodeRunEndT]] (empty)

196 """The node definition itself.""" (empty)

197 node_id: str (empty)

198 """ID of the node.""" (empty)

199 note: str | None (empty)

200 """Note about the node to render on mermaid charts.""" (empty)

201 next_node_edges: dict[str, Edge] (empty)

202 """IDs of the nodes that can be called next.""" (empty)

203 end_edge: Edge | None (empty)

204 """If node definition returns an `End` this is an Edge, indicating the node can end the run.""" (empty)

205 returns_base_node: bool (empty)

206 """The node definition returns a `BaseNode`, hence any node in the next can be called next.""" (empty)