agent_stream_kit/
spec.rs

1use std::ops::Not;
2
3use serde::{Deserialize, Serialize};
4use serde_json::Value;
5
6use crate::FnvIndexMap;
7use crate::config::AgentConfigs;
8use crate::definition::{AgentConfigSpecs, AgentDefinition};
9use crate::error::AgentError;
10use crate::id::new_id;
11
12pub type AgentStreamSpecs = FnvIndexMap<String, AgentStreamSpec>;
13
14#[derive(Clone, Debug, Default, Deserialize, Serialize)]
15pub struct AgentStreamSpec {
16    pub agents: im::Vector<AgentSpec>,
17
18    pub channels: im::Vector<ChannelSpec>,
19
20    #[serde(default, skip_serializing_if = "<&bool>::not")]
21    pub run_on_start: bool,
22
23    #[serde(flatten)]
24    pub extensions: im::HashMap<String, Value>,
25}
26
27impl AgentStreamSpec {
28    pub fn add_agent(&mut self, agent: AgentSpec) {
29        self.agents.push_back(agent);
30    }
31
32    pub fn remove_agent(&mut self, agent_id: &str) {
33        self.agents.retain(|agent| agent.id != agent_id);
34    }
35
36    pub fn add_channels(&mut self, channel: ChannelSpec) {
37        self.channels.push_back(channel);
38    }
39
40    pub fn remove_channel(&mut self, channel_id: &str) -> Option<ChannelSpec> {
41        if let Some(channel) = self
42            .channels
43            .iter()
44            .find(|channel| channel.id == channel_id)
45            .cloned()
46        {
47            self.channels.retain(|e| e.id != channel_id);
48            Some(channel)
49        } else {
50            None
51        }
52    }
53
54    pub fn to_json(&self) -> Result<String, AgentError> {
55        let json = serde_json::to_string_pretty(self)
56            .map_err(|e| AgentError::SerializationError(e.to_string()))?;
57        Ok(json)
58    }
59
60    pub fn from_json(json_str: &str) -> Result<Self, AgentError> {
61        let stream: AgentStreamSpec = serde_json::from_str(json_str)
62            .map_err(|e| AgentError::SerializationError(e.to_string()))?;
63        Ok(stream)
64    }
65}
66
67pub fn copy_sub_stream(
68    agents: &Vec<AgentSpec>,
69    channels: &Vec<ChannelSpec>,
70) -> (Vec<AgentSpec>, Vec<ChannelSpec>) {
71    let mut new_agents = Vec::new();
72    let mut agent_id_map = FnvIndexMap::default();
73    for agent in agents {
74        let new_id = new_id();
75        agent_id_map.insert(agent.id.clone(), new_id.clone());
76        let mut new_agent = agent.clone();
77        new_agent.id = new_id;
78        new_agents.push(new_agent);
79    }
80
81    let mut new_channels = Vec::new();
82    for channel in channels {
83        let Some(source) = agent_id_map.get(&channel.source) else {
84            continue;
85        };
86        let Some(target) = agent_id_map.get(&channel.target) else {
87            continue;
88        };
89        let mut new_channel = channel.clone();
90        new_channel.id = new_id();
91        new_channel.source = source.clone();
92        new_channel.target = target.clone();
93        new_channels.push(new_channel);
94    }
95
96    (new_agents, new_channels)
97}
98
99/// Information held by each agent.
100#[derive(Debug, Clone, Serialize, Deserialize)]
101pub struct AgentSpec {
102    #[serde(skip_serializing_if = "String::is_empty", default)]
103    pub id: String,
104
105    /// Name of the AgentDefinition.
106    #[serde(skip_serializing_if = "String::is_empty", default)]
107    pub def_name: String,
108
109    /// List of input pin names.
110    #[serde(skip_serializing_if = "Option::is_none", default)]
111    pub inputs: Option<Vec<String>>,
112
113    /// List of output pin names.
114    #[serde(skip_serializing_if = "Option::is_none", default)]
115    pub outputs: Option<Vec<String>>,
116
117    /// Config values.
118    #[serde(skip_serializing_if = "Option::is_none")]
119    pub configs: Option<AgentConfigs>,
120
121    /// Config specs.
122    #[serde(skip_serializing_if = "Option::is_none")]
123    pub config_specs: Option<AgentConfigSpecs>,
124
125    #[deprecated(note = "Use `disabled` instead")]
126    #[serde(default, skip_serializing_if = "<&bool>::not")]
127    pub enabled: bool,
128
129    #[serde(default, skip_serializing_if = "<&bool>::not")]
130    pub disabled: bool,
131
132    #[serde(flatten)]
133    pub extensions: FnvIndexMap<String, serde_json::Value>,
134}
135
136impl AgentSpec {
137    pub fn from_def(def: &AgentDefinition) -> Self {
138        let mut spec = def.to_spec();
139        spec.id = new_id();
140        spec
141    }
142}
143
144// ChannelSpec
145
146#[derive(Debug, Default, Serialize, Deserialize, Clone)]
147pub struct ChannelSpec {
148    pub id: String,
149    pub source: String,
150    pub source_handle: String,
151    pub target: String,
152    pub target_handle: String,
153}