4445464748495051525354555657585960616263646566

@dataclassclassNodeSnapshot(Generic[StateT,RunEndT]):"""History step describing the execution of a node in a graph."""state:StateT"""The state of the graph before the node is run."""node:Annotated[BaseNode[StateT,Any,RunEndT],_utils.CustomNodeSchema()]"""The node to run next."""start_ts:datetime|None=None"""The timestamp when the node started running, `None` until the run starts."""duration:float|None=None"""The duration of the node run in seconds, if the node has been run."""status:SnapshotStatus='created'"""The status of the snapshot."""kind:Literal['node']='node'"""The kind of history step, can be used as a discriminator when deserializing history."""id:str=UNSET_SNAPSHOT_ID"""Unique ID of the snapshot."""def__post_init__(self)->None:ifself.id==UNSET_SNAPSHOT_ID:self.id=self.node.get_snapshot_id()

697071727374757677787980818283848586878889909192939495

@dataclassclassEndSnapshot(Generic[StateT,RunEndT]):"""History step describing the end of a graph run."""state:StateT"""The state of the graph at the end of the run."""result:End[RunEndT]"""The result of the graph run."""ts:datetime=field(default_factory=_utils.now_utc)"""The timestamp when the graph run ended."""kind:Literal['end']='end'"""The kind of history step, can be used as a discriminator when deserializing history."""id:str=UNSET_SNAPSHOT_ID"""Unique ID of the snapshot."""def__post_init__(self)->None:ifself.id==UNSET_SNAPSHOT_ID:self.id=self.node.get_snapshot_id()@propertydefnode(self)->End[RunEndT]:"""Shim to get the [`result`][pydantic_graph.persistence.EndSnapshot.result].        Useful to allow `[snapshot.node for snapshot in persistence.history]`.        """returnself.result

106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225

classBaseStatePersistence(ABC,Generic[StateT,RunEndT]):"""Abstract base class for storing the state of a graph run.    Each instance of a `BaseStatePersistence` subclass should be used for a single graph run.    """@abstractmethodasyncdefsnapshot_node(self,state:StateT,next_node:BaseNode[StateT,Any,RunEndT])->None:"""Snapshot the state of a graph, when the next step is to run a node.        This method should add a [`NodeSnapshot`][pydantic_graph.persistence.NodeSnapshot] to persistence.        Args:            state: The state of the graph.            next_node: The next node to run.        """raiseNotImplementedError@abstractmethodasyncdefsnapshot_node_if_new(self,snapshot_id:str,state:StateT,next_node:BaseNode[StateT,Any,RunEndT])->None:"""Snapshot the state of a graph if the snapshot ID doesn't already exist in persistence.        This method will generally call [`snapshot_node`][pydantic_graph.persistence.BaseStatePersistence.snapshot_node]        but should do so in an atomic way.        Args:            snapshot_id: The ID of the snapshot to check.            state: The state of the graph.            next_node: The next node to run.        """raiseNotImplementedError@abstractmethodasyncdefsnapshot_end(self,state:StateT,end:End[RunEndT])->None:"""Snapshot the state of a graph when the graph has ended.        This method should add an [`EndSnapshot`][pydantic_graph.persistence.EndSnapshot] to persistence.        Args:            state: The state of the graph.            end: data from the end of the run.        """raiseNotImplementedError@abstractmethoddefrecord_run(self,snapshot_id:str)->AbstractAsyncContextManager[None]:"""Record the run of the node, or error if the node is already running.        Args:            snapshot_id: The ID of the snapshot to record.        Raises:            GraphNodeRunningError: if the node status it not `'created'` or `'pending'`.            LookupError: if the snapshot ID is not found in persistence.        Returns:            An async context manager that records the run of the node.        In particular this should set:        - [`NodeSnapshot.status`][pydantic_graph.persistence.NodeSnapshot.status] to `'running'` and          [`NodeSnapshot.start_ts`][pydantic_graph.persistence.NodeSnapshot.start_ts] when the run starts.        - [`NodeSnapshot.status`][pydantic_graph.persistence.NodeSnapshot.status] to `'success'` or `'error'` and          [`NodeSnapshot.duration`][pydantic_graph.persistence.NodeSnapshot.duration] when the run finishes.        """raiseNotImplementedError@abstractmethodasyncdefload_next(self)->NodeSnapshot[StateT,RunEndT]|None:"""Retrieve a node snapshot with status `'created`' and set its status to `'pending'`.        This is used by [`Graph.iter_from_persistence`][pydantic_graph.graph.Graph.iter_from_persistence]        to get the next node to run.        Returns: The snapshot, or `None` if no snapshot with status `'created`' exists.        """raiseNotImplementedError@abstractmethodasyncdefload_all(self)->list[Snapshot[StateT,RunEndT]]:"""Load the entire history of snapshots.        `load_all` is not used by pydantic-graph itself, instead it's provided to make it convenient to        get all [snapshots][pydantic_graph.persistence.Snapshot] from persistence.        Returns: The list of snapshots.        """raiseNotImplementedErrordefset_graph_types(self,graph:Graph[StateT,Any,RunEndT])->None:"""Set the types of the state and run end from a graph.        You generally won't need to customise this method, instead implement        [`set_types`][pydantic_graph.persistence.BaseStatePersistence.set_types] and        [`should_set_types`][pydantic_graph.persistence.BaseStatePersistence.should_set_types].        """ifself.should_set_types():with_utils.set_nodes_type_context(graph.get_nodes()):self.set_types(*graph.inferred_types)defshould_set_types(self)->bool:"""Whether types need to be set.        Implementations should override this method to return `True` when types have not been set if they are needed.        """returnFalsedefset_types(self,state_type:type[StateT],run_end_type:type[RunEndT])->None:"""Set the types of the state and run end.        This can be used to create [type adapters][pydantic.TypeAdapter] for serializing and deserializing snapshots,        e.g. with [`build_snapshot_list_type_adapter`][pydantic_graph.persistence.build_snapshot_list_type_adapter].        Args:            state_type: The state type.            run_end_type: The run end type.        """pass

