askit_std_agents/
array.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 im::{Vector, vector};
9use mini_moka::sync::Cache;
10
11const CATEGORY: &str = "Std/Array";
12
13const PIN_ARRAY: &str = "array";
14const PIN_IN1: &str = "in1";
15const PIN_IN2: &str = "in2";
16const PIN_T: &str = "T";
17const PIN_F: &str = "F";
18const PIN_VALUE: &str = "value";
19
20const CONFIG_N: &str = "n";
21const CONFIG_USE_CTX: &str = "use_ctx";
22const CONFIG_TTL_SEC: &str = "ttl_sec";
23const CONFIG_CAPACITY: &str = "capacity";
24
25/// Check if an input is an array.
26#[askit_agent(
27    title = "IsArray",
28    category = CATEGORY,
29    inputs = [PIN_VALUE],
30    outputs = [PIN_T, PIN_F],
31)]
32struct IsArrayAgent {
33    data: AgentData,
34}
35
36#[async_trait]
37impl AsAgent for IsArrayAgent {
38    fn new(askit: ASKit, id: String, spec: AgentSpec) -> Result<Self, AgentError> {
39        let data = AgentData::new(askit, id, spec);
40        Ok(Self { data })
41    }
42    async fn process(
43        &mut self,
44        ctx: AgentContext,
45        _pin: String,
46        value: AgentValue,
47    ) -> Result<(), AgentError> {
48        if value.is_array() {
49            self.output(ctx, PIN_T, value).await
50        } else {
51            self.output(ctx, PIN_F, value).await
52        }
53    }
54}
55
56/// Checks if an input array is empty, emitting to T or F accordingly.
57/// If the input is not an array, it is treated as non-empty.
58#[askit_agent(
59    title = "IsEmptyArray",
60    category = CATEGORY,
61    inputs = [PIN_ARRAY],
62    outputs = [PIN_T, PIN_F],
63)]
64struct IsEmptyArrayAgent {
65    data: AgentData,
66}
67
68#[async_trait]
69impl AsAgent for IsEmptyArrayAgent {
70    fn new(askit: ASKit, id: String, spec: AgentSpec) -> Result<Self, AgentError> {
71        let data = AgentData::new(askit, id, spec);
72        Ok(Self { data })
73    }
74
75    async fn process(
76        &mut self,
77        ctx: AgentContext,
78        _pin: String,
79        value: AgentValue,
80    ) -> Result<(), AgentError> {
81        let mut is_empty = false;
82        if value.is_array() {
83            let arr = value.as_array().unwrap();
84            if arr.is_empty() {
85                is_empty = true;
86            }
87        }
88        if is_empty {
89            self.output(ctx, PIN_T, value).await
90        } else {
91            self.output(ctx, PIN_F, value).await
92        }
93    }
94}
95
96/// Outputs the length of the input array.
97/// If the input is not an array, outputs 1.
98/// This is different from IsEmpty, but is designed for consistency with Map.
99#[askit_agent(
100    title = "ArrayLength",
101    category = CATEGORY,
102    inputs = [PIN_ARRAY],
103    outputs = [PIN_VALUE],
104)]
105struct ArrayLengthAgent {
106    data: AgentData,
107}
108
109#[async_trait]
110impl AsAgent for ArrayLengthAgent {
111    fn new(askit: ASKit, id: String, spec: AgentSpec) -> Result<Self, AgentError> {
112        let data = AgentData::new(askit, id, spec);
113        Ok(Self { data })
114    }
115
116    async fn process(
117        &mut self,
118        ctx: AgentContext,
119        _pin: String,
120        value: AgentValue,
121    ) -> Result<(), AgentError> {
122        let length = if value.is_array() {
123            let arr = value.as_array().unwrap();
124            arr.len() as i64
125        } else {
126            1
127        };
128        self.output(ctx, PIN_VALUE, AgentValue::integer(length)).await
129    }
130}
131
132/// Output the first item of the input array.
133/// If the input is not an array, outputs the input itself.
134/// Errors if the input array is empty.
135#[askit_agent(
136    title = "ArrayFirst",
137    category = CATEGORY,
138    inputs = [PIN_ARRAY],
139    outputs = [PIN_VALUE],
140)]
141struct ArrayFirstAgent {
142    data: AgentData,
143}
144
145#[async_trait]
146impl AsAgent for ArrayFirstAgent {
147    fn new(askit: ASKit, id: String, spec: AgentSpec) -> Result<Self, AgentError> {
148        let data = AgentData::new(askit, id, spec);
149        Ok(Self { data })
150    }
151
152    async fn process(
153        &mut self,
154        ctx: AgentContext,
155        _pin: String,
156        value: AgentValue,
157    ) -> Result<(), AgentError> {
158        match value {
159            AgentValue::Array(mut arr) => {
160                if let Some(first_item) = arr.pop_front() {
161                    self.output(ctx, PIN_VALUE, first_item).await
162                } else {
163                    Err(AgentError::InvalidValue(
164                        "Input array is empty, no first item".into(),
165                    ))
166                }
167            }
168            other => self.output(ctx, PIN_VALUE, other).await,
169        }
170    }
171}
172
173/// Output the rest of the input array after removing the first item.
174/// If the input is not an array, outputs an empty array.
175/// Output an empty array if the input array is empty.
176#[askit_agent(
177    title = "ArrayRest",
178    category = CATEGORY,
179    inputs = [PIN_ARRAY],
180    outputs = [PIN_ARRAY],
181)]
182struct ArrayRestAgent {
183    data: AgentData,
184}
185
186#[async_trait]
187impl AsAgent for ArrayRestAgent {
188    fn new(askit: ASKit, id: String, spec: AgentSpec) -> Result<Self, AgentError> {
189        let data = AgentData::new(askit, id, spec);
190        Ok(Self { data })
191    }
192
193    async fn process(
194        &mut self,
195        ctx: AgentContext,
196        _pin: String,
197        value: AgentValue,
198    ) -> Result<(), AgentError> {
199        if let Some(mut arr) = value.into_array() {
200            if arr.is_empty() {
201                return self.output(ctx, PIN_ARRAY, AgentValue::array_default()).await;
202            }
203            arr.pop_front();
204            self.output(ctx, PIN_ARRAY, AgentValue::array(arr)).await
205        } else {
206            self.output(ctx, PIN_ARRAY, AgentValue::array_default()).await
207        }
208    }
209}
210
211//// Output the last item of the input array.
212/// If the input is not an array, outputs the input itself.
213/// Errors if the input array is empty.
214#[askit_agent(
215    title = "ArrayLast",
216    category = CATEGORY,
217    inputs = [PIN_ARRAY],
218    outputs = [PIN_VALUE],
219)]
220struct ArrayLastAgent {
221    data: AgentData,
222}
223
224#[async_trait]
225impl AsAgent for ArrayLastAgent {
226    fn new(askit: ASKit, id: String, spec: AgentSpec) -> Result<Self, AgentError> {
227        let data = AgentData::new(askit, id, spec);
228        Ok(Self { data })
229    }
230
231    async fn process(
232        &mut self,
233        ctx: AgentContext,
234        _pin: String,
235        value: AgentValue,
236    ) -> Result<(), AgentError> {
237        match value {
238            AgentValue::Array(mut arr) => {
239                if let Some(last_item) = arr.pop_back() {
240                    self.output(ctx, PIN_VALUE, last_item).await
241                } else {
242                    Err(AgentError::InvalidValue(
243                        "Input array is empty, no last item".into(),
244                    ))
245                }
246            }
247            other => self.output(ctx, PIN_VALUE, other).await,
248        }
249    }
250}
251
252/// Output the nth-item of the input array.
253/// If the input is not an array, outputs the input itself if n=0, else errors.
254/// Errors if the input array is shorter than n+1.
255#[askit_agent(
256    title = "ArrayNth",
257    category = CATEGORY,
258    inputs = [PIN_ARRAY],
259    outputs = [PIN_VALUE],
260    integer_config(name = CONFIG_N, default = 0),
261)]
262struct ArrayNthAgent {
263    data: AgentData,
264}
265
266#[async_trait]
267impl AsAgent for ArrayNthAgent {
268    fn new(askit: ASKit, id: String, spec: AgentSpec) -> Result<Self, AgentError> {
269        let data = AgentData::new(askit, id, spec);
270        Ok(Self { data })
271    }
272
273    async fn process(
274        &mut self,
275        ctx: AgentContext,
276        _pin: String,
277        value: AgentValue,
278    ) -> Result<(), AgentError> {
279        let n = self
280            .data
281            .spec
282            .configs
283            .as_ref()
284            .map(|cfg| cfg.get_integer_or(CONFIG_N, 0))
285            .unwrap_or(0);
286        if n < 0 {
287            return Err(AgentError::InvalidConfig("n must be non-negative".into()));
288        }
289        let n = n as usize;
290
291        match value {
292            AgentValue::Array(arr) => {
293                if let Some(item) = arr.get(n) {
294                    self.output(ctx, PIN_VALUE, item.clone()).await
295                } else {
296                    Err(AgentError::InvalidValue(format!(
297                        "Input array length {} is less than n+1={}",
298                        arr.len(),
299                        n + 1
300                    )))
301                }
302            }
303            other => {
304                if n == 0 {
305                    self.output(ctx, PIN_VALUE, other).await
306                } else {
307                    Err(AgentError::InvalidValue(
308                        "Input is not an array and n != 0".into(),
309                    ))
310                }
311            }
312        }
313    }
314}
315
316/// Takes the first n items from the input array.
317/// If the input is not an array, outputs an array with the input as the only item.
318/// If n is greater than the array length, outputs the entire array.
319#[askit_agent(
320    title = "ArrayTake",
321    category = CATEGORY,
322    inputs = [PIN_ARRAY],
323    outputs = [PIN_ARRAY],
324    integer_config(name = CONFIG_N, default = 0),
325)]
326struct ArrayTakeAgent {
327    data: AgentData,
328}
329
330#[async_trait]
331impl AsAgent for ArrayTakeAgent {
332    fn new(askit: ASKit, id: String, spec: AgentSpec) -> Result<Self, AgentError> {
333        let data = AgentData::new(askit, id, spec);
334        Ok(Self { data })
335    }
336
337    async fn process(
338        &mut self,
339        ctx: AgentContext,
340        _pin: String,
341        value: AgentValue,
342    ) -> Result<(), AgentError> {
343        let n = self
344            .data
345            .spec
346            .configs
347            .as_ref()
348            .map(|cfg| cfg.get_integer_or(CONFIG_N, 0))
349            .unwrap_or(0);
350        if n <= 0 {
351            // output empty array
352            return self.output(ctx, PIN_ARRAY, AgentValue::array_default()).await;
353        }
354        let n = n as usize;
355
356        if value.is_array() {
357            let arr = value.as_array().unwrap();
358            if n >= arr.len() {
359                return self.output(ctx, PIN_ARRAY, value).await;
360            }
361            let taken_items = arr.take(n);
362            self.output(ctx, PIN_ARRAY, AgentValue::array(taken_items)).await
363        } else {
364            self.output(ctx, PIN_ARRAY, AgentValue::array(vector![value])).await
365        }
366    }
367}
368
369/// Maps over an input array, emitting each item individually with a `map` frame that captures the index and length.
370/// Nested maps accumulate frames to preserve lineage. If the input is not an array, it is treated as a single-item array.
371#[askit_agent(
372    title = "Map",
373    category = CATEGORY,
374    inputs = [PIN_ARRAY],
375    outputs = [PIN_VALUE],
376)]
377struct MapAgent {
378    data: AgentData,
379}
380
381#[async_trait]
382impl AsAgent for MapAgent {
383    fn new(askit: ASKit, id: String, spec: AgentSpec) -> Result<Self, AgentError> {
384        let data = AgentData::new(askit, id, spec);
385        Ok(Self { data })
386    }
387
388    async fn process(
389        &mut self,
390        ctx: AgentContext,
391        _pin: String,
392        value: AgentValue,
393    ) -> Result<(), AgentError> {
394        match value {
395            AgentValue::Array(arr) => {
396                let n = arr.len();
397                for (i, item) in arr.into_iter().enumerate() {
398                    let c = ctx.push_map_frame(i, n)?;
399                    self.output(c, PIN_VALUE, item).await?;
400                }
401            }
402            other => {
403                let c = ctx.push_map_frame(0, 1)?;
404                self.output(c, PIN_VALUE, other).await?;
405            }
406        }
407        Ok(())
408    }
409}
410
411/// Collects input values into an array.
412///
413/// Expects a `map` frame to determine the position and length for each input value.
414/// The `map` frame stores keys `i` (index) and `n` (length). Nested maps stack frames.
415/// If a `map` frame is not present, the input value is emitted directly.
416///
417/// Incomplete arrays are emitted when the context changes.
418#[askit_agent(
419    title = "Collect",
420    category = CATEGORY,
421    description = "Collects input values into an array",
422    inputs = [PIN_VALUE],
423    outputs = [PIN_ARRAY],
424)]
425struct CollectAgent {
426    data: AgentData,
427
428    // Records the context ID being processed to prevent other contexts from mixing
429    current_ctx_id: Option<usize>,
430
431    // Data buffer
432    input_values: Vec<Option<AgentValue>>,
433
434    // Expected size of the array
435    expected_size: usize,
436
437    // Number of items received (counter to avoid scanning input_values every time)
438    received_count: usize,
439}
440
441#[async_trait]
442impl AsAgent for CollectAgent {
443    fn new(askit: ASKit, id: String, spec: AgentSpec) -> Result<Self, AgentError> {
444        let data = AgentData::new(askit, id, spec);
445        Ok(Self {
446            data,
447            current_ctx_id: None,
448            input_values: Vec::new(),
449            expected_size: 0,
450            received_count: 0,
451        })
452    }
453
454    async fn process(
455        &mut self,
456        ctx: AgentContext,
457        _pin: String,
458        value: AgentValue,
459    ) -> Result<(), AgentError> {
460        // Check for map frame
461        // If not within a map, pass the value through as-is.
462        let Some((idx, n)) = ctx.current_map_frame()? else {
463            return self.output(ctx, PIN_ARRAY, value).await;
464        };
465
466        // Detect context switch and flush processing
467        // If a new context ID arrives while the previous context hasn't finished processing
468        let ctx_id = ctx.id();
469        if let Some(last_id) = &self.current_ctx_id {
470            if last_id != &ctx_id {
471                log::warn!("Context changed before collection completed. Dropping partial data.");
472                self.reset_state();
473            }
474        }
475
476        // Initialize state (when the first item of this context arrives)
477        if self.input_values.is_empty() {
478            self.current_ctx_id = Some(ctx_id);
479            self.expected_size = n;
480            // Fill with None for the required size
481            self.input_values = vec![None; n];
482            self.received_count = 0;
483        }
484
485        // Validation
486        if n != self.expected_size {
487            // Size shouldn't change within the same context ID, but check just in case
488            return Err(AgentError::InvalidValue(
489                "Map frame size mismatch within the same context".into(),
490            ));
491        }
492        if idx >= n {
493            return Err(AgentError::InvalidValue(
494                "Map frame index is out of bounds".into(),
495            ));
496        }
497
498        // Store data
499        // Check if attempting to write to a position that's already filled (duplicate index)
500        if self.input_values[idx].is_some() {
501            // If duplicate data arrives, overwrite (could also error instead).
502        } else {
503            self.received_count += 1;
504        }
505        self.input_values[idx] = Some(value);
506
507        // Check for completion
508        if self.received_count == self.expected_size {
509            // All items collected, output the result
510            let arr = self.drain_buffer_to_vector();
511
512            // Reset state
513            self.reset_state();
514
515            // Pop one map frame and output
516            let next_ctx = ctx.pop_map_frame()?;
517            self.output(next_ctx, PIN_ARRAY, AgentValue::array(arr)).await
518        } else {
519            // Not yet complete, keep waiting
520            Ok(())
521        }
522    }
523}
524
525impl CollectAgent {
526    fn reset_state(&mut self) {
527        self.current_ctx_id = None;
528        self.input_values.clear(); // Capacity is preserved for efficient reuse
529        self.expected_size = 0;
530        self.received_count = 0;
531    }
532
533    // Drain the buffer contents and convert to im::Vector
534    fn drain_buffer_to_vector(&mut self) -> Vector<AgentValue> {
535        self.input_values
536            .drain(..)
537            .map(|v| v.unwrap_or(AgentValue::Unit)) // Fill missing values with Unit
538            .collect()
539    }
540}
541
542/// Zips multiple inputs into an array.
543///
544/// The number of inputs n is specified via configuration.
545///
546/// If n=2, it takes two inputs: in1 and in2. Once all inputs are present,
547/// it emits them as [in1, in2].
548///
549/// If in2 arrives repeatedly before in1, the in2 values are queued; when in1 arrives,
550/// they’re paired in order from the head of the queue and emitted.
551///
552/// When the `use_ctx` config is true, inputs are matched by context key (including map frames)
553/// so that mapped items zip correctly even when they interleave.
554#[askit_agent(
555    title = "ZipToArray",
556    category = CATEGORY,
557    inputs = [PIN_IN1, PIN_IN2],
558    outputs = [PIN_ARRAY],
559    integer_config(name = CONFIG_N, default = 2),
560    boolean_config(name = CONFIG_USE_CTX),
561    integer_config(name = CONFIG_TTL_SEC, default = 60), 
562    integer_config(name = CONFIG_CAPACITY, default = 1000),
563)]
564struct ZipToArrayAgent {
565    data: AgentData,
566    n: usize,
567    use_ctx: bool,
568
569    ttl_sec: u64,
570    capacity: u64,
571    queues: Vec<VecDeque<AgentValue>>, // for non-ctx mode
572
573    // Context Key -> PendingZip
574    ctx_buffers: Cache<String, PendingZip>,
575}
576
577#[derive(Clone)]
578struct PendingZip {
579    values: Vec<Option<AgentValue>>,
580    count: usize,
581}
582
583impl ZipToArrayAgent {
584    fn update_spec(spec: &mut AgentSpec) -> Result<(usize, bool, u64, u64), AgentError> {
585        let mut n = spec
586            .configs
587            .as_ref()
588            .map(|cfg| cfg.get_integer_or(CONFIG_N, 2))
589            .unwrap_or(2) as usize;
590        if n < 1 {
591            n = 1;
592        }
593
594        let use_ctx = spec
595            .configs
596            .as_ref()
597            .map(|cfg| cfg.get_bool_or_default(CONFIG_USE_CTX))
598            .unwrap_or(false);
599
600        let ttl_sec = spec
601            .configs
602            .as_ref()
603            .map(|c| c.get_integer_or(CONFIG_TTL_SEC, 60))
604            .unwrap_or(60) as u64;
605
606        let capacity = spec
607            .configs
608            .as_ref()
609            .map(|c| c.get_integer_or(CONFIG_CAPACITY, 1000))
610            .unwrap_or(1000) as u64;
611
612        spec.inputs = Some((1..=n).map(|i| format!("in{}", i)).collect());
613
614        Ok((n, use_ctx, ttl_sec, capacity))
615    }
616
617    fn reset_state(&mut self) {
618        self.queues = vec![VecDeque::new(); self.n];
619        self.ctx_buffers.invalidate_all();
620    }
621}
622
623#[async_trait]
624impl AsAgent for ZipToArrayAgent {
625    fn new(askit: ASKit, id: String, mut spec: AgentSpec) -> Result<Self, AgentError> {
626        let (n, use_ctx, ttl_sec, capacity) = Self::update_spec(&mut spec)?;
627
628        let cache = Cache::builder()
629            .max_capacity(capacity) // Capacity limit (oldest entries are evicted on overflow)
630            .time_to_live(Duration::from_secs(ttl_sec)) // TTL (entries expire X seconds after write)
631            .build();
632
633        let data = AgentData::new(askit, id, spec);
634
635        Ok(Self {
636            data,
637            n,
638            use_ctx,
639            ttl_sec,
640            capacity,
641            queues: vec![VecDeque::new(); n],
642            ctx_buffers: cache,
643        })
644    }
645
646    fn configs_changed(&mut self) -> Result<(), AgentError> {
647        let (n, use_ctx, ttl_sec, capacity) = Self::update_spec(&mut self.data.spec)?;
648        let mut changed = false;
649        if n != self.n {
650            self.n = n;
651            changed = true;
652        }
653        if use_ctx != self.use_ctx {
654            self.use_ctx = use_ctx;
655            changed = true;
656        }
657        if ttl_sec != self.ttl_sec {
658            self.ttl_sec = ttl_sec;
659            changed = true;
660        }
661        if capacity != self.capacity {
662            self.capacity = capacity;
663            changed = true;
664        }
665        if changed {
666            self.reset_state();
667            // Rebuild cache with new capacity and TTL
668            self.ctx_buffers = Cache::builder()
669                .max_capacity(capacity)
670                .time_to_live(Duration::from_secs(ttl_sec))
671                .build();
672            self.emit_agent_spec_updated();
673        }
674        Ok(())
675    }
676
677    async fn stop(&mut self) -> Result<(), AgentError> {
678        self.reset_state();
679        Ok(())
680    }
681
682    async fn process(
683        &mut self,
684        ctx: AgentContext,
685        pin: String,
686        value: AgentValue,
687    ) -> Result<(), AgentError> {
688        // Parse pin number
689        let Some(idx) = pin
690            .strip_prefix("in")
691            .and_then(|s| s.parse::<usize>().ok())
692            .filter(|&i| i >= 1 && i <= self.n)
693            .map(|i| i - 1)
694        else {
695            return Err(AgentError::InvalidValue(format!(
696                "Invalid input pin: {}",
697                pin
698            )));
699        };
700
701        if self.use_ctx {
702            let ctx_key = ctx.ctx_key()?;
703
704            // Get from cache (or create new if not present)
705            let mut entry = self.ctx_buffers.get(&ctx_key).unwrap_or_else(|| PendingZip {
706                values: vec![None; self.n],
707                count: 0,
708            });
709
710            // Update
711            if entry.values[idx].is_none() {
712                entry.count += 1;
713            }
714            entry.values[idx] = Some(value);
715
716            // Check for completion
717            if entry.count == self.n {
718                // All inputs collected, remove from cache (invalidate)
719                self.ctx_buffers.invalidate(&ctx_key);
720
721                let arr: Vector<AgentValue> = entry.values
722                    .into_iter()
723                    .map(|v| v.unwrap())
724                    .collect();
725
726                return self.output(ctx, PIN_ARRAY, AgentValue::array(arr)).await;
727            }
728
729            return Ok(());
730        }
731
732        // Simple FIFO mode processing
733        self.queues[idx].push_back(value);
734
735        // Check if all queues have data
736        if self.queues.iter().all(|q| !q.is_empty()) {
737            let arr: Vector<AgentValue> = self.queues
738                .iter_mut()
739                .map(|q| q.pop_front().unwrap())
740                .collect();
741
742            self.output(ctx, PIN_ARRAY, AgentValue::array(arr)).await
743        } else {
744            Ok(())
745        }
746    }
747}
askit_std_agents/array.rs

askit_std_agents/
array.rs
