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.push((key.to_string(), into_runnable_any(wrapped)));
78        self
79    }
80
81    /// Number of mappings.
82    pub fn len(&self) -> usize {
83        self.mappings.len()
84    }
85
86    /// Whether there are no mappings.
87    pub fn is_empty(&self) -> bool {
88        self.mappings.is_empty()
89    }
90}
91
92impl Default for RunnableAssign {
93    fn default() -> Self {
94        Self::new()
95    }
96}
97
98/// Internal wrapper that serializes the output of a Runnable to Value.
99struct AssignStepWrapper<O, R>
100where
101    O: serde::Serialize + Send + Sync + 'static,
102    R: Runnable<HashMap<String, Value>, O>,
103{
104    inner: R,
105    serialize: fn(&O) -> Result<Value, serde_json::Error>,
106    _marker: std::marker::PhantomData<O>,
107}
108
109#[async_trait]
110impl<O, R> Runnable<HashMap<String, Value>, Value> for AssignStepWrapper<O, R>
111where
112    O: serde::Serialize + Send + Sync + 'static,
113    R: Runnable<HashMap<String, Value>, O>,
114    R::Error: Into<LcelError>,
115{
116    type Error = LcelError;
117
118    async fn invoke(
119        &self,
120        input: HashMap<String, Value>,
121        config: Option<RunnableConfig>,
122    ) -> Result<Value, LcelError> {
123        let result = self.inner.invoke(input, config).await.map_err(Into::into)?;
124        (self.serialize)(&result)
125            .map_err(|e| LcelError::Other(format!("assign serialization: {}", e)))
126    }
127}
128
129#[async_trait]
130impl Runnable<HashMap<String, Value>, HashMap<String, Value>> for RunnableAssign {
131    type Error = LcelError;
132
133    /// Execute all mappings and merge results into the input HashMap.
134    async fn invoke(
135        &self,
136        mut input: HashMap<String, Value>,
137        config: Option<RunnableConfig>,
138    ) -> Result<HashMap<String, Value>, LcelError> {
139        use std::any::Any;
140
141        for (key, step) in &self.mappings {
142            let boxed_input = Box::new(input.clone()) as Box<dyn Any + Send>;
143            let result = step.invoke_any(boxed_input, config.clone()).await?;
144
145            // The result should be a Value (from AssignStepWrapper)
146            let value = result
147                .downcast::<Value>()
148                .map(|b| *b)
149                .map_err(|_| LcelError::TypeMismatch(format!(
150                    "assign step output downcast: expected Value, got unknown type for key '{}'",
151                    key
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<Pin<Box<dyn Stream<Item = Result<HashMap<String, Value>, LcelError>> + Send>>, LcelError> {
166        let result = self.invoke(input, config).await?;
167        Ok(Box::pin(futures_util::stream::once(async move { Ok(result) })))
168    }
169
170    /// Batch: invoke per input.
171    async fn batch(
172        &self,
173        inputs: Vec<HashMap<String, Value>>,
174        config: Option<RunnableConfig>,
175    ) -> Result<Vec<HashMap<String, Value>>, LcelError> {
176        let mut results = Vec::with_capacity(inputs.len());
177        for input in inputs {
178            results.push(self.invoke(input, config.clone()).await?);
179        }
180        Ok(results)
181    }
182}
183
184#[cfg(test)]
185mod tests {
186    use super::*;
187    use crate::runnables::RunnableLambda;
188    use futures_util::StreamExt;
189
190    #[tokio::test]
191    async fn assign_single_field() {
192        let assign = RunnableAssign::new()
193            .with("length", RunnableLambda::new_sync(|m: HashMap<String, Value>| {
194                m.get("text")
195                    .and_then(|v| v.as_str())
196                    .map(|s| s.len() as i64)
197                    .unwrap_or(0)
198            }));
199
200        let mut input = HashMap::new();
201        input.insert("text".to_string(), Value::String("hello".to_string()));
202
203        let result = assign.invoke(input, None).await.unwrap();
204        assert_eq!(result.get("text").unwrap(), &Value::String("hello".to_string()));
205        assert_eq!(result.get("length").unwrap(), &Value::Number(serde_json::Number::from(5)));
206    }
207
208    #[tokio::test]
209    async fn assign_multiple_fields() {
210        let assign = RunnableAssign::new()
211            .with("length", RunnableLambda::new_sync(|m: HashMap<String, Value>| {
212                m.get("text")
213                    .and_then(|v| v.as_str())
214                    .map(|s| s.len() as i64)
215                    .unwrap_or(0)
216            }))
217            .with("upper", RunnableLambda::new_sync(|m: HashMap<String, Value>| {
218                m.get("text")
219                    .and_then(|v| v.as_str())
220                    .map(|s| s.to_uppercase())
221                    .unwrap_or_default()
222            }));
223
224        let mut input = HashMap::new();
225        input.insert("text".to_string(), Value::String("hello".to_string()));
226
227        let result = assign.invoke(input, None).await.unwrap();
228        assert_eq!(result.get("length").unwrap(), &Value::Number(serde_json::Number::from(5)));
229        assert_eq!(result.get("upper").unwrap(), &Value::String("HELLO".to_string()));
230        // Original field preserved
231        assert_eq!(result.get("text").unwrap(), &Value::String("hello".to_string()));
232    }
233
234    #[tokio::test]
235    async fn assign_overwrites_existing_key() {
236        let assign = RunnableAssign::new()
237            .with("text", RunnableLambda::new_sync(|_m: HashMap<String, Value>| {
238                "replaced".to_string()
239            }));
240
241        let mut input = HashMap::new();
242        input.insert("text".to_string(), Value::String("original".to_string()));
243
244        let result = assign.invoke(input, None).await.unwrap();
245        assert_eq!(result.get("text").unwrap(), &Value::String("replaced".to_string()));
246    }
247
248    #[tokio::test]
249    async fn assign_stream_works() {
250        let assign = RunnableAssign::new()
251            .with("length", RunnableLambda::new_sync(|m: HashMap<String, Value>| {
252                m.get("text")
253                    .and_then(|v| v.as_str())
254                    .map(|s| s.len() as i64)
255                    .unwrap_or(0)
256            }));
257
258        let mut input = HashMap::new();
259        input.insert("text".to_string(), Value::String("hi".to_string()));
260
261        let mut stream = assign.stream(input, None).await.unwrap();
262        let result = stream.next().await.unwrap().unwrap();
263        assert_eq!(result.get("length").unwrap(), &Value::Number(serde_json::Number::from(2)));
264    }
265
266    #[tokio::test]
267    async fn assign_batch_works() {
268        let assign = RunnableAssign::new()
269            .with("length", RunnableLambda::new_sync(|m: HashMap<String, Value>| {
270                m.get("text")
271                    .and_then(|v| v.as_str())
272                    .map(|s| s.len() as i64)
273                    .unwrap_or(0)
274            }));
275
276        let mut input1 = HashMap::new();
277        input1.insert("text".to_string(), Value::String("hi".to_string()));
278
279        let mut input2 = HashMap::new();
280        input2.insert("text".to_string(), Value::String("hello".to_string()));
281
282        let results = assign.batch(vec![input1, input2], None).await.unwrap();
283        assert_eq!(results[0].get("length").unwrap(), &Value::Number(serde_json::Number::from(2)));
284        assert_eq!(results[1].get("length").unwrap(), &Value::Number(serde_json::Number::from(5)));
285    }
286
287    #[tokio::test]
288    async fn assign_empty_works() {
289        let assign = RunnableAssign::new();
290        let mut input = HashMap::new();
291        input.insert("key".to_string(), Value::String("value".to_string()));
292
293        let result = assign.invoke(input, None).await.unwrap();
294        assert_eq!(result.len(), 1);
295    }
296}
lc_core/runnables/assign.rs

lc_core/runnables/
assign.rs
