Merge branch 'main' into wfh/_validate_more

Signed-off-by: William Fu-Hinthorn <13333726+hinthornw@users.noreply.github.com>

libs/langgraph/langgraph/graph/schema_utils.py

@@ -57,8 +57,11 @@ class SchemaCoercionMapper:
 
         self.schema = schema
         self.max_depth = max_depth
-        self.type_hints = type_hints or get_type_hints(
-            schema, localns={schema.__name__: schema}
+
+        self.type_hints = (
+            type_hints
+            if type_hints is not None
+            else get_type_hints(schema, localns={schema.__name__: schema})
         )
 
         if issubclass(schema, BaseModelV1):
@@ -71,9 +74,9 @@ class SchemaCoercionMapper:
         elif issubclass(schema, BaseModel):
             self._fields = {
                 n: self.type_hints.get(n, f.annotation)
-                for n, f in schema.model_fields.items()  # type: ignore[attr-defined]
+                for n, f in schema.model_fields.items()
             }
-            self._construct: Callable[..., Any] = schema.model_construct  # type: ignore[attr-defined,no-redef]
+            self._construct: Callable[..., Any] = schema.model_construct  # type: ignore
 
         else:
             raise TypeError("Schema is neither a Pydantic v1 nor v2 model.")
@@ -129,7 +132,7 @@ class SchemaCoercionMapper:
                 mapper = SchemaCoercionMapper(field_type, max_depth=depth - 1)
                 return lambda v, d: mapper.coerce(v, d) if isinstance(v, dict) else v
 
-        if origin in (list, set):
+        if origin is list:
             args = get_args(field_type)
             if len(args) != 1:
                 return self._passthrough
@@ -144,13 +147,18 @@ class SchemaCoercionMapper:
 
         if origin is set or field_type is set:
             args = get_args(field_type)
-            if len(args) != 1:
+            if len(args) > 1:
                 return self._passthrough
-            sub = self._build_coercer(args[0], depth - 1)
+            elif len(args) == 1:
+                sub = self._build_coercer(args[0], depth - 1)
+            else:
+                sub = None  # type: ignore
 
             def set_coercer(v: Any, d: Any) -> Any:
                 if not isinstance(v, (list, tuple, set)):
                     return v
+                if sub is None:
+                    return set(v)
                 return {sub(x, d - 1) for x in v}
 
             return set_coercer
@@ -258,6 +266,7 @@ try:
             return v
 
     try:
+        from pydantic.v1 import parse_obj_as
         from pydantic.v1.main import create_model
     except ImportError:
         create_model = None  # type: ignore
@@ -268,9 +277,8 @@ try:
                 parser = create_model(
                     f"ParsingModel[{tp}]",
                     __root__=(tp, ...),
-                    __config__={"arbitrary_types_allowed": True},
                 )
-                return lambda v: parser(__root__=v).__root__
+                return lambda v: parser(__root__=v).__root__  # type: ignore
             except RuntimeError:
                 return lambda v: v
         return lambda v: parse_obj_as(tp, v)
@@ -298,7 +306,7 @@ except ImportError:
                 f"ParsingModel[{tp}]",
                 __root__=(tp, ...),
             )
-            return lambda v: parser(__root__=v).__root__
+            return lambda v: parser(__root__=v).__root__  # type: ignore
         except RuntimeError:
             return lambda v: v
 

libs/langgraph/langgraph/graph/state.py

@@ -777,10 +777,10 @@ class CompiledStateGraph(CompiledGraph):
             elif (t := type(input)) and get_type_hints(t):
                 # Pydantic v2
                 if isinstance(input, BaseModelV1):
-                    keep = input.__fields_set__
+                    keep: Optional[set[str]] = input.__fields_set__
                     defaults = {k: v.default for k, v in t.__fields__.items()}
                 elif isinstance(input, BaseModel):
-                    keep: Optional[set[str]] = input.model_fields_set
+                    keep = input.model_fields_set
                     defaults = {k: v.default for k, v in input.model_fields.items()}
                 # Pydantic v1
                 else:

libs/langgraph/langgraph/pregel/remote.py

@@ -654,9 +654,10 @@ class RemoteGraph(PregelProtocol):
             # raise interrupt or errors
             if chunk.event.startswith("updates"):
                 if isinstance(chunk.data, dict) and INTERRUPT in chunk.data:
