awaken-contract 0.4.0

Core types, traits, and state model for the Awaken AI agent runtime
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

use serde::{Deserialize, Serialize, de::DeserializeOwned};

use crate::error::StateError;
use crate::state::StateKey;

use super::{JsonValue, Phase, decode_json, encode_json};

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ScheduledAction {
    pub phase: Phase,
    pub key: String,
    pub payload: JsonValue,
}

impl ScheduledAction {
    pub fn new(phase: Phase, key: impl Into<String>, payload: JsonValue) -> Self {
        Self {
            phase,
            key: key.into(),
            payload,
        }
    }
}

pub trait ScheduledActionSpec: 'static + Send + Sync {
    const KEY: &'static str;
    const PHASE: Phase;

    type Payload: Serialize + DeserializeOwned + Send + Sync + 'static;

    fn encode_payload(payload: &Self::Payload) -> Result<JsonValue, StateError> {
        encode_json(Self::KEY, payload)
    }

    fn decode_payload(payload: JsonValue) -> Result<Self::Payload, StateError> {
        decode_json(Self::KEY, payload)
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct ScheduledActionEnvelope {
    pub id: u64,
    pub action: ScheduledAction,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
pub struct FailedScheduledAction {
    pub id: u64,
    pub action: ScheduledAction,
    pub error: String,
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "op", rename_all = "snake_case")]
pub enum ScheduledActionQueueUpdate {
    Push(ScheduledActionEnvelope),
    Remove { id: u64 },
}

pub struct PendingScheduledActions;

impl StateKey for PendingScheduledActions {
    const KEY: &'static str = "__runtime.pending_scheduled_actions";

    type Value = Vec<ScheduledActionEnvelope>;
    type Update = ScheduledActionQueueUpdate;

    fn apply(value: &mut Self::Value, update: Self::Update) {
        match update {
            ScheduledActionQueueUpdate::Push(entry) => value.push(entry),
            ScheduledActionQueueUpdate::Remove { id } => value.retain(|entry| entry.id != id),
        }
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "op", rename_all = "snake_case")]
pub enum FailedScheduledActionUpdate {
    Push(FailedScheduledAction),
    Remove { id: u64 },
}

pub struct FailedScheduledActions;

impl StateKey for FailedScheduledActions {
    const KEY: &'static str = "__runtime.failed_scheduled_actions";

    type Value = Vec<FailedScheduledAction>;
    type Update = FailedScheduledActionUpdate;

    fn apply(value: &mut Self::Value, update: Self::Update) {
        match update {
            FailedScheduledActionUpdate::Push(entry) => value.push(entry),
            FailedScheduledActionUpdate::Remove { id } => value.retain(|entry| entry.id != id),
        }
    }
}

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

    struct TestAction;

    impl ScheduledActionSpec for TestAction {
        const KEY: &'static str = "test.action";
        const PHASE: Phase = Phase::BeforeInference;
        type Payload = String;
    }

    #[test]
    fn scheduled_action_spec_round_trip_works() {
        let payload = "hello".to_string();
        let encoded = TestAction::encode_payload(&payload).expect("encode should succeed");
        let decoded = TestAction::decode_payload(encoded).expect("decode should succeed");

        assert_eq!(decoded, payload);
    }
}