diff --git a/libs/langgraph/langgraph/constants.py b/libs/langgraph/langgraph/constants.py index fd5f8758c..c37a48fa0 100644 --- a/libs/langgraph/langgraph/constants.py +++ b/libs/langgraph/langgraph/constants.py @@ -23,6 +23,7 @@ END = sys.intern("__end__") """The last (maybe virtual) node in graph-style Pregel.""" SELF = sys.intern("__self__") """The implicit branch that handles each node's Control values.""" +PREVIOUS = sys.intern("__previous__") # --- Reserved write keys --- INPUT = sys.intern("__input__") @@ -78,7 +79,7 @@ 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") +CONFIG_KEY_PREVIOUS = sys.intern("__pregel_previous") # holds the previous return value from a stateful Pregel graph. # --- Other constants --- diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index 1ca426e8f..f73c63724 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -4,12 +4,17 @@ import functools import inspect import types from collections.abc import Iterator +from dataclasses import dataclass from typing import ( Any, Awaitable, Callable, + Generic, Optional, + TypeVar, Union, + get_args, + get_origin, overload, ) @@ -20,14 +25,14 @@ from langchain_core.runnables.graph import Graph, Node from langgraph.channels.ephemeral_value import EphemeralValue from langgraph.channels.last_value import LastValue from langgraph.checkpoint.base import BaseCheckpointSaver -from langgraph.constants import END, START, TAG_HIDDEN +from langgraph.constants import END, PREVIOUS, START, TAG_HIDDEN from langgraph.pregel import Pregel from langgraph.pregel.call import P, T, call, get_runnable_for_entrypoint from langgraph.pregel.protocol import PregelProtocol from langgraph.pregel.read import PregelNode from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry from langgraph.store.base import BaseStore -from langgraph.types import RetryPolicy, StreamMode, StreamWriter +from langgraph.types import _DC_KWARGS, RetryPolicy, StreamMode, StreamWriter @overload @@ -140,12 +145,15 @@ def task( return decorator -def entrypoint( - *, - checkpointer: Optional[BaseCheckpointSaver] = None, - store: Optional[BaseStore] = None, - config_schema: Optional[type[Any]] = None, -) -> Callable[[types.FunctionType], Pregel]: +R = TypeVar("R") +S = TypeVar("S") + + +# The decorator was wrapped in a class to support the `final` attribute. +# In this form, the `final` attribute should play nicely with IDE autocompletion, +# and type checking tools. +# In addition, we'll be able to surface this information in the API Reference. +class entrypoint: """Define a LangGraph workflow using the `entrypoint` decorator. !!! warning "Experimental" @@ -178,9 +186,6 @@ def entrypoint( config_schema: Specifies the schema for the configuration object that will be passed to the workflow. - Returns: - A decorator that converts a function into a Pregel graph. - Example: Using entrypoint and tasks ```python import time @@ -266,7 +271,34 @@ def entrypoint( ``` """ - def _imp(func: types.FunctionType) -> Pregel: + def __init__( + self, + checkpointer: Optional[BaseCheckpointSaver] = None, + store: Optional[BaseStore] = None, + config_schema: Optional[type[Any]] = None, + ) -> None: + """Initialize the entrypoint decorator.""" + self.checkpointer = checkpointer + self.store = store + self.config_schema = config_schema + + @dataclass(**_DC_KWARGS) + class final(Generic[R, S]): + """A primitive that can be returned from an entrypoint. + + This primitive allows to save a value to the checkpointer distinct from the + return value from the entrypoint. + """ + + value: R + """Value to return. A value will always be returned even if it is None.""" + save: S + """The value for the state for the next checkpoint. + + A value will always be saved even if it is None. + """ + + def __call__(self, func: types.FunctionType) -> Pregel: """Convert a function into a Pregel graph. Args: @@ -286,22 +318,53 @@ def entrypoint( @functools.wraps(func) def gen_wrapper(*args: Any, writer: StreamWriter, **kwargs: Any) -> Any: + final_: Optional[entrypoint.final] = None chunks = [] for chunk in func(*args, writer=writer, **kwargs): - writer(chunk) - chunks.append(chunk) - return chunks + if isinstance(chunk, entrypoint.final): + if final_ is not None: + raise RuntimeError( + "Yielding multiple entrypoint.final " + "objects is not allowed." + ) + else: + final_ = chunk + else: + if final_ is not None: + raise RuntimeError( + "Yielding a value after a entrypoint.final " + "object is not allowed." + ) + writer(chunk) + chunks.append(chunk) + + return final_ if final_ else chunks else: @functools.wraps(func) def gen_wrapper(*args: Any, writer: StreamWriter, **kwargs: Any) -> Any: + final_: Optional[entrypoint.final] = None 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 + if isinstance(chunk, entrypoint.final): + if final_ is not None: + raise RuntimeError( + "Yielding multiple entrypoint.final " + "objects is not allowed." + ) + else: + final_ = chunk + else: + if final_ is not None: + raise RuntimeError( + "Yielding a value after a entrypoint.final " + "object is not allowed." + ) + writer(chunk) + chunks.append(chunk) + return final_ if final_ else chunks # Create a new parameter for the writer argument extra_param = inspect.Parameter( @@ -329,22 +392,50 @@ def entrypoint( async def agen_wrapper( *args: Any, writer: StreamWriter, **kwargs: Any ) -> Any: + final_: Optional[entrypoint.final] = None chunks = [] async for chunk in func(*args, writer=writer, **kwargs): - writer(chunk) - chunks.append(chunk) - return chunks + if isinstance(chunk, entrypoint.final): + if final_ is not None: + raise RuntimeError( + "Yielding multiple entrypoint.final objects is not allowed." + ) + else: + final_ = chunk + else: + if final_ is not None: + raise RuntimeError( + "Yielding a value after a entrypoint.final object is not allowed." + ) + writer(chunk) + chunks.append(chunk) + + return final_ if final_ else chunks else: @functools.wraps(func) async def agen_wrapper( *args: Any, writer: StreamWriter, **kwargs: Any ) -> Any: + final_: Optional[entrypoint.final] = None chunks = [] async for chunk in func(*args, **kwargs): - writer(chunk) - chunks.append(chunk) - return chunks + if isinstance(chunk, entrypoint.final): + if final_ is not None: + raise RuntimeError( + "Yielding multiple entrypoint.final objects is not allowed." + ) + else: + final_ = chunk + else: + if final_ is not None: + raise RuntimeError( + "Yielding a value after a entrypoint.final object is not allowed." + ) + writer(chunk) + chunks.append(chunk) + + return final_ if final_ else chunks # Create a new parameter for the writer argument extra_param = inspect.Parameter( @@ -376,11 +467,36 @@ def entrypoint( is not inspect.Signature.empty else Any ) - output_type = ( - sig.return_annotation - if sig.return_annotation is not inspect.Signature.empty - else Any - ) + + def _pluck_return_value(value: Any) -> Any: + """Extract the return_ value the entrypoint.final object or passthrough.""" + return value.value if isinstance(value, entrypoint.final) else value + + def _pluck_save_value(value: Any) -> Any: + """Get save value from the entrypoint.final object or passthrough.""" + return value.save if isinstance(value, entrypoint.final) else value + + output_type, save_type = Any, Any + if sig.return_annotation is not inspect.Signature.empty: + # User does not parameterize entrypoint.final properly + if ( + sig.return_annotation is entrypoint.final + ): # Un-parameterized entrypoint.final + output_type = save_type = Any + else: + origin = get_origin(sig.return_annotation) + if origin is entrypoint.final: + type_annotations = get_args(sig.return_annotation) + if len(type_annotations) != 2: + raise TypeError( + "Please an annotation for both the return_ and " + "the save values." + "For example, `-> entrypoint.final[int, str]` would assign a " + "return_ a type of `int` and save the type `str`." + ) + output_type, save_type = get_args(sig.return_annotation) + else: + output_type = save_type = sig.return_annotation return EntrypointPregel( nodes={ @@ -388,25 +504,32 @@ def entrypoint( bound=bound, triggers=[START], channels=[START], - writers=[ChannelWrite([ChannelWriteEntry(END)], tags=[TAG_HIDDEN])], + writers=[ + ChannelWrite( + [ + ChannelWriteEntry(END, mapper=_pluck_return_value), + ChannelWriteEntry(PREVIOUS, mapper=_pluck_save_value), + ], + tags=[TAG_HIDDEN], + ) + ], ) }, channels={ START: EphemeralValue(input_type), END: LastValue(output_type, END), + PREVIOUS: LastValue(save_type, PREVIOUS), }, input_channels=START, output_channels=END, stream_channels=END, stream_mode=stream_mode, stream_eager=True, - checkpointer=checkpointer, - store=store, - config_type=config_schema, + checkpointer=self.checkpointer, + store=self.store, + config_type=self.config_schema, ) - return _imp - class EntrypointPregel(Pregel): def get_graph( diff --git a/libs/langgraph/langgraph/pregel/algo.py b/libs/langgraph/langgraph/pregel/algo.py index eff62e0fb..7a78c196a 100644 --- a/libs/langgraph/langgraph/pregel/algo.py +++ b/libs/langgraph/langgraph/pregel/algo.py @@ -38,7 +38,7 @@ from langgraph.constants import ( CONFIG_KEY_CHECKPOINT_MAP, CONFIG_KEY_CHECKPOINT_NS, CONFIG_KEY_CHECKPOINTER, - CONFIG_KEY_END, + CONFIG_KEY_PREVIOUS, CONFIG_KEY_READ, CONFIG_KEY_SCRATCHPAD, CONFIG_KEY_SEND, @@ -51,6 +51,7 @@ from langgraph.constants import ( NS_END, NS_SEP, NULL_TASK_ID, + PREVIOUS, PULL, PUSH, RESERVED, @@ -508,9 +509,6 @@ def prepare_single_task( pending_writes, task_id, ), - CONFIG_KEY_END: checkpoint["channel_values"].get( - "__end__", None - ), }, ), triggers, @@ -620,8 +618,8 @@ def prepare_single_task( pending_writes, task_id, ), - CONFIG_KEY_END: checkpoint["channel_values"].get( - "__end__", None + CONFIG_KEY_PREVIOUS: checkpoint["channel_values"].get( + PREVIOUS, None ), }, ), @@ -744,8 +742,8 @@ def prepare_single_task( pending_writes, task_id, ), - CONFIG_KEY_END: checkpoint["channel_values"].get( - "__end__", None + CONFIG_KEY_PREVIOUS: checkpoint["channel_values"].get( + PREVIOUS, None ), }, ), diff --git a/libs/langgraph/langgraph/utils/runnable.py b/libs/langgraph/langgraph/utils/runnable.py index fbdb2502b..8416efc72 100644 --- a/libs/langgraph/langgraph/utils/runnable.py +++ b/libs/langgraph/langgraph/utils/runnable.py @@ -36,7 +36,7 @@ from typing_extensions import TypeGuard from langgraph.constants import ( CONF, - CONFIG_KEY_END, + CONFIG_KEY_PREVIOUS, CONFIG_KEY_STORE, CONFIG_KEY_STREAM_WRITER, ) @@ -85,7 +85,7 @@ KWARGS_CONFIG_KEYS: tuple[tuple[str, tuple[Any, ...], str, Any], ...] = ( ( sys.intern("previous"), (ANY_TYPE,), - CONFIG_KEY_END, + CONFIG_KEY_PREVIOUS, inspect.Parameter.empty, ), ) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index b48ce9d8d..3b0d33fd6 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -6108,3 +6108,140 @@ def test_multiple_subgraphs_mixed_checkpointer( ), ((), {"parent_node": {"parent_counter": 7}}), ] + + +def test_entrypoint_output_schema_with_return_and_save() -> None: + """Test output schema inference with entrypoint.final.""" + + # Un-parameterized entrypoint.final is interpreted as entrypoint.final[Any, Any] + @entrypoint() + def foo2(inputs, *, previous: Any) -> entrypoint.final: + return entrypoint.final(value="foo", save=1) + + assert foo2.get_output_schema().model_json_schema() == { + "title": "LangGraphOutput", + } + + @entrypoint() + def foo(inputs, *, previous: Any) -> entrypoint.final[str, int]: + return entrypoint.final(value="foo", save=1) + + assert foo.get_output_schema().model_json_schema() == { + "title": "LangGraphOutput", + "type": "string", + } + + with pytest.raises(TypeError): + # Raise an exception on an improperly parameterized entrypoint.final + # User is attempting to parameterize in this case, so we'll offer + # a bit of help if it's not done correctly. + @entrypoint() + def foo(inputs, *, previous: Any) -> entrypoint.final[int]: + return entrypoint.final(value=1, save=1) # type: ignore + + @entrypoint() + def foo(inputs, *, previous: Any) -> Generator[int, None, None]: + yield 1 + + assert foo.get_output_schema().model_json_schema() == { + "items": { + "type": "integer", + }, + "title": "LangGraphOutput", + "type": "array", + } + + +def test_entrypoint_with_return_and_save() -> None: + """Test entrypoint with return and save.""" + previous_ = None + + @entrypoint(checkpointer=MemorySaver()) + def foo(msg: str, *, previous: Any) -> entrypoint.final[int, list[str]]: + nonlocal previous_ + previous_ = previous + previous = previous or [] + return entrypoint.final(value=len(previous), save=previous + [msg]) + + assert foo.get_output_schema().model_json_schema() == { + "title": "LangGraphOutput", + "type": "integer", + } + + config = {"configurable": {"thread_id": "1"}} + assert foo.invoke("hello", config) == 0 + assert previous_ is None + assert foo.invoke("goodbye", config) == 1 + assert previous_ == ["hello"] + assert foo.invoke("definitely", config) == 2 + assert previous_ == ["hello", "goodbye"] + + +def test_entrypoint_generator_with_return_and_save() -> None: + """Verify that generators produce expected results.""" + previous_ = None + + @entrypoint(checkpointer=MemorySaver()) + def workflow(inputs: dict, *, previous: Any): + nonlocal previous_ + previous_ = previous + + yield "hello" + yield "world" + yield entrypoint.final(value="!", save="saved value") + + assert list(workflow.stream({}, {"configurable": {"thread_id": "0"}})) == [ + "hello", + "world", + ] + assert list( + workflow.stream({}, {"configurable": {"thread_id": "0"}}, stream_mode="updates") + ) == [ + { + "workflow": "!", + } + ] + + assert workflow.invoke({}, {"configurable": {"thread_id": "1"}}) == "!" + assert previous_ is None + + # 2nd time around previous is set + assert workflow.invoke({}, {"configurable": {"thread_id": "1"}}) == "!" + assert previous_ == "saved value" + + # Test with another thread + assert workflow.invoke({}, {"configurable": {"thread_id": "2"}}) == "!" + assert previous_ is None + + +async def test_entrypoint_async_generator_with_return_and_save() -> None: + """Verify that generators produce expected results.""" + previous_ = None + + @entrypoint(checkpointer=MemorySaver()) + async def workflow(inputs: dict, *, previous: Any): + nonlocal previous_ + previous_ = previous + + yield "hello" + yield "world" + yield entrypoint.final(value="!", save="saved value") + + assert [ + c async for c in workflow.astream({}, {"configurable": {"thread_id": "0"}}) + ] == [ + "hello", + "world", + ] + + assert await workflow.ainvoke({}, {"configurable": {"thread_id": "1"}}) == "!" + + assert previous_ is None + + # 2nd time around previous is set + assert await workflow.ainvoke({}, {"configurable": {"thread_id": "1"}}) == "!" + assert previous_ == "saved value" + + # Test with another thread + assert await workflow.ainvoke({}, {"configurable": {"thread_id": "2"}}) == "!" + assert previous_ is None