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;
21use std::sync::Arc;
22
23/// Type-erased Runnable with unified `LcelError`.
24///
25/// This trait is the runtime representation of a `Runnable<I, O>` step
26/// inside a `RunnableSequence`. All generic types are erased to `Box<dyn Any + Send>`.
27#[async_trait]
28pub trait RunnableAny: Send + Sync {
29    /// Type-erased invoke: `Box<dyn Any + Send>` → `Box<dyn Any + Send>`.
30    async fn invoke_any(
31        &self,
32        input: Box<dyn Any + Send>,
33        config: Option<RunnableConfig>,
34    ) -> Result<Box<dyn Any + Send>, LcelError>;
35
36    /// Type-erased stream: `Box<dyn Any + Send>` → Stream of `Box<dyn Any + Send>`.
37    async fn stream_any(
38        &self,
39        input: Box<dyn Any + Send>,
40        config: Option<RunnableConfig>,
41    ) -> Result<Pin<Box<dyn Stream<Item = Result<Box<dyn Any + Send>, LcelError>> + Send>>, LcelError>;
42
43    /// Type-erased transform: Stream of `Box<dyn Any + Send>` → Stream of `Box<dyn Any + Send>`.
44    ///
45    /// This is the core of LCEL streaming: each step takes an input stream
46    /// and produces an output stream, enabling pipeline streaming without
47    /// buffering intermediate results.
48    async fn transform_any(
49        &self,
50        input: Pin<Box<dyn Stream<Item = Result<Box<dyn Any + Send>, LcelError>> + Send>>,
51        config: Option<RunnableConfig>,
52    ) -> Result<Pin<Box<dyn Stream<Item = Result<Box<dyn Any + Send>, LcelError>> + Send>>, LcelError>;
53
54    /// Type-erased batch: `Vec<Box<dyn Any + Send>>` → `Vec<Box<dyn Any + Send>>`.
55    async fn batch_any(
56        &self,
57        inputs: Vec<Box<dyn Any + Send>>,
58        config: Option<RunnableConfig>,
59    ) -> Result<Vec<Box<dyn Any + Send>>, LcelError>;
60}
61
62/// Wrapper that implements `RunnableAny` for any `Runnable<I, O>`.
63///
64/// We use a concrete wrapper struct instead of a blanket impl because
65/// Rust's orphan rules and type parameter constraints make a blanket
66/// `impl<I, O, R> RunnableAny for R` impossible (I and O would be
67/// unconstrained).
68pub struct RunnableAnyWrapper<I, O, R>
69where
70    I: Send + Sync + 'static,
71    O: Send + Sync + 'static,
72    R: super::Runnable<I, O>,
73{
74    // `Arc` so that a lazy `transform_any` stream can own a clone of the inner
75    // runnable and keep calling `stream` per item **after** `transform_any`
76    // has returned — no `&self` borrow in the returned stream (which would
77    // otherwise force a `+ '_` lifetime through every caller).
78    inner: Arc<R>,
79    _marker: std::marker::PhantomData<(I, O)>,
80}
81
82impl<I, O, R> RunnableAnyWrapper<I, O, R>
83where
84    I: Send + Sync + 'static,
85    O: Send + Sync + 'static,
86    R: super::Runnable<I, O>,
87{
88    /// Create a new wrapper.
89    pub fn new(runnable: R) -> Self {
90        Self {
91            inner: Arc::new(runnable),
92            _marker: std::marker::PhantomData,
93        }
94    }
95}
96
97#[async_trait]
98impl<I, O, R> RunnableAny for RunnableAnyWrapper<I, O, R>
99where
100    I: Send + Sync + 'static,
101    O: Send + Sync + 'static,
102    R: super::Runnable<I, O> + 'static,
103    R::Error: Into<LcelError>,
104{
105    async fn invoke_any(
106        &self,
107        input: Box<dyn Any + Send>,
108        config: Option<RunnableConfig>,
109    ) -> Result<Box<dyn Any + Send>, LcelError> {
110        let typed_input = input.downcast::<I>().map_err(|_| {
111            LcelError::TypeMismatch(format!(
112                "invoke_any: expected {}, got unknown type",
113                std::any::type_name::<I>()
114            ))
115        })?;
116        let result = self
117            .inner
118            .invoke(*typed_input, config)
119            .await
120            .map_err(Into::into)?;
121        Ok(Box::new(result) as Box<dyn Any + Send>)
122    }
123
124    async fn stream_any(
125        &self,
126        input: Box<dyn Any + Send>,
127        config: Option<RunnableConfig>,
128    ) -> Result<Pin<Box<dyn Stream<Item = Result<Box<dyn Any + Send>, LcelError>> + Send>>, LcelError>
129    {
130        let typed_input = input.downcast::<I>().map_err(|_| {
131            LcelError::TypeMismatch(format!(
132                "stream_any: expected {}, got unknown type",
133                std::any::type_name::<I>()
134            ))
135        })?;
136        let stream = self
137            .inner
138            .stream(*typed_input, config)
139            .await
140            .map_err(Into::into)?;
141        let any_stream = stream.map(|result| {
142            result
143                .map(|output| Box::new(output) as Box<dyn Any + Send>)
144                .map_err(Into::into)
145        });
146        Ok(Box::pin(any_stream))
147    }
148
149    async fn transform_any(
150        &self,
151        input: Pin<Box<dyn Stream<Item = Result<Box<dyn Any + Send>, LcelError>> + Send>>,
152        config: Option<RunnableConfig>,
153    ) -> Result<Pin<Box<dyn Stream<Item = Result<Box<dyn Any + Send>, LcelError>> + Send>>, LcelError>
154    {
155        // We can't forward the input stream to `self.inner.transform()` directly
156        // because the error types don't match (LcelError vs `R::Error` and
157        // `R::Error: From<LcelError>` does not hold for e.g. `Infallible`).
158        //
159        // Instead we drive each input item through `inner.stream(item, ..)`
160        // **lazily**: as soon as an input item arrives it is immediately run
161        // through `stream` and its output yielded, before pulling the next input
162        // item. This matches `Runnable::transform` default semantics — a step
163        // that overrides `stream` (e.g. an LLM) yields a real token stream per
164        // item, while a step using the default `stream` (single-element) degrades
165        // to elementwise `invoke`. Downstream receives output incrementally
166        // instead of waiting for the whole input stream, and an infinite/long-lived
167        // upstream never accumulates unboundedly in memory.
168        //
169        // The stream owns a cloned `Arc` of the inner runnable, so it can keep
170        // calling `stream` per item **after** this method returns — no `&self`
171        // borrow escapes into the returned `'static` stream.
172        //
173        // Note: `?` is not usable inside `async_stream::stream!`, so errors are
174        // yielded and the stream terminates after them.
175        use futures_util::StreamExt;
176
177        let inner = Arc::clone(&self.inner);
178        let config = config.clone();
179        let out = async_stream::stream! {
180            let mut input = input;
181            loop {
182                let boxed = match input.next().await {
183                    Some(item) => item,
184                    None => return,
185                };
186                let boxed = match boxed {
187                    Ok(b) => b,
188                    Err(e) => {
189                        yield Err(e);
190                        return;
191                    }
192                };
193                let typed = match boxed.downcast::<I>() {
194                    Ok(t) => *t,
195                    Err(_) => {
196                        yield Err(LcelError::TypeMismatch(format!(
197                            "transform_any input: expected {}",
198                            std::any::type_name::<I>()
199                        )));
200                        return;
201                    }
202                };
203                // 驱动当前元素经过 `stream`,把输出全部放出去后再拉下一个输入。
204                let item_stream = match inner.stream(typed, config.clone()).await {
205                    Ok(s) => s,
206                    Err(e) => {
207                        yield Err(e.into());
208                        return;
209                    }
210                };
211                let mut any_stream = item_stream.map(|result| {
212                    result
213                        .map(|output| Box::new(output) as Box<dyn Any + Send>)
214                        .map_err(Into::into)
215                });
216                while let Some(res) = any_stream.next().await {
217                    yield res;
218                }
219            }
220        };
221        Ok(Box::pin(out))
222    }
223
224    async fn batch_any(
225        &self,
226        inputs: Vec<Box<dyn Any + Send>>,
227        config: Option<RunnableConfig>,
228    ) -> Result<Vec<Box<dyn Any + Send>>, LcelError> {
229        let typed_inputs: Vec<I> = inputs
230            .into_iter()
231            .map(|boxed| {
232                boxed.downcast::<I>().map(|b| *b).map_err(|_| {
233                    LcelError::TypeMismatch(format!(
234                        "batch_any: expected {}",
235                        std::any::type_name::<I>()
236                    ))
237                })
238            })
239            .collect::<Result<Vec<I>, LcelError>>()?;
240        let results = self
241            .inner
242            .batch(typed_inputs, config)
243            .await
244            .map_err(Into::into)?;
245        Ok(results
246            .into_iter()
247            .map(|r| Box::new(r) as Box<dyn Any + Send>)
248            .collect())
249    }
250}
251
252/// Helper function to convert any `Runnable` into `Box<dyn RunnableAny>`.
253///
254/// This is used internally by `RunnableSequence` and `RunnableExt`
255/// to wrap typed runnables into type-erased boxes.
256pub fn into_runnable_any<I, O, R>(runnable: R) -> Box<dyn RunnableAny>
257where
258    I: Send + Sync + 'static,
259    O: Send + Sync + 'static,
260    R: super::Runnable<I, O> + 'static,
261    R::Error: Into<LcelError>,
262{
263    Box::new(RunnableAnyWrapper::new(runnable))
264}
265
266#[cfg(test)]
267mod tests {
268    use super::*;
269    use futures_util::StreamExt;
270
271    struct AddOne;
272
273    #[async_trait]
274    impl super::super::Runnable<i32, i32> for AddOne {
275        type Error = std::convert::Infallible;
276
277        async fn invoke(
278            &self,
279            input: i32,
280            _config: Option<RunnableConfig>,
281        ) -> Result<i32, Self::Error> {
282            Ok(input + 1)
283        }
284    }
285
286    #[tokio::test]
287    async fn invoke_any_works() {
288        let wrapper = RunnableAnyWrapper::new(AddOne);
289        let input: Box<dyn Any + Send> = Box::new(41i32);
290        let result = wrapper.invoke_any(input, None).await.unwrap();
291        let output: i32 = *result.downcast::<i32>().unwrap();
292        assert_eq!(output, 42);
293    }
294
295    #[tokio::test]
296    async fn batch_any_works() {
297        let wrapper = RunnableAnyWrapper::new(AddOne);
298        let inputs: Vec<Box<dyn Any + Send>> = vec![Box::new(1i32), Box::new(2i32), Box::new(3i32)];
299        let results = wrapper.batch_any(inputs, None).await.unwrap();
300        let outputs: Vec<i32> = results
301            .into_iter()
302            .map(|b| *b.downcast::<i32>().unwrap())
303            .collect();
304        assert_eq!(outputs, vec![2, 3, 4]);
305    }
306
307    #[tokio::test]
308    async fn stream_any_works() {
309        let wrapper = RunnableAnyWrapper::new(AddOne);
310        let input: Box<dyn Any + Send> = Box::new(9i32);
311        let mut stream = wrapper.stream_any(input, None).await.unwrap();
312        let result = stream.next().await.unwrap().unwrap();
313        let output: i32 = *result.downcast::<i32>().unwrap();
314        assert_eq!(output, 10);
315    }
316
317    #[tokio::test]
318    async fn invoke_any_type_mismatch() {
319        let wrapper = RunnableAnyWrapper::new(AddOne);
320        let wrong_input: Box<dyn Any + Send> = Box::new("not an i32");
321        let result = wrapper.invoke_any(wrong_input, None).await;
322        assert!(result.is_err());
323        let err = result.unwrap_err();
324        assert!(matches!(err, LcelError::TypeMismatch(_)));
325    }
326
327    #[tokio::test]
328    async fn into_runnable_any_works() {
329        let boxed: Box<dyn RunnableAny> = into_runnable_any::<i32, i32, _>(AddOne);
330        let input: Box<dyn Any + Send> = Box::new(5i32);
331        let result = boxed.invoke_any(input, None).await.unwrap();
332        let output: i32 = *result.downcast::<i32>().unwrap();
333        assert_eq!(output, 6);
334    }
335
336    /// `transform_any` 必须惰性:下游收到第一条输出时,上游流尚未产完。
337    /// 这是 `llm.pipe(parser)` 这类链能"边生成边输出"的关键。
338    #[tokio::test]
339    async fn transform_any_is_lazy_incremental() {
340        use std::sync::atomic::{AtomicBool, Ordering};
341
342        let produced_last = Arc::new(AtomicBool::new(false));
343        let flag = Arc::clone(&produced_last);
344
345        let src = async_stream::stream! {
346            yield Ok::<Box<dyn Any + Send>, LcelError>(Box::new(1i32));
347            yield Ok::<Box<dyn Any + Send>, LcelError>(Box::new(2i32));
348            yield Ok::<Box<dyn Any + Send>, LcelError>(Box::new(3i32));
349            flag.store(true, Ordering::SeqCst);
350        };
351        let input_stream = Box::pin(src)
352            as Pin<Box<dyn Stream<Item = Result<Box<dyn Any + Send>, LcelError>> + Send>>;
353
354        let wrapper = RunnableAnyWrapper::new(AddOne);
355        let mut output = wrapper.transform_any(input_stream, None).await.unwrap();
356
357        let first = output.next().await.unwrap().unwrap();
358        let v: i32 = *first.downcast::<i32>().unwrap();
359        assert_eq!(v, 2);
360        assert!(
361            !produced_last.load(Ordering::SeqCst),
362            "transform_any 不应在上游流结束前就攒齐整条输入"
363        );
364
365        while output.next().await.is_some() {}
366        assert!(produced_last.load(Ordering::SeqCst));
367    }
368}
lc_core/runnables/any.rs

lc_core/runnables/
any.rs
