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;
155
156 let mut items = Vec::new();
157 let mut input = input;
158 while let Some(item) = input.next().await {
159 let boxed = item?;
160 let typed = boxed.downcast::<I>().map_err(|_| {
161 LcelError::TypeMismatch(format!(
162 "transform_any input: expected {}",
163 std::any::type_name::<I>()
164 ))
165 })?;
166 items.push(*typed);
167 }
168
169 if let Some(last) = items.into_iter().last() {
170 let result = self.inner.invoke(last, config).await.map_err(Into::into)?;
171 Ok(Box::pin(futures_util::stream::once(async move {
172 Ok(Box::new(result) as Box<dyn Any + Send>)
173 })))
174 } else {
175 Ok(Box::pin(futures_util::stream::empty()))
176 }
177 }
178
179 async fn batch_any(
180 &self,
181 inputs: Vec<Box<dyn Any + Send>>,
182 config: Option<RunnableConfig>,
183 ) -> Result<Vec<Box<dyn Any + Send>>, LcelError> {
184 let typed_inputs: Vec<I> = inputs
185 .into_iter()
186 .map(|boxed| {
187 boxed.downcast::<I>().map(|b| *b).map_err(|_| {
188 LcelError::TypeMismatch(format!(
189 "batch_any: expected {}",
190 std::any::type_name::<I>()
191 ))
192 })
193 })
194 .collect::<Result<Vec<I>, LcelError>>()?;
195 let results = self
196 .inner
197 .batch(typed_inputs, config)
198 .await
199 .map_err(Into::into)?;
200 Ok(results
201 .into_iter()
202 .map(|r| Box::new(r) as Box<dyn Any + Send>)
203 .collect())
204 }
205}
206
207pub fn into_runnable_any<I, O, R>(runnable: R) -> Box<dyn RunnableAny>
212where
213 I: Send + Sync + 'static,
214 O: Send + Sync + 'static,
215 R: super::Runnable<I, O> + 'static,
216 R::Error: Into<LcelError>,
217{
218 Box::new(RunnableAnyWrapper::new(runnable))
219}
220
221#[cfg(test)]
222mod tests {
223 use super::*;
224 use futures_util::StreamExt;
225
226 struct AddOne;
227
228 #[async_trait]
229 impl super::super::Runnable<i32, i32> for AddOne {
230 type Error = std::convert::Infallible;
231
232 async fn invoke(
233 &self,
234 input: i32,
235 _config: Option<RunnableConfig>,
236 ) -> Result<i32, Self::Error> {
237 Ok(input + 1)
238 }
239 }
240
241 #[tokio::test]
242 async fn invoke_any_works() {
243 let wrapper = RunnableAnyWrapper::new(AddOne);
244 let input: Box<dyn Any + Send> = Box::new(41i32);
245 let result = wrapper.invoke_any(input, None).await.unwrap();
246 let output: i32 = *result.downcast::<i32>().unwrap();
247 assert_eq!(output, 42);
248 }
249
250 #[tokio::test]
251 async fn batch_any_works() {
252 let wrapper = RunnableAnyWrapper::new(AddOne);
253 let inputs: Vec<Box<dyn Any + Send>> = vec![Box::new(1i32), Box::new(2i32), Box::new(3i32)];
254 let results = wrapper.batch_any(inputs, None).await.unwrap();
255 let outputs: Vec<i32> = results
256 .into_iter()
257 .map(|b| *b.downcast::<i32>().unwrap())
258 .collect();
259 assert_eq!(outputs, vec![2, 3, 4]);
260 }
261
262 #[tokio::test]
263 async fn stream_any_works() {
264 let wrapper = RunnableAnyWrapper::new(AddOne);
265 let input: Box<dyn Any + Send> = Box::new(9i32);
266 let mut stream = wrapper.stream_any(input, None).await.unwrap();
267 let result = stream.next().await.unwrap().unwrap();
268 let output: i32 = *result.downcast::<i32>().unwrap();
269 assert_eq!(output, 10);
270 }
271
272 #[tokio::test]
273 async fn invoke_any_type_mismatch() {
274 let wrapper = RunnableAnyWrapper::new(AddOne);
275 let wrong_input: Box<dyn Any + Send> = Box::new("not an i32");
276 let result = wrapper.invoke_any(wrong_input, None).await;
277 assert!(result.is_err());
278 let err = result.unwrap_err();
279 assert!(matches!(err, LcelError::TypeMismatch(_)));
280 }
281
282 #[tokio::test]
283 async fn into_runnable_any_works() {
284 let boxed: Box<dyn RunnableAny> = into_runnable_any::<i32, i32, _>(AddOne);
285 let input: Box<dyn Any + Send> = Box::new(5i32);
286 let result = boxed.invoke_any(input, None).await.unwrap();
287 let output: i32 = *result.downcast::<i32>().unwrap();
288 assert_eq!(output, 6);
289 }
290}