diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index 86c5403f9..c7b4d1bb4 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -2,6 +2,7 @@ import asyncio import concurrent import concurrent.futures import functools +import inspect import types from typing import ( Any, @@ -111,8 +112,13 @@ def entrypoint( store: Optional[BaseStore] = None, ) -> Callable[[types.FunctionType], Pregel]: def _imp(func: types.FunctionType) -> Pregel: - bound = get_runnable_for_func(func) - stream_mode: StreamMode = "updates" + if inspect.isgeneratorfunction(func): + raise TypeError("@entrypoint does not support generator functions.") + elif inspect.isasyncgenfunction(func): + raise TypeError("@entrypoint does not support async generator functions.") + else: + bound = get_runnable_for_func(func) + stream_mode: StreamMode = "updates" return Pregel( nodes={ diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 8cf2e6b4f..432eea069 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -15,6 +15,7 @@ from typing import ( Any, Dict, Generator, + Iterable, Iterator, List, Literal, @@ -5302,9 +5303,32 @@ def test_multiple_updates() -> None: ] -def test_version_1_of_entrypoint() -> None: - from langgraph.func import entrypoint +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} + + +def test_entrypoint_stateful() -> None: + """Test stateful entrypoint invoke.""" + + # Test invoke states = [] # In this version reducers do not work @@ -5315,11 +5339,49 @@ def test_version_1_of_entrypoint() -> None: config = {"configurable": {"thread_id": "1"}} - foo.invoke({"a": "1"}, config) - foo.invoke({"a": "2"}, config) - foo.invoke({"a": "3"}, config) + 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}}] + + +async def test_entrypoint_from_generator() -> None: + """@entrypoint does not support sync generators.""" + + with pytest.raises(TypeError): + + @entrypoint(checkpointer=MemorySaver()) + def foo(inputs: Any) -> Iterable[dict]: + yield "a" + + +async def test_entrypoint_from_async_generator() -> None: + """@entrypoint does not support async generators.""" + + with pytest.raises(TypeError): + + @entrypoint(checkpointer=MemorySaver()) + def foo(inputs: Any) -> Iterable[dict]: + yield "a"