use serde::{de::Error as DeError, Deserialize, Deserializer, Serialize};
use std::collections::HashMap;
#[derive(Debug, Clone, PartialEq, Default, Serialize, Deserialize)]
pub struct StateSchema {
#[serde(default)]
pub fields: Vec<StateField>,
}
impl StateSchema {
pub fn new() -> Self {
Self::default()
}
pub fn add_field(&mut self, field: StateField) {
self.fields.push(field);
}
pub fn get_field(&self, name: &str) -> Option<&StateField> {
self.fields.iter().find(|f| f.name == name)
}
pub fn has_field(&self, name: &str) -> bool {
self.get_field(name).is_some()
}
pub fn field_names(&self) -> Vec<&str> {
self.fields.iter().map(|f| f.name.as_str()).collect()
}
pub fn validate_required(&self, value: &serde_yaml::Value) -> crate::errors::Result<()> {
if let serde_yaml::Value::Mapping(map) = value {
for field in &self.fields {
if field.default.is_none() {
let key = serde_yaml::Value::String(field.name.clone());
if !map.contains_key(&key) {
return Err(crate::errors::SelfwareError::Internal(format!(
"Required field '{}' is missing",
field.name
)));
}
}
}
}
Ok(())
}
pub fn apply_defaults(&self, value: &mut serde_yaml::Value) {
if let serde_yaml::Value::Mapping(map) = value {
for field in &self.fields {
let key = serde_yaml::Value::String(field.name.clone());
if !map.contains_key(&key) {
if let Some(ref default) = field.default {
map.insert(key, default.clone());
}
}
}
}
}
}
#[derive(Debug, Clone, PartialEq, Serialize)]
pub struct StateField {
pub name: String,
#[serde(rename = "type")]
pub field_type: FieldType,
#[serde(default)]
pub default: Option<serde_yaml::Value>,
#[serde(default)]
pub description: Option<String>,
}
impl<'de> Deserialize<'de> for StateField {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
#[derive(Deserialize)]
struct RawStateField {
name: String,
#[serde(rename = "type")]
field_type: FieldType,
#[serde(default)]
element_type: Option<FieldType>,
#[serde(default)]
default: Option<serde_yaml::Value>,
#[serde(default)]
description: Option<String>,
}
let raw = RawStateField::deserialize(deserializer)?;
let field_type = match (raw.field_type, raw.element_type) {
(FieldType::Array(_), Some(element_type)) => FieldType::Array(Box::new(element_type)),
(FieldType::Array(_), None) => FieldType::Array(Box::new(FieldType::String)),
(FieldType::Object(_), _) => FieldType::Object(HashMap::new()),
(field_type, _) => field_type,
};
Ok(Self {
name: raw.name,
field_type,
default: raw.default,
description: raw.description,
})
}
}
#[derive(Debug, Clone, PartialEq, Default, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum FieldType {
#[default]
String,
Integer,
Float,
Boolean,
Array(Box<FieldType>),
Object(HashMap<String, FieldType>),
AgentRef,
ToolRef,
}
impl<'de> Deserialize<'de> for FieldType {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
let raw = String::deserialize(deserializer)?;
match raw.as_str() {
"string" => Ok(Self::String),
"integer" => Ok(Self::Integer),
"float" => Ok(Self::Float),
"boolean" => Ok(Self::Boolean),
"array" => Ok(Self::Array(Box::new(Self::String))),
"object" => Ok(Self::Object(HashMap::new())),
"agent_ref" => Ok(Self::AgentRef),
"tool_ref" => Ok(Self::ToolRef),
other => Err(D::Error::unknown_variant(
other,
&[
"string",
"integer",
"float",
"boolean",
"array",
"object",
"agent_ref",
"tool_ref",
],
)),
}
}
}
impl FieldType {
pub fn type_name(&self) -> &'static str {
match self {
FieldType::String => "string",
FieldType::Integer => "integer",
FieldType::Float => "float",
FieldType::Boolean => "boolean",
FieldType::Array(_) => "array",
FieldType::Object(_) => "object",
FieldType::AgentRef => "agent_ref",
FieldType::ToolRef => "tool_ref",
}
}
pub fn is_scalar(&self) -> bool {
matches!(
self,
FieldType::String | FieldType::Integer | FieldType::Float | FieldType::Boolean
)
}
pub fn is_collection(&self) -> bool {
matches!(self, FieldType::Array(_) | FieldType::Object(_))
}
pub fn is_reference(&self) -> bool {
matches!(self, FieldType::AgentRef | FieldType::ToolRef)
}
}
#[derive(Debug, Default)]
pub struct SchemaBuilder {
fields: Vec<StateField>,
}
impl SchemaBuilder {
pub fn new() -> Self {
Self::default()
}
pub fn string(mut self, name: &str) -> Self {
self.fields.push(StateField {
name: name.to_string(),
field_type: FieldType::String,
default: None,
description: None,
});
self
}
pub fn string_with_default(mut self, name: &str, default: &str) -> Self {
self.fields.push(StateField {
name: name.to_string(),
field_type: FieldType::String,
default: Some(serde_yaml::Value::String(default.to_string())),
description: None,
});
self
}
pub fn integer(mut self, name: &str) -> Self {
self.fields.push(StateField {
name: name.to_string(),
field_type: FieldType::Integer,
default: None,
description: None,
});
self
}
pub fn integer_with_default(mut self, name: &str, default: i64) -> Self {
self.fields.push(StateField {
name: name.to_string(),
field_type: FieldType::Integer,
default: Some(serde_yaml::Value::Number(default.into())),
description: None,
});
self
}
pub fn float(mut self, name: &str) -> Self {
self.fields.push(StateField {
name: name.to_string(),
field_type: FieldType::Float,
default: None,
description: None,
});
self
}
pub fn boolean(mut self, name: &str) -> Self {
self.fields.push(StateField {
name: name.to_string(),
field_type: FieldType::Boolean,
default: None,
description: None,
});
self
}
pub fn boolean_with_default(mut self, name: &str, default: bool) -> Self {
self.fields.push(StateField {
name: name.to_string(),
field_type: FieldType::Boolean,
default: Some(serde_yaml::Value::Bool(default)),
description: None,
});
self
}
pub fn array(mut self, name: &str, element_type: FieldType) -> Self {
self.fields.push(StateField {
name: name.to_string(),
field_type: FieldType::Array(Box::new(element_type)),
default: Some(serde_yaml::Value::Sequence(vec![])),
description: None,
});
self
}
pub fn agent_ref(mut self, name: &str) -> Self {
self.fields.push(StateField {
name: name.to_string(),
field_type: FieldType::AgentRef,
default: None,
description: None,
});
self
}
pub fn tool_ref(mut self, name: &str) -> Self {
self.fields.push(StateField {
name: name.to_string(),
field_type: FieldType::ToolRef,
default: None,
description: None,
});
self
}
pub fn with_description(mut self, description: &str) -> Self {
if let Some(field) = self.fields.last_mut() {
field.description = Some(description.to_string());
}
self
}
pub fn build(self) -> StateSchema {
StateSchema {
fields: self.fields,
}
}
}
#[cfg(test)]
#[path = "../../../tests/unit/swl/types/schema/schema_test.rs"]
mod tests;