Skip to main content

lc_core/runnables/
ext.rs

1// lc-core/src/runnables/ext.rs
2//! Extension trait for Runnable providing LCEL composition methods.
3//!
4//! `RunnableExt` adds the `pipe()` method to any `Runnable` whose `Error`
5//! type can be converted into `LcelError`. This enables the core LCEL
6//! composition pattern:
7//!
8//! ```rust,ignore
9//! let chain = prompt.pipe(llm).pipe(parser);
10//! ```
11
12use super::error::LcelError;
13use super::runnable_trait::Runnable;
14use super::sequence::RunnableSequence;
15
16/// Extension trait that provides LCEL composition methods for `Runnable`.
17///
18/// Automatically implemented for all `Runnable<I, O>` types where
19/// `Self::Error: Into<LcelError>`.
20pub trait RunnableExt<Input: Send + Sync + 'static, Output: Send + Sync + 'static>:
21    Runnable<Input, Output>
22where
23    Self::Error: Into<LcelError>,
24    Self: Sized + 'static,
25{
26    /// Pipe the output of this runnable into another runnable.
27    ///
28    /// This is the core LCEL composition operator. It creates a
29    /// `RunnableSequence` that executes `self` first, then passes
30    /// the output to `other`.
31    ///
32    /// # Type Safety
33    ///
34    /// The compiler ensures that the output type of `self` matches
35    /// the input type of `other`. At runtime, intermediate types
36    /// are erased via `RunnableAny`, but the type relationship
37    /// was already proven at compile time.
38    ///
39    /// # Example
40    ///
41    /// ```rust,ignore
42    /// let chain = prompt.pipe(llm).pipe(parser);
43    /// let result = chain.invoke("What is Rust?".to_string(), None).await?;
44    /// ```
45    fn pipe<O2, R2>(self, other: R2) -> RunnableSequence<Input, O2>
46    where
47        O2: Send + Sync + 'static,
48        R2: Runnable<Output, O2> + Send + Sync + 'static,
49        R2::Error: Into<LcelError>,
50    {
51        RunnableSequence::from_pair(self, other)
52    }
53
54    /// Create a `RunnableSequence` from this runnable as a single step.
55    ///
56    /// Useful when you want to start building a pipeline and add
57    /// steps later via `pipe()`.
58    fn into_sequence(self) -> RunnableSequence<Input, Output> {
59        RunnableSequence::from_single(self)
60    }
61}
62
63// Blanket implementation: all Runnables with compatible Error get RunnableExt
64impl<I, O, R> RunnableExt<I, O> for R
65where
66    I: Send + Sync + 'static,
67    O: Send + Sync + 'static,
68    R: Runnable<I, O>,
69    R::Error: Into<LcelError>,
70    R: Sized + 'static,
71{
72}
73
74#[cfg(test)]
75mod tests {
76    use super::*;
77    use crate::RunnableConfig;
78    use async_trait::async_trait;
79    use futures_util::StreamExt;
80
81    struct Double;
82
83    #[async_trait]
84    impl Runnable<i32, i32> for Double {
85        type Error = std::convert::Infallible;
86
87        async fn invoke(&self, input: i32, _config: Option<RunnableConfig>) -> Result<i32, Self::Error> {
88            Ok(input * 2)
89        }
90    }
91
92    struct AddSuffix;
93
94    #[async_trait]
95    impl Runnable<i32, String> for AddSuffix {
96        type Error = std::convert::Infallible;
97
98        async fn invoke(&self, input: i32, _config: Option<RunnableConfig>) -> Result<String, Self::Error> {
99            Ok(format!("result: {}", input))
100        }
101    }
102
103    #[tokio::test]
104    async fn pipe_creates_sequence() {
105        let chain = Double.pipe(AddSuffix);
106        let result = chain.invoke(5, None).await.unwrap();
107        assert_eq!(result, "result: 10");
108    }
109
110    #[tokio::test]
111    async fn pipe_chain_multiple() {
112        // Double → Double → AddSuffix
113        let chain = Double.pipe(Double).pipe(AddSuffix);
114        let result = chain.invoke(3, None).await.unwrap();
115        assert_eq!(result, "result: 12"); // 3 * 2 * 2 = 12
116    }
117
118    #[tokio::test]
119    async fn pipe_stream_works() {
120        let chain = Double.pipe(AddSuffix);
121        let mut stream = chain.stream(5, None).await.unwrap();
122        let result = stream.next().await.unwrap().unwrap();
123        assert_eq!(result, "result: 10");
124    }
125}