From b4dc3a851f2f6f54593efbdbff5db5e54f70fdf9 Mon Sep 17 00:00:00 2001 From: Eugene Yurtsev Date: Mon, 27 Jan 2025 19:42:35 -0500 Subject: [PATCH] functional api: remove generator support (#3220) Remove generator support in entrypoint. --- libs/langgraph/langgraph/func/__init__.py | 206 +--------------------- libs/langgraph/tests/test_pregel.py | 131 +------------- libs/langgraph/tests/test_pregel_async.py | 19 +- 3 files changed, 20 insertions(+), 336 deletions(-) diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index 8e392c3ea..603a1164b 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -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) diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 4ce241086..58ac09b65 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -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: diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 85c2945b1..932330dc7 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/libs/langgraph/tests/test_pregel_async.py @@ -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