lc_chains/
base.rs

1// lc-chains/src/base.rs
2//! Chain base trait.
3
4use async_trait::async_trait;
5use futures_util::Stream;
6use serde_json::Value;
7use std::collections::HashMap;
8use std::pin::Pin;
9
10/// Chain error type.
11#[derive(Debug, thiserror::Error)]
12pub enum ChainError {
13    /// Missing input.
14    #[error("Missing input: {0}")]
15    MissingInput(String),
16
17    /// Output error.
18    #[error("Output error: {0}")]
19    OutputError(String),
20
21    /// Execution error.
22    #[error("Execution error: {0}")]
23    ExecutionError(String),
24
25    /// Stream error.
26    #[error("Stream error: {0}")]
27    StreamError(String),
28
29    /// Other error.
30    #[error("Chain error: {0}")]
31    Other(String),
32}
33
34/// Chain execution result.
35pub type ChainResult = HashMap<String, Value>;
36
37/// Stream output item: token-by-token output.
38#[derive(Debug, Clone)]
39pub struct StreamToken {
40    /// Token text.
41    pub token: String,
42    /// Whether this is the final token.
43    pub is_final: bool,
44}
45
46/// Chain stream output type.
47pub type ChainStream = Pin<Box<dyn Stream<Item = Result<StreamToken, ChainError>> + Send>>;
48
49/// Base Chain trait.
50///
51/// Chain is LangChain's core abstraction, representing a sequence of operations.
52#[async_trait]
53pub trait BaseChain: Send + Sync {
54    /// Get input keys.
55    fn input_keys(&self) -> Vec<&str>;
56
57    /// Get output keys.
58    fn output_keys(&self) -> Vec<&str>;
59
60    /// Execute the Chain.
61    ///
62    /// # Arguments
63    /// * `inputs` - Input parameter dictionary
64    ///
65    /// # Returns
66    /// Output result dictionary
67    async fn invoke(&self, inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError>;
68
69    /// Stream execute the Chain -- token by token output.
70    ///
71    /// Default implementation wraps the invoke result as a single-element stream.
72    /// Chains that support LLM streaming (LLMChain / ConversationChain) should
73    /// override this method, calling `BaseChatModel::stream_chat` internally.
74    ///
75    /// # Arguments
76    /// * `inputs` - Input parameter dictionary
77    ///
78    /// # Returns
79    /// Token stream
80    async fn stream(&self, inputs: HashMap<String, Value>) -> Result<ChainStream, ChainError> {
81        // Default: wrap invoke result as single-element stream
82        let result = self.invoke(inputs).await?;
83        let output_text = result
84            .values()
85            .next()
86            .and_then(|v| v.as_str())
87            .unwrap_or("")
88            .to_string();
89        let stream = futures_util::stream::once(async move {
90            Ok(StreamToken {
91                token: output_text,
92                is_final: true,
93            })
94        });
95        Ok(Box::pin(stream))
96    }
97
98    /// Validate inputs.
99    fn validate_inputs(&self, inputs: &HashMap<String, Value>) -> Result<(), ChainError> {
100        for key in self.input_keys() {
101            if !inputs.contains_key(key) {
102                return Err(ChainError::MissingInput(key.to_string()));
103            }
104        }
105        Ok(())
106    }
107
108    /// Get Chain name.
109    fn name(&self) -> &str {
110        "chain"
111    }
112}
113
114#[cfg(test)]
115mod tests {
116    use super::*;
117
118    #[test]
119    fn test_chain_error_display() {
120        let error = ChainError::MissingInput("test".to_string());
121        assert!(error.to_string().contains("Missing input"));
122
123        let error = ChainError::ExecutionError("test".to_string());
124        assert!(error.to_string().contains("Execution error"));
125    }
126}
lc_chains/base.rs

lc_chains/
base.rs
