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