yoagent_state/
schema.rs

1use crate::{Node, PackId, Relation, StateError};
2use serde::{Deserialize, Serialize};
3use serde_json::Value as JsonValue;
4use std::collections::{BTreeMap, BTreeSet};
5
6#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
7pub struct Pack {
8    pub id: PackId,
9    pub name: String,
10    pub version: String,
11    pub object_types: BTreeMap<String, ObjectType>,
12    pub relation_types: BTreeMap<String, RelationType>,
13    pub policies: Vec<String>,
14    pub prompts: Vec<String>,
15    pub settings: JsonValue,
16}
17
18impl Pack {
19    pub fn new(id: PackId, name: impl Into<String>, version: impl Into<String>) -> Self {
20        Self {
21            id,
22            name: name.into(),
23            version: version.into(),
24            object_types: BTreeMap::new(),
25            relation_types: BTreeMap::new(),
26            policies: Vec::new(),
27            prompts: Vec::new(),
28            settings: JsonValue::Object(Default::default()),
29        }
30    }
31
32    pub fn add_object_type(mut self, object_type: ObjectType) -> Self {
33        self.object_types
34            .insert(object_type.kind.clone(), object_type);
35        self
36    }
37
38    pub fn add_relation_type(mut self, relation_type: RelationType) -> Self {
39        self.relation_types
40            .insert(relation_type.rel.clone(), relation_type);
41        self
42    }
43
44    pub fn validate_node(&self, node: &Node) -> Result<(), StateError> {
45        let Some(object_type) = self.object_types.get(&node.kind) else {
46            return Ok(());
47        };
48
49        for required in &object_type.required_props {
50            if node.props.get(required).is_none_or(|value| value.is_null()) {
51                return Err(StateError::Validation(format!(
52                    "node {} of kind {} is missing required prop {}",
53                    node.id, node.kind, required
54                )));
55            }
56        }
57
58        Ok(())
59    }
60
61    pub fn validate_relation(
62        &self,
63        relation: &Relation,
64        from: &Node,
65        to: &Node,
66    ) -> Result<(), StateError> {
67        let Some(relation_type) = self.relation_types.get(&relation.rel) else {
68            return Ok(());
69        };
70
71        if !relation_type.from_kinds.is_empty() && !relation_type.from_kinds.contains(&from.kind) {
72            return Err(StateError::Validation(format!(
73                "relation {} cannot start from kind {}",
74                relation.rel, from.kind
75            )));
76        }
77
78        if !relation_type.to_kinds.is_empty() && !relation_type.to_kinds.contains(&to.kind) {
79            return Err(StateError::Validation(format!(
80                "relation {} cannot point to kind {}",
81                relation.rel, to.kind
82            )));
83        }
84
85        Ok(())
86    }
87}
88
89#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
90pub struct ObjectType {
91    pub kind: String,
92    pub required_props: BTreeSet<String>,
93    pub description: Option<String>,
94}
95
96impl ObjectType {
97    pub fn new(kind: impl Into<String>) -> Self {
98        Self {
99            kind: kind.into(),
100            required_props: BTreeSet::new(),
101            description: None,
102        }
103    }
104
105    pub fn require(mut self, prop: impl Into<String>) -> Self {
106        self.required_props.insert(prop.into());
107        self
108    }
109}
110
111#[derive(Debug, Clone, PartialEq, Eq, Serialize, Deserialize)]
112pub struct RelationType {
113    pub rel: String,
114    pub from_kinds: BTreeSet<String>,
115    pub to_kinds: BTreeSet<String>,
116    pub description: Option<String>,
117}
118
119impl RelationType {
120    pub fn new(rel: impl Into<String>) -> Self {
121        Self {
122            rel: rel.into(),
123            from_kinds: BTreeSet::new(),
124            to_kinds: BTreeSet::new(),
125            description: None,
126        }
127    }
128
129    pub fn from_kind(mut self, kind: impl Into<String>) -> Self {
130        self.from_kinds.insert(kind.into());
131        self
132    }
133
134    pub fn to_kind(mut self, kind: impl Into<String>) -> Self {
135        self.to_kinds.insert(kind.into());
136        self
137    }
138}
yoagent_state/schema.rs

yoagent_state/
schema.rs
