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        // Buffer all input, downcast to I, take the last, invoke, return as stream.
151        // We can't call self.inner.transform() directly because the error types
152        // don't match (LcelError vs Self::Error). Instead, we use the default
153        // transform semantics: buffer → last → invoke.
154        use futures_util::StreamExt;
155
156        let mut items = Vec::new();
157        let mut input = input;
158        while let Some(item) = input.next().await {
159            let boxed = item?;
160            let typed = boxed.downcast::<I>().map_err(|_| {
161                LcelError::TypeMismatch(format!(
162                    "transform_any input: expected {}",
163                    std::any::type_name::<I>()
164                ))
165            })?;
166            items.push(*typed);
167        }
168
169        if let Some(last) = items.into_iter().last() {
170            let result = self.inner.invoke(last, config).await.map_err(Into::into)?;
171            Ok(Box::pin(futures_util::stream::once(async move {
172                Ok(Box::new(result) as Box<dyn Any + Send>)
173            })))
174        } else {
175            Ok(Box::pin(futures_util::stream::empty()))
176        }
177    }
178
179    async fn batch_any(
180        &self,
181        inputs: Vec<Box<dyn Any + Send>>,
182        config: Option<RunnableConfig>,
183    ) -> Result<Vec<Box<dyn Any + Send>>, LcelError> {
184        let typed_inputs: Vec<I> = inputs
185            .into_iter()
186            .map(|boxed| {
187                boxed.downcast::<I>().map(|b| *b).map_err(|_| {
188                    LcelError::TypeMismatch(format!(
189                        "batch_any: expected {}",
190                        std::any::type_name::<I>()
191                    ))
192                })
193            })
194            .collect::<Result<Vec<I>, LcelError>>()?;
195        let results = self
196            .inner
197            .batch(typed_inputs, config)
198            .await
199            .map_err(Into::into)?;
200        Ok(results
201            .into_iter()
202            .map(|r| Box::new(r) as Box<dyn Any + Send>)
203            .collect())
204    }
205}
206
207/// Helper function to convert any `Runnable` into `Box<dyn RunnableAny>`.
208///
209/// This is used internally by `RunnableSequence` and `RunnableExt`
210/// to wrap typed runnables into type-erased boxes.
211pub fn into_runnable_any<I, O, R>(runnable: R) -> Box<dyn RunnableAny>
212where
213    I: Send + Sync + 'static,
214    O: Send + Sync + 'static,
215    R: super::Runnable<I, O> + 'static,
216    R::Error: Into<LcelError>,
217{
218    Box::new(RunnableAnyWrapper::new(runnable))
219}
220
221#[cfg(test)]
222mod tests {
223    use super::*;
224    use futures_util::StreamExt;
225
226    struct AddOne;
227
228    #[async_trait]
229    impl super::super::Runnable<i32, i32> for AddOne {
230        type Error = std::convert::Infallible;
231
232        async fn invoke(
233            &self,
234            input: i32,
235            _config: Option<RunnableConfig>,
236        ) -> Result<i32, Self::Error> {
237            Ok(input + 1)
238        }
239    }
240
241    #[tokio::test]
242    async fn invoke_any_works() {
243        let wrapper = RunnableAnyWrapper::new(AddOne);
244        let input: Box<dyn Any + Send> = Box::new(41i32);
245        let result = wrapper.invoke_any(input, None).await.unwrap();
246        let output: i32 = *result.downcast::<i32>().unwrap();
247        assert_eq!(output, 42);
248    }
249
250    #[tokio::test]
251    async fn batch_any_works() {
252        let wrapper = RunnableAnyWrapper::new(AddOne);
253        let inputs: Vec<Box<dyn Any + Send>> = vec![Box::new(1i32), Box::new(2i32), Box::new(3i32)];
254        let results = wrapper.batch_any(inputs, None).await.unwrap();
255        let outputs: Vec<i32> = results
256            .into_iter()
257            .map(|b| *b.downcast::<i32>().unwrap())
258            .collect();
259        assert_eq!(outputs, vec![2, 3, 4]);
260    }
261
262    #[tokio::test]
263    async fn stream_any_works() {
264        let wrapper = RunnableAnyWrapper::new(AddOne);
265        let input: Box<dyn Any + Send> = Box::new(9i32);
266        let mut stream = wrapper.stream_any(input, None).await.unwrap();
267        let result = stream.next().await.unwrap().unwrap();
268        let output: i32 = *result.downcast::<i32>().unwrap();
269        assert_eq!(output, 10);
270    }
271
272    #[tokio::test]
273    async fn invoke_any_type_mismatch() {
274        let wrapper = RunnableAnyWrapper::new(AddOne);
275        let wrong_input: Box<dyn Any + Send> = Box::new("not an i32");
276        let result = wrapper.invoke_any(wrong_input, None).await;
277        assert!(result.is_err());
278        let err = result.unwrap_err();
279        assert!(matches!(err, LcelError::TypeMismatch(_)));
280    }
281
282    #[tokio::test]
283    async fn into_runnable_any_works() {
284        let boxed: Box<dyn RunnableAny> = into_runnable_any::<i32, i32, _>(AddOne);
285        let input: Box<dyn Any + Send> = Box::new(5i32);
286        let result = boxed.invoke_any(input, None).await.unwrap();
287        let output: i32 = *result.downcast::<i32>().unwrap();
288        assert_eq!(output, 6);
289    }
290}