Try to improve coverage

Author: dmontagu

pydantic_graph/pydantic_graph/beta/graph.py

@@ -512,7 +512,7 @@ def __post_init__(self):
     @property
     def task_group(self) -> TaskGroup:
         if self._task_group is None:
-            raise RuntimeError("This graph iterator hasn't been started")
+            raise RuntimeError("This graph iterator hasn't been started")  # pragma: no cover
         return self._task_group
 
     async def iter_graph(  # noqa C901
@@ -528,7 +528,7 @@ async def iter_graph(  # noqa C901
                     # Handle task results
                     async with self.iter_stream_receiver:
                         while self.active_tasks or self.active_reducers:
-                            async for task_result in self.iter_stream_receiver:
+                            async for task_result in self.iter_stream_receiver:  # pragma: no branch
                                 # If we encounter a mock task, add it to the active tasks to ensure we don't proceed until everything downstream is handled
                                 if (
                                     not task_result.source_is_finished
@@ -653,11 +653,10 @@ async def iter_graph(  # noqa C901
             'Graph run completed, but no result was produced. This is either a bug in the graph or a bug in the graph runner.'
         )
 
-    async def _finish_task(self, task_id: TaskID, keep_cancel_scope: bool = False) -> None:
-        if not keep_cancel_scope:
-            scope = self.cancel_scopes.pop(task_id, None)
-            if scope is not None:
-                scope.cancel()
+    async def _finish_task(self, task_id: TaskID) -> None:
+        scope = self.cancel_scopes.pop(task_id, None)
+        if scope is not None:
+            scope.cancel()
         self.active_tasks.pop(task_id, None)
 
     def _handle_execution_request(self, request: Sequence[GraphTask]) -> None:

pydantic_graph/pydantic_graph/beta/graph_builder.py

@@ -9,7 +9,7 @@
 
 import inspect
 from collections import Counter, defaultdict
-from collections.abc import Callable, Iterable
+from collections.abc import AsyncIterable, Callable, Iterable
 from dataclasses import dataclass, replace
 from types import NoneType
 from typing import Any, Generic, Literal, cast, get_origin, get_type_hints, overload
@@ -44,7 +44,7 @@
     Path,
     PathBuilder,
 )
-from pydantic_graph.beta.step import NodeStep, Step, StepFunction, StepNode
+from pydantic_graph.beta.step import NodeStep, Step, StepAsyncIteratorFunction, StepContext, StepFunction, StepNode
 from pydantic_graph.beta.util import TypeOrTypeExpression, get_callable_name, unpack_type_expression
 from pydantic_graph.exceptions import GraphBuildingError
 from pydantic_graph.nodes import BaseNode, End
@@ -162,21 +162,21 @@ def end_node(self) -> EndNode[GraphOutputT]:
         return self._end_node
 
     @overload
-    def _step(
+    def step(
         self,
         *,
         node_id: str | None = None,
         label: str | None = None,
     ) -> Callable[[StepFunction[StateT, DepsT, InputT, OutputT]], Step[StateT, DepsT, InputT, OutputT]]: ...
     @overload
-    def _step(
+    def step(
         self,
         call: StepFunction[StateT, DepsT, InputT, OutputT],
         *,
         node_id: str | None = None,
         label: str | None = None,
     ) -> Step[StateT, DepsT, InputT, OutputT]: ...
-    def _step(
+    def step(
         self,
         call: StepFunction[StateT, DepsT, InputT, OutputT] | None = None,
         *,
@@ -186,10 +186,10 @@ def _step(
         Step[StateT, DepsT, InputT, OutputT]
         | Callable[[StepFunction[StateT, DepsT, InputT, OutputT]], Step[StateT, DepsT, InputT, OutputT]]
     ):
-        """Create a step from a step function (internal implementation).
+        """Create a step from a step function.
 
-        This internal method handles the actual step creation logic and
-        automatic edge inference from type hints.
+        This method can be used as a decorator or called directly to create
+        a step node from an async function.
 
         Args:
             call: The step function to wrap
@@ -204,7 +204,7 @@ def _step(
             def decorator(
                 func: StepFunction[StateT, DepsT, InputT, OutputT],
             ) -> Step[StateT, DepsT, InputT, OutputT]:
-                return self._step(call=func, node_id=node_id, label=label)
+                return self.step(call=func, node_id=node_id, label=label)
 
             return decorator
 
@@ -215,29 +215,48 @@ def decorator(
         return step
 
     @overload
-    def step(
+    def step_async_iterable(
         self,
         *,
         node_id: str | None = None,
         label: str | None = None,
-    ) -> Callable[[StepFunction[StateT, DepsT, InputT, OutputT]], Step[StateT, DepsT, InputT, OutputT]]: ...
+    ) -> Callable[
+        [StepAsyncIteratorFunction[StateT, DepsT, InputT, OutputT]], Step[StateT, DepsT, InputT, AsyncIterable[OutputT]]
+    ]: ...
     @overload
-    def step(
+    def step_async_iterable(
         self,
-        call: StepFunction[StateT, DepsT, InputT, OutputT],
+        call: StepAsyncIteratorFunction[StateT, DepsT, InputT, OutputT],
         *,
         node_id: str | None = None,
         label: str | None = None,
-    ) -> Step[StateT, DepsT, InputT, OutputT]: ...
-    def step(
+    ) -> Step[StateT, DepsT, InputT, AsyncIterable[OutputT]]: ...
+    @overload
+    def step_async_iterable(
         self,
-        call: StepFunction[StateT, DepsT, InputT, OutputT] | None = None,
+        call: StepAsyncIteratorFunction[StateT, DepsT, InputT, OutputT] | None = None,
         *,
         node_id: str | None = None,
         label: str | None = None,
     ) -> (
-        Step[StateT, DepsT, InputT, OutputT]
-        | Callable[[StepFunction[StateT, DepsT, InputT, OutputT]], Step[StateT, DepsT, InputT, OutputT]]
+        Step[StateT, DepsT, InputT, AsyncIterable[OutputT]]
+        | Callable[
+            [StepAsyncIteratorFunction[StateT, DepsT, InputT, OutputT]],
+            Step[StateT, DepsT, InputT, AsyncIterable[OutputT]],
+        ]
+    ): ...
+    def step_async_iterable(
+        self,
+        call: StepAsyncIteratorFunction[StateT, DepsT, InputT, OutputT] | None = None,
+        *,
+        node_id: str | None = None,
+        label: str | None = None,
+    ) -> (
+        Step[StateT, DepsT, InputT, AsyncIterable[OutputT]]
+        | Callable[
+            [StepAsyncIteratorFunction[StateT, DepsT, InputT, OutputT]],
+            Step[StateT, DepsT, InputT, AsyncIterable[OutputT]],
+        ]
     ):
         """Create a step from a step function.
 
@@ -253,9 +272,19 @@ def step(
             Either a Step instance or a decorator function
         """
         if call is None:
-            return self._step(node_id=node_id, label=label)
-        else:
-            return self._step(call=call, node_id=node_id, label=label)
+
+            def decorator(
+                func: StepAsyncIteratorFunction[StateT, DepsT, InputT, OutputT],
+            ) -> Step[StateT, DepsT, InputT, AsyncIterable[OutputT]]:
+                return self.step_async_iterable(call=func, node_id=node_id, label=label)
+
+            return decorator
+
+        # We need to wrap the call so that we can call `await` even though the result is an async iterator
+        async def wrapper(ctx: StepContext[StateT, DepsT, InputT]):
+            return call(ctx)
+
+        return self.step(call=wrapper, node_id=node_id, label=label)
 
     @overload
     def join(

pydantic_graph/pydantic_graph/beta/join.py

@@ -143,8 +143,6 @@ class ReduceFirstValue(Generic[T]):
 
     def __call__(self, ctx: ReducerContext[object, object], current: T, inputs: T) -> T:
         """The reducer function."""
-        if ctx.cancelled_sibling_tasks:
-            return current
         ctx.cancel_sibling_tasks()
         return inputs
 

pydantic_graph/pydantic_graph/beta/step.py

@@ -7,7 +7,7 @@
 
 from __future__ import annotations
 
-from collections.abc import Awaitable
+from collections.abc import AsyncIterator, Awaitable
 from dataclasses import dataclass
 from typing import Any, Generic, Protocol, cast, get_origin, overload
 
@@ -90,6 +90,33 @@ def __call__(self, ctx: StepContext[StateT, DepsT, InputT]) -> Awaitable[OutputT
         raise NotImplementedError
 
 
+class StepAsyncIteratorFunction(Protocol[StateT, DepsT, InputT, OutputT]):
+    """Protocol for step functions that can be executed in the graph.
+
+    Step functions are async callables that receive a step context and return
+    a result. This protocol enables serialization and deserialization of step
+    calls similar to how evaluators work.
+
+    Type Parameters:
+        StateT: The type of the graph state
+        DepsT: The type of the dependencies
+        InputT: The type of the input data
+        OutputT: The type of the output data
+    """
+
+    def __call__(self, ctx: StepContext[StateT, DepsT, InputT]) -> AsyncIterator[OutputT]:
+        """Execute the step function with the given context.
+
+        Args:
+            ctx: The step context containing state, dependencies, and inputs
+
+        Returns:
+            An awaitable that resolves to the step's output
+        """
+        raise NotImplementedError
+        yield
+
+
 AnyStepFunction = StepFunction[Any, Any, Any, Any]
 """Type alias for a step function with any type parameters."""
 

tests/graph/beta/test_graph_edge_cases.py

@@ -205,12 +205,11 @@ def get_early_stopping_reducer():
 
         def reduce(ctx: ReducerContext[EarlyStopState, object], current: int, inputs: int) -> int:
             nonlocal count
-            if not ctx.cancelled_sibling_tasks:
-                count += 1
-                current += inputs
-                if count >= 2:
-                    ctx.state.stopped = True  # update the state so we can assert on it later
-                    ctx.cancel_sibling_tasks()
+            count += 1
+            current += inputs
+            if count >= 2:
+                ctx.state.stopped = True  # update the state so we can assert on it later
+                ctx.cancel_sibling_tasks()
             return current
 
         return reduce