Skip to main content

lc_core/runnables/
any.rs

1// lc-core/src/runnables/any.rs
2//! Type-erased Runnable trait for LCEL pipeline internals.
3//!
4//! `RunnableAny` erases the generic `Input`/`Output` types of `Runnable`
5//! into `Box<dyn Any + Send>`, allowing heterogeneous steps to be stored
6//! in a single `RunnableSequence`.
7//!
8//! # Safety Guarantee
9//!
10//! Type safety is maintained at the `pipe()` boundary: the compiler
11//! ensures that `A: Runnable<I, M>` and `B: Runnable<M, O>` have matching
12//! intermediate types. The `Any` downcast only happens internally within
13//! `RunnableSequence`, where the type relationship is already proven.
14
15use super::config::RunnableConfig;
16use super::error::LcelError;
17use async_trait::async_trait;
18use futures_util::{Stream, StreamExt};
19use std::any::Any;
20use std::pin::Pin;
21
22/// Type-erased Runnable with unified `LcelError`.
23///
24/// This trait is the runtime representation of a `Runnable<I, O>` step
25/// inside a `RunnableSequence`. All generic types are erased to `Box<dyn Any + Send>`.
26#[async_trait]
27pub trait RunnableAny: Send + Sync {
28    /// Type-erased invoke: `Box<dyn Any + Send>` → `Box<dyn Any + Send>`.
29    async fn invoke_any(
30        &self,
31        input: Box<dyn Any + Send>,
32        config: Option<RunnableConfig>,
33    ) -> Result<Box<dyn Any + Send>, LcelError>;
34
35    /// Type-erased stream: `Box<dyn Any + Send>` → Stream of `Box<dyn Any + Send>`.
36    async fn stream_any(
37        &self,
38        input: Box<dyn Any + Send>,
39        config: Option<RunnableConfig>,
40    ) -> Result<Pin<Box<dyn Stream<Item = Result<Box<dyn Any + Send>, LcelError>> + Send>>, LcelError>;
41
42    /// Type-erased transform: Stream of `Box<dyn Any + Send>` → Stream of `Box<dyn Any + Send>`.
43    ///
44    /// This is the core of LCEL streaming: each step takes an input stream
45    /// and produces an output stream, enabling pipeline streaming without
46    /// buffering intermediate results.
47    async fn transform_any(
48        &self,
49        input: Pin<Box<dyn Stream<Item = Result<Box<dyn Any + Send>, LcelError>> + Send>>,
50        config: Option<RunnableConfig>,
51    ) -> Result<Pin<Box<dyn Stream<Item = Result<Box<dyn Any + Send>, LcelError>> + Send>>, LcelError>;
52
53    /// Type-erased batch: `Vec<Box<dyn Any + Send>>` → `Vec<Box<dyn Any + Send>>`.
54    async fn batch_any(
55        &self,
56        inputs: Vec<Box<dyn Any + Send>>,
57        config: Option<RunnableConfig>,
58    ) -> Result<Vec<Box<dyn Any + Send>>, LcelError>;
59}
60
61/// Wrapper that implements `RunnableAny` for any `Runnable<I, O>`.
62///
63/// We use a concrete wrapper struct instead of a blanket impl because
64/// Rust's orphan rules and type parameter constraints make a blanket
65/// `impl<I, O, R> RunnableAny for R` impossible (I and O would be
66/// unconstrained).
67pub struct RunnableAnyWrapper<I, O, R>
68where
69    I: Send + Sync + 'static,
70    O: Send + Sync + 'static,
71    R: super::Runnable<I, O>,
72{
73    inner: R,
74    _marker: std::marker::PhantomData<(I, O)>,
75}
76
77impl<I, O, R> RunnableAnyWrapper<I, O, R>
78where
79    I: Send + Sync + 'static,
80    O: Send + Sync + 'static,
81    R: super::Runnable<I, O>,
82{
83    /// Create a new wrapper.
84    pub fn new(runnable: R) -> Self {
85        Self {
86            inner: runnable,
87            _marker: std::marker::PhantomData,
88        }
89    }
90}
91
92#[async_trait]
93impl<I, O, R> RunnableAny for RunnableAnyWrapper<I, O, R>
94where
95    I: Send + Sync + 'static,
96    O: Send + Sync + 'static,
97    R: super::Runnable<I, O> + 'static,
98    R::Error: Into<LcelError>,
99{
100    async fn invoke_any(
101        &self,
102        input: Box<dyn Any + Send>,
103        config: Option<RunnableConfig>,
104    ) -> Result<Box<dyn Any + Send>, LcelError> {
105        let typed_input = input.downcast::<I>().map_err(|_| {
106            LcelError::TypeMismatch(format!(
107                "invoke_any: expected {}, got unknown type",
108                std::any::type_name::<I>()
109            ))
110        })?;
111        let result = self
112            .inner
113            .invoke(*typed_input, config)
114            .await
115            .map_err(Into::into)?;
116        Ok(Box::new(result) as Box<dyn Any + Send>)
117    }
118
119    async fn stream_any(
120        &self,
121        input: Box<dyn Any + Send>,
122        config: Option<RunnableConfig>,
123    ) -> Result<Pin<Box<dyn Stream<Item = Result<Box<dyn Any + Send>, LcelError>> + Send>>, LcelError>
124    {
125        let typed_input = input.downcast::<I>().map_err(|_| {
126            LcelError::TypeMismatch(format!(
127                "stream_any: expected {}, got unknown type",
128                std::any::type_name::<I>()
129            ))
130        })?;
131        let stream = self
132            .inner
133            .stream(*typed_input, config)
134            .await
135            .map_err(Into::into)?;
136        let any_stream = stream.map(|result| {
137            result
138                .map(|output| Box::new(output) as Box<dyn Any + Send>)
139                .map_err(Into::into)
140        });
141        Ok(Box::pin(any_stream))
142    }
143
144    async fn transform_any(
145        &self,
146        input: Pin<Box<dyn Stream<Item = Result<Box<dyn Any + Send>, LcelError>> + Send>>,
147        config: Option<RunnableConfig>,
148    ) -> Result<Pin<Box<dyn Stream<Item = Result<Box<dyn Any + Send>, LcelError>> + Send>>, LcelError>
149    {
150        // We can't forward the input stream to `self.inner.transform()` directly
151        // because the error types don't match (LcelError vs `R::Error` and
152        // `R::Error: From<LcelError>` does not hold for e.g. `Infallible`).
153        //
154        // Instead we drive each input item through `self.inner.stream(item, ..)`
155        // and concatenate the per-item streams in order. This is the LangChain
156        // default `transform` semantics: a step that overrides `stream` (e.g. an
157        // LLM) yields a real token stream per item, while a step using the
158        // default `stream` (single-element) degrades to elementwise `invoke`.
159        use futures_util::StreamExt;
160
161        let mut items = Vec::new();
162        let mut input = input;
163        while let Some(item) = input.next().await {
164            let boxed = item?;
165            let typed = boxed.downcast::<I>().map_err(|_| {
166                LcelError::TypeMismatch(format!(
167                    "transform_any input: expected {}",
168                    std::any::type_name::<I>()
169                ))
170            })?;
171            items.push(*typed);
172        }
173
174        let mut per_item_streams = Vec::with_capacity(items.len());
175        for item in items {
176            let stream = self
177                .inner
178                .stream(item, config.clone())
179                .await
180                .map_err(Into::into)?;
181            let any_stream = stream.map(|result| {
182                result
183                    .map(|output| Box::new(output) as Box<dyn Any + Send>)
184                    .map_err(Into::into)
185            });
186            per_item_streams.push(any_stream);
187        }
188
189        let flattened = futures_util::stream::iter(per_item_streams).flatten();
190        Ok(Box::pin(flattened))
191    }
192
193    async fn batch_any(
194        &self,
195        inputs: Vec<Box<dyn Any + Send>>,
196        config: Option<RunnableConfig>,
197    ) -> Result<Vec<Box<dyn Any + Send>>, LcelError> {
198        let typed_inputs: Vec<I> = inputs
199            .into_iter()
200            .map(|boxed| {
201                boxed.downcast::<I>().map(|b| *b).map_err(|_| {
202                    LcelError::TypeMismatch(format!(
203                        "batch_any: expected {}",
204                        std::any::type_name::<I>()
205                    ))
206                })
207            })
208            .collect::<Result<Vec<I>, LcelError>>()?;
209        let results = self
210            .inner
211            .batch(typed_inputs, config)
212            .await
213            .map_err(Into::into)?;
214        Ok(results
215            .into_iter()
216            .map(|r| Box::new(r) as Box<dyn Any + Send>)
217            .collect())
218    }
219}
220
221/// Helper function to convert any `Runnable` into `Box<dyn RunnableAny>`.
222///
223/// This is used internally by `RunnableSequence` and `RunnableExt`
224/// to wrap typed runnables into type-erased boxes.
225pub fn into_runnable_any<I, O, R>(runnable: R) -> Box<dyn RunnableAny>
226where
227    I: Send + Sync + 'static,
228    O: Send + Sync + 'static,
229    R: super::Runnable<I, O> + 'static,
230    R::Error: Into<LcelError>,
231{
232    Box::new(RunnableAnyWrapper::new(runnable))
233}
234
235#[cfg(test)]
236mod tests {
237    use super::*;
238    use futures_util::StreamExt;
239
240    struct AddOne;
241
242    #[async_trait]
243    impl super::super::Runnable<i32, i32> for AddOne {
244        type Error = std::convert::Infallible;
245
246        async fn invoke(
247            &self,
248            input: i32,
249            _config: Option<RunnableConfig>,
250        ) -> Result<i32, Self::Error> {
251            Ok(input + 1)
252        }
253    }
254
255    #[tokio::test]
256    async fn invoke_any_works() {
257        let wrapper = RunnableAnyWrapper::new(AddOne);
258        let input: Box<dyn Any + Send> = Box::new(41i32);
259        let result = wrapper.invoke_any(input, None).await.unwrap();
260        let output: i32 = *result.downcast::<i32>().unwrap();
261        assert_eq!(output, 42);
262    }
263
264    #[tokio::test]
265    async fn batch_any_works() {
266        let wrapper = RunnableAnyWrapper::new(AddOne);
267        let inputs: Vec<Box<dyn Any + Send>> = vec![Box::new(1i32), Box::new(2i32), Box::new(3i32)];
268        let results = wrapper.batch_any(inputs, None).await.unwrap();
269        let outputs: Vec<i32> = results
270            .into_iter()
271            .map(|b| *b.downcast::<i32>().unwrap())
272            .collect();
273        assert_eq!(outputs, vec![2, 3, 4]);
274    }
275
276    #[tokio::test]
277    async fn stream_any_works() {
278        let wrapper = RunnableAnyWrapper::new(AddOne);
279        let input: Box<dyn Any + Send> = Box::new(9i32);
280        let mut stream = wrapper.stream_any(input, None).await.unwrap();
281        let result = stream.next().await.unwrap().unwrap();
282        let output: i32 = *result.downcast::<i32>().unwrap();
283        assert_eq!(output, 10);
284    }
285
286    #[tokio::test]
287    async fn invoke_any_type_mismatch() {
288        let wrapper = RunnableAnyWrapper::new(AddOne);
289        let wrong_input: Box<dyn Any + Send> = Box::new("not an i32");
290        let result = wrapper.invoke_any(wrong_input, None).await;
291        assert!(result.is_err());
292        let err = result.unwrap_err();
293        assert!(matches!(err, LcelError::TypeMismatch(_)));
294    }
295
296    #[tokio::test]
297    async fn into_runnable_any_works() {
298        let boxed: Box<dyn RunnableAny> = into_runnable_any::<i32, i32, _>(AddOne);
299        let input: Box<dyn Any + Send> = Box::new(5i32);
300        let result = boxed.invoke_any(input, None).await.unwrap();
301        let output: i32 = *result.downcast::<i32>().unwrap();
302        assert_eq!(output, 6);
303    }
304}