pydantic-graph 1.13.0__tar.gz → 1.31.0__tar.gz

{pydantic_graph-1.13.0 → pydantic_graph-1.31.0}/.gitignore

@@ -21,3 +21,5 @@ node_modules/
 /test_tmp/
 .mcp.json
 .claude/
+/.cursor/
+/.devcontainer/

{pydantic_graph-1.13.0 → pydantic_graph-1.31.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-graph
-Version: 1.13.0
+Version: 1.31.0
 Summary: Graph and state machine library
 Project-URL: Homepage, https://ai.pydantic.dev/graph/tree/main/pydantic_graph
 Project-URL: Source, https://github.com/pydantic/pydantic-ai

{pydantic_graph-1.13.0 → pydantic_graph-1.31.0}/pydantic_graph/beta/graph.py

@@ -8,13 +8,12 @@ the graph-based workflow system.
 from __future__ import annotations as _annotations
 
 import sys
-import uuid
-from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Iterable, Sequence
+from collections.abc import AsyncGenerator, AsyncIterable, AsyncIterator, Callable, Iterable, Sequence
 from contextlib import AbstractContextManager, AsyncExitStack, ExitStack, asynccontextmanager, contextmanager
 from dataclasses import dataclass, field
 from typing import TYPE_CHECKING, Any, Generic, Literal, TypeGuard, cast, get_args, get_origin, overload
 
-from anyio import CancelScope, create_memory_object_stream, create_task_group
+from anyio import BrokenResourceError, CancelScope, create_memory_object_stream, create_task_group
 from anyio.abc import TaskGroup
 from anyio.streams.memory import MemoryObjectReceiveStream, MemoryObjectSendStream
 from typing_extensions import TypeVar, assert_never
@@ -22,7 +21,7 @@ from typing_extensions import TypeVar, assert_never
 from pydantic_graph import exceptions
 from pydantic_graph._utils import AbstractSpan, get_traceparent, infer_obj_name, logfire_span
 from pydantic_graph.beta.decision import Decision
-from pydantic_graph.beta.id_types import ForkID, ForkStack, ForkStackItem, GraphRunID, JoinID, NodeID, NodeRunID, TaskID
+from pydantic_graph.beta.id_types import ForkID, ForkStack, ForkStackItem, JoinID, NodeID, NodeRunID, TaskID
 from pydantic_graph.beta.join import Join, JoinNode, JoinState, ReducerContext
 from pydantic_graph.beta.node import (
     EndNode,
@@ -44,9 +43,9 @@ from pydantic_graph.beta.util import unpack_type_expression
 from pydantic_graph.nodes import BaseNode, End
 
 if sys.version_info < (3, 11):
-    from exceptiongroup import ExceptionGroup as ExceptionGroup  # pragma: lax no cover
+    from exceptiongroup import BaseExceptionGroup as BaseExceptionGroup  # pragma: lax no cover
 else:
-    ExceptionGroup = ExceptionGroup  # pragma: lax no cover
+    BaseExceptionGroup = BaseExceptionGroup  # pragma: lax no cover
 
 if TYPE_CHECKING:
     from pydantic_graph.beta.mermaid import StateDiagramDirection
@@ -306,14 +305,13 @@ class Graph(Generic[StateT, DepsT, InputT, OutputT]):
 
 
 @dataclass
-class GraphTask:
-    """A single task representing the execution of a node in the graph.
+class GraphTaskRequest:
+    """A request to run a task representing the execution of a node in the graph.
 
-    GraphTask encapsulates all the information needed to execute a specific
+    GraphTaskRequest encapsulates all the information needed to execute a specific
     node, including its inputs and the fork context it's executing within.
     """
 
-    # With our current BaseNode thing, next_node_id and next_node_inputs are merged into `next_node` itself
     node_id: NodeID
     """The ID of the node to execute."""
 
@@ -326,9 +324,26 @@ class GraphTask:
     Used by the GraphRun to decide when to proceed through joins.
     """
 
-    task_id: TaskID = field(default_factory=lambda: TaskID(str(uuid.uuid4())), repr=False)
+
+@dataclass
+class GraphTask(GraphTaskRequest):
+    """A task representing the execution of a node in the graph.
+
+    GraphTask encapsulates all the information needed to execute a specific
+    node, including its inputs and the fork context it's executing within,
+    and has a unique ID to identify the task within the graph run.
+    """
+
+    task_id: TaskID = field(repr=False)
     """Unique identifier for this task."""
 
+    @staticmethod
+    def from_request(request: GraphTaskRequest, get_task_id: Callable[[], TaskID]) -> GraphTask:
+        # Don't call the get_task_id callable, this is already a task
+        if isinstance(request, GraphTask):
+            return request
+        return GraphTask(request.node_id, request.inputs, request.fork_stack, get_task_id())
+
 
 class GraphRun(Generic[StateT, DepsT, OutputT]):
     """A single execution instance of a graph.
@@ -378,12 +393,20 @@ class GraphRun(Generic[StateT, DepsT, OutputT]):
         self._next: EndMarker[OutputT] | Sequence[GraphTask] | None = None
         """The next item to be processed."""
 
-        run_id = GraphRunID(str(uuid.uuid4()))
-        initial_fork_stack: ForkStack = (ForkStackItem(StartNode.id, NodeRunID(run_id), 0),)
-        self._first_task = GraphTask(node_id=StartNode.id, inputs=inputs, fork_stack=initial_fork_stack)
+        self._next_task_id = 0
+        self._next_node_run_id = 0
+        initial_fork_stack: ForkStack = (ForkStackItem(StartNode.id, self._get_next_node_run_id(), 0),)
+        self._first_task = GraphTask(
+            node_id=StartNode.id, inputs=inputs, fork_stack=initial_fork_stack, task_id=self._get_next_task_id()
+        )
         self._iterator_task_group = create_task_group()
         self._iterator_instance = _GraphIterator[StateT, DepsT, OutputT](
-            self.graph, self.state, self.deps, self._iterator_task_group
+            self.graph,
+            self.state,
+            self.deps,
+            self._iterator_task_group,
+            self._get_next_node_run_id,
+            self._get_next_task_id,
         )
         self._iterator = self._iterator_instance.iter_graph(self._first_task)
 
@@ -449,7 +472,7 @@ class GraphRun(Generic[StateT, DepsT, OutputT]):
         return self._next
 
     async def next(
-        self, value: EndMarker[OutputT] | Sequence[GraphTask] | None = None
+        self, value: EndMarker[OutputT] | Sequence[GraphTaskRequest] | None = None
     ) -> EndMarker[OutputT] | Sequence[GraphTask]:
         """Advance the graph execution by one step.
 
@@ -467,7 +490,10 @@ class GraphRun(Generic[StateT, DepsT, OutputT]):
             # if `next` is called before the `first_node` has run.
             await anext(self)
         if value is not None:
-            self._next = value
+            if isinstance(value, EndMarker):
+                self._next = value
+            else:
+                self._next = [GraphTask.from_request(gtr, self._get_next_task_id) for gtr in value]
         return await anext(self)
 
     @property
@@ -490,6 +516,16 @@ class GraphRun(Generic[StateT, DepsT, OutputT]):
             return self._next.value
         return None
 
+    def _get_next_task_id(self) -> TaskID:
+        next_id = TaskID(f'task:{self._next_task_id}')
+        self._next_task_id += 1
+        return next_id
+
+    def _get_next_node_run_id(self) -> NodeRunID:
+        next_id = NodeRunID(f'task:{self._next_node_run_id}')
+        self._next_node_run_id += 1
+        return next_id
+
 
 @dataclass
 class _GraphTaskAsyncIterable:
@@ -510,6 +546,8 @@ class _GraphIterator(Generic[StateT, DepsT, OutputT]):
     state: StateT
     deps: DepsT
     task_group: TaskGroup
+    get_next_node_run_id: Callable[[], NodeRunID]
+    get_next_task_id: Callable[[], TaskID]
 
     cancel_scopes: dict[TaskID, CancelScope] = field(init=False)
     active_tasks: dict[TaskID, GraphTask] = field(init=False)
@@ -522,8 +560,9 @@ class _GraphIterator(Generic[StateT, DepsT, OutputT]):
         self.active_tasks = {}
         self.active_reducers = {}
         self.iter_stream_sender, self.iter_stream_receiver = create_memory_object_stream[_GraphTaskResult]()
+        self._next_node_run_id = 1
 
-    async def iter_graph(  # noqa C901
+    async def iter_graph(  # noqa: C901
         self, first_task: GraphTask
     ) -> AsyncGenerator[EndMarker[OutputT] | Sequence[GraphTask], EndMarker[OutputT] | Sequence[GraphTask]]:
         async with self.iter_stream_sender:
@@ -709,12 +748,15 @@ class _GraphIterator(Generic[StateT, DepsT, OutputT]):
         with CancelScope() as scope:
             self.cancel_scopes[t_.task_id] = scope
             result = await self._run_task(t_)
-            if isinstance(result, _GraphTaskAsyncIterable):
-                async for new_tasks in result.iterable:
-                    await self.iter_stream_sender.send(_GraphTaskResult(t_, new_tasks, False))
-                await self.iter_stream_sender.send(_GraphTaskResult(t_, []))
-            else:
-                await self.iter_stream_sender.send(_GraphTaskResult(t_, result))
+            try:
+                if isinstance(result, _GraphTaskAsyncIterable):
+                    async for new_tasks in result.iterable:
+                        await self.iter_stream_sender.send(_GraphTaskResult(t_, new_tasks, False))
+                    await self.iter_stream_sender.send(_GraphTaskResult(t_, []))
+                else:
+                    await self.iter_stream_sender.send(_GraphTaskResult(t_, result))
+            except BrokenResourceError:
+                pass  # pragma: no cover # This can happen in difficult-to-reproduce circumstances when cancelling an asyncio task
 
     async def _run_task(
         self,
@@ -782,12 +824,12 @@ class _GraphIterator(Generic[StateT, DepsT, OutputT]):
         fork_stack: ForkStack,
     ) -> Sequence[GraphTask] | JoinItem | EndMarker[OutputT]:
         if isinstance(next_node, StepNode):
-            return [GraphTask(next_node.step.id, next_node.inputs, fork_stack)]
+            return [GraphTask(next_node.step.id, next_node.inputs, fork_stack, self.get_next_task_id())]
         elif isinstance(next_node, JoinNode):
             return JoinItem(next_node.join.id, next_node.inputs, fork_stack)
         elif isinstance(next_node, BaseNode):
             node_step = NodeStep(next_node.__class__)
-            return [GraphTask(node_step.id, next_node, fork_stack)]
+            return [GraphTask(node_step.id, next_node, fork_stack, self.get_next_task_id())]
         elif isinstance(next_node, End):
             return EndMarker(next_node.data)
         else:
@@ -821,7 +863,7 @@ class _GraphIterator(Generic[StateT, DepsT, OutputT]):
             'These markers should be removed from paths during graph building'
         )
         if isinstance(item, DestinationMarker):
-            return [GraphTask(item.destination_id, inputs, fork_stack)]
+            return [GraphTask(item.destination_id, inputs, fork_stack, self.get_next_task_id())]
         elif isinstance(item, TransformMarker):
             inputs = item.transform(StepContext(state=self.state, deps=self.deps, inputs=inputs))
             return self._handle_path(path.next_path, inputs, fork_stack)
@@ -853,7 +895,7 @@ class _GraphIterator(Generic[StateT, DepsT, OutputT]):
         )  # this should have already been ensured during graph building
 
         new_tasks: list[GraphTask] = []
-        node_run_id = NodeRunID(str(uuid.uuid4()))
+        node_run_id = self.get_next_node_run_id()
         if node.is_map:
             # If the map specifies a downstream join id, eagerly create a join state for it
             if (join_id := node.downstream_join_id) is not None:
@@ -931,7 +973,7 @@ def _unwrap_exception_groups():
     else:
         try:
             yield
-        except ExceptionGroup as e:
+        except BaseExceptionGroup as e:
             exception = e.exceptions[0]
             if exception.__cause__ is None:
                 # bizarrely, this prevents recursion errors when formatting the exception for logfire

{pydantic_graph-1.13.0 → pydantic_graph-1.31.0}/pydantic_graph/beta/graph_builder.py

@@ -284,6 +284,8 @@ class GraphBuilder(Generic[StateT, DepsT, GraphInputT, GraphOutputT]):
         async def wrapper(ctx: StepContext[StateT, DepsT, InputT]):
             return call(ctx)
 
+        node_id = node_id or get_callable_name(call)
+
         return self.step(call=wrapper, node_id=node_id, label=label)
 
     @overload
@@ -318,7 +320,7 @@ class GraphBuilder(Generic[StateT, DepsT, GraphInputT, GraphOutputT]):
         preferred_parent_fork: Literal['farthest', 'closest'] = 'farthest',
     ) -> Join[StateT, DepsT, InputT, OutputT]:
         if initial_factory is UNSET:
