llm 1.3.8

A Rust library unifying multiple LLM backends.
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
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106

use std::collections::HashMap;

use crate::{
    chat::{ChatMessage, ChatRole, MessageType},
    completion::CompletionRequest,
    error::LLMError,
    LLMProvider,
};

use super::registry::LLMRegistry;
use super::step::{MultiChainStep, MultiChainStepMode, ResponseTransform};

/// The multi-backend chain.
pub struct MultiPromptChain<'a> {
    registry: &'a LLMRegistry,
    steps: Vec<MultiChainStep>,
    memory: HashMap<String, String>,
}

impl<'a> MultiPromptChain<'a> {
    pub fn new(registry: &'a LLMRegistry) -> Self {
        Self {
            registry,
            steps: vec![],
            memory: HashMap::new(),
        }
    }

    /// Adds a step.
    pub fn step(mut self, step: MultiChainStep) -> Self {
        self.steps.push(step);
        self
    }

    /// Adds multiple steps at once.
    pub fn chain(mut self, steps: Vec<MultiChainStep>) -> Self {
        self.steps.extend(steps);
        self
    }

    /// Executes all steps.
    pub async fn run(mut self) -> Result<HashMap<String, String>, LLMError> {
        for step in &self.steps {
            let response = self.run_step(step).await?;
            self.memory.insert(step.id.clone(), response);
        }
        Ok(self.memory)
    }

    async fn run_step(&self, step: &MultiChainStep) -> Result<String, LLMError> {
        let prompt_text = self.replace_template(&step.template);
        let llm = self.provider(step)?;
        let response = match step.mode {
            MultiChainStepMode::Chat => run_chat(llm, prompt_text).await?,
            MultiChainStepMode::Completion => run_completion(llm, step, prompt_text).await?,
            MultiChainStepMode::SpeechToText => llm.transcribe_file(&prompt_text).await?,
        };
        Ok(apply_transform(response, step.response_transform.as_ref()))
    }

    fn provider(&self, step: &MultiChainStep) -> Result<&dyn LLMProvider, LLMError> {
        self.registry.get(&step.provider_id).ok_or_else(|| {
            LLMError::InvalidRequest(format!(
                "No provider with id '{}' found in registry",
                step.provider_id
            ))
        })
    }

    fn replace_template(&self, input: &str) -> String {
        let mut out = input.to_string();
        for (k, v) in &self.memory {
            let pattern = format!("{{{{{k}}}}}");
            out = out.replace(&pattern, v);
        }
        out
    }
}

async fn run_chat(llm: &dyn LLMProvider, prompt_text: String) -> Result<String, LLMError> {
    let messages = vec![ChatMessage {
        role: ChatRole::User,
        message_type: MessageType::Text,
        content: prompt_text,
    }];
    Ok(llm.chat(&messages).await?.text().unwrap_or_default())
}

async fn run_completion(
    llm: &dyn LLMProvider,
    step: &MultiChainStep,
    prompt_text: String,
) -> Result<String, LLMError> {
    let mut req = CompletionRequest::new(prompt_text);
    req.temperature = step.temperature;
    req.max_tokens = step.max_tokens;
    let response = llm.complete(&req).await?;
    Ok(response.text.to_string())
}

fn apply_transform(response: String, transform: Option<&ResponseTransform>) -> String {
    match transform {
        Some(transform) => transform(response),
        None => response,
    }
}