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

1from __future__ import annotations as _annotations (empty)

2 

3from abc import ABC, abstractmethod (empty)

4from dataclasses import dataclass, is_dataclass (empty)

5from functools import cache (empty)

6from typing import TYPE_CHECKING, Any, ClassVar, Generic, get_origin, get_type_hints (empty)

7 

8from typing_extensions import Never, TypeVar (empty)

9 

10from . import _utils, exceptions (empty)

11 

12if TYPE_CHECKING: (empty)

13 from .state import StateT 

14else: 

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

16 

17__all__ = 'GraphRunContext', 'BaseNode', 'End', 'Edge', 'NodeDef', 'DepsT' (empty)

18 

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 state: StateT (empty)

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

33 deps: DepsT (empty)

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

35 

36 

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

38 """Base class for a node.""" 

39 

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

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

42 

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

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

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

46 """ 

47 

48 @abstractmethod (empty)

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

50 """Run the node. 

51 

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

53 

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

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

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

57 and enforced when running the graph. 

58 

59 Args: 

60 ctx: The graph context. 

61 

62 Returns: 

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

64 """ 

65 ... 

66 

67 @classmethod (empty)

68 @cache (empty)

69 def get_id(cls) -> str: (empty)

70 """Get the ID of the node.""" 

71 return cls.__name__ (empty)

72 

73 @classmethod (empty)

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

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

76 

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

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

79 """ 

80 if not cls.docstring_notes: (empty)

81 return None (empty)

82 docstring = cls.__doc__ (empty)

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

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

85 docstring = None 

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

87 # remove indentation from docstring 

88 import inspect (empty)

89 

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

91 return docstring (empty)

92 

93 @classmethod (empty)

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

95 """Get the node definition.""" 

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

97 try: (empty)

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

99 except KeyError as e: (empty)

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

101 

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

103 end_edge: Edge | None = None (empty)

104 returns_base_node: bool = False (empty)

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

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

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

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

109 if return_type_origin is End: (empty)

110 end_edge = edge (empty)

111 elif return_type_origin is BaseNode: (empty)

112 # TODO: Should we disallow this? 

113 returns_base_node = True (empty)

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

115 next_node_edges[return_type.get_id()] = edge (empty)

116 else: 

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

118 

119 return NodeDef( (empty)

120 cls, 

121 cls.get_id(), 

122 cls.get_note(), 

123 next_node_edges, 

124 end_edge, 

125 returns_base_node, 

126 ) 

127 

128 

129@dataclass (empty)

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

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

132 

133 data: RunEndT (empty)

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

135 

136 

137@dataclass (empty)

138class Edge: (empty)

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

140 

141 label: str | None (empty)

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

143 

144 

145@dataclass (empty)

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

147 """Definition of a node. 

148 

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

150 

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

152 mermaid graphs. 

153 """ 

154 

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

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

157 node_id: str (empty)

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

159 note: str | None (empty)

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

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

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

163 end_edge: Edge | None (empty)

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

165 returns_base_node: bool (empty)

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