lc_chains/
adapter.rs

1// lc-chains/src/adapter.rs
2//! ChainRunnable adapter - bridges BaseChain to the Runnable trait.
3//!
4//! This allows chains to participate in LCEL pipelines via `pipe()`.
5
6use async_trait::async_trait;
7use futures_util::{Stream, StreamExt};
8use lc_core::runnables::{LcelError, Runnable, RunnableConfig};
9use serde_json::Value;
10use std::collections::HashMap;
11use std::pin::Pin;
12use std::sync::Arc;
13
14use crate::base::BaseChain;
15
16/// Adapter that wraps a `BaseChain` as a `Runnable<HashMap<String, Value>, HashMap<String, Value>>`.
17///
18/// This enables chains to participate in LCEL pipelines:
19///
20/// ```rust,ignore
21/// let chain_runnable = ChainRunnable::new(Arc::new(my_chain));
22/// let pipeline = chain_runnable.pipe(parser);
23/// ```
24pub struct ChainRunnable {
25    chain: Arc<dyn BaseChain>,
26}
27
28impl ChainRunnable {
29    /// Create a new adapter wrapping the given chain.
30    pub fn new(chain: Arc<dyn BaseChain>) -> Self {
31        Self { chain }
32    }
33}
34
35#[async_trait]
36impl Runnable<HashMap<String, Value>, HashMap<String, Value>> for ChainRunnable {
37    type Error = LcelError;
38
39    async fn invoke(
40        &self,
41        input: HashMap<String, Value>,
42        _config: Option<RunnableConfig>,
43    ) -> Result<HashMap<String, Value>, LcelError> {
44        self.chain.invoke(input).await.map_err(|e| LcelError::Chain(e.to_string()))
45    }
46
47    async fn stream(
48        &self,
49        input: HashMap<String, Value>,
50        _config: Option<RunnableConfig>,
51    ) -> Result<Pin<Box<dyn Stream<Item = Result<HashMap<String, Value>, LcelError>> + Send>>, LcelError> {
52        let chain_stream = self.chain.stream(input).await.map_err(|e| LcelError::Stream(e.to_string()))?;
53
54        // Convert StreamToken stream to HashMap stream
55        let mapped = chain_stream.map(|result| {
56            result
57                .map(|token| {
58                    let mut map = HashMap::new();
59                    map.insert("text".to_string(), Value::String(token.token));
60                    map
61                })
62                .map_err(|e| LcelError::Stream(e.to_string()))
63        });
64
65        Ok(Box::pin(mapped))
66    }
67
68    // batch and transform use default implementations
69}
70
71/// Allow `ChainError` to convert into `LcelError`.
72impl From<crate::base::ChainError> for LcelError {
73    fn from(err: crate::base::ChainError) -> Self {
74        LcelError::Chain(err.to_string())
75    }
76}
77
78#[cfg(test)]
79mod tests {
80    use super::*;
81    use lc_core::runnables::Runnable;
82
83    struct TestChain;
84
85    #[async_trait]
86    impl BaseChain for TestChain {
87        fn input_keys(&self) -> Vec<&str> {
88            vec!["input"]
89        }
90
91        fn output_keys(&self) -> Vec<&str> {
92            vec!["output"]
93        }
94
95        async fn invoke(
96            &self,
97            inputs: HashMap<String, Value>,
98        ) -> Result<HashMap<String, Value>, crate::base::ChainError> {
99            let input = inputs
100                .get("input")
101                .and_then(|v| v.as_str())
102                .unwrap_or("");
103            let mut result = HashMap::new();
104            result.insert("output".to_string(), Value::String(format!("echo: {}", input)));
105            Ok(result)
106        }
107    }
108
109    #[tokio::test]
110    async fn chain_runnable_invoke() {
111        let chain = ChainRunnable::new(Arc::new(TestChain));
112        let mut input = HashMap::new();
113        input.insert("input".to_string(), Value::String("hello".to_string()));
114        let result = chain.invoke(input, None).await.unwrap();
115        assert_eq!(result.get("output").unwrap(), &Value::String("echo: hello".to_string()));
116    }
117}
lc_chains/adapter.rs

lc_chains/
adapter.rs
