Merge pull request #1641 from langchain-ai/wfh/output_schema

Support optional fields in the output schema when defining as Dataclass or TypedDict

libs/langgraph/langgraph/graph/state.py

@@ -481,45 +481,22 @@ class CompiledStateGraph(CompiledGraph):
     def get_input_schema(
         self, config: Optional[RunnableConfig] = None
     ) -> type[BaseModel]:
-        if isclass(self.builder.input) and issubclass(
-            self.builder.input, (BaseModel, BaseModelV1)
-        ):
-            return self.builder.input
-        else:
-            keys = list(self.builder.schemas[self.builder.input].keys())
-            if len(keys) == 1 and keys[0] == "__root__":
-                return create_model(  # type: ignore[call-overload]
-                    self.get_name("Input"),
-                    __root__=(self.channels[keys[0]].UpdateType, None),
-                )
-            else:
-                return create_model(  # type: ignore[call-overload]
-                    self.get_name("Input"),
-                    **{
-                        k: (
-                            self.channels[k].UpdateType,
-                            (
-                                get_field_default(
-                                    k,
-                                    self.channels[k].UpdateType,
-                                    self.builder.input,
-                                )
-                            ),
-                        )
-                        for k in self.builder.schemas[self.builder.input]
-                        if isinstance(self.channels[k], BaseChannel)
-                    },
-                )
+        return _get_schema(
+            typ=self.builder.input,
+            schemas=self.builder.schemas,
+            channels=self.builder.channels,
+            name=self.get_name("Input"),
+        )
 
     def get_output_schema(
         self, config: Optional[RunnableConfig] = None
     ) -> type[BaseModel]:
-        if isclass(self.builder.output) and issubclass(
-            self.builder.output, (BaseModel, BaseModelV1)
-        ):
-            return self.builder.output
-
-        return super().get_output_schema(config)
+        return _get_schema(
+            typ=self.builder.output,
+            schemas=self.builder.schemas,
+            channels=self.builder.channels,
+            name=self.get_name("Output"),
+        )
 
     def attach_node(self, key: str, node: Optional[StateNodeSpec]) -> None:
         if key == START:
@@ -779,3 +756,38 @@ def _is_field_managed_value(name: str, typ: Type[Any]) -> Optional[ManagedValueS
                 return decoration
 
     return None
+
+
+def _get_schema(
+    typ: Type,
+    schemas: dict,
+    channels: dict,
+    name: str,
+) -> type[BaseModel]:
+    if isclass(typ) and issubclass(typ, (BaseModel, BaseModelV1)):
+        return typ
+    else:
+        keys = list(schemas[typ].keys())
+        if len(keys) == 1 and keys[0] == "__root__":
+            return create_model(  # type: ignore[call-overload]
+                name,
+                __root__=(channels[keys[0]].UpdateType, None),
+            )
+        else:
+            return create_model(  # type: ignore[call-overload]
+                name,
+                **{
+                    k: (
+                        channels[k].UpdateType,
+                        (
+                            get_field_default(
+                                k,
+                                channels[k].UpdateType,
+                                typ,
+                            )
+                        ),
+                    )
+                    for k in schemas[typ]
+                    if k in channels and isinstance(channels[k], BaseChannel)
+                },
+            )

libs/langgraph/tests/test_state.py

@@ -107,14 +107,25 @@ def test_state_schema_optional_values(total_: bool):
         val5: Annotated[Required[str], "foo"]
         val6: Annotated[NotRequired[str], "bar"]
 
+    class OutputState(SomeParentState, total=total_):  # type: ignore
+        out_val1: str
+        out_val2: Optional[str]
+        out_val3: Required[str]
+        out_val4: NotRequired[dict]
+        out_val5: Annotated[Required[str], "foo"]
+        out_val6: Annotated[NotRequired[str], "bar"]
+
     class State(InputState):  # this would be ignored
         val4: dict
+        some_shared_channel: Annotated[str, SharedValue.on("assistant_id")] = field(
+            default="foo"
+        )
 
-    builder = StateGraph(State, input=InputState)
+    builder = StateGraph(State, input=InputState, output=OutputState)
     builder.add_node("n", lambda x: x)
     builder.add_edge("__start__", "n")
     graph = builder.compile()
-    model = graph.input_schema
+    model = graph.get_input_schema()
     json_schema = model.schema()
 
     if total_ is False:
@@ -134,6 +145,23 @@ def test_state_schema_optional_values(total_: bool):
         set(json_schema["properties"].keys()) == expected_required | expected_optional
     )
 
+    # Check output schema. Should be the same process
+    output_schema = graph.get_output_schema().schema()
+    if total_ is False:
+        expected_required = set()
+        expected_optional = {"out_val2", "out_val1"}
+    else:
+        expected_required = {"out_val1"}
+        expected_optional = {"out_val2"}
+
+    expected_required |= {"val0a", "out_val3", "out_val5"}
+    expected_optional |= {"val0b", "out_val4", "out_val6"}
+
+    assert set(output_schema.get("required", set())) == expected_required
+    assert (
+        set(output_schema["properties"].keys()) == expected_required | expected_optional
+    )
+
 
 @pytest.mark.parametrize("kw_only_", [False, True])
 def test_state_schema_default_values(kw_only_: bool):
@@ -156,26 +184,29 @@ def test_state_schema_default_values(kw_only_: bool):
         val11: Annotated[list[str], "annotated list"] = field(
             default_factory=lambda: ["a", "b"]
         )
+        some_shared_channel: Annotated[str, SharedValue.on("assistant_id")] = field(
+            default="foo"
+        )
 
     builder = StateGraph(InputState)
     builder.add_node("n", lambda x: x)
     builder.add_edge("__start__", "n")
     graph = builder.compile()
-    model = graph.input_schema
-    json_schema = model.schema()
+    for model in [graph.get_input_schema(), graph.get_output_schema()]:
+        json_schema = model.schema()
 
-    expected_required = {"val1", "val7"}
-    expected_optional = {
-        "val2",
-        "val3",
-        "val4",
-        "val5",
-        "val6",
-        "val8",
-        "val9",
-        "val10",
-        "val11",
-    }
+        expected_required = {"val1", "val7"}
+        expected_optional = {
+            "val2",
+            "val3",
+            "val4",
+            "val5",
+            "val6",
+            "val8",
+            "val9",
+            "val10",
+            "val11",
+        }
 
     assert set(json_schema.get("required", set())) == expected_required
     assert (