Rename IDs and update some docs

Author: dmontagu

docs/graph/beta/index.md

@@ -1,7 +1,8 @@
 # Beta Graph API
 
 !!! warning "Beta API"
-    This is the new beta graph API. It provides enhanced capabilities for parallel execution, conditional branching, and complex workflows. The original graph API is still available and documented in the [main graph documentation](../../graph.md).
+    This is the new beta graph API. It provides enhanced capabilities for parallel execution, conditional branching, and complex workflows.
+The original graph API is still available (and compatible of interop with the new beta API) and is documented in the [main graph documentation](../../graph.md).
 
 ## Overview
 
@@ -13,7 +14,7 @@ The beta graph API in `pydantic-graph` provides a powerful builder pattern for c
 - **Broadcast operations** for sending the same data to multiple parallel paths
 - **Join nodes and Reducers** for aggregating results from parallel execution
 
-This API is designed for advanced workflows where you need explicit control over parallelism, routing, and data aggregation.
+This API is designed for advanced workflows where you want declarative control over parallelism, routing, and data aggregation.
 
 ## Installation
 

pydantic_graph/pydantic_graph/beta/decision.py

@@ -13,7 +13,7 @@
 
 from typing_extensions import Never, Self, TypeVar
 
-from pydantic_graph.beta.id_types import ForkId, JoinId, NodeId
+from pydantic_graph.beta.id_types import ForkID, JoinID, NodeID
 from pydantic_graph.beta.paths import Path, PathBuilder
 from pydantic_graph.beta.step import StepFunction
 from pydantic_graph.beta.util import TypeOrTypeExpression
@@ -42,7 +42,7 @@ class Decision(Generic[StateT, DepsT, HandledT]):
     branches based on the input data type or custom matching logic.
     """
 
-    id: NodeId
+    id: NodeID
     """Unique identifier for this decision node."""
 
     branches: list[DecisionBranch[Any]]
@@ -145,7 +145,7 @@ class DecisionBranchBuilder(Generic[StateT, DepsT, OutputT, SourceT, HandledT]):
     """Builder for the execution path."""
 
     @property
-    def last_fork_id(self) -> ForkId | None:
+    def last_fork_id(self) -> ForkID | None:
         """Get the ID of the last fork in the path.
 
         Returns:
