diff --git a/libs/langgraph/langgraph/pregel/utils.py b/libs/langgraph/langgraph/pregel/utils.py index f484664a0..fd62e345f 100644 --- a/libs/langgraph/langgraph/pregel/utils.py +++ b/libs/langgraph/langgraph/pregel/utils.py @@ -1,7 +1,10 @@ -from typing import Optional +import ast +import inspect +import textwrap +from typing import Any, Callable, Optional from langchain_core.runnables import RunnableLambda, RunnableSequence -from langchain_core.runnables.utils import get_function_nonlocals +from typing_extensions import override from langgraph.checkpoint.base import ChannelVersions from langgraph.pregel.protocol import PregelProtocol @@ -55,3 +58,152 @@ def find_subgraph_pregel(candidate: Runnable) -> Optional[PregelProtocol]: ) return None + + +def get_function_nonlocals(func: Callable) -> list[Any]: + """Get the nonlocal variables accessed by a function. + + Args: + func: The function to check. + + Returns: + List[Any]: The nonlocal variables accessed by the function. + """ + try: + code = inspect.getsource(func) + tree = ast.parse(textwrap.dedent(code)) + visitor = FunctionNonLocals() + visitor.visit(tree) + values: list[Any] = [] + closure = ( + inspect.getclosurevars(func.__wrapped__) + if hasattr(func, "__wrapped__") and callable(func.__wrapped__) + else inspect.getclosurevars(func) + ) + candidates = {**closure.globals, **closure.nonlocals} + for k, v in candidates.items(): + if k in visitor.nonlocals: + values.append(v) + for kk in visitor.nonlocals: + if "." in kk and kk.startswith(k): + vv = v + for part in kk.split(".")[1:]: + if vv is None: + break + else: + try: + vv = getattr(vv, part) + except AttributeError: + break + else: + values.append(vv) + except (SyntaxError, TypeError, OSError, SystemError): + return [] + + return values + + +class FunctionNonLocals(ast.NodeVisitor): + """Get the nonlocal variables accessed of a function.""" + + def __init__(self) -> None: + self.nonlocals: set[str] = set() + + @override + def visit_FunctionDef(self, node: ast.FunctionDef) -> Any: + """Visit a function definition. + + Args: + node: The node to visit. + + Returns: + Any: The result of the visit. + """ + visitor = NonLocals() + visitor.visit(node) + self.nonlocals.update(visitor.loads - visitor.stores) + + @override + def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef) -> Any: + """Visit an async function definition. + + Args: + node: The node to visit. + + Returns: + Any: The result of the visit. + """ + visitor = NonLocals() + visitor.visit(node) + self.nonlocals.update(visitor.loads - visitor.stores) + + @override + def visit_Lambda(self, node: ast.Lambda) -> Any: + """Visit a lambda function. + + Args: + node: The node to visit. + + Returns: + Any: The result of the visit. + """ + visitor = NonLocals() + visitor.visit(node) + self.nonlocals.update(visitor.loads - visitor.stores) + + +class NonLocals(ast.NodeVisitor): + """Get nonlocal variables accessed.""" + + def __init__(self) -> None: + self.loads: set[str] = set() + self.stores: set[str] = set() + + @override + def visit_Name(self, node: ast.Name) -> Any: + """Visit a name node. + + Args: + node: The node to visit. + + Returns: + Any: The result of the visit. + """ + if isinstance(node.ctx, ast.Load): + self.loads.add(node.id) + elif isinstance(node.ctx, ast.Store): + self.stores.add(node.id) + + @override + def visit_Attribute(self, node: ast.Attribute) -> Any: + """Visit an attribute node. + + Args: + node: The node to visit. + + Returns: + Any: The result of the visit. + """ + if isinstance(node.ctx, ast.Load): + parent = node.value + attr_expr = node.attr + while isinstance(parent, ast.Attribute): + attr_expr = parent.attr + "." + attr_expr + parent = parent.value + if isinstance(parent, ast.Name): + self.loads.add(parent.id + "." + attr_expr) + self.loads.discard(parent.id) + elif isinstance(parent, ast.Call): + if isinstance(parent.func, ast.Name): + self.loads.add(parent.func.id) + else: + parent = parent.func + attr_expr = "" + while isinstance(parent, ast.Attribute): + if attr_expr: + attr_expr = parent.attr + "." + attr_expr + else: + attr_expr = parent.attr + parent = parent.value + if isinstance(parent, ast.Name): + self.loads.add(parent.id + "." + attr_expr)