diff --git a/libs/langgraph/langgraph/func/__init__.py b/libs/langgraph/langgraph/func/__init__.py index 0e3cd5793..76dec76d2 100644 --- a/libs/langgraph/langgraph/func/__init__.py +++ b/libs/langgraph/langgraph/func/__init__.py @@ -110,6 +110,7 @@ def entrypoint( *, checkpointer: Optional[BaseCheckpointSaver] = None, store: Optional[BaseStore] = None, + config_schema: Optional[type[Any]] = None, ) -> Callable[[types.FunctionType], Pregel]: def _imp(func: types.FunctionType) -> Pregel: # wrap generators in a function that writes to StreamWriter @@ -172,6 +173,7 @@ def entrypoint( stream_eager=True, checkpointer=checkpointer, store=store, + config_type=config_schema, ) return _imp diff --git a/libs/langgraph/tests/test_pregel.py b/libs/langgraph/tests/test_pregel.py index 905b58b2e..4dec2ac05 100644 --- a/libs/langgraph/tests/test_pregel.py +++ b/libs/langgraph/tests/test_pregel.py @@ -1436,6 +1436,9 @@ def test_imp_task(request: pytest.FixtureRequest, checkpointer_name: str) -> Non checkpointer = request.getfixturevalue(f"checkpointer_{checkpointer_name}") mapper_calls = 0 + class Config: + model: str + @task() def mapper(input: int) -> str: nonlocal mapper_calls @@ -1443,7 +1446,7 @@ def test_imp_task(request: pytest.FixtureRequest, checkpointer_name: str) -> Non time.sleep(input / 100) return str(input) * 2 - @entrypoint(checkpointer=checkpointer) + @entrypoint(checkpointer=checkpointer, config_schema=Config) def graph(input: list[int]) -> list[str]: futures = [mapper(i) for i in input] mapped = [f.result() for f in futures] @@ -1460,6 +1463,39 @@ def test_imp_task(request: pytest.FixtureRequest, checkpointer_name: str) -> Non "items": {"type": "string"}, "title": "LangGraphOutput", } + assert graph.get_config_jsonschema() == { + "$defs": { + "Configurable": { + "properties": { + "model": {"default": None, "title": "Model", "type": "string"}, + "checkpoint_id": { + "anyOf": [{"type": "string"}, {"type": "null"}], + "default": None, + "description": "Pass to fetch a past checkpoint. If None, fetches the latest checkpoint.", + "title": "Checkpoint ID", + }, + "checkpoint_ns": { + "default": "", + "description": 'Checkpoint namespace. Denotes the path to the subgraph node the checkpoint originates from, separated by `|` character, e.g. `"child|grandchild"`. Defaults to "" (root graph).', + "title": "Checkpoint NS", + "type": "string", + }, + "thread_id": { + "default": "", + "title": "Thread ID", + "type": "string", + }, + }, + "title": "Configurable", + "type": "object", + } + }, + "properties": { + "configurable": {"$ref": "#/$defs/Configurable", "default": None} + }, + "title": "LangGraphConfig", + "type": "object", + } thread1 = {"configurable": {"thread_id": "1"}} assert [*graph.stream([0, 1], thread1)] == [