1use 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
16pub struct ChainRunnable {
25 chain: Arc<dyn BaseChain>,
26}
27
28impl ChainRunnable {
29 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 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 let chain_stream = self.chain.stream_with_config(input, config).await?;
64
65 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 }
81
82impl 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}