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

agent_stream_kit/
agent.rs
