1use super::config::RunnableConfig;
16use super::error::LcelError;
17use async_trait::async_trait;
18use futures_util::{Stream, StreamExt};
19use std::any::Any;
20use std::pin::Pin;
21
22#[async_trait]
27pub trait RunnableAny: Send + Sync {
28 async fn invoke_any(
30 &self,
31 input: Box<dyn Any + Send>,
32 config: Option<RunnableConfig>,
33 ) -> Result<Box<dyn Any + Send>, LcelError>;
34
35 async fn stream_any(
37 &self,
38 input: Box<dyn Any + Send>,
39 config: Option<RunnableConfig>,
40 ) -> Result<Pin<Box<dyn Stream<Item = Result<Box<dyn Any + Send>, LcelError>> + Send>>, LcelError>;
41
42 async fn transform_any(
48 &self,
49 input: Pin<Box<dyn Stream<Item = Result<Box<dyn Any + Send>, LcelError>> + Send>>,
50 config: Option<RunnableConfig>,
51 ) -> Result<Pin<Box<dyn Stream<Item = Result<Box<dyn Any + Send>, LcelError>> + Send>>, LcelError>;
52
53 async fn batch_any(
55 &self,
56 inputs: Vec<Box<dyn Any + Send>>,
57 config: Option<RunnableConfig>,
58 ) -> Result<Vec<Box<dyn Any + Send>>, LcelError>;
59}
60
61pub struct RunnableAnyWrapper<I, O, R>
68where
69 I: Send + Sync + 'static,
70 O: Send + Sync + 'static,
71 R: super::Runnable<I, O>,
72{
73 inner: R,
74 _marker: std::marker::PhantomData<(I, O)>,
75}
76
77impl<I, O, R> RunnableAnyWrapper<I, O, R>
78where
79 I: Send + Sync + 'static,
80 O: Send + Sync + 'static,
81 R: super::Runnable<I, O>,
82{
83 pub fn new(runnable: R) -> Self {
85 Self {
86 inner: runnable,
87 _marker: std::marker::PhantomData,
88 }
89 }
90}
91
92#[async_trait]
93impl<I, O, R> RunnableAny for RunnableAnyWrapper<I, O, R>
94where
95 I: Send + Sync + 'static,
96 O: Send + Sync + 'static,
97 R: super::Runnable<I, O> + 'static,
98 R::Error: Into<LcelError>,
99{
100 async fn invoke_any(
101 &self,
102 input: Box<dyn Any + Send>,
103 config: Option<RunnableConfig>,
104 ) -> Result<Box<dyn Any + Send>, LcelError> {
105 let typed_input = input.downcast::<I>().map_err(|_| {
106 LcelError::TypeMismatch(format!(
107 "invoke_any: expected {}, got unknown type",
108 std::any::type_name::<I>()
109 ))
110 })?;
111 let result = self.inner.invoke(*typed_input, config).await.map_err(Into::into)?;
112 Ok(Box::new(result) as Box<dyn Any + Send>)
113 }
114
115 async fn stream_any(
116 &self,
117 input: Box<dyn Any + Send>,
118 config: Option<RunnableConfig>,
119 ) -> Result<Pin<Box<dyn Stream<Item = Result<Box<dyn Any + Send>, LcelError>> + Send>>, LcelError> {
120 let typed_input = input.downcast::<I>().map_err(|_| {
121 LcelError::TypeMismatch(format!(
122 "stream_any: expected {}, got unknown type",
123 std::any::type_name::<I>()
124 ))
125 })?;
126 let stream = self.inner.stream(*typed_input, config).await.map_err(Into::into)?;
127 let any_stream = stream.map(|result| {
128 result
129 .map(|output| Box::new(output) as Box<dyn Any + Send>)
130 .map_err(Into::into)
131 });
132 Ok(Box::pin(any_stream))
133 }
134
135 async fn transform_any(
136 &self,
137 input: Pin<Box<dyn Stream<Item = Result<Box<dyn Any + Send>, LcelError>> + Send>>,
138 config: Option<RunnableConfig>,
139 ) -> Result<Pin<Box<dyn Stream<Item = Result<Box<dyn Any + Send>, LcelError>> + Send>>, LcelError> {
140 use futures_util::StreamExt;
145
146 let mut items = Vec::new();
147 let mut input = input;
148 while let Some(item) = input.next().await {
149 let boxed = item?;
150 let typed = boxed.downcast::<I>().map_err(|_| {
151 LcelError::TypeMismatch(format!(
152 "transform_any input: expected {}",
153 std::any::type_name::<I>()
154 ))
155 })?;
156 items.push(*typed);
157 }
158
159 if let Some(last) = items.into_iter().last() {
160 let result = self.inner.invoke(last, config).await.map_err(Into::into)?;
161 Ok(Box::pin(futures_util::stream::once(async move {
162 Ok(Box::new(result) as Box<dyn Any + Send>)
163 })))
164 } else {
165 Ok(Box::pin(futures_util::stream::empty()))
166 }
167 }
168
169 async fn batch_any(
170 &self,
171 inputs: Vec<Box<dyn Any + Send>>,
172 config: Option<RunnableConfig>,
173 ) -> Result<Vec<Box<dyn Any + Send>>, LcelError> {
174 let typed_inputs: Vec<I> = inputs
175 .into_iter()
176 .map(|boxed| {
177 boxed.downcast::<I>().map(|b| *b).map_err(|_| {
178 LcelError::TypeMismatch(format!(
179 "batch_any: expected {}",
180 std::any::type_name::<I>()
181 ))
182 })
183 })
184 .collect::<Result<Vec<I>, LcelError>>()?;
185 let results = self.inner.batch(typed_inputs, config).await.map_err(Into::into)?;
186 Ok(results.into_iter().map(|r| Box::new(r) as Box<dyn Any + Send>).collect())
187 }
188}
189
190pub fn into_runnable_any<I, O, R>(runnable: R) -> Box<dyn RunnableAny>
195where
196 I: Send + Sync + 'static,
197 O: Send + Sync + 'static,
198 R: super::Runnable<I, O> + 'static,
199 R::Error: Into<LcelError>,
200{
201 Box::new(RunnableAnyWrapper::new(runnable))
202}
203
204#[cfg(test)]
205mod tests {
206 use super::*;
207 use futures_util::StreamExt;
208
209 struct AddOne;
210
211 #[async_trait]
212 impl super::super::Runnable<i32, i32> for AddOne {
213 type Error = std::convert::Infallible;
214
215 async fn invoke(
216 &self,
217 input: i32,
218 _config: Option<RunnableConfig>,
219 ) -> Result<i32, Self::Error> {
220 Ok(input + 1)
221 }
222 }
223
224 #[tokio::test]
225 async fn invoke_any_works() {
226 let wrapper = RunnableAnyWrapper::new(AddOne);
227 let input: Box<dyn Any + Send> = Box::new(41i32);
228 let result = wrapper.invoke_any(input, None).await.unwrap();
229 let output: i32 = *result.downcast::<i32>().unwrap();
230 assert_eq!(output, 42);
231 }
232
233 #[tokio::test]
234 async fn batch_any_works() {
235 let wrapper = RunnableAnyWrapper::new(AddOne);
236 let inputs: Vec<Box<dyn Any + Send>> = vec![Box::new(1i32), Box::new(2i32), Box::new(3i32)];
237 let results = wrapper.batch_any(inputs, None).await.unwrap();
238 let outputs: Vec<i32> = results.into_iter().map(|b| *b.downcast::<i32>().unwrap()).collect();
239 assert_eq!(outputs, vec![2, 3, 4]);
240 }
241
242 #[tokio::test]
243 async fn stream_any_works() {
244 let wrapper = RunnableAnyWrapper::new(AddOne);
245 let input: Box<dyn Any + Send> = Box::new(9i32);
246 let mut stream = wrapper.stream_any(input, None).await.unwrap();
247 let result = stream.next().await.unwrap().unwrap();
248 let output: i32 = *result.downcast::<i32>().unwrap();
249 assert_eq!(output, 10);
250 }
251
252 #[tokio::test]
253 async fn invoke_any_type_mismatch() {
254 let wrapper = RunnableAnyWrapper::new(AddOne);
255 let wrong_input: Box<dyn Any + Send> = Box::new("not an i32");
256 let result = wrapper.invoke_any(wrong_input, None).await;
257 assert!(result.is_err());
258 let err = result.unwrap_err();
259 assert!(matches!(err, LcelError::TypeMismatch(_)));
260 }
261
262 #[tokio::test]
263 async fn into_runnable_any_works() {
264 let boxed: Box<dyn RunnableAny> = into_runnable_any::<i32, i32, _>(AddOne);
265 let input: Box<dyn Any + Send> = Box::new(5i32);
266 let result = boxed.invoke_any(input, None).await.unwrap();
267 let output: i32 = *result.downcast::<i32>().unwrap();
268 assert_eq!(output, 6);
269 }
270}