lc_chains/

sequential_chain.rs

// lc-chains/src/sequential_chain.rs
//! Sequential Chain
//!
//! Execute multiple chains sequentially.

use async_trait::async_trait;
use serde_json::Value;
use std::collections::HashMap;
use std::sync::Arc;

use crate::base::{BaseChain, ChainError, ChainResult, ChainStream};

/// Sequential Chain
///
/// Executes multiple chains sequentially, where the output of one chain
/// can be used as the input of the next.
pub struct SequentialChain {
    /// Chain list.
    chains: Vec<ChainStep>,

    /// Chain name.
    name: String,
}

/// Chain step.
struct ChainStep {
    /// Chain instance.
    chain: Arc<dyn BaseChain>,

    /// Input mapping (from global input or previous output).
    input_mapping: HashMap<String, String>,

    /// Output mapping (to global result).
    output_mapping: HashMap<String, String>,
}

impl SequentialChain {
    /// Create an empty SequentialChain.
    pub fn new() -> Self {
        Self {
            chains: Vec::new(),
            name: "sequential_chain".to_string(),
        }
    }

    /// Set name.
    pub fn with_name(mut self, name: impl Into<String>) -> Self {
        self.name = name.into();
        self
    }

    /// Add a Chain.
    ///
    /// # Arguments
    /// * `chain` - Chain to add
    /// * `input_keys` - Input keys (from global input)
    /// * `output_keys` - Output keys (to global result)
    pub fn add_chain(
        mut self,
        chain: Arc<dyn BaseChain>,
        input_keys: Vec<&str>,
        output_keys: Vec<&str>,
    ) -> Self {
        let input_mapping = input_keys
            .into_iter()
            .map(|k| (k.to_string(), k.to_string()))
            .collect();

        let output_mapping = output_keys
            .into_iter()
            .map(|k| (k.to_string(), k.to_string()))
            .collect();

        self.chains.push(ChainStep {
            chain,
            input_mapping,
            output_mapping,
        });

        self
    }

    /// Add a Chain with mapping.
    ///
    /// # Arguments
    /// * `chain` - Chain to add
    /// * `input_mapping` - Input mapping {chain_input_key: global_key}
    /// * `output_mapping` - Output mapping {chain_output_key: global_key}
    pub fn add_chain_with_mapping(
        mut self,
        chain: Arc<dyn BaseChain>,
        input_mapping: HashMap<String, String>,
        output_mapping: HashMap<String, String>,
    ) -> Self {
        self.chains.push(ChainStep {
            chain,
            input_mapping,
            output_mapping,
        });

        self
    }
}

impl Default for SequentialChain {
    fn default() -> Self {
        Self::new()
    }
}

#[async_trait]
impl BaseChain for SequentialChain {
    fn input_keys(&self) -> Vec<&str> {
        if let Some(first) = self.chains.first() {
            first.input_mapping.values().map(|s| s.as_str()).collect()
        } else {
            vec![]
        }
    }

    fn output_keys(&self) -> Vec<&str> {
        if let Some(last) = self.chains.last() {
            last.output_mapping.values().map(|s| s.as_str()).collect()
        } else {
            vec![]
        }
    }

    async fn invoke(&self, inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError> {
        let mut current_state = inputs.clone();
        let mut final_output = HashMap::new();

        for (step_index, step) in self.chains.iter().enumerate() {
            let mut chain_inputs = HashMap::new();
            for (chain_key, global_key) in &step.input_mapping {
                if let Some(value) = current_state.get(global_key) {
                    chain_inputs.insert(chain_key.clone(), value.clone());
                } else {
                    return Err(ChainError::MissingInput(format!(
                        "Step {}: missing input '{}' (mapped from '{}')",
                        step_index, chain_key, global_key
                    )));
                }
            }

            let chain_output = step.chain.invoke(chain_inputs).await.map_err(|e| {
                ChainError::ExecutionError(format!(
                    "Step {} ({}) execution failed: {}",
                    step_index,
                    step.chain.name(),
                    e
                ))
            })?;

            for (chain_key, global_key) in &step.output_mapping {
                if let Some(value) = chain_output.get(chain_key) {
                    current_state.insert(global_key.clone(), value.clone());
                    final_output.insert(global_key.clone(), value.clone());
                } else {
                    return Err(ChainError::OutputError(format!(
                        "Step {} ({}) did not produce expected output key '{}' (mapped to '{}')",
                        step_index,
                        step.chain.name(),
                        chain_key,
                        global_key
                    )));
                }
            }
        }

        Ok(final_output)
    }

    /// Stream execution for SequentialChain.
    ///
    /// Runs all chains except the last via invoke (since their output feeds
    /// into subsequent chains). The last chain's output is streamed token
    /// by token by delegating to its `stream()` method.
    async fn stream(&self, inputs: HashMap<String, Value>) -> Result<ChainStream, ChainError> {
        if self.chains.is_empty() {
            return Err(ChainError::ExecutionError(
                "SequentialChain has no chains".to_string(),
            ));
        }

        let mut current_state = inputs.clone();

        // Run all chains except the last via invoke
        let last_idx = self.chains.len() - 1;
        for (step_index, step) in self.chains[..last_idx].iter().enumerate() {
            let mut chain_inputs = HashMap::new();
            for (chain_key, global_key) in &step.input_mapping {
                if let Some(value) = current_state.get(global_key) {
                    chain_inputs.insert(chain_key.clone(), value.clone());
                } else {
                    return Err(ChainError::MissingInput(format!(
                        "Step {}: missing input '{}' (mapped from '{}')",
                        step_index, chain_key, global_key
                    )));
                }
            }

            let chain_output = step.chain.invoke(chain_inputs).await.map_err(|e| {
                ChainError::ExecutionError(format!(
                    "Step {} ({}) execution failed: {}",
                    step_index,
                    step.chain.name(),
                    e
                ))
            })?;

            for (chain_key, global_key) in &step.output_mapping {
                if let Some(value) = chain_output.get(chain_key) {
                    current_state.insert(global_key.clone(), value.clone());
                } else {
                    return Err(ChainError::OutputError(format!(
                        "Step {} ({}) did not produce expected output key '{}'",
                        step_index,
                        step.chain.name(),
                        chain_key,
                    )));
                }
            }
        }

        // Stream the last chain
        let last_step = &self.chains[last_idx];
        let mut chain_inputs = HashMap::new();
        for (chain_key, global_key) in &last_step.input_mapping {
            if let Some(value) = current_state.get(global_key) {
                chain_inputs.insert(chain_key.clone(), value.clone());
            } else {
                return Err(ChainError::MissingInput(format!(
                    "Last step: missing input '{}' (mapped from '{}')",
                    chain_key, global_key
                )));
            }
        }

        last_step.chain.stream(chain_inputs).await
    }

    fn name(&self) -> &str {
        &self.name
    }
}

impl std::fmt::Debug for SequentialChain {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("SequentialChain")
            .field("steps", &self.chains.len())
            .field("name", &self.name)
            .finish()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use async_trait::async_trait;
    use serde_json::json;

