From 7e9cf02922b80a832e02609d4d4ff9d62a9f9820 Mon Sep 17 00:00:00 2001 From: Tat Dat Duong Date: Tue, 8 Oct 2024 17:08:08 +0200 Subject: [PATCH] Use casting instead --- libs/langgraph/langgraph/pregel/__init__.py | 2 +- libs/langgraph/langgraph/pregel/utils.py | 7 ++----- 2 files changed, 3 insertions(+), 6 deletions(-) diff --git a/libs/langgraph/langgraph/pregel/__init__.py b/libs/langgraph/langgraph/pregel/__init__.py index 268c0241b..98d82712b 100644 --- a/libs/langgraph/langgraph/pregel/__init__.py +++ b/libs/langgraph/langgraph/pregel/__init__.py @@ -390,7 +390,7 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]): continue # find the subgraph, if any - graph: Optional[Pregel] = find_subgraph_pregel(node.bound) + graph = cast(Optional[Pregel], find_subgraph_pregel(node.bound)) # if found, yield recursively if graph: diff --git a/libs/langgraph/langgraph/pregel/utils.py b/libs/langgraph/langgraph/pregel/utils.py index 2f9b9d411..3f5fe54d1 100644 --- a/libs/langgraph/langgraph/pregel/utils.py +++ b/libs/langgraph/langgraph/pregel/utils.py @@ -1,4 +1,4 @@ -from typing import TYPE_CHECKING, Optional +from typing import Optional from langchain_core.runnables import RunnableLambda, RunnableSequence from langchain_core.runnables.utils import get_function_nonlocals @@ -6,9 +6,6 @@ from langchain_core.runnables.utils import get_function_nonlocals from langgraph.checkpoint.base import ChannelVersions from langgraph.utils.runnable import Runnable, RunnableCallable, RunnableSeq -if TYPE_CHECKING: - from langgraph.pregel import Pregel - def get_new_channel_versions( previous_versions: ChannelVersions, current_versions: ChannelVersions @@ -28,7 +25,7 @@ def get_new_channel_versions( return new_versions -def find_subgraph_pregel(candidate: Runnable) -> Optional[Pregel]: +def find_subgraph_pregel(candidate: Runnable) -> Optional[Runnable]: candidates: list[Runnable] = [candidate] for c in candidates: