pydantic-graph 1.18.0__py3-none-any.whl → 1.20.0__py3-none-any.whl

pydantic_graph/beta/graph.py

@@ -8,8 +8,7 @@ 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
@@ -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,
@@ -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,6 +560,7 @@ 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
         self, first_task: GraphTask
@@ -782,12 +821,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 +860,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 +892,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:

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.18.0.dist-info → pydantic_graph-1.20.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pydantic-graph
-Version: 1.18.0
+Version: 1.20.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.18.0.dist-info → pydantic_graph-1.20.0.dist-info}/RECORD

@@ -7,9 +7,9 @@ pydantic_graph/nodes.py,sha256=CkY3lrC6jqZtzwhSRjFzmM69TdFFFrr58XSDU4THKHA,7450
 pydantic_graph/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 pydantic_graph/beta/__init__.py,sha256=VVmbEFaCSXYHwXqS4pANg4B3cn_c86tT62tW_EXcuyw,751
 pydantic_graph/beta/decision.py,sha256=x-Ta549b-j5hyBPUWFdwRQDRaJqnBHF1pfBP9L8I3vI,11239
-pydantic_graph/beta/graph.py,sha256=O8MRVyOLwk2HNq7Q9GSrLTMIjyGmkflQJ664AlIlW1E,41349
+pydantic_graph/beta/graph.py,sha256=CBCqWTJ1u7o0k2Rud193IJq2NffydEOJ5UMl0A_tYJs,42790
 pydantic_graph/beta/graph_builder.py,sha256=dCw4LePreagujGNtTdCVZfRVWkCs35MpoPEtAncLo5U,43326
-pydantic_graph/beta/id_types.py,sha256=mIhS3HYvmTWWfkGZmt5UEyedn5Ave424FyzUalP9nsU,2915
+pydantic_graph/beta/id_types.py,sha256=FZ3rYSubF6g_Ocv0faL3yJsy1lNN9AGZl9f_izvORUg,2814
 pydantic_graph/beta/join.py,sha256=rzCumDX_YgaU_a5bisfbjbbOuI3IwSZsCZs9TC0T9E4,8002
 pydantic_graph/beta/mermaid.py,sha256=Bj8a3CODPcojwT7BnrYqLBKTp0AbA1T3XsTmK2St3v4,7127
 pydantic_graph/beta/node.py,sha256=cTEGKiT3Lutg-PWxBbZDihpnBTVoPMSyCbfB50fjKeY,3071
@@ -22,7 +22,7 @@ pydantic_graph/persistence/__init__.py,sha256=NLBGvUWhem23EdMHHxtX0XgTS2vyixmuWt
 pydantic_graph/persistence/_utils.py,sha256=6ySxCc1lFz7bbLUwDLkoZWNqi8VNLBVU4xxJbKI23fQ,2264
 pydantic_graph/persistence/file.py,sha256=XZy295cGc86HfUl_KuB-e7cECZW3bubiEdyJMVQ1OD0,6906
 pydantic_graph/persistence/in_mem.py,sha256=MmahaVpdzmDB30Dm3ZfSCZBqgmx6vH4HXdBaWwVF0K0,6799
-pydantic_graph-1.18.0.dist-info/METADATA,sha256=SplmDrAtzc7Ir_55iKDdgKUnAXYOnnElO8o_8ri4O0U,3895
-pydantic_graph-1.18.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
-pydantic_graph-1.18.0.dist-info/licenses/LICENSE,sha256=vA6Jc482lEyBBuGUfD1pYx-cM7jxvLYOxPidZ30t_PQ,1100
-pydantic_graph-1.18.0.dist-info/RECORD,,
+pydantic_graph-1.20.0.dist-info/METADATA,sha256=Iob0hAz2T1R6EGWDsSesCp03KNFZh6E9RZgS24pweSA,3895
+pydantic_graph-1.20.0.dist-info/WHEEL,sha256=qtCwoSJWgHk21S1Kb4ihdzI2rlJ1ZKaIurTj_ngOhyQ,87
+pydantic_graph-1.20.0.dist-info/licenses/LICENSE,sha256=vA6Jc482lEyBBuGUfD1pYx-cM7jxvLYOxPidZ30t_PQ,1100
+pydantic_graph-1.20.0.dist-info/RECORD,,