diff --git a/libs/langgraph/langgraph/graph/schema_utils.py b/libs/langgraph/langgraph/graph/schema_utils.py index 774ab6967..35ccada7b 100644 --- a/libs/langgraph/langgraph/graph/schema_utils.py +++ b/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 diff --git a/libs/langgraph/langgraph/graph/state.py b/libs/langgraph/langgraph/graph/state.py index c9f503621..093faf5d4 100644 --- a/libs/langgraph/langgraph/graph/state.py +++ b/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: diff --git a/libs/langgraph/langgraph/pregel/remote.py b/libs/langgraph/langgraph/pregel/remote.py index f58d0a4d9..07d44cb66 100644 --- a/libs/langgraph/langgraph/pregel/remote.py +++ b/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 diff --git a/libs/langgraph/tests/conftest.py b/libs/langgraph/tests/conftest.py index 67ed40ee7..689ef6ab9 100644 --- a/libs/langgraph/tests/conftest.py +++ b/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, diff --git a/libs/langgraph/tests/test_pregel_async.py b/libs/langgraph/tests/test_pregel_async.py index 327236d7f..375beb261 100644 --- a/libs/langgraph/tests/test_pregel_async.py +++ b/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"}}, diff --git a/libs/langgraph/tests/test_remote_graph.py b/libs/langgraph/tests/test_remote_graph.py index 4aea1fc32..878f5e6f1 100644 --- a/libs/langgraph/tests/test_remote_graph.py +++ b/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__": ()}), ]