agent_stream_kit/
agent.rs

1use async_trait::async_trait;
2
3use super::askit::ASKit;
4use super::config::AgentConfig;
5use super::context::AgentContext;
6use super::data::AgentData;
7use super::error::AgentError;
8
9#[derive(Debug, Default, Clone, PartialEq)]
10pub enum AgentStatus {
11    #[default]
12    Init,
13    Start,
14    Stop,
15}
16
17pub enum AgentMessage {
18    Input { ctx: AgentContext, data: AgentData },
19    Config { config: AgentConfig },
20    Stop,
21}
22
23#[async_trait]
24pub trait Agent {
25    fn new(
26        askit: ASKit,
27        id: String,
28        def_name: String,
29        config: Option<AgentConfig>,
30    ) -> Result<Self, AgentError>
31    where
32        Self: Sized;
33
34    fn askit(&self) -> &ASKit;
35
36    fn id(&self) -> &str;
37
38    fn status(&self) -> &AgentStatus;
39
40    fn def_name(&self) -> &str;
41
42    fn config(&self) -> Option<&AgentConfig>;
43
44    fn set_config(&mut self, config: AgentConfig) -> Result<(), AgentError>;
45
46    fn get_global_config(&self) -> Option<AgentConfig> {
47        self.askit().get_global_config(self.def_name())
48    }
49
50    fn start(&mut self) -> Result<(), AgentError>;
51
52    fn stop(&mut self) -> Result<(), AgentError>;
53
54    async fn process(&mut self, ctx: AgentContext, data: AgentData) -> Result<(), AgentError>;
55}
56
57pub struct AsAgentData {
58    pub askit: ASKit,
59
60    pub id: String,
61    pub status: AgentStatus,
62    pub def_name: String,
63    pub config: Option<AgentConfig>,
64}
65
66impl AsAgentData {
67    pub fn new(askit: ASKit, id: String, def_name: String, config: Option<AgentConfig>) -> Self {
68        Self {
69            askit,
70            id,
71            status: AgentStatus::Init,
72            def_name,
73            config,
74        }
75    }
76}
77
78#[async_trait]
79pub trait AsAgent {
80    fn new(
81        askit: ASKit,
82        id: String,
83        def_name: String,
84        config: Option<AgentConfig>,
85    ) -> Result<Self, AgentError>
86    where
87        Self: Sized + Send + Sync;
88
89    fn data(&self) -> &AsAgentData;
90
91    fn mut_data(&mut self) -> &mut AsAgentData;
92
93    fn set_config(&mut self, _config: AgentConfig) -> Result<(), AgentError> {
94        Ok(())
95    }
96
97    fn start(&mut self) -> Result<(), AgentError> {
98        Ok(())
99    }
100
101    fn stop(&mut self) -> Result<(), AgentError> {
102        Ok(())
103    }
104
105    async fn process(&mut self, _ctx: AgentContext, _data: AgentData) -> Result<(), AgentError> {
106        Ok(())
107    }
108}
109
110#[async_trait]
111impl<T: AsAgent + Send + Sync> Agent for T {
112    fn new(
113        askit: ASKit,
114        id: String,
115        def_name: String,
116        config: Option<AgentConfig>,
117    ) -> Result<Self, AgentError> {
118        let mut agent = T::new(askit, id, def_name, config)?;
119        agent.mut_data().status = AgentStatus::Init;
120        Ok(agent)
121    }
122
123    fn askit(&self) -> &ASKit {
124        &self.data().askit
125    }
126
127    fn id(&self) -> &str {
128        &self.data().id
129    }
130
131    fn status(&self) -> &AgentStatus {
132        &self.data().status
133    }
134
135    fn def_name(&self) -> &str {
136        self.data().def_name.as_str()
137    }
138
139    fn config(&self) -> Option<&AgentConfig> {
140        self.data().config.as_ref()
141    }
142
143    fn set_config(&mut self, config: AgentConfig) -> Result<(), AgentError> {
144        self.mut_data().config = Some(config.clone());
145        self.set_config(config)
146    }
147
148    fn start(&mut self) -> Result<(), AgentError> {
149        self.mut_data().status = AgentStatus::Start;
150
151        if let Err(e) = self.start() {
152            self.askit()
153                .emit_error(self.id().to_string(), e.to_string());
154            return Err(e);
155        }
156
157        Ok(())
158    }
159
160    fn stop(&mut self) -> Result<(), AgentError> {
161        self.mut_data().status = AgentStatus::Stop;
162        self.stop()?;
163        self.mut_data().status = AgentStatus::Init;
164        Ok(())
165    }
166
167    async fn process(&mut self, ctx: AgentContext, data: AgentData) -> Result<(), AgentError> {
168        if let Err(e) = self.process(ctx, data).await {
169            self.askit()
170                .emit_error(self.id().to_string(), e.to_string());
171            return Err(e);
172        }
173        Ok(())
174    }
175
176    fn get_global_config(&self) -> Option<AgentConfig> {
177        self.askit().get_global_config(self.def_name())
178    }
179}
180
181pub fn new_boxed<T: Agent + Send + Sync + 'static>(
182    askit: ASKit,
183    id: String,
184    def_name: String,
185    config: Option<AgentConfig>,
186) -> Result<Box<dyn Agent + Send + Sync>, AgentError> {
187    Ok(Box::new(T::new(askit, id, def_name, config)?))
188}
189
190pub fn agent_new(
191    askit: ASKit,
192    agent_id: String,
193    def_name: &str,
194    config: Option<AgentConfig>,
195) -> Result<Box<dyn Agent + Send + Sync>, AgentError> {
196    let def;
197    {
198        let defs = askit.defs.lock().unwrap();
199        def = defs
200            .get(def_name)
201            .ok_or_else(|| AgentError::UnknownDefName(def_name.to_string()))?
202            .clone();
203    }
204
205    if let Some(new_boxed) = def.new_boxed {
206        return new_boxed(askit, agent_id, def_name.to_string(), config);
207    }
208
209    match def.kind.as_str() {
210        // "Command" => {
211        //     return new_boxed::<super::builtins::CommandAgent>(
212        //         askit,
213        //         agent_id,
214        //         def_name.to_string(),
215        //         config,
216        //     );
217        // }
218        _ => return Err(AgentError::UnknownDefKind(def.kind.to_string()).into()),
219    }
220}
agent_stream_kit/agent.rs

agent_stream_kit/
agent.rs
