diff --git a/libs/langgraph/langgraph/graph/branch.py b/libs/langgraph/langgraph/graph/branch.py new file mode 100644 index 000000000..8e90d9097 --- /dev/null +++ b/libs/langgraph/langgraph/graph/branch.py @@ -0,0 +1,215 @@ +import asyncio +from inspect import ( + isfunction, + ismethod, + signature, +) +from types import FunctionType +from typing import ( + Any, + Awaitable, + Callable, + Hashable, + Literal, + NamedTuple, + Optional, + Sequence, + Type, + Union, + cast, + get_args, + get_origin, + get_type_hints, +) + +from langchain_core.runnables import ( + Runnable, + RunnableConfig, + RunnableLambda, +) + +from langgraph.constants import END, START +from langgraph.errors import InvalidUpdateError +from langgraph.pregel.write import ChannelWrite +from langgraph.types import Send +from langgraph.utils.runnable import ( + RunnableCallable, +) + + +def _get_branch_path_input_schema( + path: Union[ + Callable[..., Union[Hashable, list[Hashable]]], + Callable[..., Awaitable[Union[Hashable, list[Hashable]]]], + Runnable[Any, Union[Hashable, list[Hashable]]], + ], +) -> Optional[Type[Any]]: + input = None + # detect input schema annotation in the branch callable + try: + callable_: Optional[ + Union[ + Callable[..., Union[Hashable, list[Hashable]]], + Callable[..., Awaitable[Union[Hashable, list[Hashable]]]], + ] + ] = None + if isinstance(path, (RunnableCallable, RunnableLambda)): + if isfunction(path.func) or ismethod(path.func): + callable_ = path.func + elif (callable_method := getattr(path.func, "__call__", None)) and ismethod( + callable_method + ): + callable_ = callable_method + elif isfunction(path.afunc) or ismethod(path.afunc): + callable_ = path.afunc + elif ( + callable_method := getattr(path.afunc, "__call__", None) + ) and ismethod(callable_method): + callable_ = callable_method + elif callable(path): + callable_ = path + + if callable_ is not None and (hints := get_type_hints(callable_)): + first_parameter_name = next( + iter(signature(cast(FunctionType, callable_)).parameters.keys()) + ) + if input_hint := hints.get(first_parameter_name): + if isinstance(input_hint, type) and get_type_hints(input_hint): + input = input_hint + except (TypeError, StopIteration): + pass + + return input + + +class Branch(NamedTuple): + path: Runnable[Any, Union[Hashable, list[Hashable]]] + ends: Optional[dict[Hashable, str]] + then: Optional[str] = None + input_schema: Optional[Type[Any]] = None + + @classmethod + def from_path( + cls, + path: Runnable[Any, Union[Hashable, list[Hashable]]], + path_map: Optional[Union[dict[Hashable, str], list[str]]], + then: Optional[str] = None, + infer_schema: bool = False, + ) -> "Branch": + # coerce path_map to a dictionary + path_map_: Optional[dict[Hashable, str]] = None + try: + if isinstance(path_map, dict): + path_map_ = path_map.copy() + elif isinstance(path_map, list): + path_map_ = {name: name for name in path_map} + else: + # find func + func: Optional[Callable] = None + if isinstance(path, (RunnableCallable, RunnableLambda)): + func = path.func or path.afunc + if func is not None: + # find callable method + if (cal := getattr(path, "__call__", None)) and ismethod(cal): + func = cal + # get the return type + if rtn_type := get_type_hints(func).get("return"): + if get_origin(rtn_type) is Literal: + path_map_ = {name: name for name in get_args(rtn_type)} + except Exception: + pass + # infer input schema + input_schema = _get_branch_path_input_schema(path) if infer_schema else None + # create branch + return cls(path=path, ends=path_map_, then=then, input_schema=input_schema) + + def run( + self, + writer: Callable[ + [Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite] + ], + reader: Optional[Callable[[RunnableConfig], Any]] = None, + ) -> RunnableCallable: + return ChannelWrite.register_writer( + RunnableCallable( + func=self._route, + afunc=self._aroute, + writer=writer, + reader=reader, + name=None, + trace=False, + ) + ) + + def _route( + self, + input: Any, + config: RunnableConfig, + *, + reader: Optional[Callable[[RunnableConfig], Any]], + writer: Callable[ + [Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite] + ], + ) -> Runnable: + if reader: + value = reader(config) + # passthrough additional keys from node to branch + # only doable when using dict states + if ( + isinstance(value, dict) + and isinstance(input, dict) + and self.input_schema is None + ): + value = {**input, **value} + else: + value = input + result = self.path.invoke(value, config) + return self._finish(writer, input, result, config) + + async def _aroute( + self, + input: Any, + config: RunnableConfig, + *, + reader: Optional[Callable[[RunnableConfig], Any]], + writer: Callable[ + [Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite] + ], + ) -> Runnable: + if reader: + value = await asyncio.to_thread(reader, config) + # passthrough additional keys from node to branch + # only doable when using dict states + if ( + isinstance(value, dict) + and isinstance(input, dict) + and self.input_schema is None + ): + value = {**input, **value} + else: + value = input + result = await self.path.ainvoke(value, config) + return self._finish(writer, input, result, config) + + def _finish( + self, + writer: Callable[ + [Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite] + ], + input: Any, + result: Any, + config: RunnableConfig, + ) -> Union[Runnable, Any]: + if not isinstance(result, (list, tuple)): + result = [result] + if self.ends: + destinations: Sequence[Union[Send, str]] = [ + r if isinstance(r, Send) else self.ends[r] for r in result + ] + else: + destinations = cast(Sequence[Union[Send, str]], result) + if any(dest is None or dest == START for dest in destinations): + raise ValueError("Branch did not return a valid destination") + if any(p.node == END for p in destinations if isinstance(p, Send)): + raise InvalidUpdateError("Cannot send a packet to the END node") + return writer(destinations, config) or input diff --git a/libs/langgraph/langgraph/graph/graph.py b/libs/langgraph/langgraph/graph/graph.py index 4bda590c3..ba2c8df72 100644 --- a/libs/langgraph/langgraph/graph/graph.py +++ b/libs/langgraph/langgraph/graph/graph.py @@ -1,4 +1,3 @@ -import asyncio import logging from collections import defaultdict from typing import ( @@ -6,15 +5,11 @@ from typing import ( Awaitable, Callable, Hashable, - Literal, NamedTuple, Optional, Sequence, Union, cast, - get_args, - get_origin, - get_type_hints, overload, ) @@ -34,12 +29,12 @@ from langgraph.constants import ( TAG_HIDDEN, Send, ) -from langgraph.errors import InvalidUpdateError +from langgraph.graph.branch import Branch from langgraph.pregel import Channel, Pregel from langgraph.pregel.read import PregelNode from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry from langgraph.types import All, Checkpointer -from langgraph.utils.runnable import RunnableCallable, RunnableLike, coerce_to_runnable +from langgraph.utils.runnable import RunnableLike, coerce_to_runnable logger = logging.getLogger(__name__) @@ -50,95 +45,6 @@ class NodeSpec(NamedTuple): ends: Optional[Union[tuple[str, ...], dict[str, str]]] = EMPTY_SEQ -class Branch(NamedTuple): - path: Runnable[Any, Union[Hashable, list[Hashable]]] - ends: Optional[dict[Hashable, str]] - then: Optional[str] = None - - def run( - self, - writer: Callable[ - [Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite] - ], - reader: Optional[Callable[[RunnableConfig], Any]] = None, - ) -> RunnableCallable: - return ChannelWrite.register_writer( - RunnableCallable( - func=self._route, - afunc=self._aroute, - writer=writer, - reader=reader, - name=None, - trace=False, - ) - ) - - def _route( - self, - input: Any, - config: RunnableConfig, - *, - reader: Optional[Callable[[RunnableConfig], Any]], - writer: Callable[ - [Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite] - ], - ) -> Runnable: - if reader: - value = reader(config) - # passthrough additional keys from node to branch - # only doable when using dict states - if isinstance(value, dict) and isinstance(input, dict): - value = {**input, **value} - else: - value = input - result = self.path.invoke(value, config) - return self._finish(writer, input, result, config) - - async def _aroute( - self, - input: Any, - config: RunnableConfig, - *, - reader: Optional[Callable[[RunnableConfig], Any]], - writer: Callable[ - [Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite] - ], - ) -> Runnable: - if reader: - value = await asyncio.to_thread(reader, config) - # passthrough additional keys from node to branch - # only doable when using dict states - if isinstance(value, dict) and isinstance(input, dict): - value = {**input, **value} - else: - value = input - result = await self.path.ainvoke(value, config) - return self._finish(writer, input, result, config) - - def _finish( - self, - writer: Callable[ - [Sequence[Union[str, Send]], RunnableConfig], Optional[ChannelWrite] - ], - input: Any, - result: Any, - config: RunnableConfig, - ) -> Union[Runnable, Any]: - if not isinstance(result, (list, tuple)): - result = [result] - if self.ends: - destinations: Sequence[Union[Send, str]] = [ - r if isinstance(r, Send) else self.ends[r] for r in result - ] - else: - destinations = cast(Sequence[Union[Send, str]], result) - if any(dest is None or dest == START for dest in destinations): - raise ValueError("Branch did not return a valid destination") - if any(p.node == END for p in destinations if isinstance(p, Send)): - raise InvalidUpdateError("Cannot send a packet to the END node") - return writer(destinations, config) or input - - class Graph: def __init__(self) -> None: self.nodes: dict[str, NodeSpec] = {} @@ -267,25 +173,7 @@ class Graph: "Adding an edge to a graph that has already been compiled. This will " "not be reflected in the compiled graph." ) - # coerce path_map to a dictionary - try: - if isinstance(path_map, dict): - path_map_ = path_map.copy() - elif isinstance(path_map, list): - path_map_ = {name: name for name in path_map} - elif isinstance(path, Runnable): - path_map_ = None - elif rtn_type := get_type_hints(path.__call__).get( # type: ignore[operator] - "return" - ) or get_type_hints(path).get("return"): - if get_origin(rtn_type) is Literal: - path_map_ = {name: name for name in get_args(rtn_type)} - else: - path_map_ = None - else: - path_map_ = None - except Exception: - path_map_ = None + # find a name for the condition path = coerce_to_runnable(path, name=None, trace=True) name = path.name or "condition" @@ -295,7 +183,7 @@ class Graph: f"Branch with name `{path.name}` already exists for node " f"`{source}`" ) # save it - self.branches[source][name] = Branch(path, path_map_, then) + self.branches[source][name] = Branch.from_path(path, path_map, then, False) return self def set_entry_point(self, key: str) -> Self: diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index a48356df5..48267c5df 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/libs/langgraph/langgraph/graph/state.py @@ -7,7 +7,9 @@ from inspect import isclass, isfunction, ismethod, signature from types import FunctionType from typing import ( Any, + Awaitable, Callable, + Hashable, Literal, NamedTuple, Optional, @@ -40,7 +42,14 @@ from langgraph.errors import ( ParentCommand, create_error_message, ) -from langgraph.graph.graph import END, START, Branch, CompiledGraph, Graph, Send +from langgraph.graph.branch import Branch +from langgraph.graph.graph import ( + END, + START, + CompiledGraph, + Graph, + Send, +) from langgraph.managed.base import ( ChannelKeyPlaceholder, ChannelTypePlaceholder, @@ -461,6 +470,57 @@ class StateGraph(Graph): self.waiting_edges.add((tuple(start_key), end_key)) return self + def add_conditional_edges( + self, + source: str, + path: Union[ + Callable[..., Union[Hashable, list[Hashable]]], + Callable[..., Awaitable[Union[Hashable, list[Hashable]]]], + Runnable[Any, Union[Hashable, list[Hashable]]], + ], + path_map: Optional[Union[dict[Hashable, str], list[str]]] = None, + then: Optional[str] = None, + ) -> Self: + """Add a conditional edge from the starting node to any number of destination nodes. + + Args: + source (str): The starting node. This conditional edge will run when + exiting this node. + path (Union[Callable, Runnable]): The callable that determines the next + node or nodes. If not specifying `path_map` it should return one or + more nodes. If it returns END, the graph will stop execution. + path_map (Optional[dict[Hashable, str]]): Optional mapping of paths to node + names. If omitted the paths returned by `path` should be node names. + then (Optional[str]): The name of a node to execute after the nodes + selected by `path`. + + Returns: + Self: The instance of the graph, allowing for method chaining. + + Note: Without typehints on the `path` function's return value (e.g., `-> Literal["foo", "__end__"]:`) + or a path_map, the graph visualization assumes the edge could transition to any node in the graph. + + """ # noqa: E501 + if self.compiled: + logger.warning( + "Adding an edge to a graph that has already been compiled. This will " + "not be reflected in the compiled graph." + ) + + # find a name for the condition + path = coerce_to_runnable(path, name=None, trace=True) + name = path.name or "condition" + # validate the condition + if name in self.branches[source]: + raise ValueError( + f"Branch with name `{path.name}` already exists for node " f"`{source}`" + ) + # save it + self.branches[source][name] = Branch.from_path(path, path_map, then, True) + if schema := self.branches[source][name].input_schema: + self._add_schema(schema) + return self + def add_sequence( self, nodes: Sequence[Union[RunnableLike, tuple[str, RunnableLike]]], @@ -826,12 +886,12 @@ class CompiledStateGraph(CompiledGraph): config, cast(Sequence[Union[Send, ChannelWriteEntry]], writes) ) - # attach branch publisher - schema = ( + schema = branch.input_schema or ( self.builder.nodes[start].input if start in self.builder.nodes else self.builder.schema ) + # attach branch publisher self.nodes[start] |= branch.run( branch_writer, _get_state_reader(self.builder, schema) if with_reader else None, diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index f562720b8..f049e4e81 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -2827,7 +2827,7 @@ def test_state_graph_packets( } # Define decision-making logic - def should_continue(data: AgentState) -> str: + def should_continue(data: dict) -> str: assert isinstance(data["session"], httpx.Client) assert ( data["something_extra"] == "hi there" diff --git a/libs/langgraph/tests/test_state.py b/libs/langgraph/tests/test_state.py index 5af406669..d85ccef94 100644 --- a/libs/langgraph/tests/test_state.py +++ b/libs/langgraph/tests/test_state.py @@ -1,4 +1,5 @@ import inspect +import operator import warnings from dataclasses import dataclass, field from typing import Annotated as Annotated2 @@ -328,3 +329,74 @@ def test__get_node_name() -> None: # class method assert _get_node_name(MyClass().class_method) == "class_method" + + +def test_input_schema_conditional_edge(): + class OverallState(TypedDict): + foo: Annotated[int, operator.add] + bar: str + + class PrivateState(TypedDict): + baz: str + + builder = StateGraph(OverallState) + + def node_1(state: OverallState): + return {"foo": 1, "baz": "bar"} + + def node_2(state: PrivateState): + return {"foo": 1, "bar": state["baz"], "something_else": "meow"} + + def node_3(state: OverallState): + return {"foo": 1} + + def router(state: OverallState): + assert state == {"foo": 2, "bar": "bar"} + if state["foo"] == 2: + return "node_3" + else: + return "__end__" + + builder.add_node(node_1) + builder.add_node(node_2) + builder.add_node(node_3) + builder.add_conditional_edges("node_2", router) + builder.add_edge("__start__", "node_1") + builder.add_edge("node_1", "node_2") + graph = builder.compile() + assert graph.invoke({"foo": 0}) == {"foo": 3, "bar": "bar"} + + +def test_private_input_schema_conditional_edge(): + class OverallState(TypedDict): + foo: Annotated[int, operator.add] + bar: str + + class RouterState(TypedDict): + baz: str + + class Node2State(TypedDict): + foo: Annotated[int, operator.add] + baz: str + + builder = StateGraph(OverallState) + + def node_1(state: OverallState): + return {"foo": 1, "baz": "meow"} + + def node_2(state: Node2State): + return {"foo": 1, "bar": state["baz"]} + + def router(state: RouterState): + assert state == {"baz": "meow"} + if state["baz"] == "meow": + return "node_2" + else: + return "__end__" + + builder.add_node(node_1) + builder.add_node(node_2) + builder.add_conditional_edges("node_1", router) + builder.add_edge("__start__", "node_1") + graph = builder.compile() + assert graph.invoke({"foo": 0}) == {"foo": 2, "bar": "meow"}