-            initial_factory = lambda: initial  # pyright: ignore[reportAssignmentType]  # noqa E731
+            initial_factory = lambda: initial  # pyright: ignore[reportAssignmentType]  # noqa: E731
 
         return Join[StateT, DepsT, InputT, OutputT](
             id=JoinID(NodeID(node_id or generate_placeholder_node_id(get_callable_name(reducer)))),
@@ -329,7 +331,7 @@ class GraphBuilder(Generic[StateT, DepsT, GraphInputT, GraphOutputT]):
         )
 
     # Edge building
-    def add(self, *edges: EdgePath[StateT, DepsT]) -> None:  # noqa C901
+    def add(self, *edges: EdgePath[StateT, DepsT]) -> None:  # noqa: C901
         """Add one or more edge paths to the graph.
 
         This method processes edge paths and automatically creates any necessary
@@ -674,7 +676,7 @@ class GraphBuilder(Generic[StateT, DepsT, GraphInputT, GraphOutputT]):
         )
 
 
-def _validate_graph_structure(  # noqa C901
+def _validate_graph_structure(  # noqa: C901
     nodes: dict[NodeID, AnyNode],
     edges_by_source: dict[NodeID, list[Path]],
 ) -> None:

{pydantic_graph-1.13.0 → pydantic_graph-1.31.0}/pydantic_graph/beta/id_types.py

@@ -24,9 +24,6 @@ JoinID = NodeID
 ForkID = NodeID
 """Alias for NodeId when referring to fork nodes."""
 
-GraphRunID = NewType('GraphRunID', str)
-"""Unique identifier for a complete graph execution run."""
-
 TaskID = NewType('TaskID', str)
 """Unique identifier for a task within the graph execution."""
 

{pydantic_graph-1.13.0 → pydantic_graph-1.31.0}/pydantic_graph/beta/mermaid.py

@@ -49,7 +49,7 @@ class MermaidEdge:
     label: str | None
 
 
-def build_mermaid_graph(  # noqa C901
+def build_mermaid_graph(  # noqa: C901
     graph_nodes: dict[NodeID, AnyNode], graph_edges_by_source: dict[NodeID, list[Path]]
 ) -> MermaidGraph:
     """Build a mermaid graph."""