diff --git a/langgraph/pregel/__init__.py b/langgraph/pregel/__init__.py index 2521708ac..46baeb6d4 100644 --- a/langgraph/pregel/__init__.py +++ b/langgraph/pregel/__init__.py @@ -162,6 +162,8 @@ class Pregel( saver: Optional[BaseCheckpointSaver] = None + name: str = "LangGraph" + class Config: arbitrary_types_allowed = True @@ -191,7 +193,7 @@ class Pregel( return super().get_input_schema(config) else: return create_model( # type: ignore[call-overload] - "PregelInput", + self.get_name("Input"), **{ k: (self.channels[k].UpdateType, None) for k in self.input or self.channels.keys() @@ -210,7 +212,7 @@ class Pregel( return super().get_output_schema(config) else: return create_model( # type: ignore[call-overload] - "PregelOutput", + self.get_name("Output"), **{k: (self.channels[k].ValueType, None) for k in self.output}, ) diff --git a/tests/test_pregel.py b/tests/test_pregel.py index a56e641f2..ded25a5f5 100644 --- a/tests/test_pregel.py +++ b/tests/test_pregel.py @@ -41,8 +41,8 @@ def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: graph.set_finish_point("add_one") gapp = graph.compile() - assert app.input_schema.schema() == {"title": "PregelInput", "type": "integer"} - assert app.output_schema.schema() == {"title": "PregelOutput", "type": "integer"} + assert app.input_schema.schema() == {"title": "LangGraphInput", "type": "integer"} + assert app.output_schema.schema() == {"title": "LangGraphOutput", "type": "integer"} assert app.invoke(2) == 3 assert app.invoke(2, output=["output"]) == {"output": 3} assert repr(app), "does not raise recursion error" @@ -56,8 +56,8 @@ def test_invoke_single_process_in_out_implicit_channels(mocker: MockerFixture) - app = Pregel(nodes={"one": chain}) - assert app.input_schema.schema() == {"title": "PregelInput"} - assert app.output_schema.schema() == {"title": "PregelOutput"} + assert app.input_schema.schema() == {"title": "LangGraphInput"} + assert app.output_schema.schema() == {"title": "LangGraphOutput"} assert app.invoke(2) == 3 @@ -71,9 +71,9 @@ def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> None: app = Pregel(nodes={"one": chain}, output=["output", "fixed", "output_plus_one"]) - assert app.input_schema.schema() == {"title": "PregelInput"} + assert app.input_schema.schema() == {"title": "LangGraphInput"} assert app.output_schema.schema() == { - "title": "PregelOutput", + "title": "LangGraphOutput", "type": "object", "properties": { "output": {"title": "Output"}, @@ -95,8 +95,8 @@ def test_invoke_single_process_in_out_reserved_is_last(mocker: MockerFixture) -> app = Pregel(nodes={"one": chain}) - assert app.input_schema.schema() == {"title": "PregelInput"} - assert app.output_schema.schema() == {"title": "PregelOutput"} + assert app.input_schema.schema() == {"title": "LangGraphInput"} + assert app.output_schema.schema() == {"title": "LangGraphOutput"} assert app.invoke(2) == {"input": 3, "is_last_step": False} assert app.invoke(2, {"recursion_limit": 1}) == {"input": 3, "is_last_step": True} @@ -112,9 +112,9 @@ def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: output=["output"], ) - assert app.input_schema.schema() == {"title": "PregelInput"} + assert app.input_schema.schema() == {"title": "LangGraphInput"} assert app.output_schema.schema() == { - "title": "PregelOutput", + "title": "LangGraphOutput", "type": "object", "properties": {"output": {"title": "Output"}}, } @@ -134,12 +134,12 @@ def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> None: ) assert app.input_schema.schema() == { - "title": "PregelInput", + "title": "LangGraphInput", "type": "object", "properties": {"input": {"title": "Input"}}, } assert app.output_schema.schema() == { - "title": "PregelOutput", + "title": "LangGraphOutput", "type": "object", "properties": {"output": {"title": "Output"}}, } diff --git a/tests/test_pregel_async.py b/tests/test_pregel_async.py index 96adce78f..3f94cf86d 100644 --- a/tests/test_pregel_async.py +++ b/tests/test_pregel_async.py @@ -39,8 +39,8 @@ async def test_invoke_single_process_in_out(mocker: MockerFixture) -> None: graph.set_finish_point("add_one") gapp = graph.compile() - assert app.input_schema.schema() == {"title": "PregelInput", "type": "integer"} - assert app.output_schema.schema() == {"title": "PregelOutput", "type": "integer"} + assert app.input_schema.schema() == {"title": "LangGraphInput", "type": "integer"} + assert app.output_schema.schema() == {"title": "LangGraphOutput", "type": "integer"} assert await app.ainvoke(2) == 3 assert await app.ainvoke(2, output=["output"]) == {"output": 3} @@ -55,8 +55,8 @@ async def test_invoke_single_process_in_out_implicit_channels( app = Pregel(nodes={"one": chain}) - assert app.input_schema.schema() == {"title": "PregelInput"} - assert app.output_schema.schema() == {"title": "PregelOutput"} + assert app.input_schema.schema() == {"title": "LangGraphInput"} + assert app.output_schema.schema() == {"title": "LangGraphOutput"} assert await app.ainvoke(2) == 3 @@ -70,9 +70,9 @@ async def test_invoke_single_process_in_write_kwargs(mocker: MockerFixture) -> N app = Pregel(nodes={"one": chain}, output=["output", "fixed", "output_plus_one"]) - assert app.input_schema.schema() == {"title": "PregelInput"} + assert app.input_schema.schema() == {"title": "LangGraphInput"} assert app.output_schema.schema() == { - "title": "PregelOutput", + "title": "LangGraphOutput", "type": "object", "properties": { "output": {"title": "Output"}, @@ -96,8 +96,8 @@ async def test_invoke_single_process_in_out_reserved_is_last( app = Pregel(nodes={"one": chain}) - assert app.input_schema.schema() == {"title": "PregelInput"} - assert app.output_schema.schema() == {"title": "PregelOutput"} + assert app.input_schema.schema() == {"title": "LangGraphInput"} + assert app.output_schema.schema() == {"title": "LangGraphOutput"} assert await app.ainvoke(2) == {"input": 3, "is_last_step": False} assert await app.ainvoke(2, {"recursion_limit": 1}) == { "input": 3, @@ -114,9 +114,9 @@ async def test_invoke_single_process_in_out_dict(mocker: MockerFixture) -> None: output=["output"], ) - assert app.input_schema.schema() == {"title": "PregelInput"} + assert app.input_schema.schema() == {"title": "LangGraphInput"} assert app.output_schema.schema() == { - "title": "PregelOutput", + "title": "LangGraphOutput", "type": "object", "properties": {"output": {"title": "Output"}}, } @@ -136,12 +136,12 @@ async def test_invoke_single_process_in_dict_out_dict(mocker: MockerFixture) -> ) assert app.input_schema.schema() == { - "title": "PregelInput", + "title": "LangGraphInput", "type": "object", "properties": {"input": {"title": "Input"}}, } assert app.output_schema.schema() == { - "title": "PregelOutput", + "title": "LangGraphOutput", "type": "object", "properties": {"output": {"title": "Output"}}, }