langgraph-checkpoint 0.1.0

Checkpoint saver trait and in-memory implementation for durable graph execution.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50

use std::collections::{BTreeMap, HashMap};
use std::sync::RwLock;

use crate::traits::{CheckpointError, CheckpointSaver};
use crate::types::{Checkpoint, CheckpointId, ThreadId};

#[derive(Debug, Default)]
pub struct InMemorySaver {
    inner: RwLock<HashMap<ThreadId, BTreeMap<CheckpointId, Checkpoint>>>,
}

impl InMemorySaver {
    #[must_use]
    pub fn new() -> Self {
        Self::default()
    }
}

impl CheckpointSaver for InMemorySaver {
    fn put(&self, checkpoint: Checkpoint) -> Result<(), CheckpointError> {
        let mut guard = self.inner.write().map_err(|e| CheckpointError::Storage(e.to_string()))?;
        let entries = guard.entry(checkpoint.thread_id.clone()).or_insert_with(BTreeMap::new);
        if entries.contains_key(&checkpoint.checkpoint_id) {
            return Err(CheckpointError::Conflict(format!(
                "thread `{}` already has checkpoint `{}`",
                checkpoint.thread_id, checkpoint.checkpoint_id
            )));
        }
        entries.insert(checkpoint.checkpoint_id.clone(), checkpoint);
        Ok(())
    }

    fn get(
        &self,
        thread_id: &str,
        checkpoint_id: &str,
    ) -> Result<Option<Checkpoint>, CheckpointError> {
        let guard = self.inner.read().map_err(|e| CheckpointError::Storage(e.to_string()))?;
        Ok(guard.get(thread_id).and_then(|checkpoints| checkpoints.get(checkpoint_id).cloned()))
    }

    fn list(&self, thread_id: &str) -> Result<Vec<Checkpoint>, CheckpointError> {
        let guard = self.inner.read().map_err(|e| CheckpointError::Storage(e.to_string()))?;
        let result = guard
            .get(thread_id)
            .map(|checkpoints| checkpoints.values().cloned().collect())
            .unwrap_or_default();
        Ok(result)
    }
}