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)) if n.as_f64().is_some() => {
89                    base = base.with_temperature(n.as_f64().unwrap() as f32);
90                }
91                ("max_tokens", Value::Number(n)) if n.as_u64().is_some() => {
92                    base = base.with_max_tokens(n.as_u64().unwrap() as usize);
93                }
94                _ => base = base.with_metadata(key.clone(), value.clone()),
95            }
96        }
97
98        // Merge with invocation config
99        if let Some(inv) = invocation_config {
100            base.merge(inv)
101        } else {
102            base
103        }
104    }
105}
106
107#[async_trait]
108impl<I: Send + Sync + 'static, O: Send + Sync + 'static> Runnable<I, O> for RunnableBinding<I, O> {
109    type Error = LcelError;
110
111    async fn invoke(&self, input: I, config: Option<RunnableConfig>) -> Result<O, LcelError> {
112        let merged = self.merged_config(config);
113        let result = self
114            .bound
115            .invoke_any(Box::new(input) as Box<dyn Any + Send>, Some(merged))
116            .await?;
117        result.downcast::<O>().map(|b| *b).map_err(|_| {
118            LcelError::TypeMismatch(format!(
119                "binding output downcast: expected {}",
120                std::any::type_name::<O>()
121            ))
122        })
123    }
124
125    async fn batch(
126        &self,
127        inputs: Vec<I>,
128        config: Option<RunnableConfig>,
129    ) -> Result<Vec<O>, LcelError> {
130        let merged = self.merged_config(config);
131        let boxed_inputs: Vec<Box<dyn Any + Send>> = inputs
132            .into_iter()
133            .map(|i| Box::new(i) as Box<dyn Any + Send>)
134            .collect();
135        let results = self.bound.batch_any(boxed_inputs, Some(merged)).await?;
136        results
137            .into_iter()
138            .map(|boxed| {
139                boxed.downcast::<O>().map(|b| *b).map_err(|_| {
140                    LcelError::TypeMismatch(format!(
141                        "binding batch downcast: expected {}",
142                        std::any::type_name::<O>()
143                    ))
144                })
145            })
146            .collect()
147    }
148
149    async fn stream(
150        &self,
151        input: I,
152        config: Option<RunnableConfig>,
153    ) -> Result<Pin<Box<dyn Stream<Item = Result<O, LcelError>> + Send>>, LcelError> {
154        let merged = self.merged_config(config);
155        let stream = self
156            .bound
157            .stream_any(Box::new(input) as Box<dyn Any + Send>, Some(merged))
158            .await?;
159        let output_stream = stream.map(|result| {
160            result.and_then(|boxed| {
161                boxed.downcast::<O>().map(|b| *b).map_err(|_| {
162                    LcelError::TypeMismatch(format!(
163                        "binding stream downcast: expected {}",
164                        std::any::type_name::<O>()
165                    ))
166                })
167            })
168        });
169        Ok(Box::pin(output_stream))
170    }
171
172    async fn transform(
173        &self,
174        input: Pin<Box<dyn Stream<Item = Result<I, LcelError>> + Send>>,
175        config: Option<RunnableConfig>,
176    ) -> Result<Pin<Box<dyn Stream<Item = Result<O, LcelError>> + Send>>, LcelError> {
177        let merged = self.merged_config(config);
178        let any_input: Pin<Box<dyn Stream<Item = Result<Box<dyn Any + Send>, LcelError>> + Send>> =
179            Box::pin(input.map(|result| result.map(|item| Box::new(item) as Box<dyn Any + Send>)));
180
181        let output_stream = self.bound.transform_any(any_input, Some(merged)).await?;
182
183        let typed_output = output_stream.map(|result| {
184            result.and_then(|boxed| {
185                boxed.downcast::<O>().map(|b| *b).map_err(|_| {
186                    LcelError::TypeMismatch(format!(
187                        "binding transform downcast: expected {}",
188                        std::any::type_name::<O>()
189                    ))
190                })
191            })
192        });
193        Ok(Box::pin(typed_output))
194    }
195}
196
197#[cfg(test)]
198mod tests {
199    use super::*;
200    use futures_util::StreamExt;
201
202    struct EchoRunnable;
203
204    #[async_trait]
205    impl Runnable<String, String> for EchoRunnable {
206        type Error = std::convert::Infallible;
207
208        async fn invoke(
209            &self,
210            input: String,
211            config: Option<RunnableConfig>,
212        ) -> Result<String, Self::Error> {
213            let tags = config.map(|c| c.tags.join(",")).unwrap_or_default();
214            if tags.is_empty() {
215                Ok(input)
216            } else {
217                Ok(format!("[{}] {}", tags, input))
218            }
219        }
220    }
221
222    #[tokio::test]
223    async fn binding_with_config() {
224        let binding = RunnableBinding::new(EchoRunnable)
225            .with_config(RunnableConfig::default().with_tag("prod"));
226
227        let result = binding.invoke("hello".to_string(), None).await.unwrap();
228        assert_eq!(result, "[prod] hello");
229    }
230
231    #[tokio::test]
232    async fn binding_with_kwargs() {
233        let binding =
234            RunnableBinding::new(EchoRunnable).bind("stop", Value::String("\n".to_string()));
235
236        // kwargs are stored in metadata
237        let result = binding.invoke("test".to_string(), None).await.unwrap();
238        assert_eq!(result, "test"); // EchoRunnable doesn't use metadata
239    }
240
241    #[tokio::test]
242    async fn binding_stream_works() {
243        let binding = RunnableBinding::new(EchoRunnable)
244            .with_config(RunnableConfig::default().with_tag("stream"));
245
246        let mut stream = binding.stream("hello".to_string(), None).await.unwrap();
247        let result = stream.next().await.unwrap().unwrap();
248        assert_eq!(result, "[stream] hello");
249    }
250
251    /// 探针:invoke 时把收到的 config.temperature / config.max_tokens 打出来。
252    struct ConfigProbe;
253
254    #[async_trait]
255    impl Runnable<(), String> for ConfigProbe {
256        type Error = std::convert::Infallible;
257
258        async fn invoke(
259            &self,
260            _input: (),
261            config: Option<RunnableConfig>,
262        ) -> Result<String, Self::Error> {
263            Ok(format!(
264                "temp={:?},max={:?}",
265                config.as_ref().and_then(|c| c.temperature),
266                config.as_ref().and_then(|c| c.max_tokens)
267            ))
268        }
269    }
270
271    #[tokio::test]
272    async fn binding_temperature_kwarg_affects_sampling() {
273        // `llm.bind("temperature", 0.5)` 应真正改采样温度(进 typed config 字段),
274        // 而不是只进 metadata —— 对齐 Python `.bind(**kwargs)` 语义。
275        let binding = RunnableBinding::new(ConfigProbe).bind("temperature", Value::from(0.5));
276        let result = binding.invoke((), None).await.unwrap();
277        assert_eq!(result, "temp=Some(0.5),max=None");
278    }
279
280    #[tokio::test]
281    async fn binding_max_tokens_kwarg_affects_sampling() {
282        let binding = RunnableBinding::new(ConfigProbe).bind("max_tokens", Value::from(128));
283        let result = binding.invoke((), None).await.unwrap();
284        assert_eq!(result, "temp=None,max=Some(128)");
285    }
286
287    #[tokio::test]
288    async fn binding_unknown_kwarg_stays_in_metadata() {
289        let binding = RunnableBinding::new(ConfigProbe).bind("stop", Value::String("\n".to_string()));
290        let result = binding.invoke((), None).await.unwrap();
291        // stop 不识别为采样字段 → 只进 metadata,typed 字段不受影响
292        assert_eq!(result, "temp=None,max=None");
293    }
294}
lc_core/runnables/binding.rs

lc_core/runnables/
binding.rs
