daimon 0.16.0

A Rust-native AI agent framework
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
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
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136

//! In-memory checkpoint backend (ephemeral, useful for testing).

use std::collections::HashMap;

use tokio::sync::Mutex;

use crate::error::Result;

use super::traits::Checkpoint;
use super::types::CheckpointState;

/// Stores checkpoints in memory. Lost on process exit.
///
/// Thread-safe via internal `Mutex`.
pub struct InMemoryCheckpoint {
    store: Mutex<HashMap<String, CheckpointState>>,
}

impl InMemoryCheckpoint {
    /// Creates an empty in-memory checkpoint store.
    pub fn new() -> Self {
        Self {
            store: Mutex::new(HashMap::new()),
        }
    }
}

impl Default for InMemoryCheckpoint {
    fn default() -> Self {
        Self::new()
    }
}

impl Checkpoint for InMemoryCheckpoint {
    async fn save(&self, state: &CheckpointState) -> Result<()> {
        let mut store = self.store.lock().await;
        store.insert(state.run_id.clone(), state.clone());
        Ok(())
    }

    async fn load(&self, run_id: &str) -> Result<Option<CheckpointState>> {
        let store = self.store.lock().await;
        Ok(store.get(run_id).cloned())
    }

    async fn list_runs(&self) -> Result<Vec<String>> {
        let store = self.store.lock().await;
        Ok(store.keys().cloned().collect())
    }

    async fn delete(&self, run_id: &str) -> Result<()> {
        let mut store = self.store.lock().await;
        store.remove(run_id);
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::model::types::Message;

    #[tokio::test]
    async fn test_save_and_load() {
        let cp = InMemoryCheckpoint::new();
        let state = CheckpointState::new(
            "run-1",
            vec![Message::user("hello"), Message::assistant("hi")],
            1,
        );

        cp.save(&state).await.unwrap();
        let loaded = cp.load("run-1").await.unwrap().unwrap();
        assert_eq!(loaded.run_id, "run-1");
        assert_eq!(loaded.iteration, 1);
        assert_eq!(loaded.messages.len(), 2);
    }

    #[tokio::test]
    async fn test_load_nonexistent() {
        let cp = InMemoryCheckpoint::new();
        assert!(cp.load("nope").await.unwrap().is_none());
    }

    #[tokio::test]
    async fn test_overwrite() {
        let cp = InMemoryCheckpoint::new();
        let s1 = CheckpointState::new("run-1", vec![Message::user("a")], 1);
        let s2 = CheckpointState::new("run-1", vec![Message::user("a"), Message::user("b")], 2);

        cp.save(&s1).await.unwrap();
        cp.save(&s2).await.unwrap();

        let loaded = cp.load("run-1").await.unwrap().unwrap();
        assert_eq!(loaded.iteration, 2);
        assert_eq!(loaded.messages.len(), 2);
    }

    #[tokio::test]
    async fn test_list_runs() {
        let cp = InMemoryCheckpoint::new();
        cp.save(&CheckpointState::new("a", vec![], 0))
            .await
            .unwrap();
        cp.save(&CheckpointState::new("b", vec![], 0))
            .await
            .unwrap();

        let mut runs = cp.list_runs().await.unwrap();
        runs.sort();
        assert_eq!(runs, vec!["a", "b"]);
    }

    #[tokio::test]
    async fn test_delete() {
        let cp = InMemoryCheckpoint::new();
        cp.save(&CheckpointState::new("run-1", vec![], 0))
            .await
            .unwrap();
        cp.delete("run-1").await.unwrap();
        assert!(cp.load("run-1").await.unwrap().is_none());
    }

    #[tokio::test]
    async fn test_completed_metadata() {
        let cp = InMemoryCheckpoint::new();
        let state = CheckpointState::new("run-1", vec![], 3)
            .mark_completed()
            .with_metadata("final_answer", serde_json::json!("42"));

        cp.save(&state).await.unwrap();
        let loaded = cp.load("run-1").await.unwrap().unwrap();
        assert!(loaded.completed);
        assert_eq!(loaded.metadata["final_answer"], serde_json::json!("42"));
    }
}