    /// A simple mock chain that echoes input to output.
    struct EchoChain {
        input_key: String,
        output_key: String,
    }

    impl EchoChain {
        fn new(input_key: &str, output_key: &str) -> Self {
            Self {
                input_key: input_key.to_string(),
                output_key: output_key.to_string(),
            }
        }
    }

    #[async_trait]
    impl BaseChain for EchoChain {
        fn input_keys(&self) -> Vec<&str> {
            vec![&self.input_key]
        }
        fn output_keys(&self) -> Vec<&str> {
            vec![&self.output_key]
        }
        async fn invoke(&self, inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError> {
            let mut result = HashMap::new();
            if let Some(v) = inputs.get(&self.input_key) {
                result.insert(self.output_key.clone(), v.clone());
            }
            Ok(result)
        }
    }

    /// A mock chain that transforms input (uppercases).
    struct UppercaseChain {
        input_key: String,
        output_key: String,
    }

    #[async_trait]
    impl BaseChain for UppercaseChain {
        fn input_keys(&self) -> Vec<&str> { vec![&self.input_key] }
        fn output_keys(&self) -> Vec<&str> { vec![&self.output_key] }
        async fn invoke(&self, inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError> {
            let mut result = HashMap::new();
            if let Some(Value::String(s)) = inputs.get(&self.input_key) {
                result.insert(self.output_key.clone(), Value::String(s.to_uppercase()));
            }
            Ok(result)
        }
    }

    #[tokio::test]
    async fn test_sequential_chain_single_step() {
        let chain = SequentialChain::new()
            .add_chain(Arc::new(EchoChain::new("text", "result")), vec!["text"], vec!["result"]);

        let mut inputs = HashMap::new();
        inputs.insert("text".to_string(), json!("hello"));

        let result = chain.invoke(inputs).await.unwrap();
        assert_eq!(result.get("result").unwrap(), &json!("hello"));
    }

    #[tokio::test]
    async fn test_sequential_chain_two_steps() {
        let chain = SequentialChain::new()
            .add_chain(Arc::new(EchoChain::new("text", "intermediate")), vec!["text"], vec!["intermediate"])
            .add_chain(Arc::new(UppercaseChain { input_key: "intermediate".to_string(), output_key: "result".to_string() }), vec!["intermediate"], vec!["result"]);

        let mut inputs = HashMap::new();
        inputs.insert("text".to_string(), json!("hello"));

        let result = chain.invoke(inputs).await.unwrap();
        assert_eq!(result.get("result").unwrap(), &json!("HELLO"));
    }

    #[tokio::test]
    async fn test_sequential_chain_missing_input() {
        let chain = SequentialChain::new()
            .add_chain(Arc::new(EchoChain::new("text", "result")), vec!["text"], vec!["result"]);

        let inputs = HashMap::new();
        let result = chain.invoke(inputs).await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_sequential_chain_with_name() {
        let chain = SequentialChain::new().with_name("my_chain");
        assert_eq!(chain.name(), "my_chain");
    }

    #[tokio::test]
    async fn test_sequential_chain_default() {
        let chain = SequentialChain::default();
        assert_eq!(chain.name(), "sequential_chain");
    }

    #[tokio::test]
    async fn test_sequential_chain_debug() {
        let chain = SequentialChain::new().with_name("test_chain");
        let debug_str = format!("{:?}", chain);
        assert!(debug_str.contains("test_chain"));
        assert!(debug_str.contains("0")); // 0 steps
    }

    #[tokio::test]
    async fn test_sequential_chain_input_keys() {
        let chain = SequentialChain::new()
            .add_chain(Arc::new(EchoChain::new("query", "result")), vec!["query"], vec!["result"]);
        assert_eq!(chain.input_keys(), vec!["query"]);
    }

    #[tokio::test]
    async fn test_sequential_chain_output_keys() {
        let chain = SequentialChain::new()
            .add_chain(Arc::new(EchoChain::new("query", "answer")), vec!["query"], vec!["answer"]);
        assert_eq!(chain.output_keys(), vec!["answer"]);
    }
}