diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index 85c4a890c..fd5f8758c 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -78,6 +78,8 @@ CONFIG_KEY_NODE_FINISHED = sys.intern("__pregel_node_finished") # holds a callback to be called when a node is finished CONFIG_KEY_SCRATCHPAD = sys.intern("__pregel_scratchpad") # holds a mutable dict for temporary storage scoped to the current task +CONFIG_KEY_END = sys.intern("__pregel_previous") +# holds the previous return value from a stateful Pregel graph. # --- Other constants --- PUSH = sys.intern("__pregel_push") diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index 76dec76d2..a051e465e 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -113,22 +113,96 @@ def entrypoint( config_schema: Optional[type[Any]] = None, ) -> Callable[[types.FunctionType], Pregel]: def _imp(func: types.FunctionType) -> Pregel: + """Convert a function into a Pregel graph. + + Args: + func: The function to convert. Support both sync and async functions, as well + as generator and async generator functions. + + Returns: + A Pregel graph. + """ # wrap generators in a function that writes to StreamWriter if inspect.isgeneratorfunction(func): + original_sig = inspect.signature(func) + # Check if original signature has a writer argument with a matching type. + # If not, we'll inject it into the decorator, but not pass it + # to the wrapped function. + if "writer" in original_sig.parameters: - def gen_wrapper(*args: Any, writer: StreamWriter, **kwargs: Any) -> Any: - for chunk in func(*args, **kwargs): - writer(chunk) + @functools.wraps(func) + def gen_wrapper(*args: Any, writer: StreamWriter, **kwargs: Any) -> Any: + chunks = [] + for chunk in func(*args, writer=writer, **kwargs): + writer(chunk) + chunks.append(chunk) + return chunks + else: + @functools.wraps(func) + def gen_wrapper(*args: Any, writer: StreamWriter, **kwargs: Any) -> Any: + chunks = [] + # Do not pass the writer argument to the wrapped function + # as it does not have a matching parameter + for chunk in func(*args, **kwargs): + writer(chunk) + chunks.append(chunk) + return chunks + + # Create a new parameter for the writer argument + extra_param = inspect.Parameter( + "writer", + inspect.Parameter.KEYWORD_ONLY, + # The extra argument is a keyword-only argument + default=lambda _: None, + ) + # Update the function's signature to include the extra argument + new_params = list(original_sig.parameters.values()) + [extra_param] + new_sig = original_sig.replace(parameters=new_params) + # Update the signature of the wrapper function + gen_wrapper.__signature__ = new_sig # type: ignore bound = get_runnable_for_func(gen_wrapper) stream_mode: StreamMode = "custom" elif inspect.isasyncgenfunction(func): + original_sig = inspect.signature(func) + # Check if original signature has a writer argument with a matching type. + # If not, we'll inject it into the decorator, but not pass it + # to the wrapped function. + if "writer" in original_sig.parameters: - async def agen_wrapper( - *args: Any, writer: StreamWriter, **kwargs: Any - ) -> Any: - async for chunk in func(*args, **kwargs): - writer(chunk) + @functools.wraps(func) + async def agen_wrapper( + *args: Any, writer: StreamWriter, **kwargs: Any + ) -> Any: + chunks = [] + async for chunk in func(*args, writer=writer, **kwargs): + writer(chunk) + chunks.append(chunk) + return chunks + else: + + @functools.wraps(func) + async def agen_wrapper( + *args: Any, writer: StreamWriter, **kwargs: Any + ) -> Any: + chunks = [] + async for chunk in func(*args, **kwargs): + writer(chunk) + chunks.append(chunk) + return chunks + + # Create a new parameter for the writer argument + extra_param = inspect.Parameter( + "writer", + inspect.Parameter.KEYWORD_ONLY, + # The extra argument is a keyword-only argument + default=lambda _: None, + ) + # Update the function's signature to include the extra argument + new_params = list(original_sig.parameters.values()) + [extra_param] + new_sig = original_sig.replace(parameters=new_params) + # Update the signature of the wrapper function + agen_wrapper.__signature__ = new_sig # type: ignore bound = get_runnable_for_func(agen_wrapper) stream_mode = "custom" diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index 39bc9462d..d1247dbc4 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -37,6 +37,7 @@ from langgraph.constants import ( CONFIG_KEY_CHECKPOINT_MAP, CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_CHECKPOINTER, + CONFIG_KEY_END, CONFIG_KEY_READ, CONFIG_KEY_SCRATCHPAD, CONFIG_KEY_SEND, @@ -507,6 +508,9 @@ def prepare_single_task( pending_writes, task_id, ), + CONFIG_KEY_END: checkpoint["channel_values"].get( + "__end__", None + ), }, ), triggers, @@ -616,6 +620,9 @@ def prepare_single_task( pending_writes, task_id, ), + CONFIG_KEY_END: checkpoint["channel_values"].get( + "__end__", None + ), }, ), triggers, @@ -737,6 +744,9 @@ def prepare_single_task( pending_writes, task_id, ), + CONFIG_KEY_END: checkpoint["channel_values"].get( + "__end__", None + ), }, ), triggers, diff --git a/libs/langgraph/langgraph/utils/runnable.py b/libs/langgraph/langgraph/utils/runnable.py index 9f1a7ef8d..19a70e4cf 100644 --- a/libs/langgraph/langgraph/utils/runnable.py +++ b/libs/langgraph/langgraph/utils/runnable.py @@ -34,7 +34,12 @@ from langchain_core.runnables.utils import Input from langchain_core.tracers._streaming import _StreamingCallbackHandler from typing_extensions import TypeGuard -from langgraph.constants import CONF, CONFIG_KEY_STORE, CONFIG_KEY_STREAM_WRITER +from langgraph.constants import ( + CONF, + CONFIG_KEY_END, + CONFIG_KEY_STORE, + CONFIG_KEY_STREAM_WRITER, +) from langgraph.store.base import BaseStore from langgraph.types import StreamWriter from langgraph.utils.config import ( @@ -58,6 +63,10 @@ class StrEnum(str, enum.Enum): """A string enum.""" +# Special type to denote any type is accepted +ANY_TYPE = object() + + ASYNCIO_ACCEPTS_CONTEXT = sys.version_info >= (3, 11) KWARGS_CONFIG_KEYS: tuple[tuple[str, tuple[Any, ...], str, Any], ...] = ( @@ -73,6 +82,12 @@ KWARGS_CONFIG_KEYS: tuple[tuple[str, tuple[Any, ...], str, Any], ...] = ( CONFIG_KEY_STORE, inspect.Parameter.empty, ), + ( + sys.intern("previous"), + (ANY_TYPE,), + CONFIG_KEY_END, + inspect.Parameter.empty, + ), ) """List of kwargs that can be passed to functions, and their corresponding config keys, default values and type annotations. @@ -135,9 +150,12 @@ class RunnableCallable(Runnable): self.func_accepts: dict[str, bool] = {} for kw, typ, _, _ in KWARGS_CONFIG_KEYS: p = params.get(kw) - self.func_accepts[kw] = ( - p is not None and p.annotation in typ and p.kind in VALID_KINDS - ) + if typ == (ANY_TYPE,): + self.func_accepts[kw] = p is not None and p.kind in VALID_KINDS + else: + self.func_accepts[kw] = ( + p is not None and p.annotation in typ and p.kind in VALID_KINDS + ) def __repr__(self) -> str: repr_args = { @@ -162,16 +180,20 @@ class RunnableCallable(Runnable): if self.func_accepts_config: kwargs["config"] = config _conf = config[CONF] - for kw, _, ck, defv in KWARGS_CONFIG_KEYS: + for kw, _, config_key, default_value in KWARGS_CONFIG_KEYS: if not self.func_accepts[kw]: continue - if defv is inspect.Parameter.empty and kw not in kwargs and ck not in _conf: + if ( + default_value is inspect.Parameter.empty + and kw not in kwargs + and config_key not in _conf + ): raise ValueError( - f"Missing required config key '{ck}' for '{self.name}'." + f"Missing required config key '{config_key}' for '{self.name}'." ) elif kwargs.get(kw) is None: - kwargs[kw] = _conf.get(ck, defv) + kwargs[kw] = _conf.get(config_key, default_value) context = copy_context() if self.trace: @@ -210,16 +232,20 @@ class RunnableCallable(Runnable): if self.func_accepts_config: kwargs["config"] = config _conf = config[CONF] - for kw, _, ck, defv in KWARGS_CONFIG_KEYS: + for kw, _, config_key, default_value in KWARGS_CONFIG_KEYS: if not self.func_accepts[kw]: continue - if defv is inspect.Parameter.empty and kw not in kwargs and ck not in _conf: + if ( + default_value is inspect.Parameter.empty + and kw not in kwargs + and config_key not in _conf + ): raise ValueError( - f"Missing required config key '{ck}' for '{self.name}'." + f"Missing required config key '{config_key}' for '{self.name}'." ) elif kwargs.get(kw) is None: - kwargs[kw] = _conf.get(ck, defv) + kwargs[kw] = _conf.get(config_key, default_value) context = copy_context() if self.trace: callback_manager = get_async_callback_manager_for_config(config, self.tags) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 430a98dd9..a30a66cba 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -5636,3 +5636,164 @@ def test_sync_streaming_with_functional_api() -> None: delta = arrival_times[1] - arrival_times[0] # Delta cannot be less than 10 ms if it is streaming as results are generated. assert delta > time_delay + + +def test_entrypoint_without_checkpointer() -> None: + """Test no checkpointer.""" + states = [] + config = {"configurable": {"thread_id": "1"}} + + # Test without previous + @entrypoint() + def foo(inputs: Any) -> Any: + states.append(inputs) + return inputs + + assert foo.invoke({"a": "1"}, config) == {"a": "1"} + + @entrypoint() + def foo(inputs: Any, *, previous: Any) -> Any: + states.append(previous) + return {"previous": previous, "current": inputs} + + assert foo.invoke({"a": "1"}, config) == {"current": {"a": "1"}, "previous": None} + assert foo.invoke({"a": "1"}, config) == {"current": {"a": "1"}, "previous": None} + + +async def test_async_entrypoint_without_checkpointer() -> None: + """Test no checkpointer.""" + states = [] + config = {"configurable": {"thread_id": "1"}} + + # Test without previous + @entrypoint() + async def foo(inputs: Any) -> Any: + states.append(inputs) + return inputs + + assert (await foo.ainvoke({"a": "1"}, config)) == {"a": "1"} + + @entrypoint() + async def foo(inputs: Any, *, previous: Any) -> Any: + states.append(previous) + return {"previous": previous, "current": inputs} + + assert (await foo.ainvoke({"a": "1"}, config)) == { + "current": {"a": "1"}, + "previous": None, + } + assert (await foo.ainvoke({"a": "1"}, config)) == { + "current": {"a": "1"}, + "previous": None, + } + + +def test_entrypoint_stateful() -> None: + """Test stateful entrypoint invoke.""" + + # Test invoke + states = [] + + @entrypoint(checkpointer=MemorySaver()) + def foo(inputs, *, previous: Any) -> Any: + states.append(previous) + return {"previous": previous, "current": inputs} + + config = {"configurable": {"thread_id": "1"}} + + assert foo.invoke({"a": "1"}, config) == {"current": {"a": "1"}, "previous": None} + assert foo.invoke({"a": "2"}, config) == { + "current": {"a": "2"}, + "previous": {"current": {"a": "1"}, "previous": None}, + } + assert foo.invoke({"a": "3"}, config) == { + "current": {"a": "3"}, + "previous": { + "current": {"a": "2"}, + "previous": {"current": {"a": "1"}, "previous": None}, + }, + } + assert states == [ + None, + {"current": {"a": "1"}, "previous": None}, + {"current": {"a": "2"}, "previous": {"current": {"a": "1"}, "previous": None}}, + ] + + # Test stream + @entrypoint(checkpointer=MemorySaver()) + def foo(inputs, *, previous: Any) -> Any: + return {"previous": previous, "current": inputs} + + config = {"configurable": {"thread_id": "1"}} + items = [item for item in foo.stream({"a": "1"}, config)] + assert items == [{"foo": {"current": {"a": "1"}, "previous": None}}] + + +def test_entrypoint_from_sync_generator() -> None: + """@entrypoint does not support sync generators.""" + previous_return_values = [] + + @entrypoint(checkpointer=MemorySaver()) + def foo(inputs, previous=None) -> Any: + previous_return_values.append(previous) + yield "a" + yield "b" + + config = {"configurable": {"thread_id": "1"}} + + assert foo.invoke({"a": "1"}, config) == ["a", "b"] + assert previous_return_values == [None] + assert foo.invoke({"a": "2"}, config) == ["a", "b"] + assert previous_return_values == [None, ["a", "b"]] + + +def test_entrypoint_request_stream_writer() -> None: + """Test using a stream writer with an entrypoint.""" + + @entrypoint(checkpointer=MemorySaver()) + def foo(inputs, writer: StreamWriter) -> Any: + writer("a") + yield "b" + + config = {"configurable": {"thread_id": "1"}} + + # Different invocations + # Are any of these confusing or unexpected? + assert list(foo.invoke({}, config)) == ["b"] + assert list(foo.stream({}, config)) == ["a", "b"] + + # Stream modes + assert list(foo.stream({}, config, stream_mode=["updates"])) == [ + ("updates", {"foo": ["b"]}) + ] + assert list(foo.stream({}, config, stream_mode=["values"])) == [("values", ["b"])] + assert list(foo.stream({}, config, stream_mode=["custom"])) == [ + ( + "custom", + "a", + ), + ( + "custom", + "b", + ), + ] + + +async def test_entrypoint_from_async_generator() -> None: + """@entrypoint does not support sync generators.""" + # Test invoke + previous_return_values = [] + + # In this version reducers do not work + @entrypoint(checkpointer=MemorySaver()) + async def foo(inputs, previous=None) -> Any: + previous_return_values.append(previous) + yield "a" + yield "b" + + config = {"configurable": {"thread_id": "1"}} + + assert list(await foo.ainvoke({"a": "1"}, config)) == ["a", "b"] + assert previous_return_values == [None] + assert list(foo.invoke({"a": "2"}, config)) == ["a", "b"] + assert previous_return_values == [None, ["a", "b"]] diff --git a/libs/scheduler-kafka/tests/test_subgraph.py b/libs/scheduler-kafka/tests/test_subgraph.py index 053eaedd0..2a6c9992a 100644 --- a/libs/scheduler-kafka/tests/test_subgraph.py +++ b/libs/scheduler-kafka/tests/test_subgraph.py @@ -195,6 +195,7 @@ async def test_subgraph_w_interrupt( "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": False, + "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, "__pregel_scratchpad": { @@ -267,6 +268,7 @@ async def test_subgraph_w_interrupt( "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": False, + "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, "__pregel_scratchpad": { @@ -369,6 +371,7 @@ async def test_subgraph_w_interrupt( "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": False, + "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, "__pregel_scratchpad": { @@ -481,6 +484,7 @@ async def test_subgraph_w_interrupt( "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": True, + "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, "__pregel_scratchpad": { @@ -548,6 +552,7 @@ async def test_subgraph_w_interrupt( "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": True, + "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, "__pregel_scratchpad": { @@ -671,6 +676,7 @@ async def test_subgraph_w_interrupt( "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": True, + "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, "__pregel_scratchpad": { diff --git a/libs/scheduler-kafka/tests/test_subgraph_sync.py b/libs/scheduler-kafka/tests/test_subgraph_sync.py index 7d5de920c..c2c9a8fc1 100644 --- a/libs/scheduler-kafka/tests/test_subgraph_sync.py +++ b/libs/scheduler-kafka/tests/test_subgraph_sync.py @@ -194,6 +194,7 @@ def test_subgraph_w_interrupt( "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": False, + "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[0].tasks[0].id, "__pregel_scratchpad": { @@ -267,6 +268,7 @@ def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_dedupe_tasks": True, "__pregel_resuming": False, + "__pregel_previous": None, "__pregel_task_id": history[0].tasks[0].id, "__pregel_scratchpad": { "subgraph_counter": AnyInt(), @@ -370,6 +372,7 @@ def test_subgraph_w_interrupt( "__pregel_store": None, "__pregel_resuming": False, "__pregel_task_id": history[0].tasks[0].id, + "__pregel_previous": None, "__pregel_scratchpad": { "subgraph_counter": AnyInt(), "call_counter": 0, @@ -480,6 +483,7 @@ def test_subgraph_w_interrupt( "__pregel_dedupe_tasks": True, "__pregel_store": None, "__pregel_resuming": True, + "__pregel_previous": None, "__pregel_task_id": history[1].tasks[0].id, "__pregel_scratchpad": { "subgraph_counter": AnyInt(), @@ -547,6 +551,7 @@ def test_subgraph_w_interrupt( "__pregel_dedupe_tasks": True, "__pregel_store": None, "__pregel_resuming": True, + "__pregel_previous": None, "__pregel_task_id": history[1].tasks[0].id, "__pregel_scratchpad": { "subgraph_counter": AnyInt(), @@ -669,6 +674,7 @@ def test_subgraph_w_interrupt( "__pregel_ensure_latest": True, "__pregel_dedupe_tasks": True, "__pregel_resuming": True, + "__pregel_previous": None, "__pregel_store": None, "__pregel_task_id": history[1].tasks[0].id, "__pregel_scratchpad": {