Skip to main content

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