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        // Propagate config (including callbacks) through to the chain
45        self.chain
46            .invoke_with_config(input, config)
47            .await
48            .map_err(|e| LcelError::Chain(e.to_string()))
49    }
50
51    async fn stream(
52        &self,
53        input: HashMap<String, Value>,
54        config: Option<RunnableConfig>,
55    ) -> Result<Pin<Box<dyn Stream<Item = Result<HashMap<String, Value>, LcelError>> + Send>>, LcelError> {
56        // Propagate config (including callbacks) through to the chain
57        let chain_stream = self.chain
58            .stream_with_config(input, config)
59            .await
60            .map_err(|e| LcelError::Stream(e.to_string()))?;
61
62        // Convert StreamToken stream to HashMap stream
63        let mapped = chain_stream.map(|result| {
64            result
65                .map(|token| {
66                    let mut map = HashMap::new();
67                    map.insert("text".to_string(), Value::String(token.token));
68                    map
69                })
70                .map_err(|e| LcelError::Stream(e.to_string()))
71        });
72
73        Ok(Box::pin(mapped))
74    }
75
76    // batch and transform use default implementations
77}
78
79/// Allow `ChainError` to convert into `LcelError`.
80impl From<crate::base::ChainError> for LcelError {
81    fn from(err: crate::base::ChainError) -> Self {
82        LcelError::Chain(err.to_string())
83    }
84}
85
86#[cfg(test)]
87mod tests {
88    use super::*;
89    use lc_core::runnables::Runnable;
90
91    struct TestChain;
92
93    #[async_trait]
94    impl BaseChain for TestChain {
95        fn input_keys(&self) -> Vec<&str> {
96            vec!["input"]
97        }
98
99        fn output_keys(&self) -> Vec<&str> {
100            vec!["output"]
101        }
102
103        async fn invoke(
104            &self,
105            inputs: HashMap<String, Value>,
106        ) -> Result<HashMap<String, Value>, crate::base::ChainError> {
107            let input = inputs
108                .get("input")
109                .and_then(|v| v.as_str())
110                .unwrap_or("");
111            let mut result = HashMap::new();
112            result.insert("output".to_string(), Value::String(format!("echo: {}", input)));
113            Ok(result)
114        }
115    }
116
117    #[tokio::test]
118    async fn chain_runnable_invoke() {
119        let chain = ChainRunnable::new(Arc::new(TestChain));
120        let mut input = HashMap::new();
121        input.insert("input".to_string(), Value::String("hello".to_string()));
122        let result = chain.invoke(input, None).await.unwrap();
123        assert_eq!(result.get("output").unwrap(), &Value::String("echo: hello".to_string()));
124    }
125}
lc_chains/adapter.rs

lc_chains/
adapter.rs
