mirror of
https://github.com/langchain-ai/langgraph.git
synced 2026-08-25 17:12:26 +02:00
functional api: remove generator support (#3220)
Remove generator support in entrypoint.
This commit is contained in:
Eugene Yurtsev
@@ -31,7 +31,7 @@ from langgraph.pregel.call import (
|
||||
from langgraph.pregel.read import PregelNode
|
||||
from langgraph.pregel.write import ChannelWrite, ChannelWriteEntry
|
||||
from langgraph.store.base import BaseStore
|
||||
from langgraph.types import _DC_KWARGS, RetryPolicy, StreamMode, StreamWriter
|
||||
from langgraph.types import _DC_KWARGS, RetryPolicy, StreamMode
|
||||
|
||||
|
||||
@overload
|
||||
@@ -175,8 +175,7 @@ class entrypoint:
|
||||
| **`config`** | A configuration object (aka RunnableConfig) that holds run-time configuration values. |
|
||||
| **`previous`** | The previous return value for the given thread (available only when a checkpointer is provided). |
|
||||
|
||||
The entrypoint decorator can be applied to sync functions, async functions,
|
||||
generator functions, and async generator functions.
|
||||
The entrypoint decorator can be applied to sync functions or async functions.
|
||||
|
||||
### State management
|
||||
|
||||
@@ -188,24 +187,6 @@ class entrypoint:
|
||||
`entrypoint.final` object to return a value while saving a different value to the
|
||||
checkpoint.
|
||||
|
||||
### Generator functions
|
||||
|
||||
In generator functions, `yield` is used as a shorthand for writing
|
||||
to the `custom` channel using the `writer` parameter (i.e., `writer(chunk)`).
|
||||
|
||||
The value of `previous` will be the list of the values yielded during the previous
|
||||
run for the given thread id, unless an `entrypoint.final` was yielded.
|
||||
|
||||
If an `entrypoint.final` object is yielded, the value of `previous` will be the
|
||||
value the `save` attribute of the `entrypoint.final` object.
|
||||
|
||||
When executing an entrypoint created from a generator function, expect the following
|
||||
behavior:
|
||||
- stream_mode is set to 'custom' by default, and streaming will not stream the
|
||||
return value
|
||||
- add a `values` or `updates` stream_mode to stream the return value (if needed)
|
||||
- using `invoke` will return the return value of the entrypoint
|
||||
|
||||
Args:
|
||||
checkpointer: Specify a checkpointer to create a workflow that can persist
|
||||
its state across runs.
|
||||
@@ -324,35 +305,6 @@ class entrypoint:
|
||||
my_workflow.invoke(3, config) # 0 (previous was None)
|
||||
my_workflow.invoke(1, config) # 6 (previous was 3 * 2 from the previous invocation)
|
||||
```
|
||||
|
||||
Example: Using a generator entrypoint
|
||||
You can decorate a generator function with the `entrypoint` decorator.
|
||||
|
||||
```python
|
||||
from langgraph.checkpoint.memory import MemorySaver
|
||||
from langgraph.func import entrypoint
|
||||
|
||||
@entrypoint(checkpointer=MemorySaver())
|
||||
def workflow(inputs: dict):
|
||||
yield "hello"
|
||||
yield "world"
|
||||
|
||||
config = {
|
||||
"configurable": {
|
||||
"thread_id": "1"
|
||||
}
|
||||
}
|
||||
|
||||
for result in workflow.stream({}, config):
|
||||
print(result)
|
||||
```
|
||||
|
||||
This will print:
|
||||
```pycon
|
||||
hello
|
||||
world
|
||||
```
|
||||
|
||||
"""
|
||||
|
||||
def __init__(
|
||||
@@ -409,159 +361,19 @@ class entrypoint:
|
||||
"""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.
|
||||
func: The function to convert. Support both sync and async 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:
|
||||
if inspect.isgeneratorfunction(func) or inspect.isasyncgenfunction(func):
|
||||
raise NotImplementedError(
|
||||
"Generators are not supported in the Functional API."
|
||||
)
|
||||
|
||||
@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):
|
||||
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):
|
||||
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(
|
||||
"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_entrypoint(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:
|
||||
|
||||
@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, writer=writer, **kwargs):
|
||||
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):
|
||||
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(
|
||||
"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_entrypoint(agen_wrapper)
|
||||
stream_mode = "custom"
|
||||
else:
|
||||
bound = get_runnable_for_entrypoint(func)
|
||||
stream_mode = "updates"
|
||||
bound = get_runnable_for_entrypoint(func)
|
||||
stream_mode: StreamMode = "updates"
|
||||
|
||||
# get input and output types
|
||||
sig = inspect.signature(func)
|
||||
|
||||
@@ -5824,50 +5824,13 @@ 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"
|
||||
with pytest.raises(NotImplementedError):
|
||||
|
||||
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",
|
||||
),
|
||||
]
|
||||
@entrypoint(checkpointer=MemorySaver())
|
||||
def foo(inputs, previous=None) -> Any:
|
||||
previous_return_values.append(previous)
|
||||
yield "a"
|
||||
yield "b"
|
||||
|
||||
|
||||
@pytest.mark.parametrize("checkpointer_name", ALL_CHECKPOINTERS_SYNC)
|
||||
@@ -6231,18 +6194,6 @@ def test_entrypoint_output_schema_with_return_and_save() -> None:
|
||||
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."""
|
||||
@@ -6269,76 +6220,6 @@ def test_entrypoint_with_return_and_save() -> None:
|
||||
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
|
||||
|
||||
|
||||
def test_named_tasks_functional() -> None:
|
||||
class Foo:
|
||||
def foo(self, value: str) -> dict:
|
||||
|
||||
@@ -7386,23 +7386,14 @@ async def test_async_entrypoint_without_checkpointer() -> None:
|
||||
}
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
async def test_entrypoint_from_async_generator() -> None:
|
||||
"""@entrypoint does not support sync generators."""
|
||||
# Test invoke
|
||||
previous_return_values = []
|
||||
with pytest.raises(NotImplementedError):
|
||||
|
||||
# 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]
|
||||
@entrypoint(checkpointer=MemorySaver())
|
||||
async def foo(inputs) -> Any:
|
||||
yield "a"
|
||||
yield "b"
|
||||
|
||||
|
||||
@NEEDS_CONTEXTVARS
|
||||
|
||||
Reference in New Issue
Block a user