From 333ad65cac9500b8784f15e99fc3981fb2fec382 Mon Sep 17 00:00:00 2001 From: Bagatur Date: Tue, 23 Jan 2024 13:29:23 -0800 Subject: [PATCH] patch: make conditional_edge_mapping optional --- langgraph/graph/graph.py | 21 +++++++++++++++------ 1 file changed, 15 insertions(+), 6 deletions(-) diff --git a/langgraph/graph/graph.py b/langgraph/graph/graph.py index 4ebd8144e..7fc239b43 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(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))