1use super::any::{into_runnable_any, RunnableAny};
9use super::config::RunnableConfig;
10use super::error::LcelError;
11use super::runnable_trait::Runnable;
12use async_trait::async_trait;
13use futures_util::{Stream, StreamExt};
14use serde_json::Value;
15use std::any::Any;
16use std::collections::HashMap;
17use std::pin::Pin;
18
19pub struct RunnableBinding<I: Send + Sync + 'static, O: Send + Sync + 'static> {
30 bound: Box<dyn RunnableAny>,
31 kwargs: HashMap<String, Value>,
32 config: RunnableConfig,
33 _marker: std::marker::PhantomData<(I, O)>,
34}
35
36impl<I: Send + Sync + 'static, O: Send + Sync + 'static> std::fmt::Debug for RunnableBinding<I, O> {
37 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
38 f.debug_struct("RunnableBinding")
39 .field("kwargs", &self.kwargs)
40 .field("input", &std::any::type_name::<I>())
41 .field("output", &std::any::type_name::<O>())
42 .finish()
43 }
44}
45
46impl<I: Send + Sync + 'static, O: Send + Sync + 'static> RunnableBinding<I, O> {
47 pub fn new<R>(runnable: R) -> Self
49 where
50 R: Runnable<I, O> + 'static,
51 R::Error: Into<LcelError>,
52 {
53 Self {
54 bound: into_runnable_any(runnable),
55 kwargs: HashMap::new(),
56 config: RunnableConfig::default(),
57 _marker: std::marker::PhantomData,
58 }
59 }
60
61 pub fn bind(mut self, key: impl Into<String>, value: Value) -> Self {
66 self.kwargs.insert(key.into(), value);
67 self
68 }
69
70 pub fn with_config(mut self, config: RunnableConfig) -> Self {
74 self.config = config;
75 self
76 }
77
78 fn merged_config(&self, invocation_config: Option<RunnableConfig>) -> RunnableConfig {
80 let mut base = self.config.clone();
81
82 for (key, value) in &self.kwargs {
87 match (key.as_str(), value) {
88 ("temperature", Value::Number(n)) => match n.as_f64() {
89 Some(f) => base = base.with_temperature(f as f32),
90 None => base = base.with_metadata(key.clone(), value.clone()),
91 },
92 ("max_tokens", Value::Number(n)) => match n.as_u64() {
93 Some(u) => base = base.with_max_tokens(u as usize),
94 None => base = base.with_metadata(key.clone(), value.clone()),
95 },
96 _ => base = base.with_metadata(key.clone(), value.clone()),
97 }
98 }
99
100 if let Some(inv) = invocation_config {
102 base.merge(inv)
103 } else {
104 base
105 }
106 }
107}
108
109#[async_trait]
110impl<I: Send + Sync + 'static, O: Send + Sync + 'static> Runnable<I, O> for RunnableBinding<I, O> {
111 type Error = LcelError;
112
113 async fn invoke(&self, input: I, config: Option<RunnableConfig>) -> Result<O, LcelError> {
114 let merged = self.merged_config(config);
115 let result = self
116 .bound
117 .invoke_any(Box::new(input) as Box<dyn Any + Send>, Some(merged))
118 .await?;
119 result.downcast::<O>().map(|b| *b).map_err(|_| {
120 LcelError::TypeMismatch(format!(
121 "binding output downcast: expected {}",
122 std::any::type_name::<O>()
123 ))
124 })
125 }
126
127 async fn batch(
128 &self,
129 inputs: Vec<I>,
130 config: Option<RunnableConfig>,
131 ) -> Result<Vec<O>, LcelError> {
132 let merged = self.merged_config(config);
133 let boxed_inputs: Vec<Box<dyn Any + Send>> = inputs
134 .into_iter()
135 .map(|i| Box::new(i) as Box<dyn Any + Send>)
136 .collect();
137 let results = self.bound.batch_any(boxed_inputs, Some(merged)).await?;
138 results
139 .into_iter()
140 .map(|boxed| {
141 boxed.downcast::<O>().map(|b| *b).map_err(|_| {
142 LcelError::TypeMismatch(format!(
143 "binding batch downcast: expected {}",
144 std::any::type_name::<O>()
145 ))
146 })
147 })
148 .collect()
149 }
150
151 async fn stream(
152 &self,
153 input: I,
154 config: Option<RunnableConfig>,
155 ) -> Result<Pin<Box<dyn Stream<Item = Result<O, LcelError>> + Send>>, LcelError> {
156 let merged = self.merged_config(config);
157 let stream = self
158 .bound
159 .stream_any(Box::new(input) as Box<dyn Any + Send>, Some(merged))
160 .await?;
161 let output_stream = stream.map(|result| {
162 result.and_then(|boxed| {
163 boxed.downcast::<O>().map(|b| *b).map_err(|_| {
164 LcelError::TypeMismatch(format!(
165 "binding stream downcast: expected {}",
166 std::any::type_name::<O>()
167 ))
168 })
169 })
170 });
171 Ok(Box::pin(output_stream))
172 }
173
174 async fn transform(
175 &self,
176 input: Pin<Box<dyn Stream<Item = Result<I, LcelError>> + Send>>,
177 config: Option<RunnableConfig>,
178 ) -> Result<Pin<Box<dyn Stream<Item = Result<O, LcelError>> + Send + '_>>, LcelError> {
179 let merged = self.merged_config(config);
180 let any_input: Pin<Box<dyn Stream<Item = Result<Box<dyn Any + Send>, LcelError>> + Send>> =
181 Box::pin(input.map(|result| result.map(|item| Box::new(item) as Box<dyn Any + Send>)));
182
183 let output_stream = self.bound.transform_any(any_input, Some(merged)).await?;
184
185 let typed_output = output_stream.map(|result| {
186 result.and_then(|boxed| {
187 boxed.downcast::<O>().map(|b| *b).map_err(|_| {
188 LcelError::TypeMismatch(format!(
189 "binding transform downcast: expected {}",
190 std::any::type_name::<O>()
191 ))
192 })
193 })
194 });
195 Ok(Box::pin(typed_output))
196 }
197}
198
199#[cfg(test)]
200mod tests {
201 use super::*;
202 use futures_util::StreamExt;
203
204 struct EchoRunnable;
205
206 #[async_trait]
207 impl Runnable<String, String> for EchoRunnable {
208 type Error = std::convert::Infallible;
209
210 async fn invoke(
211 &self,
212 input: String,
213 config: Option<RunnableConfig>,
214 ) -> Result<String, Self::Error> {
215 let tags = config.map(|c| c.tags.join(",")).unwrap_or_default();
216 if tags.is_empty() {
217 Ok(input)
218 } else {
219 Ok(format!("[{}] {}", tags, input))
220 }
221 }
222 }
223
224 #[tokio::test]
225 async fn binding_with_config() {
226 let binding = RunnableBinding::new(EchoRunnable)
227 .with_config(RunnableConfig::default().with_tag("prod"));
228
229 let result = binding.invoke("hello".to_string(), None).await.unwrap();
230 assert_eq!(result, "[prod] hello");
231 }
232
233 #[tokio::test]
234 async fn binding_with_kwargs() {
235 let binding =
236 RunnableBinding::new(EchoRunnable).bind("stop", Value::String("\n".to_string()));
237
238 let result = binding.invoke("test".to_string(), None).await.unwrap();
240 assert_eq!(result, "test"); }
242
243 #[tokio::test]
244 async fn binding_stream_works() {
245 let binding = RunnableBinding::new(EchoRunnable)
246 .with_config(RunnableConfig::default().with_tag("stream"));
247
248 let mut stream = binding.stream("hello".to_string(), None).await.unwrap();
249 let result = stream.next().await.unwrap().unwrap();
250 assert_eq!(result, "[stream] hello");
251 }
252
253 struct ConfigProbe;
255
256 #[async_trait]
257 impl Runnable<(), String> for ConfigProbe {
258 type Error = std::convert::Infallible;
259
260 async fn invoke(
261 &self,
262 _input: (),
263 config: Option<RunnableConfig>,
264 ) -> Result<String, Self::Error> {
265 Ok(format!(
266 "temp={:?},max={:?}",
267 config.as_ref().and_then(|c| c.temperature),
268 config.as_ref().and_then(|c| c.max_tokens)
269 ))
270 }
271 }
272
273 #[tokio::test]
274 async fn binding_temperature_kwarg_affects_sampling() {
275 let binding = RunnableBinding::new(ConfigProbe).bind("temperature", Value::from(0.5));
278 let result = binding.invoke((), None).await.unwrap();
279 assert_eq!(result, "temp=Some(0.5),max=None");
280 }
281
282 #[tokio::test]
283 async fn binding_max_tokens_kwarg_affects_sampling() {
284 let binding = RunnableBinding::new(ConfigProbe).bind("max_tokens", Value::from(128));
285 let result = binding.invoke((), None).await.unwrap();
286 assert_eq!(result, "temp=None,max=Some(128)");
287 }
288
289 #[tokio::test]
290 async fn binding_unknown_kwarg_stays_in_metadata() {
291 let binding =
292 RunnableBinding::new(ConfigProbe).bind("stop", Value::String("\n".to_string()));
293 let result = binding.invoke((), None).await.unwrap();
294 assert_eq!(result, "temp=None,max=None");
296 }
297}