lc_core/runnables/
assign.rs

1// lc-core/src/runnables/assign.rs
2//! RunnableAssign - inject new fields into a HashMap pipeline.
3//!
4//! `RunnableAssign` adds new key-value pairs to a `HashMap<String, Value>`,
5//! enabling RAG pipelines where the context is injected alongside the
6//! original question.
7//!
8//! # Example
9//!
10//! ```rust,ignore
11//! let chain = RunnableParallel::<String>::new()
12//!     .assign("context", retriever.pipe(format_docs))
13//!     .assign("question", RunnablePassthrough)
14//!     .pipe(prompt_template)
15//!     .pipe(llm);
16//! ```
17
18use super::any::RunnableAny;
19use super::config::RunnableConfig;
20use super::error::LcelError;
21use super::runnable_trait::Runnable;
22use async_trait::async_trait;
23use futures_util::Stream;
24use serde_json::Value;
25use std::collections::HashMap;
26use std::pin::Pin;
27
28/// A `Runnable` that adds new key-value pairs to a `HashMap<String, Value>`.
29///
30/// Each mapping runs a `Runnable` on the input HashMap and merges the
31/// result back into the HashMap under the specified key.
32///
33/// This is the LCEL equivalent of Python's `RunnableAssign`.
34pub struct RunnableAssign {
35    mappings: Vec<(String, Box<dyn RunnableAny>)>,
36}
37
38impl std::fmt::Debug for RunnableAssign {
39    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
40        let keys: Vec<&str> = self.mappings.iter().map(|(k, _)| k.as_str()).collect();
41        f.debug_struct("RunnableAssign")
42            .field("mappings", &keys)
43            .finish()
44    }
45}
46
47impl RunnableAssign {
48    /// Create an empty assign runnable.
49    pub fn new() -> Self {
50        Self {
51            mappings: Vec::new(),
52        }
53    }
54
55    /// Add a mapping that runs the given runnable and stores the result
56    /// under the specified key.
57    ///
58    /// The runnable takes `HashMap<String, Value>` as input and produces
59    /// any output that implements `serde::Serialize`. The output is
60    /// serialized to `serde_json::Value` and merged into the HashMap.
61    pub fn with<O, R>(mut self, key: &str, runnable: R) -> Self
62    where
63        O: serde::Serialize + Send + Sync + 'static,
64        R: Runnable<HashMap<String, Value>, O> + Send + Sync + 'static,
65        R::Error: Into<LcelError>,
66    {
67        use super::any::into_runnable_any;
68
69        // Wrap the runnable so that its output is serialized to Value
70        // and then we can store it in the HashMap
71        let wrapped = AssignStepWrapper {
72            inner: runnable,
73            serialize: |output: &O| serde_json::to_value(output),
74            _marker: std::marker::PhantomData,
75        };
76
77        self.mappings
78            .push((key.to_string(), into_runnable_any(wrapped)));
79        self
80    }
81
82    /// Number of mappings.
83    pub fn len(&self) -> usize {
84        self.mappings.len()
85    }
86
87    /// Whether there are no mappings.
88    pub fn is_empty(&self) -> bool {
89        self.mappings.is_empty()
90    }
91}
92
93impl Default for RunnableAssign {
94    fn default() -> Self {
95        Self::new()
96    }
97}
98
99/// Internal wrapper that serializes the output of a Runnable to Value.
100struct AssignStepWrapper<O, R>
101where
102    O: serde::Serialize + Send + Sync + 'static,
103    R: Runnable<HashMap<String, Value>, O>,
104{
105    inner: R,
106    serialize: fn(&O) -> Result<Value, serde_json::Error>,
107    _marker: std::marker::PhantomData<O>,
108}
109
110#[async_trait]
111impl<O, R> Runnable<HashMap<String, Value>, Value> for AssignStepWrapper<O, R>
112where
113    O: serde::Serialize + Send + Sync + 'static,
114    R: Runnable<HashMap<String, Value>, O>,
115    R::Error: Into<LcelError>,
116{
117    type Error = LcelError;
118
119    async fn invoke(
120        &self,
121        input: HashMap<String, Value>,
122        config: Option<RunnableConfig>,
123    ) -> Result<Value, LcelError> {
124        let result = self.inner.invoke(input, config).await.map_err(Into::into)?;
125        (self.serialize)(&result)
126            .map_err(|e| LcelError::Other(format!("assign serialization: {}", e)))
127    }
128}
129
130#[async_trait]
131impl Runnable<HashMap<String, Value>, HashMap<String, Value>> for RunnableAssign {
132    type Error = LcelError;
133
134    /// Execute all mappings and merge results into the input HashMap.
135    async fn invoke(
136        &self,
137        mut input: HashMap<String, Value>,
138        config: Option<RunnableConfig>,
139    ) -> Result<HashMap<String, Value>, LcelError> {
140        use std::any::Any;
141
142        for (key, step) in &self.mappings {
143            let boxed_input = Box::new(input.clone()) as Box<dyn Any + Send>;
144            let result = step.invoke_any(boxed_input, config.clone()).await?;
145
146            // The result should be a Value (from AssignStepWrapper)
147            let value = result.downcast::<Value>().map(|b| *b).map_err(|_| {
148                LcelError::TypeMismatch(format!(
149                    "assign step output downcast: expected Value, got unknown type for key '{}'",
150                    key
151                ))
152            })?;
153
154            input.insert(key.clone(), value);
155        }
156
157        Ok(input)
158    }
159
160    /// Stream: invoke and return single-element stream.
161    async fn stream(
162        &self,
163        input: HashMap<String, Value>,
164        config: Option<RunnableConfig>,
165    ) -> Result<
166        Pin<Box<dyn Stream<Item = Result<HashMap<String, Value>, LcelError>> + Send>>,
167        LcelError,
168    > {
169        let result = self.invoke(input, config).await?;
170        Ok(Box::pin(futures_util::stream::once(
171            async move { Ok(result) },
172        )))
173    }
174
175    /// Batch: invoke per input.
176    async fn batch(
177        &self,
178        inputs: Vec<HashMap<String, Value>>,
179        config: Option<RunnableConfig>,
180    ) -> Result<Vec<HashMap<String, Value>>, LcelError> {
181        let mut results = Vec::with_capacity(inputs.len());
182        for input in inputs {
183            results.push(self.invoke(input, config.clone()).await?);
184        }
185        Ok(results)
186    }
187}
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192    use crate::runnables::RunnableLambda;
193    use futures_util::StreamExt;
194
195    #[tokio::test]
196    async fn assign_single_field() {
197        let assign = RunnableAssign::new().with(
198            "length",
199            RunnableLambda::new_sync(|m: HashMap<String, Value>| {
200                m.get("text")
201                    .and_then(|v| v.as_str())
202                    .map(|s| s.len() as i64)
203                    .unwrap_or(0)
204            }),
205        );
206
207        let mut input = HashMap::new();
208        input.insert("text".to_string(), Value::String("hello".to_string()));
209
210        let result = assign.invoke(input, None).await.unwrap();
211        assert_eq!(
212            result.get("text").unwrap(),
213            &Value::String("hello".to_string())
214        );
215        assert_eq!(
216            result.get("length").unwrap(),
217            &Value::Number(serde_json::Number::from(5))
218        );
219    }
220
221    #[tokio::test]
222    async fn assign_multiple_fields() {
223        let assign = RunnableAssign::new()
224            .with(
225                "length",
226                RunnableLambda::new_sync(|m: HashMap<String, Value>| {
227                    m.get("text")
228                        .and_then(|v| v.as_str())
229                        .map(|s| s.len() as i64)
230                        .unwrap_or(0)
231                }),
232            )
233            .with(
234                "upper",
235                RunnableLambda::new_sync(|m: HashMap<String, Value>| {
236                    m.get("text")
237                        .and_then(|v| v.as_str())
238                        .map(|s| s.to_uppercase())
239                        .unwrap_or_default()
240                }),
241            );
242
243        let mut input = HashMap::new();
244        input.insert("text".to_string(), Value::String("hello".to_string()));
245
246        let result = assign.invoke(input, None).await.unwrap();
247        assert_eq!(
248            result.get("length").unwrap(),
249            &Value::Number(serde_json::Number::from(5))
250        );
251        assert_eq!(
252            result.get("upper").unwrap(),
253            &Value::String("HELLO".to_string())
254        );
255        // Original field preserved
256        assert_eq!(
257            result.get("text").unwrap(),
258            &Value::String("hello".to_string())
259        );
260    }
261
262    #[tokio::test]
263    async fn assign_overwrites_existing_key() {
264        let assign = RunnableAssign::new().with(
265            "text",
266            RunnableLambda::new_sync(|_m: HashMap<String, Value>| "replaced".to_string()),
267        );
268
269        let mut input = HashMap::new();
270        input.insert("text".to_string(), Value::String("original".to_string()));
271
272        let result = assign.invoke(input, None).await.unwrap();
273        assert_eq!(
274            result.get("text").unwrap(),
275            &Value::String("replaced".to_string())
276        );
277    }
278
279    #[tokio::test]
280    async fn assign_stream_works() {
281        let assign = RunnableAssign::new().with(
282            "length",
283            RunnableLambda::new_sync(|m: HashMap<String, Value>| {
284                m.get("text")
285                    .and_then(|v| v.as_str())
286                    .map(|s| s.len() as i64)
287                    .unwrap_or(0)
288            }),
289        );
290
291        let mut input = HashMap::new();
292        input.insert("text".to_string(), Value::String("hi".to_string()));
293
294        let mut stream = assign.stream(input, None).await.unwrap();
295        let result = stream.next().await.unwrap().unwrap();
296        assert_eq!(
297            result.get("length").unwrap(),
298            &Value::Number(serde_json::Number::from(2))
299        );
300    }
301
302    #[tokio::test]
303    async fn assign_batch_works() {
304        let assign = RunnableAssign::new().with(
305            "length",
306            RunnableLambda::new_sync(|m: HashMap<String, Value>| {
307                m.get("text")
308                    .and_then(|v| v.as_str())
309                    .map(|s| s.len() as i64)
310                    .unwrap_or(0)
311            }),
312        );
313
314        let mut input1 = HashMap::new();
315        input1.insert("text".to_string(), Value::String("hi".to_string()));
316
317        let mut input2 = HashMap::new();
318        input2.insert("text".to_string(), Value::String("hello".to_string()));
319
320        let results = assign.batch(vec![input1, input2], None).await.unwrap();
321        assert_eq!(
322            results[0].get("length").unwrap(),
323            &Value::Number(serde_json::Number::from(2))
324        );
325        assert_eq!(
326            results[1].get("length").unwrap(),
327            &Value::Number(serde_json::Number::from(5))
328        );
329    }
330
331    #[tokio::test]
332    async fn assign_empty_works() {
333        let assign = RunnableAssign::new();
334        let mut input = HashMap::new();
335        input.insert("key".to_string(), Value::String("value".to_string()));
336
337        let result = assign.invoke(input, None).await.unwrap();
338        assert_eq!(result.len(), 1);
339    }
340}
lc_core/runnables/assign.rs

lc_core/runnables/
assign.rs