228229230231232233234235236237238

defbuild_snapshot_list_type_adapter(state_t:type[StateT],run_end_t:type[RunEndT])->pydantic.TypeAdapter[list[Snapshot[StateT,RunEndT]]]:"""Build a type adapter for a list of snapshots.    This method should be called from within    [`set_types`][pydantic_graph.persistence.BaseStatePersistence.set_types]    where context variables will be set such that Pydantic can create a schema for    [`NodeSnapshot.node`][pydantic_graph.persistence.NodeSnapshot.node].    """returnpydantic.TypeAdapter(list[Annotated[Snapshot[state_t,run_end_t],pydantic.Discriminator('kind')]])

31323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182

@dataclassclassSimpleStatePersistence(BaseStatePersistence[StateT,RunEndT]):"""Simple in memory state persistence that just hold the latest snapshot.    If no state persistence implementation is provided when running a graph, this is used by default.    """last_snapshot:Snapshot[StateT,RunEndT]|None=None"""The last snapshot."""asyncdefsnapshot_node(self,state:StateT,next_node:BaseNode[StateT,Any,RunEndT])->None:self.last_snapshot=NodeSnapshot(state=state,node=next_node)asyncdefsnapshot_node_if_new(self,snapshot_id:str,state:StateT,next_node:BaseNode[StateT,Any,RunEndT])->None:ifself.last_snapshotandself.last_snapshot.id==snapshot_id:return# pragma: no coverelse:awaitself.snapshot_node(state,next_node)asyncdefsnapshot_end(self,state:StateT,end:End[RunEndT])->None:self.last_snapshot=EndSnapshot(state=state,result=end)@asynccontextmanagerasyncdefrecord_run(self,snapshot_id:str)->AsyncIterator[None]:ifself.last_snapshotisNoneorsnapshot_id!=self.last_snapshot.id:raiseLookupError(f'No snapshot found with id={snapshot_id!r}')assertisinstance(self.last_snapshot,NodeSnapshot),'Only NodeSnapshot can be recorded'exceptions.GraphNodeStatusError.check(self.last_snapshot.status)self.last_snapshot.status='running'self.last_snapshot.start_ts=_utils.now_utc()start=perf_counter()try:yieldexceptException:self.last_snapshot.duration=perf_counter()-startself.last_snapshot.status='error'raiseelse:self.last_snapshot.duration=perf_counter()-startself.last_snapshot.status='success'asyncdefload_next(self)->NodeSnapshot[StateT,RunEndT]|None:ifisinstance(self.last_snapshot,NodeSnapshot)andself.last_snapshot.status=='created':self.last_snapshot.status='pending'returnself.last_snapshotasyncdefload_all(self)->list[Snapshot[StateT,RunEndT]]:raiseNotImplementedError('load is not supported for SimpleStatePersistence')

 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172

