functional api: Add ability to request previous output (#3025)

1. The inputs into foo do not affect any state behavior
2. `previous` always reflects the previous return value from the
function
3. Anything can be returned and that will be the new state for the
function on the next iteration
4. This API is not meant to support reducers in the inputs/state

```python
  from langgraph.func import entrypoint

  states = []

  # In this version reducers do not work
  @entrypoint(checkpointer=MemorySaver())
  def foo(inputs, *, previous: Any) -> Any:
      states.append(previous)
      return {"previous": previous, "current": inputs}

  config = {"configurable": {"thread_id": "1"}}

  foo.invoke({"a": "1"}, config)
  foo.invoke({"a": "2"}, config)
  foo.invoke({"a": "3"}, config)
  assert states == [
      None,
      {"current": {"a": "1"}, "previous": None},
      {"current": {"a": "2"}, "previous": {"current": {"a": "1"}, "previous": None}},
  ]
```

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")

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"

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,

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)

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"]]

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": {

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": {