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
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152

//! Redis checkpoint backend.
//!
//! Stores checkpoints in Redis hashes, providing fast, shared checkpoint
//! storage accessible from multiple processes. Enable with `feature = "redis"`.
//!
//! ```ignore
//! use daimon::checkpoint::RedisCheckpoint;
//!
//! let cp = RedisCheckpoint::new("redis://127.0.0.1/", "daimon:checkpoints").await?;
//! cp.save(&state).await?;
//! ```

use crate::error::{DaimonError, Result};

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

/// Checkpoint backend using Redis.
///
/// Each checkpoint is stored as a field in a Redis hash (`{prefix}:data`),
/// keyed by run ID with the `CheckpointState` serialized as JSON.
/// A secondary hash (`{prefix}:runs`) tracks all known run IDs.
pub struct RedisCheckpoint {
    client: redis::Client,
    prefix: String,
}

impl RedisCheckpoint {
    /// Connects to Redis and creates a new checkpoint backend.
    ///
    /// * `url` — Redis connection URL (e.g. `redis://127.0.0.1/`)
    /// * `prefix` — key prefix for all Redis keys (e.g. `daimon:checkpoints`)
    pub async fn new(url: &str, prefix: impl Into<String>) -> Result<Self> {
        let client = redis::Client::open(url)
            .map_err(|e| DaimonError::Other(format!("redis checkpoint connection: {e}")))?;

        let mut conn = client
            .get_multiplexed_async_connection()
            .await
            .map_err(|e| DaimonError::Other(format!("redis checkpoint connect: {e}")))?;

        redis::cmd("PING")
            .query_async::<String>(&mut conn)
            .await
            .map_err(|e| DaimonError::Other(format!("redis checkpoint ping: {e}")))?;

        Ok(Self {
            client,
            prefix: prefix.into(),
        })
    }

    fn data_key(&self) -> String {
        format!("{}:data", self.prefix)
    }

    async fn conn(&self) -> Result<redis::aio::MultiplexedConnection> {
        self.client
            .get_multiplexed_async_connection()
            .await
            .map_err(|e| DaimonError::Other(format!("redis checkpoint conn: {e}")))
    }
}

impl Checkpoint for RedisCheckpoint {
    async fn save(&self, state: &CheckpointState) -> Result<()> {
        use redis::AsyncCommands;

        let json = serde_json::to_string(state)?;
        let mut conn = self.conn().await?;

        conn.hset::<_, _, _, ()>(&self.data_key(), &state.run_id, &json)
            .await
            .map_err(|e| DaimonError::Other(format!("redis checkpoint save: {e}")))?;

        Ok(())
    }

    async fn load(&self, run_id: &str) -> Result<Option<CheckpointState>> {
        use redis::AsyncCommands;

        let mut conn = self.conn().await?;

        let json: Option<String> = conn
            .hget(&self.data_key(), run_id)
            .await
            .map_err(|e| DaimonError::Other(format!("redis checkpoint load: {e}")))?;

        match json {
            Some(j) => {
                let state: CheckpointState = serde_json::from_str(&j)
                    .map_err(|e| DaimonError::Other(format!("redis checkpoint deserialize: {e}")))?;
                Ok(Some(state))
            }
            None => Ok(None),
        }
    }

    async fn list_runs(&self) -> Result<Vec<String>> {
        use redis::AsyncCommands;

        let mut conn = self.conn().await?;

        let keys: Vec<String> = conn
            .hkeys(&self.data_key())
            .await
            .map_err(|e| DaimonError::Other(format!("redis checkpoint list: {e}")))?;

        Ok(keys)
    }

    async fn delete(&self, run_id: &str) -> Result<()> {
        use redis::AsyncCommands;

        let mut conn = self.conn().await?;

        conn.hdel::<_, _, ()>(&self.data_key(), run_id)
            .await
            .map_err(|e| DaimonError::Other(format!("redis checkpoint delete: {e}")))?;

        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_key_format() {
        let prefix = "daimon:cp";
        assert_eq!(format!("{prefix}:data"), "daimon:cp:data");
    }

    #[test]
    fn test_state_serialization_roundtrip() {
        use crate::model::types::Message;

        let state = CheckpointState::new(
            "run-redis",
            vec![Message::user("test")],
            3,
        )
        .with_metadata("model", serde_json::json!("gpt-4o"));

        let json = serde_json::to_string(&state).unwrap();
        let deser: CheckpointState = serde_json::from_str(&json).unwrap();
        assert_eq!(deser.run_id, "run-redis");
        assert_eq!(deser.iteration, 3);
        assert_eq!(deser.metadata["model"], serde_json::json!("gpt-4o"));
    }
}