rune_chain_core/
chain.rs

1use std::{collections::HashMap, pin::Pin};
2
3use async_trait::async_trait;
4use futures::Stream;
5use serde_json::Value;
6
7use crate::{ChainError, GenerateResult, PromptArgs, StreamData};
8
9const DEFAULT_OUTPUT_KEY: &str = "output";
10const DEFAULT_RESULT_KEY: &str = "generate_result";
11
12/// A composable unit of LLM work that maps [`PromptArgs`] to a [`GenerateResult`].
13///
14/// Chains are the core abstraction of the `rune-chain` ecosystem. A chain may wrap
15/// a single LLM call, a retrieval step, a sequence of sub-chains, or any other
16/// unit of work that accepts named variables and produces generated text.
17///
18/// Implementors must provide [`Chain::call`]; all other methods have default
19/// implementations built on top of it.
20///
21/// # Example
22///
23/// ```rust,ignore
24/// use rune_chain_core::{Chain, ChainError, GenerateResult, PromptArgs, prompt_args};
25/// use async_trait::async_trait;
26///
27/// struct EchoChain;
28///
29/// #[async_trait]
30/// impl Chain for EchoChain {
31///     async fn call(&self, input: PromptArgs) -> Result<GenerateResult, ChainError> {
32///         let text = input["input"].as_str().unwrap_or("").to_string();
33///         Ok(GenerateResult::from_text(text))
34///     }
35/// }
36///
37/// # tokio_test::block_on(async {
38/// let chain = EchoChain;
39/// let result = chain.invoke(prompt_args! { "input" => "hello" }).await.unwrap();
40/// assert_eq!(result, "hello");
41/// # });
42/// ```
43#[async_trait]
44pub trait Chain: Sync + Send {
45    /// Run the chain and return the full [`GenerateResult`] including token usage.
46    async fn call(&self, input_variables: PromptArgs) -> Result<GenerateResult, ChainError>;
47
48    /// Run the chain and return the generated text string only.
49    ///
50    /// Convenience wrapper around [`Chain::call`] that discards token usage.
51    async fn invoke(&self, input_variables: PromptArgs) -> Result<String, ChainError> {
52        self.call(input_variables).await.map(|r| r.generation)
53    }
54
55    /// Run the chain and return a named-output map ready to pipe into the next step.
56    ///
57    /// The returned map contains:
58    /// - the key from [`Chain::output_keys`] (default: `"output"`) → generated text
59    /// - `"generate_result"` → the full serialised [`GenerateResult`]
60    async fn execute(
61        &self,
62        input_variables: PromptArgs,
63    ) -> Result<HashMap<String, Value>, ChainError> {
64        let result = self.call(input_variables).await?;
65        let output_key = self
66            .output_keys()
67            .into_iter()
68            .next()
69            .unwrap_or_else(|| DEFAULT_OUTPUT_KEY.to_string());
70        let mut map = HashMap::new();
71        map.insert(output_key, Value::String(result.generation.clone()));
72        map.insert(
73            DEFAULT_RESULT_KEY.to_string(),
74            serde_json::to_value(&result).unwrap_or(Value::Null),
75        );
76        Ok(map)
77    }
78
79    /// Stream tokens as they are produced rather than waiting for the full completion.
80    ///
81    /// Returns a [`Stream`] of [`StreamData`] chunks. Chains that do not support
82    /// streaming return [`ChainError::Other`] by default.
83    ///
84    /// > [!NOTE]
85    /// > Memory layers cannot be updated automatically during a stream because the
86    /// > complete output is only known when the stream ends. Callers are responsible
87    /// > for persisting the accumulated output after draining the stream.
88    async fn stream(
89        &self,
90        _input_variables: PromptArgs,
91    ) -> Result<Pin<Box<dyn Stream<Item = Result<StreamData, ChainError>> + Send>>, ChainError>
92    {
93        Err(ChainError::Other(
94            "streaming is not implemented for this chain".into(),
95        ))
96    }
97
98    /// The variable names this chain reads from its [`PromptArgs`] input.
99    fn input_keys(&self) -> Vec<String> {
100        vec![]
101    }
102
103    /// The keys this chain writes into the map returned by [`Chain::execute`].
104    fn output_keys(&self) -> Vec<String> {
105        vec![
106            DEFAULT_OUTPUT_KEY.to_string(),
107            DEFAULT_RESULT_KEY.to_string(),
108        ]
109    }
110}
111
112impl<C: Chain + 'static> From<C> for Box<dyn Chain> {
113    fn from(chain: C) -> Self {
114        Box::new(chain)
115    }
116}
rune_chain_core/chain.rs

rune_chain_core/
chain.rs
