WEBDIENSTE/langgraph
1
0
Fork 0
mirror of https://github.com/langchain-ai/langgraph.git synced 2026-08-26 01:22:24 +02:00
Code Issues Packages Projects Releases Wiki Activity

Handle total=false

Browse Source
This commit is contained in:
William Fu-Hinthorn
2024-08-30 10:21:57 -07:00
parent 65b64c1331
commit a07ec9d8c8
2 changed files with 19 additions and 6 deletions
Download Patch File Download Diff File Expand all files Collapse all files
libs/langgraph/langgraph/graph/state.py
+6 -1
View File
@@ -486,6 +486,10 @@ class CompiledStateGraph(CompiledGraph):
__root__=(self.channels[keys[0]].UpdateType, None),
)
else:
is_total_false = (
hasattr(self.builder.input, "__total__")
and self.builder.input.__total__ is False
)
return create_model( # type: ignore[call-overload]
self.get_name("Input"),
**{
@@ -493,7 +497,8 @@ class CompiledStateGraph(CompiledGraph):
self.channels[k].UpdateType,
(
None
if is_optional_type(self.channels[k].UpdateType)
if is_total_false
or is_optional_type(self.channels[k].UpdateType)
else ...
),
)
libs/langgraph/tests/test_state.py
+13 -5
View File
@@ -88,8 +88,9 @@ def test_state_schema_with_type_hint():
assert c[node_name] == output_state
def test_state_schema_optional_values():
class InputState(TypedDict):
@pytest.mark.parametrize("total_", [True, False])
def test_state_schema_optional_values(total_: bool):
class InputState(TypedDict, total=total_): # type: ignore
val1: str
val2: Optional[str]
@@ -102,9 +103,16 @@ def test_state_schema_optional_values():
graph = builder.compile()
model = graph.input_schema
json_schema = model.schema()
expected_required = {"val1"}
expected_optional = {"val2"}
assert set(json_schema["required"]) == expected_required
if total_ is False:
expected_required = set()
expected_optional = {"val2", "val1"}
else:
expected_required = {"val1"}
expected_optional = {"val2"}
assert set(json_schema.get("required", set())) == expected_required
assert (
set(json_schema["properties"].keys()) == expected_required | expected_optional
)