Merge pull request #62 from langchain-ai/bagatur/optional_conditional_edge_mapping

patch: make conditional_edge_mapping optional

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()