@@ -214,8 +214,8 @@ def transform(
     def spread(
         self: DecisionBranchBuilder[StateT, DepsT, Iterable[T], SourceT, HandledT],
         *,
-        fork_id: ForkId | None = None,
-        downstream_join_id: JoinId | None = None,
+        fork_id: ForkID | None = None,
+        downstream_join_id: JoinID | None = None,
     ) -> DecisionBranchBuilder[StateT, DepsT, T, SourceT, HandledT]:
         """Spread the branch's output.
 

pydantic_graph/pydantic_graph/beta/graph.py

@@ -21,7 +21,7 @@
 from pydantic_graph import exceptions
 from pydantic_graph._utils import AbstractSpan, get_traceparent, logfire_span
 from pydantic_graph.beta.decision import Decision
-from pydantic_graph.beta.id_types import ForkStack, ForkStackItem, GraphRunId, JoinId, NodeId, NodeRunId, TaskId
+from pydantic_graph.beta.id_types import ForkStack, ForkStackItem, GraphRunID, JoinID, NodeID, NodeRunID, TaskID
 from pydantic_graph.beta.join import Join, JoinNode, Reducer
 from pydantic_graph.beta.node import (
     EndNode,
@@ -82,7 +82,7 @@ class JoinItem:
     node, along with metadata about which execution 'fork' it originated from.
     """
 
-    join_id: JoinId
+    join_id: JoinID
     """The ID of the join node this item is targeting."""
 
     inputs: Any
@@ -125,16 +125,16 @@ class Graph(Generic[StateT, DepsT, InputT, OutputT]):
     auto_instrument: bool
     """Whether to automatically create instrumentation spans."""
 
-    nodes: dict[NodeId, AnyNode]
+    nodes: dict[NodeID, AnyNode]
     """All nodes in the graph indexed by their ID."""
 
-    edges_by_source: dict[NodeId, list[Path]]
+    edges_by_source: dict[NodeID, list[Path]]
     """Outgoing paths from each source node."""
 
-    parent_forks: dict[JoinId, ParentFork[NodeId]]
+    parent_forks: dict[JoinID, ParentFork[NodeID]]
     """Parent fork information for each join node."""
 
-    def get_parent_fork(self, join_id: JoinId) -> ParentFork[NodeId]:
+    def get_parent_fork(self, join_id: JoinID) -> ParentFork[NodeID]:
         """Get the parent fork information for a join node.
 
         Args:
@@ -288,7 +288,7 @@ class GraphTask:
     """
 
     # With our current BaseNode thing, next_node_id and next_node_inputs are merged into `next_node` itself
-    node_id: NodeId
+    node_id: NodeID
     """The ID of the node to execute."""
 
     inputs: Any
@@ -300,7 +300,7 @@ class GraphTask:
     Used by the GraphRun to decide when to proceed through joins.
     """
 
-    task_id: TaskId = field(default_factory=lambda: TaskId(str(uuid.uuid4())))
+    task_id: TaskID = field(default_factory=lambda: TaskID(str(uuid.uuid4())))
     """Unique identifier for this task."""
 
 
@@ -346,14 +346,14 @@ def __init__(
         self.inputs = inputs
         """The initial input data."""
 
-        self._active_reducers: dict[tuple[JoinId, NodeRunId], tuple[Reducer[Any, Any, Any, Any], ForkStack]] = {}
+        self._active_reducers: dict[tuple[JoinID, NodeRunID], tuple[Reducer[Any, Any, Any, Any], ForkStack]] = {}
         """Active reducers for join operations."""
 
         self._next: EndMarker[OutputT] | JoinItem | Sequence[GraphTask] | None = None
         """The next item to be processed."""
 
-        run_id = GraphRunId(str(uuid.uuid4()))
-        initial_fork_stack: ForkStack = (ForkStackItem(StartNode.id, NodeRunId(run_id), 0),)
+        run_id = GraphRunID(str(uuid.uuid4()))
+        initial_fork_stack: ForkStack = (ForkStackItem(StartNode.id, NodeRunID(run_id), 0),)
         self._first_task = GraphTask(node_id=StartNode.id, inputs=inputs, fork_stack=initial_fork_stack)
         self._iterator = self._iter_graph()
 
@@ -446,7 +446,7 @@ async def _iter_graph(  # noqa C901
     ) -> AsyncGenerator[
         EndMarker[OutputT] | JoinItem | Sequence[GraphTask], EndMarker[OutputT] | JoinItem | Sequence[GraphTask]
     ]:
-        tasks_by_id: dict[TaskId, GraphTask] = {}
+        tasks_by_id: dict[TaskID, GraphTask] = {}
         pending: set[asyncio.Task[EndMarker[OutputT] | JoinItem | Sequence[GraphTask]]] = set()
 
         def _start_task(t_: GraphTask) -> None:
@@ -490,7 +490,7 @@ def _handle_result(result: EndMarker[OutputT] | JoinItem | Sequence[GraphTask])
                     reducer.reduce(StepContext(self.state, self.deps, result.inputs))
                 except StopIteration:
                     # cancel all concurrently running tasks with the same fork_run_id of the parent fork
-                    task_ids_to_cancel = set[TaskId]()
+                    task_ids_to_cancel = set[TaskID]()
                     for task_id, t in tasks_by_id.items():
                         for item in t.fork_stack:
                             if item.fork_id == parent_fork_id and item.node_run_id == fork_run_id:
@@ -510,7 +510,7 @@ def _handle_result(result: EndMarker[OutputT] | JoinItem | Sequence[GraphTask])
                 done, pending = await asyncio.wait(pending, return_when=asyncio.FIRST_COMPLETED)
                 for task in done:
                     task_result = task.result()
-                    source_task = tasks_by_id.pop(TaskId(task.get_name()))
+                    source_task = tasks_by_id.pop(TaskID(task.get_name()))
                     maybe_overridden_result = yield task_result
                     if _handle_result(maybe_overridden_result):
                         return
@@ -632,8 +632,8 @@ def _get_completed_fork_runs(
         self,
         t: GraphTask,
         active_tasks: Iterable[GraphTask],
-    ) -> list[tuple[JoinId, NodeRunId]]:
-        completed_fork_runs: list[tuple[JoinId, NodeRunId]] = []
+    ) -> list[tuple[JoinID, NodeRunID]]:
+        completed_fork_runs: list[tuple[JoinID, NodeRunID]] = []
 
         fork_run_indices = {fsi.node_run_id: i for i, fsi in enumerate(t.fork_stack)}
         for join_id, fork_run_id in self._active_reducers.keys():
@@ -661,7 +661,7 @@ def _handle_path(self, path: Path, inputs: Any, fork_stack: ForkStack) -> Sequen
             except TypeError:
                 raise RuntimeError(f'Cannot spread non-iterable value: {inputs!r}')
 
-            node_run_id = NodeRunId(str(uuid.uuid4()))
+            node_run_id = NodeRunID(str(uuid.uuid4()))
 
             # If the spread specifies a downstream join id, eagerly create a reducer for it
             if item.downstream_join_id is not None:
@@ -698,7 +698,7 @@ def _handle_edges(self, node: AnyNode, inputs: Any, fork_stack: ForkStack) -> Se
             new_tasks.extend(self._handle_path(path, inputs, fork_stack))
         return new_tasks
 
-    def _is_fork_run_completed(self, tasks: Iterable[GraphTask], join_id: JoinId, fork_run_id: NodeRunId) -> bool:
+    def _is_fork_run_completed(self, tasks: Iterable[GraphTask], join_id: JoinID, fork_run_id: NodeRunID) -> bool:
         # Check if any of the tasks in the graph have this fork_run_id in their fork_stack
         # If this is the case, then the fork run is not yet completed
         parent_fork = self.graph.get_parent_fork(join_id)

pydantic_graph/pydantic_graph/beta/graph_builder.py

@@ -19,7 +19,7 @@
 from pydantic_graph import _utils, exceptions
 from pydantic_graph.beta.decision import Decision, DecisionBranch, DecisionBranchBuilder
 from pydantic_graph.beta.graph import Graph
-from pydantic_graph.beta.id_types import ForkId, JoinId, NodeId
+from pydantic_graph.beta.id_types import ForkID, JoinID, NodeID
 from pydantic_graph.beta.join import Join, JoinNode, Reducer
 from pydantic_graph.beta.node import (
     EndNode,
@@ -99,7 +99,7 @@ def decorator(
     node_id = node_id or get_callable_name(reducer_type)
 
     return Join[StateT, DepsT, Any, Any](
-        id=JoinId(NodeId(node_id)),
+        id=JoinID(NodeID(node_id)),
         reducer_type=reducer_type,
     )
 
@@ -137,10 +137,10 @@ class GraphBuilder(Generic[StateT, DepsT, GraphInputT, GraphOutputT]):
     auto_instrument: bool
     """Whether to automatically create instrumentation spans."""
 
-    _nodes: dict[NodeId, AnyNode]
+    _nodes: dict[NodeID, AnyNode]
     """Internal storage for nodes in the graph."""
 
-    _edges_by_source: dict[NodeId, list[Path]]
+    _edges_by_source: dict[NodeID, list[Path]]
     """Internal storage for edges by source node."""
 
     _decision_index: int
@@ -253,7 +253,7 @@ def decorator(
 
         node_id = node_id or get_callable_name(call)
 
-        step = Step[StateT, DepsT, InputT, OutputT](id=NodeId(node_id), call=call, user_label=label)
+        step = Step[StateT, DepsT, InputT, OutputT](id=NodeID(node_id), call=call, user_label=label)
 
         return step
 
@@ -414,8 +414,8 @@ def add_spreading_edge(
         *,
         pre_spread_label: str | None = None,
         post_spread_label: str | None = None,
-        fork_id: ForkId | None = None,
-        downstream_join_id: JoinId | None = None,
+        fork_id: ForkID | None = None,
+        downstream_join_id: JoinID | None = None,
     ) -> None:
         """Add an edge that spreads iterable data across parallel paths.
 
@@ -461,7 +461,7 @@ def decision(self, *, note: str | None = None) -> Decision[StateT, DepsT, Never]
         Returns:
             A new Decision node with no branches
         """
-        return Decision(id=NodeId(self._get_new_decision_id()), branches=[], note=note)
+        return Decision(id=NodeID(self._get_new_decision_id()), branches=[], note=note)
 
     def match(
         self,
@@ -478,7 +478,7 @@ def match(
         Returns:
             A DecisionBranchBuilder for constructing the branch
         """
-        node_id = NodeId(self._get_new_decision_id())
+        node_id = NodeID(self._get_new_decision_id())
         decision = Decision[StateT, DepsT, Never](node_id, branches=[], note=None)
         new_path_builder = PathBuilder[StateT, DepsT, SourceT](working_items=[])
         return DecisionBranchBuilder(decision=decision, source=source, matches=matches, path_builder=new_path_builder)
@@ -721,8 +721,8 @@ def build(self) -> Graph[StateT, DepsT, GraphInputT, GraphOutputT]:
 
 
 def _normalize_forks(
-    nodes: dict[NodeId, AnyNode], edges: dict[NodeId, list[Path]]
-) -> tuple[dict[NodeId, AnyNode], dict[NodeId, list[Path]]]:
+    nodes: dict[NodeID, AnyNode], edges: dict[NodeID, list[Path]]
+) -> tuple[dict[NodeID, AnyNode], dict[NodeID, list[Path]]]:
     """Normalize the graph structure so only broadcast forks have multiple outgoing edges.
 
     This function ensures that any node with multiple outgoing edges is converted
@@ -736,7 +736,7 @@ def _normalize_forks(
         A tuple of normalized nodes and edges
     """
     new_nodes = nodes.copy()
