flip

2026-08-22 23:52:23 +02:00 · 2026-03-03 14:23:04 -08:00
parent dec15d13dc
commit 727a9a1f6c
3 changed files with 140 additions and 140 deletions
@@ -2454,9 +2454,9 @@ class Pregel(
        durability: Durability | None = None,
        subgraphs: bool = False,
        debug: bool | None = None,
-        stream_version: Literal["v2"],
+        stream_version: Literal["v1"],
        **kwargs: Unpack[DeprecatedKwargs],
-    ) -> Iterator[StreamPart[OutputT, StateT]]: ...
+    ) -> Iterator[dict[str, Any] | Any]: ...

    @overload
    def stream(
@@ -2473,9 +2473,9 @@ class Pregel(
        durability: Durability | None = None,
        subgraphs: bool = False,
        debug: bool | None = None,
-        stream_version: Literal["v1"] = ...,
+        stream_version: Literal["v2"] = ...,
        **kwargs: Unpack[DeprecatedKwargs],
-    ) -> Iterator[dict[str, Any] | Any]: ...
+    ) -> Iterator[StreamPart[OutputT, StateT]]: ...

    def stream(
        self,
@@ -2491,7 +2491,7 @@ class Pregel(
        durability: Durability | None = None,
        subgraphs: bool = False,
        debug: bool | None = None,
-        stream_version: Literal["v1", "v2"] = "v1",
+        stream_version: Literal["v1", "v2"] = "v2",
        **kwargs: Unpack[DeprecatedKwargs],
    ) -> Iterator[dict[str, Any] | Any]:
        """Stream graph steps for a single input.
@@ -2785,9 +2785,9 @@ class Pregel(
        durability: Durability | None = None,
        subgraphs: bool = False,
        debug: bool | None = None,
-        stream_version: Literal["v2"],
+        stream_version: Literal["v1"],
        **kwargs: Unpack[DeprecatedKwargs],
-    ) -> AsyncIterator[StreamPart[OutputT, StateT]]: ...
+    ) -> AsyncIterator[dict[str, Any] | Any]: ...

    @overload
    def astream(
@@ -2804,9 +2804,9 @@ class Pregel(
        durability: Durability | None = None,
        subgraphs: bool = False,
        debug: bool | None = None,
-        stream_version: Literal["v1"] = ...,
+        stream_version: Literal["v2"] = ...,
        **kwargs: Unpack[DeprecatedKwargs],
-    ) -> AsyncIterator[dict[str, Any] | Any]: ...
+    ) -> AsyncIterator[StreamPart[OutputT, StateT]]: ...

    async def astream(
        self,
@@ -2822,7 +2822,7 @@ class Pregel(
        durability: Durability | None = None,
        subgraphs: bool = False,
        debug: bool | None = None,
-        stream_version: Literal["v1", "v2"] = "v1",
+        stream_version: Literal["v1", "v2"] = "v2",
        **kwargs: Unpack[DeprecatedKwargs],
    ) -> AsyncIterator[dict[str, Any] | Any]:
        """Asynchronously stream graph steps for a single input.
@@ -3162,6 +3162,23 @@ class Pregel(
            await asyncio.shield(run_manager.on_chain_error(e))
            raise

+    @overload
+    def invoke(
+        self,
+        input: InputT | Command | None,
+        config: RunnableConfig | None = None,
+        *,
+        context: ContextT | None = None,
+        stream_mode: StreamMode = "values",
+        print_mode: StreamMode | Sequence[StreamMode] = (),
+        output_keys: str | Sequence[str] | None = None,
+        interrupt_before: All | Sequence[str] | None = None,
+        interrupt_after: All | Sequence[str] | None = None,
+        durability: Durability | None = None,
+        stream_version: Literal["v1"],
+        **kwargs: Any,
+    ) -> dict[str, Any] | Any: ...
+
    @overload
    def invoke(
        self,
@@ -3175,7 +3192,7 @@ class Pregel(
        interrupt_before: All | Sequence[str] | None = None,
        interrupt_after: All | Sequence[str] | None = None,
        durability: Durability | None = None,
-        stream_version: Literal["v2"],
+        stream_version: Literal["v2"] = ...,
        **kwargs: Any,
    ) -> GraphOutput[OutputT]: ...

@@ -3192,11 +3209,10 @@ class Pregel(
        interrupt_before: All | Sequence[str] | None = None,
        interrupt_after: All | Sequence[str] | None = None,
        durability: Durability | None = None,
-        stream_version: Literal["v2"],
+        stream_version: Literal["v2"] = ...,
        **kwargs: Any,
    ) -> list[StreamPart[OutputT, StateT]]: ...

-    @overload
    def invoke(
        self,
        input: InputT | Command | None,
@@ -3209,23 +3225,7 @@ class Pregel(
        interrupt_before: All | Sequence[str] | None = None,
        interrupt_after: All | Sequence[str] | None = None,
        durability: Durability | None = None,
-        stream_version: Literal["v1"] = ...,
-        **kwargs: Any,
-    ) -> dict[str, Any] | Any: ...
-
-    def invoke(
-        self,
-        input: InputT | Command | None,
-        config: RunnableConfig | None = None,
-        *,
-        context: ContextT | None = None,
-        stream_mode: StreamMode = "values",
-        print_mode: StreamMode | Sequence[StreamMode] = (),
-        output_keys: str | Sequence[str] | None = None,
-        interrupt_before: All | Sequence[str] | None = None,
-        interrupt_after: All | Sequence[str] | None = None,
-        durability: Durability | None = None,
-        stream_version: Literal["v1", "v2"] = "v1",
+        stream_version: Literal["v1", "v2"] = "v2",
        **kwargs: Any,
    ) -> dict[str, Any] | Any:
        """Run the graph with a single input and config.
@@ -3332,6 +3332,23 @@ class Pregel(
        else:
            return chunks

+    @overload
+    async def ainvoke(
+        self,
+        input: InputT | Command | None,
+        config: RunnableConfig | None = None,
+        *,
+        context: ContextT | None = None,
+        stream_mode: StreamMode = "values",
+        print_mode: StreamMode | Sequence[StreamMode] = (),
+        output_keys: str | Sequence[str] | None = None,
+        interrupt_before: All | Sequence[str] | None = None,
+        interrupt_after: All | Sequence[str] | None = None,
+        durability: Durability | None = None,
+        stream_version: Literal["v1"],
+        **kwargs: Any,
+    ) -> dict[str, Any] | Any: ...
+
    @overload
    async def ainvoke(
        self,
@@ -3345,7 +3362,7 @@ class Pregel(
        interrupt_before: All | Sequence[str] | None = None,
        interrupt_after: All | Sequence[str] | None = None,
        durability: Durability | None = None,
-        stream_version: Literal["v2"],
+        stream_version: Literal["v2"] = ...,
        **kwargs: Any,
    ) -> GraphOutput[OutputT]: ...

@@ -3362,11 +3379,10 @@ class Pregel(
        interrupt_before: All | Sequence[str] | None = None,
        interrupt_after: All | Sequence[str] | None = None,
        durability: Durability | None = None,
-        stream_version: Literal["v2"],
+        stream_version: Literal["v2"] = ...,
        **kwargs: Any,
    ) -> list[StreamPart[OutputT, StateT]]: ...

-    @overload
    async def ainvoke(
        self,
        input: InputT | Command | None,
@@ -3379,23 +3395,7 @@ class Pregel(
        interrupt_before: All | Sequence[str] | None = None,
        interrupt_after: All | Sequence[str] | None = None,
        durability: Durability | None = None,
-        stream_version: Literal["v1"] = ...,
-        **kwargs: Any,
-    ) -> dict[str, Any] | Any: ...
-
-    async def ainvoke(
-        self,
-        input: InputT | Command | None,
-        config: RunnableConfig | None = None,
-        *,
-        context: ContextT | None = None,
-        stream_mode: StreamMode = "values",
-        print_mode: StreamMode | Sequence[StreamMode] = (),
-        output_keys: str | Sequence[str] | None = None,
-        interrupt_before: All | Sequence[str] | None = None,
-        interrupt_after: All | Sequence[str] | None = None,
-        durability: Durability | None = None,
-        stream_version: Literal["v1", "v2"] = "v1",
+        stream_version: Literal["v1", "v2"] = "v2",
        **kwargs: Any,
    ) -> dict[str, Any] | Any:
        """Asynchronously run the graph with a single input and config.
@@ -3556,7 +3556,7 @@ def _output(
    stream_subgraphs: bool,
    getter: Callable[[], tuple[tuple[str, ...], str, Any]],
    empty_exc: type[Exception],
-    stream_version: Literal["v1", "v2"] = "v1",
+    stream_version: Literal["v1", "v2"] = "v2",
    output_mapper: Callable[[Any], Any] | None = None,
    state_mapper: Callable[[Any], Any] | None = None,
 ) -> Iterator:
@@ -116,12 +116,41 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou
        interrupt_before: All | Sequence[str] | None = None,
        interrupt_after: All | Sequence[str] | None = None,
        subgraphs: bool = False,
-        stream_version: Literal["v2"],
+        stream_version: Literal["v1"],
+    ) -> Iterator[dict[str, Any] | Any]: ...
+
+    @overload
+    @abstractmethod
+    def stream(
+        self,
+        input: InputT | Command | None,
+        config: RunnableConfig | None = None,
+        *,
+        context: ContextT | None = None,
+        stream_mode: StreamMode | list[StreamMode] | None = None,
+        interrupt_before: All | Sequence[str] | None = None,
+        interrupt_after: All | Sequence[str] | None = None,
+        subgraphs: bool = False,
+        stream_version: Literal["v2"] = ...,
+    ) -> Iterator[StreamPart[OutputT, StateT]]: ...
+
+    @abstractmethod
+    def stream(
+        self,
+        input: InputT | Command | None,
+        config: RunnableConfig | None = None,
+        *,
+        context: ContextT | None = None,
+        stream_mode: StreamMode | list[StreamMode] | None = None,
+        interrupt_before: All | Sequence[str] | None = None,
+        interrupt_after: All | Sequence[str] | None = None,
+        subgraphs: bool = False,
+        stream_version: Literal["v1", "v2"] = "v2",
    ) -> Iterator[StreamPart[OutputT, StateT]]: ...

    @overload
    @abstractmethod
-    def stream(
+    def astream(
        self,
        input: InputT | Command | None,
        config: RunnableConfig | None = None,
@@ -131,22 +160,8 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou
        interrupt_before: All | Sequence[str] | None = None,
        interrupt_after: All | Sequence[str] | None = None,
        subgraphs: bool = False,
-        stream_version: Literal["v1"] = ...,
-    ) -> Iterator[dict[str, Any] | Any]: ...
-
-    @abstractmethod
-    def stream(
-        self,
-        input: InputT | Command | None,
-        config: RunnableConfig | None = None,
-        *,
-        context: ContextT | None = None,
-        stream_mode: StreamMode | list[StreamMode] | None = None,
-        interrupt_before: All | Sequence[str] | None = None,
-        interrupt_after: All | Sequence[str] | None = None,
-        subgraphs: bool = False,
-        stream_version: Literal["v1", "v2"] = "v1",
-    ) -> Iterator[dict[str, Any] | Any]: ...
+        stream_version: Literal["v1"],
+    ) -> AsyncIterator[dict[str, Any] | Any]: ...

    @overload
    @abstractmethod
@@ -160,37 +175,35 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou
        interrupt_before: All | Sequence[str] | None = None,
        interrupt_after: All | Sequence[str] | None = None,
        subgraphs: bool = False,
-        stream_version: Literal["v2"],
+        stream_version: Literal["v2"] = ...,
+    ) -> AsyncIterator[StreamPart[OutputT, StateT]]: ...
+
+    @abstractmethod
+    def astream(
+        self,
+        input: InputT | Command | None,
+        config: RunnableConfig | None = None,
+        *,
+        context: ContextT | None = None,
+        stream_mode: StreamMode | list[StreamMode] | None = None,
+        interrupt_before: All | Sequence[str] | None = None,
+        interrupt_after: All | Sequence[str] | None = None,
+        subgraphs: bool = False,
+        stream_version: Literal["v1", "v2"] = "v2",
    ) -> AsyncIterator[StreamPart[OutputT, StateT]]: ...

    @overload
    @abstractmethod
-    def astream(
+    def invoke(
        self,
        input: InputT | Command | None,
        config: RunnableConfig | None = None,
        *,
        context: ContextT | None = None,
-        stream_mode: StreamMode | list[StreamMode] | None = None,
        interrupt_before: All | Sequence[str] | None = None,
        interrupt_after: All | Sequence[str] | None = None,
-        subgraphs: bool = False,
-        stream_version: Literal["v1"] = ...,
-    ) -> AsyncIterator[dict[str, Any] | Any]: ...
-
-    @abstractmethod
-    def astream(
-        self,
-        input: InputT | Command | None,
-        config: RunnableConfig | None = None,
-        *,
-        context: ContextT | None = None,
-        stream_mode: StreamMode | list[StreamMode] | None = None,
-        interrupt_before: All | Sequence[str] | None = None,
-        interrupt_after: All | Sequence[str] | None = None,
-        subgraphs: bool = False,
-        stream_version: Literal["v1", "v2"] = "v1",
-    ) -> AsyncIterator[dict[str, Any] | Any]: ...
+        stream_version: Literal["v1"],
+    ) -> dict[str, Any] | Any: ...

    @overload
    @abstractmethod
@@ -202,10 +215,9 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou
        context: ContextT | None = None,
        interrupt_before: All | Sequence[str] | None = None,
        interrupt_after: All | Sequence[str] | None = None,
-        stream_version: Literal["v2"],
+        stream_version: Literal["v2"] = ...,
    ) -> GraphOutput[OutputT]: ...

-    @overload
    @abstractmethod
    def invoke(
        self,
@@ -215,32 +227,7 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou
        context: ContextT | None = None,
        interrupt_before: All | Sequence[str] | None = None,
        interrupt_after: All | Sequence[str] | None = None,
-        stream_version: Literal["v1"] = ...,
-    ) -> dict[str, Any] | Any: ...
-
-    @abstractmethod
-    def invoke(
-        self,
-        input: InputT | Command | None,
-        config: RunnableConfig | None = None,
-        *,
-        context: ContextT | None = None,
-        interrupt_before: All | Sequence[str] | None = None,
-        interrupt_after: All | Sequence[str] | None = None,
-        stream_version: Literal["v1", "v2"] = "v1",
-    ) -> dict[str, Any] | Any: ...
-
-    @overload
-    @abstractmethod
-    async def ainvoke(
-        self,
-        input: InputT | Command | None,
-        config: RunnableConfig | None = None,
-        *,
-        context: ContextT | None = None,
-        interrupt_before: All | Sequence[str] | None = None,
-        interrupt_after: All | Sequence[str] | None = None,
-        stream_version: Literal["v2"],
+        stream_version: Literal["v1", "v2"] = "v2",
    ) -> GraphOutput[OutputT]: ...

    @overload
@@ -253,9 +240,22 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou
        context: ContextT | None = None,
        interrupt_before: All | Sequence[str] | None = None,
        interrupt_after: All | Sequence[str] | None = None,
-        stream_version: Literal["v1"] = ...,
+        stream_version: Literal["v1"],
    ) -> dict[str, Any] | Any: ...

