Fix some deprecation warnings in tests

libs/langgraph/langgraph/pregel/__init__.py

@@ -316,6 +316,15 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
                 },
             )
 
+    def get_input_jsonschema(
+        self, config: Optional[RunnableConfig] = None
+    ) -> Dict[All, Any]:
+        schema = self.get_input_schema(config)
+        if hasattr(schema, "model_json_schema"):
+            return schema.model_json_schema()
+        else:
+            return schema.schema()
+
     @property
     def OutputType(self) -> Any:
         if isinstance(self.output_channels, str):
@@ -335,6 +344,15 @@ class Pregel(Runnable[Union[dict[str, Any], Any], Union[dict[str, Any], Any]]):
                 },
             )
 
+    def get_output_jsonschema(
+        self, config: Optional[RunnableConfig] = None
+    ) -> Dict[All, Any]:
+        schema = self.get_output_schema(config)
+        if hasattr(schema, "model_json_schema"):
+            return schema.model_json_schema()
+        else:
+            return schema.schema()
+
     @property
     def stream_channels_list(self) -> Sequence[str]:
         stream_channels = self.stream_channels_asis

libs/langgraph/tests/test_pregel.py

@@ -361,7 +361,7 @@ def test_node_schemas_custom_output() -> None:
         "messages": [_AnyIdHumanMessage(content="hello")],
     }
 
-    builder = StateGraph(input=State, output=Output)
+    builder = StateGraph(State, output=Output)
     builder.add_node("a", node_a)
     builder.add_node("b", node_b)
     builder.add_node("c", node_c)
@@ -7663,10 +7663,9 @@ def test_in_one_fan_out_state_graph_waiting_edge_custom_state_class_pydantic1(
 
     app = workflow.compile()
 
-    # because it's a v1 pydantic, we're using .schema() here instead of the new methods
     assert app.get_graph().draw_mermaid(with_styles=False) == snapshot
-    assert app.get_input_schema().schema() == snapshot
-    assert app.get_output_schema().schema() == snapshot
+    assert app.get_input_jsonschema() == snapshot
+    assert app.get_output_jsonschema() == snapshot
 
     with pytest.raises(ValidationError), assert_ctx_once():
         app.invoke({"query": {}})
@@ -9868,7 +9867,7 @@ def test_send_to_nested_graphs(
         return {"subject": f"{subject} - hohoho"}
 
     # subgraph
-    subgraph = StateGraph(input=JokeState, output=OverallState)
+    subgraph = StateGraph(JokeState, output=OverallState)
     subgraph.add_node("edit", edit)
     subgraph.add_node(
         "generate", lambda state: {"jokes": [f"Joke about {state['subject']}"]}

libs/langgraph/tests/test_pregel_async.py

@@ -601,7 +601,7 @@ async def test_node_schemas_custom_output() -> None:
         "messages": [_AnyIdHumanMessage(content="hello")],
     }
 
-    builder = StateGraph(input=State, output=Output)
+    builder = StateGraph(State, output=Output)
     builder.add_node("a", node_a)
     builder.add_node("b", node_b)
     builder.add_node("c", node_c)
@@ -3179,10 +3179,7 @@ async def test_conditional_graph_state(
             setup.reset_mock()
             teardown.reset_mock()
 
-    class MyPydanticContextModel(BaseModel):
-        class Config:
-            arbitrary_types_allowed = True
-
+    class MyPydanticContextModel(BaseModel, arbitrary_types_allowed=True):
         session: httpx.AsyncClient
         something_else: str
 
@@ -8603,7 +8600,7 @@ async def test_send_to_nested_graphs(checkpointer_name: str) -> None:
         return {"subject": f"{subject} - hohoho"}
 
     # subgraph
-    subgraph = StateGraph(input=JokeState, output=OverallState)
+    subgraph = StateGraph(JokeState, output=OverallState)
     subgraph.add_node("edit", edit)
     subgraph.add_node(
         "generate", lambda state: {"jokes": [f"Joke about {state['subject']}"]}

libs/langgraph/tests/test_state.py

@@ -73,7 +73,7 @@ def test_state_schema_with_type_hint():
     def miss_all_hint(state, config):
         return {"input_state": state}
 
-    graph = StateGraph(input=InputState, output=OutputState)
+    graph = StateGraph(InputState, output=OutputState)
     actions = [complete_hint, miss_first_hint, only_return_hint, miss_all_hint]
 
     for action in actions:
@@ -125,8 +125,7 @@ def test_state_schema_optional_values(total_: bool):
     builder.add_node("n", lambda x: x)
     builder.add_edge("__start__", "n")
     graph = builder.compile()
-    model = graph.get_input_schema()
-    json_schema = model.schema()
+    json_schema = graph.get_input_jsonschema()
 
     if total_ is False:
         expected_required = set()
@@ -146,7 +145,7 @@ def test_state_schema_optional_values(total_: bool):
     )
 
     # Check output schema. Should be the same process
-    output_schema = graph.get_output_schema().schema()
+    output_schema = graph.get_output_jsonschema()
     if total_ is False:
         expected_required = set()
         expected_optional = {"out_val2", "out_val1"}
@@ -192,9 +191,7 @@ def test_state_schema_default_values(kw_only_: bool):
     builder.add_node("n", lambda x: x)
     builder.add_edge("__start__", "n")
     graph = builder.compile()
-    for model in [graph.get_input_schema(), graph.get_output_schema()]:
-        json_schema = model.schema()
-
+    for json_schema in [graph.get_input_jsonschema(), graph.get_output_jsonschema()]:
         expected_required = {"val1", "val7"}
         expected_optional = {
             "val2",
@@ -256,8 +253,6 @@ def test_raises_invalid_managed():
             StateGraph(_state, input=_inp, output=_outp)
     bad_output_examples = [
         (State, InputState, BadOutputState),
-        (None, InputState, BadOutputState),
-        (None, State, BadOutputState),
         (State, None, BadOutputState),
     ]
     for _state, _inp, _outp in bad_output_examples: