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>(
41        self,
42        keys: impl IntoIterator<Item = K>,
43    ) -> RunnableSequence<I, HashMap<String, Value>>
44    where
45        K: Into<String>;
46
47    /// Pull a single value out of the dict output.
48    ///
49    /// A missing key yields `Value::Null` (Python raises `KeyError`; `Null`
50    /// keeps the runnable total so the pipeline does not short-circuit).
51    ///
52    /// # Example
53    ///
54    /// ```rust,ignore
55    /// let val = parallel.pluck("a");   // output: Value
56    /// ```
57    fn pluck(self, key: impl Into<String>) -> RunnableSequence<I, Value>;
58}
59
60impl<I, R> RunnablePick<I> for R
61where
62    I: Send + Sync + 'static,
63    R: Runnable<I, HashMap<String, Value>> + Sized + 'static,
64    R::Error: Into<LcelError>,
65{
66    fn pick<K>(
67        self,
68        keys: impl IntoIterator<Item = K>,
69    ) -> RunnableSequence<I, HashMap<String, Value>>
70    where
71        K: Into<String>,
72    {
73        let keys: Vec<String> = keys.into_iter().map(Into::into).collect();
74        let filter = RunnableLambda::new_sync(move |m: HashMap<String, Value>| {
75            keys.iter()
76                .filter_map(|k| m.get(k).cloned().map(|v| (k.clone(), v)))
77                .collect::<HashMap<String, Value>>()
78        });
79        self.pipe(filter)
80    }
81
82    fn pluck(self, key: impl Into<String>) -> RunnableSequence<I, Value> {
83        let key = key.into();
84        let extract = RunnableLambda::new_sync(move |m: HashMap<String, Value>| {
85            m.get(&key).cloned().unwrap_or(Value::Null)
86        });
87        self.pipe(extract)
88    }
89}
90
91#[cfg(test)]
92mod tests {
93    use super::*;
94    use crate::Runnable;
95    use crate::RunnableLambda;
96    use crate::RunnableParallel;
97
98    fn parallel() -> RunnableParallel<String> {
99        RunnableParallel::<String>::new()
100            .with("len", RunnableLambda::new_sync(|s: String| s.len() as i64))
101            .with(
102                "upper",
103                RunnableLambda::new_sync(|s: String| s.to_uppercase()),
104            )
105    }
106
107    #[tokio::test]
108    async fn pick_keeps_only_selected_keys() {
109        let chain = parallel().pick(["len"]);
110        let out = chain.invoke("hello".to_string(), None).await.unwrap();
111        assert_eq!(out.len(), 1);
112        assert!(out.contains_key("len"));
113        assert!(!out.contains_key("upper"));
114    }
115
116    #[tokio::test]
117    async fn pick_multiple_keys_and_missing() {
118        let chain = parallel().pick(["len", "nope"]);
119        let out = chain.invoke("hello".to_string(), None).await.unwrap();
120        // "nope" 不存在 → 被省略,只剩 len
121        assert_eq!(out.len(), 1);
122        assert!(out.contains_key("len"));
123    }
124
125    #[tokio::test]
126    async fn pluck_returns_single_value() {
127        let chain = parallel().pluck("upper");
128        let out = chain.invoke("hello".to_string(), None).await.unwrap();
129        assert_eq!(out, Value::String("HELLO".to_string()));
130    }
131
132    #[tokio::test]
133    async fn pluck_missing_key_yields_null() {
134        let chain = parallel().pluck("missing");
135        let out = chain.invoke("hello".to_string(), None).await.unwrap();
136        assert_eq!(out, Value::Null);
137    }
138
139    #[tokio::test]
140    async fn pick_works_on_sequence_ending_in_map() {
141        // RunnableSequence<I, HashMap<String, Value>> 也自动获得 RunnablePick
142        let seq = RunnableLambda::new_sync(|s: String| s.to_uppercase())
143            .pipe(RunnableLambda::new_sync(|s: String| {
144                let mut m = HashMap::new();
145                m.insert("up".to_string(), Value::String(s));
146                m.insert("drop".to_string(), Value::Bool(true));
147                m
148            }))
149            .pick(["up"]);
150        let out = seq.invoke("hi".to_string(), None).await.unwrap();
151        assert_eq!(out.len(), 1);
152        assert_eq!(out["up"], Value::String("HI".to_string()));
153    }
154}
lc_core/runnables/pick.rs

lc_core/runnables/
pick.rs
