llm_chain/
summarization.rs

1//! Opinionated text summarization functionality
2//!
3//! This module contains the `TextSummarizer` struct, that provides an easy way to summarize text.
4
5use crate::{
6    chains::map_reduce::{self, MapReduceChainError},
7    frame::FormatAndExecuteError,
8    parameters, prompt,
9    step::Step,
10    traits,
11};
12
13/// A `TextSummarizer` takes a given text and summarizes it using an `Executor`.
14///
15/// The summarizer is built on top of a `map_reduce::Chain`, which takes care of the summarization process.
16pub struct TextSummarizer {
17    chain: map_reduce::Chain,
18}
19
20impl Default for TextSummarizer {
21    fn default() -> Self {
22        let map_prompt = Step::for_prompt_template(prompt!(
23            "You are a text summarizer. You will be given a text and you will have to summarize it",
24            "Text:\n\n{{text}}\n\nPlease write a summary of the text above. Respond only with the summary."
25        ));
26        let reduce_prompt = Step::for_prompt_template(prompt!(
27            "You are a text summarizer. You will be given a text and you will have to summarize it",
28            "Text:\n\n{{text}}\n\nPlease write a combined summary of the segment summaries above. Respond only with the summary."
29        ));
30
31        TextSummarizer {
32            chain: map_reduce::Chain::new(map_prompt, reduce_prompt),
33        }
34    }
35}
36
37/// The error type returned by the `TextSummarizer` when summarizing text.
38#[derive(thiserror::Error, Debug)]
39pub enum TextSummarizerError {
40    #[error("MapReduceChainError: {0}")]
41    MapReduceChainError(#[from] MapReduceChainError),
42    #[error("No output was produced")]
43    NoOutput,
44}
45
46impl TextSummarizer {
47    /// Summarizes the given text using the provided `Executor`.
48    ///
49    /// Returns the summarized text, or an error if the summarization process fails.
50    pub async fn summarize_text<E: traits::Executor>(
51        &self,
52        exec: &E,
53        text: &str,
54    ) -> Result<String, TextSummarizerError> {
55        let params = parameters! {
56            "text" => text,
57        };
58        let chain_output = self.chain.run(vec![params], parameters!(), exec).await?;
59        chain_output
60            .to_immediate()
61            .await
62            .map_err(|err| {
63                TextSummarizerError::MapReduceChainError(
64                    MapReduceChainError::FormatAndExecuteError(FormatAndExecuteError::Execute(err)),
65                )
66            })?
67            .primary_textual_output()
68            .ok_or(TextSummarizerError::NoOutput)
69    }
70}
71
72/// A convenience function to summarize text using the provided `Executor`.
73///
74/// Returns the summarized text, or an error if the summarization process fails.
75pub async fn summarize_text<E: traits::Executor>(
76    exec: &E,
77    text: &str,
78) -> Result<String, TextSummarizerError> {
79    TextSummarizer::default().summarize_text(exec, text).await
80}
llm_chain/summarization.rs

llm_chain/
summarization.rs