+    @overload
+    @abstractmethod
+    async def ainvoke(
+        self,
+        input: InputT | Command | None,
+        config: RunnableConfig | None = None,
+        *,
+        context: ContextT | None = None,
+        interrupt_before: All | Sequence[str] | None = None,
+        interrupt_after: All | Sequence[str] | None = None,
+        stream_version: Literal["v2"] = ...,
+    ) -> GraphOutput[OutputT]: ...
+
    @abstractmethod
    async def ainvoke(
        self,
@@ -265,8 +265,8 @@ class PregelProtocol(Runnable[InputT, Any], Generic[StateT, ContextT, InputT, Ou
        context: ContextT | None = None,
        interrupt_before: All | Sequence[str] | None = None,
        interrupt_after: All | Sequence[str] | None = None,
-        stream_version: Literal["v1", "v2"] = "v1",
-    ) -> dict[str, Any] | Any: ...
+        stream_version: Literal["v1", "v2"] = "v2",
+    ) -> GraphOutput[OutputT]: ...


 StreamChunk = tuple[tuple[str, ...], str, Any]
@@ -697,9 +697,9 @@ class RemoteGraph(PregelProtocol):
        subgraphs: bool = False,
        headers: dict[str, str] | None = None,
        params: QueryParamTypes | None = None,
-        stream_version: Literal["v2"],
+        stream_version: Literal["v1"],
        **kwargs: Any,
-    ) -> Iterator[StreamPart]: ...
+    ) -> Iterator[dict[str, Any] | Any]: ...

    @overload
    def stream(
@@ -713,9 +713,9 @@ class RemoteGraph(PregelProtocol):
        subgraphs: bool = False,
        headers: dict[str, str] | None = None,
        params: QueryParamTypes | None = None,
-        stream_version: Literal["v1"] = ...,
+        stream_version: Literal["v2"] = ...,
        **kwargs: Any,
-    ) -> Iterator[dict[str, Any] | Any]: ...
+    ) -> Iterator[StreamPart]: ...

    def stream(
        self,
@@ -728,7 +728,7 @@ class RemoteGraph(PregelProtocol):
        subgraphs: bool = False,
        headers: dict[str, str] | None = None,
        params: QueryParamTypes | None = None,
-        stream_version: Literal["v1", "v2"] = "v1",
+        stream_version: Literal["v1", "v2"] = "v2",
        **kwargs: Any,
    ) -> Iterator[dict[str, Any] | Any]:
        """Create a run and stream the results.
@@ -842,9 +842,9 @@ class RemoteGraph(PregelProtocol):
        subgraphs: bool = False,
        headers: dict[str, str] | None = None,
        params: QueryParamTypes | None = None,
-        stream_version: Literal["v2"],
+        stream_version: Literal["v1"],
        **kwargs: Any,
-    ) -> AsyncIterator[StreamPart]: ...
+    ) -> AsyncIterator[dict[str, Any] | Any]: ...

    @overload
    def astream(
@@ -858,9 +858,9 @@ class RemoteGraph(PregelProtocol):
        subgraphs: bool = False,
        headers: dict[str, str] | None = None,
        params: QueryParamTypes | None = None,
-        stream_version: Literal["v1"] = ...,
+        stream_version: Literal["v2"] = ...,
        **kwargs: Any,
-    ) -> AsyncIterator[dict[str, Any] | Any]: ...
+    ) -> AsyncIterator[StreamPart]: ...

    async def astream(
        self,
@@ -873,7 +873,7 @@ class RemoteGraph(PregelProtocol):
        subgraphs: bool = False,
        headers: dict[str, str] | None = None,
        params: QueryParamTypes | None = None,
-        stream_version: Literal["v1", "v2"] = "v1",
+        stream_version: Literal["v1", "v2"] = "v2",
        **kwargs: Any,
    ) -> AsyncIterator[dict[str, Any] | Any]:
        """Create a run and stream the results.
