lc_core/runnables/

fallback.rs

// lc-core/src/runnables/fallback.rs
//! RunnableWithFallbacks - fallback composition for LCEL pipelines.
//!
//! `RunnableWithFallbacks` wraps a primary `Runnable` with a list of fallback
//! runnables. If the primary fails, each fallback is tried in order until one
//! succeeds. If all fail, the primary's error is returned.
//!
//! # Example
//!
//! ```rust,ignore
//! let chain = prompt
//!     .pipe(openai_llm)
//!     .with_fallbacks(vec![anthropic_llm, ollama_llm])
//!     .pipe(parser);
//! ```

use super::any::{into_runnable_any, RunnableAny};
use super::config::RunnableConfig;
use super::error::LcelError;
use super::runnable_trait::Runnable;
use async_trait::async_trait;
use futures_util::{Stream, StreamExt};
use std::any::Any;
use std::marker::PhantomData;
use std::pin::Pin;

/// A `Runnable` that tries fallbacks when the primary fails.
///
/// If the primary runnable returns an error on `invoke`, each fallback is
/// tried in order. The first successful result is returned. If all fail,
/// the **primary's** error is returned (so the user sees the error from
/// the runnable they explicitly chose).
///
/// The input type `I` must be `Clone` so that the input can be re-boxed
/// for each fallback attempt.
pub struct RunnableWithFallbacks<I: Send + Sync + 'static, O: Send + Sync + 'static> {
    primary: Box<dyn RunnableAny>,
    fallbacks: Vec<Box<dyn RunnableAny>>,
    _marker: PhantomData<(I, O)>,
}

impl<I: Send + Sync + 'static, O: Send + Sync + 'static> std::fmt::Debug
    for RunnableWithFallbacks<I, O>
{
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("RunnableWithFallbacks")
            .field("fallbacks", &self.fallbacks.len())
            .field("input", &std::any::type_name::<I>())
            .field("output", &std::any::type_name::<O>())
            .finish()
    }
}

impl<I: Clone + Send + Sync + 'static, O: Send + Sync + 'static> RunnableWithFallbacks<I, O> {
    /// Create a new fallback runnable with the given primary and fallbacks.
    pub fn new<R>(primary: R, fallbacks: Vec<Box<dyn RunnableAny>>) -> Self
    where
        R: Runnable<I, O> + 'static,
        R::Error: Into<LcelError>,
    {
        Self {
            primary: into_runnable_any(primary),
            fallbacks,
            _marker: PhantomData,
        }
    }

    /// Number of fallback runnables.
    pub fn fallback_count(&self) -> usize {
        self.fallbacks.len()
    }

    /// Try all runnables (primary then fallbacks) on the given input.
    /// Returns the first successful result, or the primary's error if all fail.
    async fn try_all(&self, input: I, config: Option<RunnableConfig>) -> Result<O, LcelError> {
        // Try primary
        let boxed_input = Box::new(input.clone()) as Box<dyn Any + Send>;
        match self.primary.invoke_any(boxed_input, config.clone()).await {
            Ok(result) => result.downcast::<O>().map(|b| *b).map_err(|_| {
                LcelError::TypeMismatch(format!(
                    "fallback primary output downcast: expected {}",
                    std::any::type_name::<O>()
                ))
            }),
            Err(primary_error) => {
                // Try each fallback
                for fallback in &self.fallbacks {
                    let boxed_input = Box::new(input.clone()) as Box<dyn Any + Send>;
                    match fallback.invoke_any(boxed_input, config.clone()).await {
                        Ok(result) => {
                            return result.downcast::<O>().map(|b| *b).map_err(|_| {
                                LcelError::TypeMismatch(format!(
                                    "fallback output downcast: expected {}",
                                    std::any::type_name::<O>()
                                ))
                            });
                        }
                        Err(e) => {
                            log::warn!("fallback invocation failed (trying next): {e}");
                            continue;
                        }
                    }
                }
                // All failed, return the primary's error
                log::warn!("primary and all fallbacks failed; returning primary's error");
                Err(primary_error)
            }
        }
    }
}

