Skip to main content

lc_core/runnables/
binding.rs

1// lc-core/src/runnables/binding.rs
2//! RunnableBinding - bind runtime parameters and config to a Runnable.
3//!
4//! `RunnableBinding` wraps a `Runnable` and attaches runtime kwargs
5//! and/or a `RunnableConfig`, allowing pre-configuration of steps
6//! in an LCEL pipeline.
7
8use 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
19/// A `Runnable` that binds runtime parameters and config to an inner runnable.
20///
21/// # Example
22///
23/// ```rust,ignore
24/// let chain = llm
25///     .pipe(parser)
26///     .bind("stop", json!("\n"))
27///     .with_config(RunnableConfig::new().with_tag("production"));
28/// ```
29pub 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    /// Create a new binding wrapping the given runnable.
48    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    /// Bind a runtime keyword argument.
62    ///
63    /// The kwargs are merged into the `RunnableConfig` metadata
64    /// when the runnable is invoked.
65    pub fn bind(mut self, key: impl Into<String>, value: Value) -> Self {
66        self.kwargs.insert(key.into(), value);
67        self
68    }
69
70    /// Set the execution config.
71    ///
72    /// This config is merged with any config passed at invocation time.
73    pub fn with_config(mut self, config: RunnableConfig) -> Self {
74        self.config = config;
75        self
76    }
77
78    /// Merge bound kwargs and config into the invocation config.
79    fn merged_config(&self, invocation_config: Option<RunnableConfig>) -> RunnableConfig {
80        let mut base = self.config.clone();
81
82        // Bound kwargs: recognize the sampling overrides the providers already
83        // consume (`temperature` / `max_tokens`) so `llm.bind("temperature", 0.5)`
84        // actually changes sampling — matching Python's `.bind(**kwargs)` which
85        // merges into the model call. Other keys stay in metadata.
86        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        // Merge with invocation config
101        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        // kwargs are stored in metadata
239        let result = binding.invoke("test".to_string(), None).await.unwrap();
240        assert_eq!(result, "test"); // EchoRunnable doesn't use metadata
241    }
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    /// 探针:invoke 时把收到的 config.temperature / config.max_tokens 打出来。
254    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        // `llm.bind("temperature", 0.5)` 应真正改采样温度(进 typed config 字段),
276        // 而不是只进 metadata —— 对齐 Python `.bind(**kwargs)` 语义。
277        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        // stop 不识别为采样字段 → 只进 metadata,typed 字段不受影响
295        assert_eq!(result, "temp=None,max=None");
296    }
297}