@@ -1001,9 +1001,9 @@ class RemoteGraph(PregelProtocol):
        interrupt_after: All | Sequence[str] | None = None,
        headers: dict[str, str] | None = None,
        params: QueryParamTypes | None = None,
-        stream_version: Literal["v2"],
+        stream_version: Literal["v1"],
        **kwargs: Any,
-    ) -> GraphOutput[dict[str, Any]]: ...
+    ) -> dict[str, Any] | Any: ...

    @overload
    def invoke(
@@ -1015,9 +1015,9 @@ class RemoteGraph(PregelProtocol):
        interrupt_after: All | Sequence[str] | None = None,
        headers: dict[str, str] | None = None,
        params: QueryParamTypes | None = None,
-        stream_version: Literal["v1"] = ...,
+        stream_version: Literal["v2"] = ...,
        **kwargs: Any,
-    ) -> dict[str, Any] | Any: ...
+    ) -> GraphOutput[dict[str, Any]]: ...

    def invoke(
        self,
@@ -1028,7 +1028,7 @@ class RemoteGraph(PregelProtocol):
        interrupt_after: All | Sequence[str] | None = None,
        headers: dict[str, str] | None = None,
        params: QueryParamTypes | None = None,
-        stream_version: Literal["v1", "v2"] = "v1",
+        stream_version: Literal["v1", "v2"] = "v2",
        **kwargs: Any,
    ) -> dict[str, Any] | Any:
        """Create a run, wait until it finishes and return the final state.
@@ -1079,9 +1079,9 @@ class RemoteGraph(PregelProtocol):
        interrupt_after: All | Sequence[str] | None = None,
        headers: dict[str, str] | None = None,
        params: QueryParamTypes | None = None,
-        stream_version: Literal["v2"],
+        stream_version: Literal["v1"],
        **kwargs: Any,
-    ) -> GraphOutput[dict[str, Any]]: ...
+    ) -> dict[str, Any] | Any: ...

    @overload
    async def ainvoke(
@@ -1093,9 +1093,9 @@ class RemoteGraph(PregelProtocol):
        interrupt_after: All | Sequence[str] | None = None,
        headers: dict[str, str] | None = None,
        params: QueryParamTypes | None = None,
-        stream_version: Literal["v1"] = ...,
+        stream_version: Literal["v2"] = ...,
        **kwargs: Any,
-    ) -> dict[str, Any] | Any: ...
+    ) -> GraphOutput[dict[str, Any]]: ...

    async def ainvoke(
        self,
@@ -1106,7 +1106,7 @@ class RemoteGraph(PregelProtocol):
        interrupt_after: All | Sequence[str] | None = None,
        headers: dict[str, str] | None = None,
        params: QueryParamTypes | None = None,
-        stream_version: Literal["v1", "v2"] = "v1",
+        stream_version: Literal["v1", "v2"] = "v2",
        **kwargs: Any,
    ) -> dict[str, Any] | Any:
        """Create a run, wait until it finishes and return the final state.