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
112 .inner
113 .invoke(*typed_input, config)
114 .await
115 .map_err(Into::into)?;
116 Ok(Box::new(result) as Box<dyn Any + Send>)
117 }
118
119 async fn stream_any(
120 &self,
121 input: Box<dyn Any + Send>,
122 config: Option<RunnableConfig>,
123 ) -> Result<Pin<Box<dyn Stream<Item = Result<Box<dyn Any + Send>, LcelError>> + Send>>, LcelError>
124 {
125 let typed_input = input.downcast::<I>().map_err(|_| {
126 LcelError::TypeMismatch(format!(
127 "stream_any: expected {}, got unknown type",
128 std::any::type_name::<I>()
129 ))
130 })?;
131 let stream = self
132 .inner
133 .stream(*typed_input, config)
134 .await
135 .map_err(Into::into)?;
136 let any_stream = stream.map(|result| {
137 result
138 .map(|output| Box::new(output) as Box<dyn Any + Send>)
139 .map_err(Into::into)
140 });
141 Ok(Box::pin(any_stream))
142 }
143
144 async fn transform_any(
145 &self,
146 input: Pin<Box<dyn Stream<Item = Result<Box<dyn Any + Send>, LcelError>> + Send>>,
147 config: Option<RunnableConfig>,
148 ) -> Result<Pin<Box<dyn Stream<Item = Result<Box<dyn Any + Send>, LcelError>> + Send>>, LcelError>
149 {
150 use futures_util::StreamExt;
160
161 let mut items = Vec::new();
162 let mut input = input;
163 while let Some(item) = input.next().await {
164 let boxed = item?;
165 let typed = boxed.downcast::<I>().map_err(|_| {
166 LcelError::TypeMismatch(format!(
167 "transform_any input: expected {}",
168 std::any::type_name::<I>()
169 ))
170 })?;
171 items.push(*typed);
172 }
173
174 let mut per_item_streams = Vec::with_capacity(items.len());
175 for item in items {
176 let stream = self
177 .inner
178 .stream(item, config.clone())
179 .await
180 .map_err(Into::into)?;
181 let any_stream = stream.map(|result| {
182 result
183 .map(|output| Box::new(output) as Box<dyn Any + Send>)
184 .map_err(Into::into)
185 });
186 per_item_streams.push(any_stream);
187 }
188
189 let flattened = futures_util::stream::iter(per_item_streams).flatten();
190 Ok(Box::pin(flattened))
191 }
192
193 async fn batch_any(
194 &self,
195 inputs: Vec<Box<dyn Any + Send>>,
196 config: Option<RunnableConfig>,
197 ) -> Result<Vec<Box<dyn Any + Send>>, LcelError> {
198 let typed_inputs: Vec<I> = inputs
199 .into_iter()
200 .map(|boxed| {
201 boxed.downcast::<I>().map(|b| *b).map_err(|_| {
202 LcelError::TypeMismatch(format!(
203 "batch_any: expected {}",
204 std::any::type_name::<I>()
205 ))
206 })
207 })
208 .collect::<Result<Vec<I>, LcelError>>()?;
209 let results = self
210 .inner
211 .batch(typed_inputs, config)
212 .await
213 .map_err(Into::into)?;
214 Ok(results
215 .into_iter()
216 .map(|r| Box::new(r) as Box<dyn Any + Send>)
217 .collect())
218 }
219}
220
221pub fn into_runnable_any<I, O, R>(runnable: R) -> Box<dyn RunnableAny>
226where
227 I: Send + Sync + 'static,
228 O: Send + Sync + 'static,
229 R: super::Runnable<I, O> + 'static,
230 R::Error: Into<LcelError>,
231{
232 Box::new(RunnableAnyWrapper::new(runnable))
233}
234
235#[cfg(test)]
236mod tests {
237 use super::*;
238 use futures_util::StreamExt;
239
240 struct AddOne;
241
242 #[async_trait]
243 impl super::super::Runnable<i32, i32> for AddOne {
244 type Error = std::convert::Infallible;
245
246 async fn invoke(
247 &self,
248 input: i32,
249 _config: Option<RunnableConfig>,
250 ) -> Result<i32, Self::Error> {
251 Ok(input + 1)
252 }
253 }
254
255 #[tokio::test]
256 async fn invoke_any_works() {
257 let wrapper = RunnableAnyWrapper::new(AddOne);
258 let input: Box<dyn Any + Send> = Box::new(41i32);
259 let result = wrapper.invoke_any(input, None).await.unwrap();
260 let output: i32 = *result.downcast::<i32>().unwrap();
261 assert_eq!(output, 42);
262 }
263
264 #[tokio::test]
265 async fn batch_any_works() {
266 let wrapper = RunnableAnyWrapper::new(AddOne);
267 let inputs: Vec<Box<dyn Any + Send>> = vec![Box::new(1i32), Box::new(2i32), Box::new(3i32)];
268 let results = wrapper.batch_any(inputs, None).await.unwrap();
269 let outputs: Vec<i32> = results
270 .into_iter()
271 .map(|b| *b.downcast::<i32>().unwrap())
272 .collect();
273 assert_eq!(outputs, vec![2, 3, 4]);
274 }
275
276 #[tokio::test]
277 async fn stream_any_works() {
278 let wrapper = RunnableAnyWrapper::new(AddOne);
279 let input: Box<dyn Any + Send> = Box::new(9i32);
280 let mut stream = wrapper.stream_any(input, None).await.unwrap();
281 let result = stream.next().await.unwrap().unwrap();
282 let output: i32 = *result.downcast::<i32>().unwrap();
283 assert_eq!(output, 10);
284 }
285
286 #[tokio::test]
287 async fn invoke_any_type_mismatch() {
288 let wrapper = RunnableAnyWrapper::new(AddOne);
289 let wrong_input: Box<dyn Any + Send> = Box::new("not an i32");
290 let result = wrapper.invoke_any(wrong_input, None).await;
291 assert!(result.is_err());
292 let err = result.unwrap_err();
293 assert!(matches!(err, LcelError::TypeMismatch(_)));
294 }
295
296 #[tokio::test]
297 async fn into_runnable_any_works() {
298 let boxed: Box<dyn RunnableAny> = into_runnable_any::<i32, i32, _>(AddOne);
299 let input: Box<dyn Any + Send> = Box::new(5i32);
300 let result = boxed.invoke_any(input, None).await.unwrap();
301 let output: i32 = *result.downcast::<i32>().unwrap();
302 assert_eq!(output, 6);
303 }
304}