llm_chain/
step.rs

1//! Steps are individual LLM invocations in a chain. They are a combination of a prompt and a configuration.
2//!
3//! Steps are used to set the per-invocation settings for a prompt. Useful when you want to change the settings for a specific prompt in a chain.
4use crate::frame::{FormatAndExecuteError, Frame};
5use crate::options::Opt;
6use crate::options::Options;
7use crate::output::Output;
8use crate::prompt::{Prompt, StringTemplateError};
9use crate::traits::Executor;
10use crate::{chains::sequential, prompt, Parameters};
11
12use serde::Deserialize;
13use serde::Serialize;
14#[derive(derive_builder::Builder, Debug, Clone, Serialize, Deserialize)]
15/// A step in a chain of LLM invocations. It is a combination of a prompt and a configuration.
16pub struct Step {
17    pub(crate) prompt: prompt::PromptTemplate,
18    pub(crate) options: Options,
19}
20
21impl Step {
22    pub fn for_prompt_template(prompt: prompt::PromptTemplate) -> Self {
23        Self {
24            prompt,
25            options: Options::empty().clone(),
26        }
27    }
28    pub fn for_prompt_with_streaming(prompt: prompt::PromptTemplate) -> Self {
29        let mut options = Options::builder();
30        options.add_option(Opt::Stream(true));
31        let options = options.build();
32        Self { prompt, options }
33    }
34    pub fn for_prompt_and_options(prompt: prompt::PromptTemplate, options: Options) -> Self {
35        Self { prompt, options }
36    }
37    pub fn prompt(&self) -> &prompt::PromptTemplate {
38        &self.prompt
39    }
40    pub fn options(&self) -> &Options {
41        &self.options
42    }
43
44    /// Converts this step into a sequential chain with a single step.
45    ///
46    /// # Returns
47    ///
48    /// A sequential chain containing this step.
49    pub fn to_chain(self) -> sequential::Chain
50    where
51        Self: Sized,
52    {
53        crate::chains::sequential::Chain::of_one(self)
54    }
55
56    /// Formats the prompt for this step with the given parameters.
57    pub fn format(&self, parameters: &Parameters) -> Result<Prompt, StringTemplateError> {
58        self.prompt.format(parameters)
59    }
60
61    /// Executes the step with the given parameters and executor.
62    /// # Arguments
63    /// * `parameters` - A `Parameters` object containing the input parameters for the step.
64    /// * `executor` - An executor to use to execute the step.
65    /// # Returns
66    /// The output of the executor.
67    pub async fn run<E>(
68        &self,
69        parameters: &Parameters,
70        executor: &E,
71    ) -> Result<Output, FormatAndExecuteError>
72    where
73        Self: Sized,
74        E: Executor,
75    {
76        Frame::new(executor, self)
77            .format_and_execute(parameters)
78            .await
79    }
80}