agent_stream_kit/
stream.rs

1use serde::{Deserialize, Serialize};
2
3use crate::FnvIndexMap;
4use crate::askit::ASKit;
5use crate::error::AgentError;
6use crate::id::{new_id, update_ids};
7use crate::spec::AgentStreamSpec;
8
9pub type AgentStreams = FnvIndexMap<String, AgentStream>;
10
11pub struct AgentStream {
12    id: String,
13
14    name: String,
15
16    running: bool,
17
18    spec: AgentStreamSpec,
19}
20
21impl AgentStream {
22    /// Create a new agent stream with the given name and spec.
23    ///
24    /// The ids of the given spec, including agents and channels, are changed to new unique ids.
25    pub fn new(name: String, mut spec: AgentStreamSpec) -> Self {
26        let (agents, channels) = update_ids(&spec.agents, &spec.channels);
27        spec.agents = agents;
28        spec.channels = channels;
29
30        Self {
31            id: new_id(),
32            name,
33            running: false,
34            spec,
35        }
36    }
37
38    pub fn id(&self) -> &str {
39        &self.id
40    }
41
42    pub fn spec(&self) -> &AgentStreamSpec {
43        &self.spec
44    }
45
46    pub fn spec_mut(&mut self) -> &mut AgentStreamSpec {
47        &mut self.spec
48    }
49
50    pub fn name(&self) -> &str {
51        &self.name
52    }
53
54    pub fn set_name(&mut self, name: String) {
55        self.name = name;
56    }
57
58    pub fn running(&self) -> bool {
59        self.running
60    }
61
62    pub async fn start(&mut self, askit: &ASKit) -> Result<(), AgentError> {
63        if self.running {
64            // Already running
65            return Ok(());
66        }
67        self.running = true;
68
69        for agent in self.spec.agents.iter() {
70            if agent.disabled {
71                continue;
72            }
73            askit.start_agent(&agent.id).await.unwrap_or_else(|e| {
74                log::error!("Failed to start agent {}: {}", agent.id, e);
75            });
76        }
77
78        Ok(())
79    }
80
81    pub async fn stop(&mut self, askit: &ASKit) -> Result<(), AgentError> {
82        for agent in self.spec.agents.iter() {
83            askit.stop_agent(&agent.id).await.unwrap_or_else(|e| {
84                log::error!("Failed to stop agent {}: {}", agent.id, e);
85            });
86        }
87        self.running = false;
88        Ok(())
89    }
90}
91
92#[derive(Clone, Debug, Serialize, Deserialize)]
93pub struct AgentStreamInfo {
94    pub id: String,
95    pub name: String,
96    pub running: bool,
97    pub run_on_start: bool,
98}
99
100impl From<&AgentStream> for AgentStreamInfo {
101    fn from(stream: &AgentStream) -> Self {
102        Self {
103            id: stream.id.clone(),
104            name: stream.name.clone(),
105            running: stream.running,
106            run_on_start: stream.spec.run_on_start,
107        }
108    }
109}