@dataclassclassFullStatePersistence(BaseStatePersistence[StateT,RunEndT]):"""In memory state persistence that hold a list of snapshots."""deep_copy:bool=True"""Whether to deep copy the state and nodes when storing them.    Defaults to `True` so even if nodes or state are modified after the snapshot is taken,    the persistence history will record the value at the time of the snapshot.    """history:list[Snapshot[StateT,RunEndT]]=field(default_factory=list)"""List of snapshots taken during the graph run."""_snapshots_type_adapter:pydantic.TypeAdapter[list[Snapshot[StateT,RunEndT]]]|None=field(default=None,init=False,repr=False)asyncdefsnapshot_node(self,state:StateT,next_node:BaseNode[StateT,Any,RunEndT])->None:snapshot=NodeSnapshot(state=self._prep_state(state),node=next_node.deep_copy()ifself.deep_copyelsenext_node,)self.history.append(snapshot)asyncdefsnapshot_node_if_new(self,snapshot_id:str,state:StateT,next_node:BaseNode[StateT,Any,RunEndT])->None:ifnotany(s.id==snapshot_idforsinself.history):awaitself.snapshot_node(state,next_node)asyncdefsnapshot_end(self,state:StateT,end:End[RunEndT])->None:snapshot=EndSnapshot(state=self._prep_state(state),result=end.deep_copy_data()ifself.deep_copyelseend,)self.history.append(snapshot)@asynccontextmanagerasyncdefrecord_run(self,snapshot_id:str)->AsyncIterator[None]:try:snapshot=next(sforsinself.historyifs.id==snapshot_id)exceptStopIterationase:raiseLookupError(f'No snapshot found with id={snapshot_id!r}')fromeassertisinstance(snapshot,NodeSnapshot),'Only NodeSnapshot can be recorded'exceptions.GraphNodeStatusError.check(snapshot.status)snapshot.status='running'snapshot.start_ts=_utils.now_utc()start=perf_counter()try:yieldexceptException:snapshot.duration=perf_counter()-startsnapshot.status='error'raiseelse:snapshot.duration=perf_counter()-startsnapshot.status='success'asyncdefload_next(self)->NodeSnapshot[StateT,RunEndT]|None:ifsnapshot:=next((sforsinself.historyifisinstance(s,NodeSnapshot)ands.status=='created'),None):snapshot.status='pending'returnsnapshotasyncdefload_all(self)->list[Snapshot[StateT,RunEndT]]:returnself.historydefshould_set_types(self)->bool:returnself._snapshots_type_adapterisNonedefset_types(self,state_type:type[StateT],run_end_type:type[RunEndT])->None:self._snapshots_type_adapter=build_snapshot_list_type_adapter(state_type,run_end_type)defdump_json(self,*,indent:int|None=None)->bytes:"""Dump the history to JSON bytes."""assertself._snapshots_type_adapterisnotNone,'type adapter must be set to use `dump_json`'returnself._snapshots_type_adapter.dump_json(self.history,indent=indent)defload_json(self,json_data:str|bytes|bytearray)->None:"""Load the history from JSON."""assertself._snapshots_type_adapterisnotNone,'type adapter must be set to use `load_json`'self.history=self._snapshots_type_adapter.validate_json(json_data)def_prep_state(self,state:StateT)->StateT:"""Prepare state for snapshot, uses [`copy.deepcopy`][copy.deepcopy] by default."""ifnotself.deep_copyorstateisNone:returnstateelse:returncopy.deepcopy(state)

 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162

