diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index 4ebd8144e..a0ac4a172 100644 --- a/langgraph/graph/graph.py +++ b/langgraph/graph/graph.py @@ -17,11 +17,14 @@ END = "__end__" class Branch(NamedTuple): condition: Callable[..., str] - ends: dict[str, str] + ends: Optional[dict[str, str]] def runnable(self, input: Any) -> Runnable: result = self.condition(input) - destination = self.ends[result] + if self.ends: + destination = self.ends[result] + else: + destination = result return Channel.write_to(f"{destination}:inbox" if destination != END else END) @@ -59,15 +62,21 @@ class Graph: self, start_key: str, condition: Callable[..., str], - conditional_edge_mapping: Dict[str, str], + conditional_edge_mapping: Optional[Dict[str, str]] = None, ) -> None: if start_key not in self.nodes: raise ValueError(f"Need to add_node `{start_key}` first") if iscoroutinefunction(condition): raise ValueError("Condition cannot be a coroutine function") - for destination in conditional_edge_mapping.values(): - if destination not in self.nodes and destination != END: - raise ValueError(f"Need to add_node `{destination}` first") + if conditional_edge_mapping and set( + conditional_edge_mapping.values() + ).difference([END]).difference(self.nodes): + raise ValueError( + f"Missing nodes which are in conditional edge mapping. Mapping " + f"contains possible destinations: " + f"{list(conditional_edge_mapping.values())}. Possible nodes are " + f"{list(self.nodes.keys())}." + ) self.branches[start_key].append(Branch(condition, conditional_edge_mapping)) @@ -81,23 +90,30 @@ class Graph: def validate(self) -> None: all_starts = {src for src, _ in self.edges} | {src for src in self.branches} - all_ends = ( - {end for _, end in self.edges} - | { - end - for branch_list in self.branches.values() - for branch in branch_list - for end in branch.ends.values() - } - | {self.entry_point} - ) - for node in self.nodes: - if node not in all_ends: - raise ValueError(f"Node `{node}` is not reachable") if node not in all_starts: raise ValueError(f"Node `{node}` is a dead-end") + if all( + branch.ends is not None + for branch_list in self.branches.values() + for branch in branch_list + ): + all_ends = ( + {end for _, end in self.edges} + | { + end + for branch_list in self.branches.values() + for branch in branch_list + for end in branch.ends.values() + } + | {self.entry_point} + ) + + for node in self.nodes: + if node not in all_ends: + raise ValueError(f"Node `{node}` is not reachable") + def compile(self, checkpointer: Optional[BaseCheckpointSaver] = None) -> Pregel: self.validate()