Tweak automatic edge creation from step function return hints

Author: DouweM

examples/pydantic_ai_examples/temporal_graph.py

@@ -217,7 +217,9 @@ async def return_container(
     g.node(HandleStrNode),
     g.node(ReturnContainerNode),
     g.node(ForwardContainerNode),
-    g.edge_from(g.start_node).label('begin').to(begin),
+    g.edge_from(g.start_node)
+    .label('begin')
+    .to(begin),  # This also adds begin -> ChooseTypeNode
     g.edge_from(choose_type).to(
         g.decision()
         .branch(g.match(TypeExpression[Literal['int']]).to(handle_int))
@@ -236,7 +238,9 @@ async def return_container(
     g.edge_from(
         handle_int_1, handle_int_2, handle_str_1, handle_str_2, handle_field_3_item
     ).to(handle_join),
-    g.edge_from(handle_join).to(return_container),
+    g.edge_from(handle_join).to(
+        return_container
+    ),  # This also adds return_container -> ForwardContainerNode
 )
 
 graph = g.build()

pydantic_graph/pydantic_graph/v2/graph_builder.py

@@ -272,17 +272,6 @@ def decorator(
 
         step = Step[StateT, DepsT, InputT, OutputT](id=NodeId(node_id), call=call, user_label=label)
 
-        parent_namespace = _utils.get_parent_namespace(inspect.currentframe())
-        type_hints = get_type_hints(call, localns=parent_namespace, include_extras=True)
-        try:
-            return_hint = type_hints['return']
-        except KeyError:
-            pass
-        else:
-            edge = self._edge_from_return_hint(step, return_hint)
-            if edge is not None:
-                self.add(edge)
-
         return step
 
     @overload
@@ -413,15 +402,32 @@ def _handle_path(p: Path):
                 elif isinstance(item, DestinationMarker):
                     pass
 
+        destinations: list[AnyDestinationNode] = []
         for edge in edges:
             for source_node in edge.sources:
                 self._insert_node(source_node)
                 self._edges_by_source[source_node.id].append(edge.path)
             for destination_node in edge.destinations:
+                destinations.append(destination_node)
                 self._insert_node(destination_node)
 
             _handle_path(edge.path)
 
+        # Automatically create edges from step function return hints including `BaseNode`s
+        for destination in destinations:
+            if not isinstance(destination, Step) or isinstance(destination, NodeStep):
+                continue
+            parent_namespace = _utils.get_parent_namespace(inspect.currentframe())
+            type_hints = get_type_hints(destination.call, localns=parent_namespace, include_extras=True)
+            try:
+                return_hint = type_hints['return']
+            except KeyError:
+                pass
+            else:
+                edge = self._edge_from_return_hint(destination, return_hint)
+                if edge is not None:
+                    self.add(edge)
+
     def add_edge(self, source: Source[T], destination: Destination[T], *, label: str | None = None) -> None:
         """Add a simple edge between two nodes.