@dataclassclassFileStatePersistence(BaseStatePersistence[StateT,RunEndT]):"""File based state persistence that hold graph run state in a JSON file."""json_file:Path"""Path to the JSON file where the snapshots are stored.    You should use a different file for each graph run, but a single file should be reused for multiple    steps of the same run.    For example if you have a run ID of the form `run_123abc`, you might create a `FileStatePersistence` thus:    ```py    from pathlib import Path    from pydantic_graph import FullStatePersistence    run_id = 'run_123abc'    persistence = FullStatePersistence(Path('runs') / f'{run_id}.json')    ```    """_snapshots_type_adapter:pydantic.TypeAdapter[list[Snapshot[StateT,RunEndT]]]|None=field(default=None,init=False,repr=False)asyncdefsnapshot_node(self,state:StateT,next_node:BaseNode[StateT,Any,RunEndT])->None:awaitself._append_save(NodeSnapshot(state=state,node=next_node))asyncdefsnapshot_node_if_new(self,snapshot_id:str,state:StateT,next_node:BaseNode[StateT,Any,RunEndT])->None:asyncwithself._lock():snapshots=awaitself.load_all()ifnotany(s.id==snapshot_idforsinsnapshots):# pragma: no branchawaitself._append_save(NodeSnapshot(state=state,node=next_node),lock=False)asyncdefsnapshot_end(self,state:StateT,end:End[RunEndT])->None:awaitself._append_save(EndSnapshot(state=state,result=end))@asynccontextmanagerasyncdefrecord_run(self,snapshot_id:str)->AsyncIterator[None]:asyncwithself._lock():snapshots=awaitself.load_all()try:snapshot=next(sforsinsnapshotsifs.id==snapshot_id)exceptStopIterationase:raiseLookupError(f'No snapshot found with id={snapshot_id!r}')fromeassertisinstance(snapshot,NodeSnapshot),'Only NodeSnapshot can be recorded'exceptions.GraphNodeStatusError.check(snapshot.status)snapshot.status='running'snapshot.start_ts=_utils.now_utc()awaitself._save(snapshots)start=perf_counter()try:yieldexceptException:duration=perf_counter()-startasyncwithself._lock():await_graph_utils.run_in_executor(self._after_run_sync,snapshot_id,duration,'error')raiseelse:snapshot.duration=perf_counter()-startasyncwithself._lock():await_graph_utils.run_in_executor(self._after_run_sync,snapshot_id,snapshot.duration,'success')asyncdefload_next(self)->NodeSnapshot[StateT,RunEndT]|None:asyncwithself._lock():snapshots=awaitself.load_all()ifsnapshot:=next((sforsinsnapshotsifisinstance(s,NodeSnapshot)ands.status=='created'),None):snapshot.status='pending'awaitself._save(snapshots)returnsnapshotdefshould_set_types(self)->bool:"""Whether types need to be set."""returnself._snapshots_type_adapterisNonedefset_types(self,state_type:type[StateT],run_end_type:type[RunEndT])->None:self._snapshots_type_adapter=build_snapshot_list_type_adapter(state_type,run_end_type)asyncdefload_all(self)->list[Snapshot[StateT,RunEndT]]:returnawait_graph_utils.run_in_executor(self._load_sync)def_load_sync(self)->list[Snapshot[StateT,RunEndT]]:assertself._snapshots_type_adapterisnotNone,'snapshots type adapter must be set'try:content=self.json_file.read_bytes()exceptFileNotFoundError:return[]else:returnself._snapshots_type_adapter.validate_json(content)def_after_run_sync(self,snapshot_id:str,duration:float,status:SnapshotStatus)->None:snapshots=self._load_sync()snapshot=next(sforsinsnapshotsifs.id==snapshot_id)assertisinstance(snapshot,NodeSnapshot),'Only NodeSnapshot can be recorded'snapshot.duration=durationsnapshot.status=statusself._save_sync(snapshots)asyncdef_save(self,snapshots:list[Snapshot[StateT,RunEndT]])->None:await_graph_utils.run_in_executor(self._save_sync,snapshots)def_save_sync(self,snapshots:list[Snapshot[StateT,RunEndT]])->None:assertself._snapshots_type_adapterisnotNone,'snapshots type adapter must be set'self.json_file.write_bytes(self._snapshots_type_adapter.dump_json(snapshots,indent=2))asyncdef_append_save(self,snapshot:Snapshot[StateT,RunEndT],*,lock:bool=True)->None:assertself._snapshots_type_adapterisnotNone,'snapshots type adapter must be set'asyncwithAsyncExitStack()asstack:iflock:awaitstack.enter_async_context(self._lock())snapshots=awaitself.load_all()snapshots.append(snapshot)awaitself._save(snapshots)@asynccontextmanagerasyncdef_lock(self,*,timeout:float=1.0)->AsyncIterator[None]:"""Lock a file by checking and writing a `.pydantic-graph-persistence-lock` to it.        Args:            timeout: how long to wait for the lock        Returns: an async context manager that holds the lock        """lock_file=self.json_file.parent/f'{self.json_file.name}.pydantic-graph-persistence-lock'lock_id=secrets.token_urlsafe().encode()awaitasyncio.wait_for(_get_lock(lock_file,lock_id),timeout=timeout)try:yieldfinally:await_graph_utils.run_in_executor(lock_file.unlink,missing_ok=True)