-                    raise GraphInterrupt(
-                        [Interrupt(**i) for i in chunk.data[INTERRUPT]]
-                    )
+                    if caller_ns:
+                        raise GraphInterrupt(
+                            [Interrupt(**i) for i in chunk.data[INTERRUPT]]
+                        )
             elif chunk.event.startswith("error"):
                 raise RemoteException(chunk.data)
             # filter for what was actually requested
@@ -748,9 +749,10 @@ class RemoteGraph(PregelProtocol):
             # raise interrupt or errors
             if chunk.event.startswith("updates"):
                 if isinstance(chunk.data, dict) and INTERRUPT in chunk.data:
-                    raise GraphInterrupt(
-                        [Interrupt(**i) for i in chunk.data[INTERRUPT]]
-                    )
+                    if caller_ns:
+                        raise GraphInterrupt(
+                            [Interrupt(**i) for i in chunk.data[INTERRUPT]]
+                        )
             elif chunk.event.startswith("error"):
                 raise RemoteException(chunk.data)
             # filter for what was actually requested

libs/langgraph/tests/conftest.py

@@ -443,11 +443,11 @@ async def awith_store(store_name: Optional[str]) -> AsyncIterator[BaseStore]:
 SHALLOW_CHECKPOINTERS_SYNC = ["postgres_shallow"]
 REGULAR_CHECKPOINTERS_SYNC = [
     "memory",
-    # "sqlite",
-    # "postgres",
-    # "postgres_pipe",
-    # "postgres_pool",
-    # "sqlite_aes",
+    "sqlite",
+    "postgres",
+    "postgres_pipe",
+    "postgres_pool",
+    "sqlite_aes",
 ]
 ALL_CHECKPOINTERS_SYNC = [
     *REGULAR_CHECKPOINTERS_SYNC,
@@ -456,10 +456,10 @@ ALL_CHECKPOINTERS_SYNC = [
 SHALLOW_CHECKPOINTERS_ASYNC = ["postgres_aio_shallow"]
 REGULAR_CHECKPOINTERS_ASYNC = [
     "memory",
-    # "sqlite_aio",
-    # "postgres_aio",
-    # "postgres_aio_pipe",
-    # "postgres_aio_pool",
+    "sqlite_aio",
+    "postgres_aio",
+    "postgres_aio_pipe",
+    "postgres_aio_pool",
 ]
 ALL_CHECKPOINTERS_ASYNC = [
     *REGULAR_CHECKPOINTERS_ASYNC,

libs/langgraph/tests/test_pregel_async.py

@@ -4553,6 +4553,7 @@ async def test_nested_pydantic_models(version: str) -> None:
         optional_nested: Optional[NestedModel] = None
         dict_nested: dict[str, NestedModel]
         my_set: set[int]
+        another_set: set
         my_enum: MyEnum
         list_nested: Annotated[
             Union[dict, list[dict[str, NestedModel]]], lambda x, y: (x or []) + [y]
@@ -4581,6 +4582,7 @@ async def test_nested_pydantic_models(version: str) -> None:
         "nested": {"value": 42, "name": "test"},
         "optional_nested": {"value": 10, "name": "optional"},
         "my_set": [1, 2, 7],
+        "another_set": ["foo", 3],
         "my_enum": MyEnum.B,
         "my_typed_dict": {"x": 1, "my_enum": MyEnum.A},
         "dict_nested": {"a": {"value": 5, "name": "a"}},

libs/langgraph/tests/test_remote_graph.py

@@ -437,15 +437,17 @@ def test_stream():
         sync_client=mock_sync_client,
     )
 
-    # stream modes doesn't include 'updates'
-    stream_parts = []
+    # test raising graph interrupt if invoked as a subgraph
     with pytest.raises(GraphInterrupt) as exc:
         for stream_part in remote_pregel.stream(
             {"input": "data"},
-            config={"configurable": {"thread_id": "thread_1"}},
+            # pretend we invoked this as a subgraph
+            config={
+                "configurable": {"thread_id": "thread_1", "checkpoint_ns": "some_ns"}
+            },
             stream_mode="values",
         ):
-            stream_parts.append(stream_part)
+            pass
 
     assert exc.value.args[0] == [
         Interrupt(
@@ -456,6 +458,15 @@ def test_stream():
         )
     ]
 
+    # stream modes doesn't include 'updates'
+    stream_parts = []
+    for stream_part in remote_pregel.stream(
+        {"input": "data"},
+        config={"configurable": {"thread_id": "thread_1"}},
+        stream_mode="values",
+    ):
+        stream_parts.append(stream_part)
+
     assert stream_parts == [
         {"chunk": "data1"},
         {"chunk": "data2"},
@@ -470,62 +481,62 @@ def test_stream():
 
     # default stream_mode is updates
     stream_parts = []
-    with pytest.raises(GraphInterrupt):
-        for stream_part in remote_pregel.stream(
-            {"input": "data"},
-            config={"configurable": {"thread_id": "thread_1"}},
-        ):
-            stream_parts.append(stream_part)
+    for stream_part in remote_pregel.stream(
+        {"input": "data"},
+        config={"configurable": {"thread_id": "thread_1"}},
+    ):
+        stream_parts.append(stream_part)
 
     assert stream_parts == [
         {"chunk": "data3"},
         {"chunk": "data4"},
+        {"__interrupt__": ()},
     ]
 
     # list stream_mode includes mode names
     stream_parts = []
-    with pytest.raises(GraphInterrupt):
-        for stream_part in remote_pregel.stream(
-            {"input": "data"},
-            config={"configurable": {"thread_id": "thread_1"}},
-            stream_mode=["updates"],
-        ):
-            stream_parts.append(stream_part)
+    for stream_part in remote_pregel.stream(
+        {"input": "data"},
+        config={"configurable": {"thread_id": "thread_1"}},
+        stream_mode=["updates"],
+    ):
+        stream_parts.append(stream_part)
 
     assert stream_parts == [
         ("updates", {"chunk": "data3"}),
         ("updates", {"chunk": "data4"}),
+        ("updates", {"__interrupt__": ()}),
     ]
 
     # subgraphs + list modes
     stream_parts = []
-    with pytest.raises(GraphInterrupt):
-        for stream_part in remote_pregel.stream(
-            {"input": "data"},
-            config={"configurable": {"thread_id": "thread_1"}},
-            stream_mode=["updates"],
-            subgraphs=True,
-        ):
-            stream_parts.append(stream_part)
+    for stream_part in remote_pregel.stream(
+        {"input": "data"},
+        config={"configurable": {"thread_id": "thread_1"}},
+        stream_mode=["updates"],
+        subgraphs=True,
+    ):
+        stream_parts.append(stream_part)
 
     assert stream_parts == [
         ((), "updates", {"chunk": "data3"}),
         ((), "updates", {"chunk": "data4"}),
+        ((), "updates", {"__interrupt__": ()}),
     ]
 
     # subgraphs + single mode
     stream_parts = []
-    with pytest.raises(GraphInterrupt):
-        for stream_part in remote_pregel.stream(
-            {"input": "data"},
-            config={"configurable": {"thread_id": "thread_1"}},
-            subgraphs=True,
-        ):
-            stream_parts.append(stream_part)
+    for stream_part in remote_pregel.stream(
+        {"input": "data"},
+        config={"configurable": {"thread_id": "thread_1"}},
+        subgraphs=True,
+    ):
+        stream_parts.append(stream_part)
 
     assert stream_parts == [
         ((), {"chunk": "data3"}),
         ((), {"chunk": "data4"}),
+        ((), {"__interrupt__": ()}),
     ]
 
 
@@ -561,15 +572,17 @@ async def test_astream():
         client=mock_async_client,
     )
 
-    # stream modes doesn't include 'updates'
-    stream_parts = []
+    # test raising graph interrupt if invoked as a subgraph
     with pytest.raises(GraphInterrupt) as exc:
         async for stream_part in remote_pregel.astream(
             {"input": "data"},
-            config={"configurable": {"thread_id": "thread_1"}},
+            # pretend we invoked this as a subgraph
+            config={
+                "configurable": {"thread_id": "thread_1", "checkpoint_ns": "some_ns"}
+            },
             stream_mode="values",
         ):
-            stream_parts.append(stream_part)
+            pass
 
     assert exc.value.args[0] == [
         Interrupt(
@@ -580,6 +593,15 @@ async def test_astream():
         )
     ]
 
+    # stream modes doesn't include 'updates'
+    stream_parts = []
+    async for stream_part in remote_pregel.astream(
+        {"input": "data"},
+        config={"configurable": {"thread_id": "thread_1"}},
+        stream_mode="values",
+    ):
+        stream_parts.append(stream_part)
+
     assert stream_parts == [
         {"chunk": "data1"},
         {"chunk": "data2"},
@@ -596,62 +618,62 @@ async def test_astream():
 
     # default stream_mode is updates
     stream_parts = []
-    with pytest.raises(GraphInterrupt):
-        async for stream_part in remote_pregel.astream(
-            {"input": "data"},
-            config={"configurable": {"thread_id": "thread_1"}},
-        ):
-            stream_parts.append(stream_part)
+    async for stream_part in remote_pregel.astream(
+        {"input": "data"},
+        config={"configurable": {"thread_id": "thread_1"}},
+    ):
+        stream_parts.append(stream_part)
 
     assert stream_parts == [
         {"chunk": "data3"},
         {"chunk": "data4"},
+        {"__interrupt__": ()},
     ]
 
     # list stream_mode includes mode names
     stream_parts = []
-    with pytest.raises(GraphInterrupt):
-        async for stream_part in remote_pregel.astream(
-            {"input": "data"},
-            config={"configurable": {"thread_id": "thread_1"}},
-            stream_mode=["updates"],
-        ):
-            stream_parts.append(stream_part)
+    async for stream_part in remote_pregel.astream(
+        {"input": "data"},
+        config={"configurable": {"thread_id": "thread_1"}},
+        stream_mode=["updates"],
+    ):
+        stream_parts.append(stream_part)
 
     assert stream_parts == [
         ("updates", {"chunk": "data3"}),
         ("updates", {"chunk": "data4"}),
+        ("updates", {"__interrupt__": ()}),
     ]
 
     # subgraphs + list modes
     stream_parts = []
-    with pytest.raises(GraphInterrupt):
-        async for stream_part in remote_pregel.astream(
-            {"input": "data"},
-            config={"configurable": {"thread_id": "thread_1"}},
-            stream_mode=["updates"],
-            subgraphs=True,
-        ):
-            stream_parts.append(stream_part)
+    async for stream_part in remote_pregel.astream(
+        {"input": "data"},
+        config={"configurable": {"thread_id": "thread_1"}},
+        stream_mode=["updates"],
+        subgraphs=True,
+    ):
+        stream_parts.append(stream_part)
 
     assert stream_parts == [
         ((), "updates", {"chunk": "data3"}),
         ((), "updates", {"chunk": "data4"}),
+        ((), "updates", {"__interrupt__": ()}),
     ]
 
     # subgraphs + single mode
     stream_parts = []
-    with pytest.raises(GraphInterrupt):
-        async for stream_part in remote_pregel.astream(
-            {"input": "data"},
-            config={"configurable": {"thread_id": "thread_1"}},
-            subgraphs=True,
-        ):
-            stream_parts.append(stream_part)
+    async for stream_part in remote_pregel.astream(
+        {"input": "data"},
+        config={"configurable": {"thread_id": "thread_1"}},
+        subgraphs=True,
+    ):
+        stream_parts.append(stream_part)
 
     assert stream_parts == [
         ((), {"chunk": "data3"}),
         ((), {"chunk": "data4"}),
+        ((), {"__interrupt__": ()}),
     ]
 
     async_iter = MagicMock()
@@ -664,33 +686,33 @@ async def test_astream():
 
     # subgraphs + list modes
     stream_parts = []
-    with pytest.raises(GraphInterrupt):
-        async for stream_part in remote_pregel.astream(
-            {"input": "data"},
-            config={"configurable": {"thread_id": "thread_1"}},
-            stream_mode=["updates"],
-            subgraphs=True,
-        ):
-            stream_parts.append(stream_part)
+    async for stream_part in remote_pregel.astream(
+        {"input": "data"},
+        config={"configurable": {"thread_id": "thread_1"}},
+        stream_mode=["updates"],
+        subgraphs=True,
+    ):
+        stream_parts.append(stream_part)
 
     assert stream_parts == [
         (("my", "subgraph"), "updates", {"chunk": "data3"}),
         (("hello", "subgraph"), "updates", {"chunk": "data4"}),
+        (("bye", "subgraph"), "updates", {"__interrupt__": ()}),
     ]
 
     # subgraphs + single mode
     stream_parts = []
-    with pytest.raises(GraphInterrupt):
-        async for stream_part in remote_pregel.astream(
-            {"input": "data"},
-            config={"configurable": {"thread_id": "thread_1"}},
-            subgraphs=True,
-        ):
-            stream_parts.append(stream_part)
+    async for stream_part in remote_pregel.astream(
+        {"input": "data"},
+        config={"configurable": {"thread_id": "thread_1"}},
+        subgraphs=True,
+    ):
+        stream_parts.append(stream_part)
 
     assert stream_parts == [
         (("my", "subgraph"), {"chunk": "data3"}),
         (("hello", "subgraph"), {"chunk": "data4"}),
+        (("bye", "subgraph"), {"__interrupt__": ()}),
     ]