llm-chain 0.13.0

A library for running chains of LLMs (such as ChatGPT) in series to complete complex tasks, such as text summation.
Documentation 
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
76
77
78
79
80

//! 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},
    frame::FormatAndExecuteError,
    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 {
    chain: map_reduce::Chain,
}

impl Default for TextSummarizer {
    fn default() -> Self {
        let map_prompt = Step::for_prompt_template(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_template(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 {
    #[error("MapReduceChainError: {0}")]
    MapReduceChainError(#[from] MapReduceChainError),
    #[error("No output was produced")]
    NoOutput,
}

impl TextSummarizer {
    /// 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<E: traits::Executor>(
        &self,
        exec: &E,
        text: &str,
    ) -> Result<String, TextSummarizerError> {
        let params = parameters! {
            "text" => text,
        };
        let chain_output = self.chain.run(vec![params], parameters!(), exec).await?;
        chain_output
            .to_immediate()
            .await
            .map_err(|err| {
                TextSummarizerError::MapReduceChainError(
                    MapReduceChainError::FormatAndExecuteError(FormatAndExecuteError::Execute(err)),
                )
            })?
            .primary_textual_output()
            .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> {
    TextSummarizer::default().summarize_text(exec, text).await
}