WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-09-12 04:37:51 +02:00
Code Issues Packages Projects Releases Wiki Activity

Add entrypoint.final to decouple return value from save value (#3135)

Browse Source 
* Introduce `entrypoint.final` that allows decoupling what's returned
from the state update.
* moving decorator to class object w/ call to support defining `final`
as a property on it -- this should play nicely w/ IDE tooling / type
checking.


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

```
This commit is contained in:
Eugene Yurtsev
2025-01-23 10:18:23 -05:00
committed by GitHub
parent 9e066554ba
commit 5dc5853161
5 changed files with 305 additions and 46 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/constants.py
+2 -1
View File
@@ -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 ---
libs/langgraph/langgraph/func/__init__.py
+158 -35
View File
@@ -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(
libs/langgraph/langgraph/pregel/algo.py
+6 -8
View File
@@ -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
),
},
),
libs/langgraph/langgraph/utils/runnable.py
+2 -2
View File
@@ -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,
),
)
libs/langgraph/tests/test_pregel.py
+137
View File
@@ -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