#[async_trait]
impl<I: Clone + Send + Sync + 'static, O: Send + Sync + 'static> Runnable<I, O>
    for RunnableWithFallbacks<I, O>
{
    type Error = LcelError;

    /// Try the primary, then fallbacks on failure.
    async fn invoke(&self, input: I, config: Option<RunnableConfig>) -> Result<O, LcelError> {
        self.try_all(input, config).await
    }

    /// Stream: try the primary first. On failure, try fallbacks.
    /// Returns the first successful stream, or the primary's error.
    async fn stream(
        &self,
        input: I,
        config: Option<RunnableConfig>,
    ) -> Result<Pin<Box<dyn Stream<Item = Result<O, LcelError>> + Send>>, LcelError> {
        // Try primary stream
        let boxed_input = Box::new(input.clone()) as Box<dyn Any + Send>;
        match self.primary.stream_any(boxed_input, config.clone()).await {
            Ok(stream) => {
                let output_stream = stream.map(|result| {
                    result.and_then(|boxed| {
                        boxed.downcast::<O>().map(|b| *b).map_err(|_| {
                            LcelError::TypeMismatch(format!(
                                "fallback stream downcast: expected {}",
                                std::any::type_name::<O>()
                            ))
                        })
                    })
                });
                return Ok(Box::pin(output_stream));
            }
            Err(primary_error) => {
                // Try each fallback
                for fallback in &self.fallbacks {
                    let boxed_input = Box::new(input.clone()) as Box<dyn Any + Send>;
                    match fallback.stream_any(boxed_input, config.clone()).await {
                        Ok(stream) => {
                            let output_stream = stream.map(|result| {
                                result.and_then(|boxed| {
                                    boxed.downcast::<O>().map(|b| *b).map_err(|_| {
                                        LcelError::TypeMismatch(format!(
                                            "fallback stream downcast: expected {}",
                                            std::any::type_name::<O>()
                                        ))
                                    })
                                })
                            });
                            return Ok(Box::pin(output_stream));
                        }
                        Err(e) => {
                            log::warn!("fallback stream failed (trying next): {e}");
                            continue;
                        }
                    }
                }
                log::warn!("primary and all fallback streams failed; returning primary's error");
                Err(primary_error)
            }
        }
    }

    /// Batch: apply fallback logic per-item.
    async fn batch(
        &self,
        inputs: Vec<I>,
        config: Option<RunnableConfig>,
    ) -> Result<Vec<O>, LcelError> {
        let mut results = Vec::with_capacity(inputs.len());
        for input in inputs {
            results.push(self.invoke(input, config.clone()).await?);
        }
        Ok(results)
    }

    /// Transform: 逐项调用 `invoke`(带 fallback),产出与输入等长的流。
    async fn transform(
        &self,
        input: Pin<Box<dyn Stream<Item = Result<I, LcelError>> + Send>>,
        config: Option<RunnableConfig>,
    ) -> Result<Pin<Box<dyn Stream<Item = Result<O, LcelError>> + Send + '_>>, LcelError> {
        // M4: 旧实现只处理输入流最后一项,其余静默丢弃(与 batch 语义不一致)。
        // 改为逐项 elementwise:缓冲输入后对每一项调用 invoke-with-fallback,
        // 一一对应输出(与 trait 默认 transform 的缓冲模式一致,避免返回流借用 self)。
        let mut items = Vec::new();
        let mut input = input;
        while let Some(item) = input.next().await {
            items.push(item?);
        }
        let mut results = Vec::with_capacity(items.len());
        for item in items {
            results.push(self.invoke(item, config.clone()).await);
        }
        Ok(Box::pin(futures_util::stream::iter(results)))
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::runnables::{RunnableExt, RunnableLambda};
    use futures_util::StreamExt;

    #[tokio::test]
    async fn primary_succeeds_no_fallback() {
        let primary = RunnableLambda::new_sync(|x: i32| x * 2);
        let fallback = RunnableLambda::new_sync(|x: i32| x * 3);

        let with_fallbacks = primary.with_fallbacks(vec![fallback]);
        let result = with_fallbacks.invoke(5, None).await.unwrap();
        assert_eq!(result, 10); // 5 * 2, primary succeeded
    }

    #[tokio::test]
    async fn primary_fails_fallback_succeeds() {
        let primary = RunnableLambda::new_sync_fallible(|_x: i32| -> Result<i32, LcelError> {
            Err(LcelError::Other("primary failed".to_string()))
        });
        let fallback = RunnableLambda::new_sync(|x: i32| x * 3);

        let with_fallbacks = primary.with_fallbacks(vec![fallback]);
        let result = with_fallbacks.invoke(5, None).await.unwrap();
        assert_eq!(result, 15); // 5 * 3, fallback succeeded
    }

    #[tokio::test]
    async fn all_fail_returns_primary_error() {
        let primary = RunnableLambda::new_sync_fallible(|_x: i32| -> Result<i32, LcelError> {
            Err(LcelError::Provider("openai timeout".to_string()))
        });
        let fallback = RunnableLambda::new_sync_fallible(|_x: i32| -> Result<i32, LcelError> {
            Err(LcelError::Provider("anthropic timeout".to_string()))
        });

        let with_fallbacks = primary.with_fallbacks(vec![fallback]);
        let err = with_fallbacks.invoke(5, None).await.unwrap_err();
        // Should return the primary's error
        assert!(matches!(err, LcelError::Provider(msg) if msg.contains("openai")));
    }

    #[tokio::test]
    async fn multiple_fallbacks_first_wins() {
        let primary = RunnableLambda::new_sync_fallible(|_x: i32| -> Result<i32, LcelError> {
            Err(LcelError::Other("primary failed".to_string()))
        });
        let fb1 = RunnableLambda::new_sync(|x: i32| x + 100);
        let fb2 = RunnableLambda::new_sync(|x: i32| x + 200);

        let with_fallbacks = primary.with_fallbacks(vec![fb1, fb2]);
        let result = with_fallbacks.invoke(5, None).await.unwrap();
        assert_eq!(result, 105); // 5 + 100, first fallback succeeded
    }

    #[tokio::test]
    async fn first_fallback_fails_second_succeeds() {
        let primary = RunnableLambda::new_sync_fallible(|_x: i32| -> Result<i32, LcelError> {
            Err(LcelError::Other("primary failed".to_string()))
        });
        let fb1 = RunnableLambda::new_sync_fallible(|_x: i32| -> Result<i32, LcelError> {
            Err(LcelError::Other("fb1 failed".to_string()))
        });
        let fb2 = RunnableLambda::new_sync(|x: i32| x + 200);

        let with_fallbacks = primary.with_fallbacks(vec![fb1, fb2]);
        let result = with_fallbacks.invoke(5, None).await.unwrap();
        assert_eq!(result, 205); // 5 + 200, second fallback succeeded
    }

    #[tokio::test]
    async fn stream_primary_succeeds() {
        let primary = RunnableLambda::new_sync(|x: i32| x * 2);
        let with_fallbacks = primary.with_fallbacks(Vec::<RunnableLambda<i32, i32>>::new());

        let mut stream = with_fallbacks.stream(5, None).await.unwrap();
        let result = stream.next().await.unwrap().unwrap();
        assert_eq!(result, 10);
    }

    #[tokio::test]
    async fn stream_primary_fails_fallback_succeeds() {
        let primary = RunnableLambda::new_sync_fallible(|_x: i32| -> Result<i32, LcelError> {
            Err(LcelError::Other("primary failed".to_string()))
        });
        let fallback = RunnableLambda::new_sync(|x: i32| x * 3);

        let with_fallbacks = primary.with_fallbacks(vec![fallback]);
        let mut stream = with_fallbacks.stream(5, None).await.unwrap();
        let result = stream.next().await.unwrap().unwrap();
        assert_eq!(result, 15);
    }

    #[tokio::test]
    async fn batch_works() {
        let primary = RunnableLambda::new_sync(|x: i32| x * 2);
        let fallback = RunnableLambda::new_sync(|x: i32| x * 3);

        let with_fallbacks = primary.with_fallbacks(vec![fallback]);
        let results = with_fallbacks.batch(vec![1, 2, 3], None).await.unwrap();
        assert_eq!(results, vec![2, 4, 6]); // primary succeeds for all
    }

    #[tokio::test]
    async fn transform_works() {
        let primary = RunnableLambda::new_sync(|x: i32| x * 2);
        let with_fallbacks = primary.with_fallbacks(Vec::<RunnableLambda<i32, i32>>::new());

        let input = Box::pin(futures_util::stream::iter(vec![
            Ok(1i32),
            Ok(2i32),
            Ok(3i32),
        ])) as Pin<Box<dyn Stream<Item = Result<i32, LcelError>> + Send>>;

        let mut output = with_fallbacks.transform(input, None).await.unwrap();
        // M4: 逐项 elementwise 处理(与 batch 语义一致),而非只处理最后一项。
        let mut results = Vec::new();
        while let Some(item) = output.next().await {
            results.push(item.unwrap());
        }
        assert_eq!(results, vec![2, 4, 6]);
    }

    #[tokio::test]
    async fn transform_fallback_per_item() {
        // M4: 逐项失败时按项走 fallback。
        let primary = RunnableLambda::new_sync_fallible(|_x: i32| -> Result<i32, LcelError> {
            Err(LcelError::Other("primary failed".to_string()))
        });
        let fallback = RunnableLambda::new_sync(|x: i32| x * 3);

        let with_fallbacks = primary.with_fallbacks(vec![fallback]);
        let input = Box::pin(futures_util::stream::iter(vec![Ok(1i32), Ok(2i32)]))
            as Pin<Box<dyn Stream<Item = Result<i32, LcelError>> + Send>>;

        let mut output = with_fallbacks.transform(input, None).await.unwrap();
        let mut results = Vec::new();
        while let Some(item) = output.next().await {
            results.push(item.unwrap());
        }
        assert_eq!(results, vec![3, 6]);
    }

    #[tokio::test]
    async fn debug_format() {
        let primary = RunnableLambda::new_sync(|x: i32| x);
        let with_fallbacks = primary.with_fallbacks(Vec::<RunnableLambda<i32, i32>>::new());
        let debug_str = format!("{:?}", with_fallbacks);
        assert!(debug_str.contains("RunnableWithFallbacks"));
    }

    #[tokio::test]
    async fn pipe_with_fallbacks() {
        let primary = RunnableLambda::new_sync_fallible(|_x: i32| -> Result<i32, LcelError> {
            Err(LcelError::Other("fail".to_string()))
        });
        let fallback = RunnableLambda::new_sync(|x: i32| x + 10);

        // pipe: (primary | fallback) -> double
        let double = RunnableLambda::new_sync(|x: i32| x * 2);
        let chain = primary.with_fallbacks(vec![fallback]).pipe(double);

        let result = chain.invoke(5, None).await.unwrap();
        assert_eq!(result, 30); // fallback: 5+10=15, double: 15*2=30
    }
}