-    new_edges: dict[NodeId, list[Path]] = {}
+    new_edges: dict[NodeID, list[Path]] = {}
 
     paths_to_handle: list[Path] = []
 
@@ -750,7 +750,7 @@ def _normalize_forks(
         if len(edges_from_source) == 1:
             new_edges[source_id] = edges_from_source
             continue
-        new_fork = Fork[Any, Any](id=ForkId(NodeId(f'{node.id}_broadcast_fork')), is_spread=False)
+        new_fork = Fork[Any, Any](id=ForkID(NodeID(f'{node.id}_broadcast_fork')), is_spread=False)
         new_nodes[new_fork.id] = new_fork
         new_edges[source_id] = [Path(items=[BroadcastMarker(fork_id=new_fork.id, paths=edges_from_source)])]
         new_edges[new_fork.id] = edges_from_source
@@ -772,8 +772,8 @@ def _normalize_forks(
 
 
 def _collect_dominating_forks(
-    graph_nodes: dict[NodeId, AnyNode], graph_edges_by_source: dict[NodeId, list[Path]]
-) -> dict[JoinId, ParentFork[NodeId]]:
+    graph_nodes: dict[NodeID, AnyNode], graph_edges_by_source: dict[NodeID, list[Path]]
+) -> dict[JoinID, ParentFork[NodeID]]:
     """Find the dominating fork for each join node in the graph.
 
     This function analyzes the graph structure to find the parent fork that
@@ -791,10 +791,10 @@ def _collect_dominating_forks(
         ValueError: If any join node lacks a dominating fork
     """
     nodes = set(graph_nodes)
-    start_ids: set[NodeId] = {StartNode.id}
-    edges: dict[NodeId, list[NodeId]] = defaultdict(list)
+    start_ids: set[NodeID] = {StartNode.id}
+    edges: dict[NodeID, list[NodeID]] = defaultdict(list)
 
-    fork_ids: set[NodeId] = set(start_ids)
+    fork_ids: set[NodeID] = set(start_ids)
     for source_id in nodes:
         working_source_id = source_id
         node = graph_nodes.get(source_id)
@@ -803,7 +803,7 @@ def _collect_dominating_forks(
             fork_ids.add(node.id)
             continue
 
-        def _handle_path(path: Path, last_source_id: NodeId):
+        def _handle_path(path: Path, last_source_id: NodeID):
             """Process a path and collect edges and fork information.
 
             Args:
@@ -840,7 +840,7 @@ def _handle_path(path: Path, last_source_id: NodeId):
     )
 
     join_ids = {node.id for node in graph_nodes.values() if isinstance(node, Join)}
-    dominating_forks: dict[JoinId, ParentFork[NodeId]] = {}
+    dominating_forks: dict[JoinID, ParentFork[NodeID]] = {}
     for join_id in join_ids:
         dominating_fork = finder.find_parent_fork(join_id)
         if dominating_fork is None:

pydantic_graph/pydantic_graph/beta/id_types.py

@@ -9,23 +9,23 @@
 from dataclasses import dataclass
 from typing import NewType
 
-NodeId = NewType('NodeId', str)
+NodeID = NewType('NodeID', str)
 """Unique identifier for a node in the graph."""
 
-NodeRunId = NewType('NodeRunId', str)
+NodeRunID = NewType('NodeRunID', str)
 """Unique identifier for a specific execution instance of a node."""
 
 # The following aliases are just included for clarity; making them NewTypes is a hassle
-JoinId = NodeId
+JoinID = NodeID
 """Alias for NodeId when referring to join nodes."""
 
-ForkId = NodeId
+ForkID = NodeID
 """Alias for NodeId when referring to fork nodes."""
 
-GraphRunId = NewType('GraphRunId', str)
+GraphRunID = NewType('GraphRunID', str)
 """Unique identifier for a complete graph execution run."""
 
-TaskId = NewType('TaskId', str)
+TaskID = NewType('TaskID', str)
 """Unique identifier for a task within the graph execution."""
 
 
@@ -38,9 +38,9 @@ class ForkStackItem:
     and coordinate parallel branches of execution.
     """
 
-    fork_id: ForkId
+    fork_id: ForkID
     """The ID of the node that created this fork."""
-    node_run_id: NodeRunId
+    node_run_id: NodeRunID
     """The ID associated to the specific run of the node that created this fork."""
     thread_index: int
     """The index of the execution "thread" created during the node run that created this fork.

pydantic_graph/pydantic_graph/beta/join.py

@@ -14,7 +14,7 @@
 from typing_extensions import TypeVar
 
 from pydantic_graph import BaseNode, End, GraphRunContext
-from pydantic_graph.beta.id_types import ForkId, JoinId
+from pydantic_graph.beta.id_types import ForkID, JoinID
 from pydantic_graph.beta.step import StepContext
 
 StateT = TypeVar('StateT', infer_variance=True)
@@ -206,7 +206,7 @@ class Join(Generic[StateT, DepsT, InputT, OutputT]):
     """
 
     def __init__(
-        self, id: JoinId, reducer_type: type[Reducer[StateT, DepsT, InputT, OutputT]], joins: ForkId | None = None
+        self, id: JoinID, reducer_type: type[Reducer[StateT, DepsT, InputT, OutputT]], joins: ForkID | None = None
     ) -> None:
         """Initialize a join operation.