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 lc_core::runnables::RunnableConfig;
7use serde_json::Value;
8use std::collections::HashMap;
9use std::pin::Pin;
10
11/// Chain error type.
12#[derive(Debug, thiserror::Error)]
13pub enum ChainError {
14    /// Missing input.
15    #[error("Missing input: {0}")]
16    MissingInput(String),
17
18    /// Output error.
19    #[error("Output error: {0}")]
20    OutputError(String),
21
22    /// Execution error.
23    #[error("Execution error: {0}")]
24    ExecutionError(String),
25
26    /// Stream error.
27    #[error("Stream error: {0}")]
28    StreamError(String),
29
30    /// Other error.
31    #[error("Chain error: {0}")]
32    Other(String),
33}
34
35/// Chain execution result.
36pub type ChainResult = HashMap<String, Value>;
37
38/// Stream output item: token-by-token output.
39#[derive(Debug, Clone)]
40pub struct StreamToken {
41    /// Token text.
42    pub token: String,
43    /// Whether this is the final token.
44    pub is_final: bool,
45}
46
47/// Chain stream output type.
48pub type ChainStream = Pin<Box<dyn Stream<Item = Result<StreamToken, ChainError>> + Send>>;
49
50/// Base Chain trait.
51///
52/// Chain is LangChain's core abstraction, representing a sequence of operations.
53#[async_trait]
54pub trait BaseChain: Send + Sync {
55    /// Get input keys.
56    fn input_keys(&self) -> Vec<&str>;
57
58    /// Get output keys.
59    fn output_keys(&self) -> Vec<&str>;
60
61    /// Execute the Chain.
62    ///
63    /// # Arguments
64    /// * `inputs` - Input parameter dictionary
65    ///
66    /// # Returns
67    /// Output result dictionary
68    async fn invoke(&self, inputs: HashMap<String, Value>) -> Result<ChainResult, ChainError>;
69
70    /// Execute the Chain with a RunnableConfig.
71    ///
72    /// This method propagates callbacks through the chain execution.
73    /// The default implementation delegates to `invoke()` without config,
74    /// so chains that don't need callback support work automatically.
75    ///
76    /// Chains that want to propagate callbacks (on_chain_start/end, on_llm_start/end)
77    /// should override this method.
78    async fn invoke_with_config(
79        &self,
80        inputs: HashMap<String, Value>,
81        config: Option<RunnableConfig>,
82    ) -> Result<ChainResult, ChainError> {
83        // Default: fire on_chain_start/end if callbacks are present, then delegate to invoke
84        let _ = config; // suppress unused warning
85        self.invoke(inputs).await
86    }
87
88    /// Stream execute the Chain -- token by token output.
89    ///
90    /// Default implementation wraps the invoke result as a single-element stream.
91    /// Chains that support LLM streaming (LLMChain / ConversationChain) should
92    /// override this method, calling `BaseChatModel::stream_chat` internally.
93    ///
94    /// # Arguments
95    /// * `inputs` - Input parameter dictionary
96    ///
97    /// # Returns
98    /// Token stream
99    async fn stream(&self, inputs: HashMap<String, Value>) -> Result<ChainStream, ChainError> {
100        // Default: wrap invoke result as single-element stream
101        let result = self.invoke(inputs).await?;
102        let output_text = result
103            .values()
104            .next()
105            .and_then(|v| v.as_str())
106            .unwrap_or("")
107            .to_string();
108        let stream = futures_util::stream::once(async move {
109            Ok(StreamToken {
110                token: output_text,
111                is_final: true,
112            })
113        });
114        Ok(Box::pin(stream))
115    }
116
117    /// Stream execute the Chain with a RunnableConfig.
118    ///
119    /// The default implementation delegates to `stream()` without config.
120    async fn stream_with_config(
121        &self,
122        inputs: HashMap<String, Value>,
123        config: Option<RunnableConfig>,
124    ) -> Result<ChainStream, ChainError> {
125        let _ = config;
126        self.stream(inputs).await
127    }
128
129    /// Validate inputs.
130    fn validate_inputs(&self, inputs: &HashMap<String, Value>) -> Result<(), ChainError> {
131        for key in self.input_keys() {
132            if !inputs.contains_key(key) {
133                return Err(ChainError::MissingInput(key.to_string()));
134            }
135        }
136        Ok(())
137    }
138
139    /// Get Chain name.
140    fn name(&self) -> &str {
141        "chain"
142    }
143}
144
145#[cfg(test)]
146mod tests {
147    use super::*;
148
149    #[test]
150    fn test_chain_error_display() {
151        let error = ChainError::MissingInput("test".to_string());
152        assert!(error.to_string().contains("Missing input"));
153
154        let error = ChainError::ExecutionError("test".to_string());
155        assert!(error.to_string().contains("Execution error"));
156    }
157}
lc_chains/base.rs

lc_chains/
base.rs
