1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75

//! Opinionated text summarization functionality
//!
//! This module contains the `TextSummarizer` struct, that provides an easy way to summarize text.
use crate::{
    chains::map_reduce::{self, MapReduceChainError},
    output::Output,
    parameters, prompt,
    step::Step,
    traits,
};

/// A `TextSummarizer` takes a given text and summarizes it using an `Executor`.
///
/// The summarizer is built on top of a `map_reduce::Chain`, which takes care of the summarization process.
pub struct TextSummarizer<E: traits::Executor> {
    chain: map_reduce::Chain<E>,
}

impl<E: traits::Executor> Default for TextSummarizer<E> {
    fn default() -> Self {
        let map_prompt = Step::for_prompt(prompt!(
            "You are a text summarizer. You will be given a text and you will have to summarize it",
            "Text:\n\n{{text}}\n\nPlease write a summary of the text above. Respond only with the summary."
        ));
        let reduce_prompt = Step::for_prompt(prompt!(
            "You are a text summarizer. You will be given a text and you will have to summarize it",
            "Text:\n\n{{text}}\n\nPlease write a combined summary of the segment summaries above. Respond only with the summary."
        ));

        TextSummarizer {
            chain: map_reduce::Chain::new(map_prompt, reduce_prompt),
        }
    }
}

/// The error type returned by the `TextSummarizer` when summarizing text.
#[derive(thiserror::Error, Debug)]
pub enum TextSummarizerError<E: traits::ExecutorError> {
    #[error("MapReduceChainError: {0}")]
    MapReduceChainError(#[from] MapReduceChainError<E>),
    #[error("No output was produced")]
    NoOutput,
}

impl<E: traits::Executor> TextSummarizer<E> {
    /// Summarizes the given text using the provided `Executor`.
    ///
    /// Returns the summarized text, or an error if the summarization process fails.
    pub async fn summarize_text(
        &self,
        exec: &E,
        text: &str,
    ) -> Result<String, TextSummarizerError<E::Error>> {
        let params = parameters! {
            "text" => text,
        };
        let chain_output = self.chain.run(vec![params], parameters!(), exec).await?;
        chain_output
            .primary_textual_output()
            .await
            .ok_or(TextSummarizerError::NoOutput)
    }
}

/// A convenience function to summarize text using the provided `Executor`.
///
/// Returns the summarized text, or an error if the summarization process fails.
pub async fn summarize_text<E: traits::Executor>(
    exec: &E,
    text: &str,
) -> Result<String, TextSummarizerError<E::Error>> {
    TextSummarizer::<E>::default()
        .summarize_text(exec, text)
        .await
}