lc_core/runnables/
pick.rs

1// lc-core/src/runnables/pick.rs
2//! RunnablePick — extract keys / values out of a dict-returning Runnable.
3//!
4//! Rust counterpart of Python LCEL's `Runnable.pick(*keys)` and
5//! `Runnable.pluck(key)`:
6//!
7//! - `pick(["a", "b"])` — keep only the given keys of a
8//!   `HashMap<String, Value>` output, dropping everything else.
9//! - `pluck("a")` — pull a single value out of the dict output.
10//!
11//! The trait is blanket-implemented for every `Runnable` whose output is
12//! `HashMap<String, Value>` (e.g. `RunnableParallel`, or any sequence that
13//! ends in a map), so the composition is checked at compile time — a
14//! chain that does *not* produce a map simply won't compile, mirroring the
15//! dynamic check Python performs at runtime.
16
17use super::error::LcelError;
18use super::ext::RunnableExt;
19use super::lambda::RunnableLambda;
20use super::runnable_trait::Runnable;
21use super::sequence::RunnableSequence;
22use serde_json::Value;
23use std::collections::HashMap;
24
25/// Extracts keys / values from a Runnable whose output is a `HashMap<String, Value>`.
26pub trait RunnablePick<I: Send + Sync + 'static>: Sized {
27    /// Keep only the given keys of the dict output, dropping everything else.
28    ///
29    /// Missing keys are omitted from the result (Python raises `KeyError`;
30    /// Rust's map semantics make omission the closest safe equivalent).
31    ///
32    /// # Example
33    ///
34    /// ```rust,ignore
35    /// let parallel = RunnableParallel::<String>::new()
36    ///     .with("a", RunnableLambda::new_sync(|s: String| s.len() as i64))
37    ///     .with("b", RunnableLambda::new_sync(|s: String| s.to_uppercase()));
38    /// let picked = parallel.pick(["a"]);   // output: {"a": N}
39    /// ```
40    fn pick<K>(self, keys: impl IntoIterator<Item = K>) -> RunnableSequence<I, HashMap<String, Value>>
41    where
42        K: Into<String>;
43
44    /// Pull a single value out of the dict output.
45    ///
46    /// A missing key yields `Value::Null` (Python raises `KeyError`; `Null`
47    /// keeps the runnable total so the pipeline does not short-circuit).
48    ///
49    /// # Example
50    ///
51    /// ```rust,ignore
52    /// let val = parallel.pluck("a");   // output: Value
53    /// ```
54    fn pluck(self, key: impl Into<String>) -> RunnableSequence<I, Value>;
55}
56
57impl<I, R> RunnablePick<I> for R
58where
59    I: Send + Sync + 'static,
60    R: Runnable<I, HashMap<String, Value>> + Sized + 'static,
61    R::Error: Into<LcelError>,
62{
63    fn pick<K>(self, keys: impl IntoIterator<Item = K>) -> RunnableSequence<I, HashMap<String, Value>>
64    where
65        K: Into<String>,
66    {
67        let keys: Vec<String> = keys.into_iter().map(Into::into).collect();
68        let filter = RunnableLambda::new_sync(move |m: HashMap<String, Value>| {
69            keys.iter()
70                .filter_map(|k| m.get(k).cloned().map(|v| (k.clone(), v)))
71                .collect::<HashMap<String, Value>>()
72        });
73        self.pipe(filter)
74    }
75
76    fn pluck(self, key: impl Into<String>) -> RunnableSequence<I, Value> {
77        let key = key.into();
78        let extract = RunnableLambda::new_sync(move |m: HashMap<String, Value>| {
79            m.get(&key).cloned().unwrap_or(Value::Null)
80        });
81        self.pipe(extract)
82    }
83}
84
85#[cfg(test)]
86mod tests {
87    use super::*;
88    use crate::Runnable;
89    use crate::RunnableLambda;
90    use crate::RunnableParallel;
91
92    fn parallel() -> RunnableParallel<String> {
93        RunnableParallel::<String>::new()
94            .with("len", RunnableLambda::new_sync(|s: String| s.len() as i64))
95            .with("upper", RunnableLambda::new_sync(|s: String| s.to_uppercase()))
96    }
97
98    #[tokio::test]
99    async fn pick_keeps_only_selected_keys() {
100        let chain = parallel().pick(["len"]);
101        let out = chain.invoke("hello".to_string(), None).await.unwrap();
102        assert_eq!(out.len(), 1);
103        assert!(out.contains_key("len"));
104        assert!(!out.contains_key("upper"));
105    }
106
107    #[tokio::test]
108    async fn pick_multiple_keys_and_missing() {
109        let chain = parallel().pick(["len", "nope"]);
110        let out = chain.invoke("hello".to_string(), None).await.unwrap();
111        // "nope" 不存在 → 被省略,只剩 len
112        assert_eq!(out.len(), 1);
113        assert!(out.contains_key("len"));
114    }
115
116    #[tokio::test]
117    async fn pluck_returns_single_value() {
118        let chain = parallel().pluck("upper");
119        let out = chain.invoke("hello".to_string(), None).await.unwrap();
120        assert_eq!(out, Value::String("HELLO".to_string()));
121    }
122
123    #[tokio::test]
124    async fn pluck_missing_key_yields_null() {
125        let chain = parallel().pluck("missing");
126        let out = chain.invoke("hello".to_string(), None).await.unwrap();
127        assert_eq!(out, Value::Null);
128    }
129
130    #[tokio::test]
131    async fn pick_works_on_sequence_ending_in_map() {
132        // RunnableSequence<I, HashMap<String, Value>> 也自动获得 RunnablePick
133        let seq = RunnableLambda::new_sync(|s: String| s.to_uppercase())
134            .pipe(RunnableLambda::new_sync(|s: String| {
135                let mut m = HashMap::new();
136                m.insert("up".to_string(), Value::String(s));
137                m.insert("drop".to_string(), Value::Bool(true));
138                m
139            }))
140            .pick(["up"]);
141        let out = seq.invoke("hi".to_string(), None).await.unwrap();
142        assert_eq!(out.len(), 1);
143        assert_eq!(out["up"], Value::String("HI".to_string()));
144    }
145}
lc_core/runnables/pick.rs

lc_core/runnables/
pick.rs
