diff --git a/libs/langgraph/langgraph/pregel/call.py b/libs/langgraph/langgraph/pregel/call.py index 9a74eb0c3..a4bba63fe 100644 --- a/libs/langgraph/langgraph/pregel/call.py +++ b/libs/langgraph/langgraph/pregel/call.py @@ -8,6 +8,7 @@ import sys import types from typing import Any, Callable, Optional, TypeVar, Union +from langchain_core.runnables import Runnable from typing_extensions import ParamSpec from langgraph.constants import CONF, CONFIG_KEY_CALL, RETURN, TAG_HIDDEN @@ -158,14 +159,9 @@ def get_runnable_for_entrypoint(func: Callable[..., Any]) -> RunnableSeq: trace=False, recurse=False, ) - seq = RunnableSeq( - run, - ChannelWrite([ChannelWriteEntry(RETURN)], tags=[TAG_HIDDEN]), - name=func.__name__, - ) if not _lookup_module_and_qualname(func): - return seq - return CACHE.setdefault(key, seq) + return run + return CACHE.setdefault(key, run) def get_runnable_for_task(func: Callable[..., Any]) -> RunnableSeq: @@ -204,7 +200,7 @@ def get_runnable_for_task(func: Callable[..., Any]) -> RunnableSeq: return CACHE.setdefault(key, seq) -CACHE: dict[tuple[Callable[..., Any], bool], RunnableSeq] = {} +CACHE: dict[tuple[Callable[..., Any], bool], Runnable] = {} P = ParamSpec("P") diff --git a/libs/langgraph/tests/test_large_cases.py b/libs/langgraph/tests/test_large_cases.py index 6980b7b68..f562720b8 100644 --- a/libs/langgraph/tests/test_large_cases.py +++ b/libs/langgraph/tests/test_large_cases.py @@ -2829,9 +2829,9 @@ def test_state_graph_packets( # Define decision-making logic def should_continue(data: AgentState) -> str: assert isinstance(data["session"], httpx.Client) - assert data["something_extra"] == "hi there", ( - "nodes can pass extra data to their cond edges, which isn't saved in state" - ) + assert ( + data["something_extra"] == "hi there" + ), "nodes can pass extra data to their cond edges, which isn't saved in state" # Logic to decide whether to continue in the loop or exit if tool_calls := data["messages"][-1].tool_calls: return [Send("tools", tool_call) for tool_call in tool_calls]