askit_std_agents/
sequence.rs

1use std::collections::VecDeque;
2use std::time::Duration;
3
4use agent_stream_kit::{
5    ASKit, AgentContext, AgentData, AgentError, AgentOutput, AgentSpec, AgentValue, AsAgent,
6    askit_agent, async_trait,
7};
8use mini_moka::sync::Cache;
9
10const CONFIG_TTL_SEC: &str = "ttl_sec";
11const CONFIG_CAPACITY: &str = "capacity";
12
13const CATEGORY: &str = "Std/Sequence";
14
15const PIN_IN: &str = "in";
16const PIN_IN1: &str = "in1";
17const PIN_IN2: &str = "in2";
18const PIN_OUT1: &str = "out1";
19const PIN_OUT2: &str = "out2";
20
21const CONFIG_N: &str = "n";
22const CONFIG_USE_CTX: &str = "use_ctx";
23
24/// Receives an input and emits it sequentially to n outputs.
25#[askit_agent(
26    title = "Sequence",
27    category = CATEGORY,
28    inputs = [PIN_IN],
29    outputs = [PIN_OUT1, PIN_OUT2],
30    integer_config(name = CONFIG_N, default = 2),
31)]
32struct SequenceAgent {
33    data: AgentData,
34    n: usize,
35}
36
37impl SequenceAgent {
38    fn update_spec(spec: &mut AgentSpec) -> Result<usize, AgentError> {
39        let mut n = spec
40            .configs
41            .as_ref()
42            .map(|cfg| cfg.get_integer_or(CONFIG_N, 2))
43            .unwrap_or(2) as usize;
44        if n < 1 {
45            n = 1;
46        }
47
48        spec.outputs = Some((1..=n).map(|i| format!("out{}", i)).collect());
49
50        Ok(n)
51    }
52}
53
54#[async_trait]
55impl AsAgent for SequenceAgent {
56    fn new(askit: ASKit, id: String, mut spec: AgentSpec) -> Result<Self, AgentError> {
57        let n = Self::update_spec(&mut spec)?;
58        let data = AgentData::new(askit, id, spec);
59        Ok(Self { data, n })
60    }
61
62    fn configs_changed(&mut self) -> Result<(), AgentError> {
63        let n = Self::update_spec(&mut self.data.spec)?;
64        let mut changed = false;
65        if n != self.n {
66            self.n = n;
67            changed = true;
68        }
69        if changed {
70            self.emit_agent_spec_updated();
71        }
72        Ok(())
73    }
74
75    async fn process(
76        &mut self,
77        ctx: AgentContext,
78        _pin: String,
79        value: AgentValue,
80    ) -> Result<(), AgentError> {
81        for i in 0..self.n {
82            let out_pin = format!("out{}", i + 1);
83            self.output(ctx.clone(), out_pin, value.clone()).await?;
84        }
85        Ok(())
86    }
87}
88
89/// Receives inputs in any order and, once all are present, emits them sequentially.
90#[askit_agent(
91    title = "Sync",
92    category = CATEGORY,
93    inputs = [PIN_IN1, PIN_IN2],
94    outputs = [PIN_OUT1, PIN_OUT2],
95    integer_config(name = CONFIG_N, default = 2),
96    boolean_config(name = CONFIG_USE_CTX),
97    integer_config(name = CONFIG_TTL_SEC, default = 60), 
98    integer_config(name = CONFIG_CAPACITY, default = 1000),
99)]
100struct SyncAgent {
101    data: AgentData,
102    n: usize,
103    use_ctx: bool,
104        ttl_sec: u64,
105    capacity: u64,
106
107    // Optimization: Pre-generate and store output pin names ("out1", "out2"...)
108    output_pins: Vec<String>,
109
110    // For simple mode
111    queues: Vec<VecDeque<AgentValue>>,
112
113    // For use_ctx mode: Cache with TTL
114    ctx_buffers: Cache<String, PendingSync>,
115}
116
117#[derive(Clone)]
118struct PendingSync {
119    values: Vec<Option<AgentValue>>,
120    count: usize,
121}
122
123impl SyncAgent {
124    fn update_spec(spec: &mut AgentSpec) -> Result<(usize, bool, u64, u64, Vec<String>), AgentError> {
125        let n = spec.configs.as_ref()
126            .map(|cfg| cfg.get_integer_or(CONFIG_N, 2))
127            .unwrap_or(2) as usize;
128        let n = if n < 1 { 1 } else { n };
129
130        let use_ctx = spec
131            .configs
132            .as_ref()
133            .map(|cfg| cfg.get_bool_or_default(CONFIG_USE_CTX))
134            .unwrap_or(false);
135
136        let ttl_sec = spec
137            .configs
138            .as_ref()
139            .map(|c| c.get_integer_or(CONFIG_TTL_SEC, 60))
140            .unwrap_or(60) as u64;
141
142        let capacity = spec
143            .configs
144            .as_ref()
145            .map(|c| c.get_integer_or(CONFIG_CAPACITY, 1000))
146            .unwrap_or(1000) as u64;
147
148        spec.inputs = Some((1..=n).map(|i| format!("in{}", i)).collect());
149
150        let output_pins: Vec<String> = (1..=n).map(|i| format!("out{}", i)).collect();
151        spec.outputs = Some(output_pins.clone());
152
153        Ok((n, use_ctx, ttl_sec, capacity, output_pins))
154    }
155
156    fn reset_state(&mut self) {
157        self.queues = vec![VecDeque::new(); self.n];
158        self.ctx_buffers.invalidate_all();
159    }
160}
161
162#[async_trait]
163impl AsAgent for SyncAgent {
164    fn new(askit: ASKit, id: String, mut spec: AgentSpec) -> Result<Self, AgentError> {
165        let (n, use_ctx, ttl_sec, capacity, output_pins) = Self::update_spec(&mut spec)?;
166
167        let cache = Cache::builder()
168            .max_capacity(capacity)
169            .time_to_live(Duration::from_secs(ttl_sec))
170            .build();
171
172        let data = AgentData::new(askit, id, spec);
173        Ok(Self {
174            data,
175            n,
176            use_ctx,
177            ttl_sec,
178            capacity,
179            output_pins,
180            queues: vec![VecDeque::new(); n],
181            ctx_buffers: cache,
182        })
183    }
184
185    fn configs_changed(&mut self) -> Result<(), AgentError> {
186        let (n, use_ctx, ttl_sec, capacity, output_pins) = Self::update_spec(&mut self.data.spec)?;
187        let mut changed = false;
188        if n != self.n {
189            self.n = n;
190            changed = true;
191        }
192        if use_ctx != self.use_ctx {
193            self.use_ctx = use_ctx;
194            changed = true;
195        }
196        if ttl_sec != self.ttl_sec {
197            self.ttl_sec = ttl_sec;
198            changed = true;
199        }
200        if capacity != self.capacity {
201            self.capacity = capacity;
202            changed = true;
203        }
204        if changed {
205            self.reset_state();
206            self.output_pins = output_pins;
207            self.ctx_buffers = Cache::builder()
208                .max_capacity(capacity)
209                .time_to_live(Duration::from_secs(ttl_sec))
210                .build();
211            self.emit_agent_spec_updated();
212        }
213        Ok(())
214    }
215
216    async fn stop(&mut self) -> Result<(), AgentError> {
217        // Clear input queues on stop
218        self.reset_state();
219        Ok(())
220    }
221
222    async fn process(
223        &mut self,
224        ctx: AgentContext,
225        pin: String,
226        value: AgentValue,
227    ) -> Result<(), AgentError> {
228        // Parse pin number
229        let Some(idx) = pin
230            .strip_prefix("in")
231            .and_then(|s| s.parse::<usize>().ok())
232            .filter(|&i| i >= 1 && i <= self.n)
233            .map(|i| i - 1)
234        else {
235            return Err(AgentError::InvalidValue(format!("Invalid input pin: {}", pin)));
236        };
237
238        // Context Mode
239        if self.use_ctx {
240            let ctx_key = ctx.ctx_key()?;
241
242            // Get from cache or create new
243            let mut entry = self.ctx_buffers.get(&ctx_key).unwrap_or_else(|| PendingSync {
244                values: vec![None; self.n],
245                count: 0,
246            });
247
248            if entry.values[idx].is_none() {
249                entry.count += 1;
250            }
251            entry.values[idx] = Some(value);
252
253            if entry.count == self.n {
254                // All inputs collected, remove from cache
255                self.ctx_buffers.invalidate(&ctx_key);
256
257                // Output sequentially
258                for (i, val_opt) in entry.values.into_iter().enumerate() {
259                    if let Some(val) = val_opt {
260                        self.output(ctx.clone(), &self.output_pins[i], val).await?;
261                    }
262                }
263            }
264            return Ok(());
265        }
266
267        // Simple FIFO Mode
268        self.queues[idx].push_back(value);
269
270        // Check if all queues have data
271        if self.queues.iter().all(|q| !q.is_empty()) {
272            let ready_values: Vec<AgentValue> = self.queues
273                .iter_mut()
274                .map(|q| q.pop_front().unwrap())
275                .collect();
276
277            for (i, val) in ready_values.into_iter().enumerate() {
278                self.output(ctx.clone(), &self.output_pins[i], val).await?;
279            }
280        }
281
282        Ok(())
283    }
284}
askit_std_agents/sequence.rs

askit_std_agents/
sequence.rs
