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

use crate::completion::{
    Completion, CompletionError, CompletionModel, CompletionRequestBuilder, CompletionResponse,
    Message, ModelChoice, Prompt, PromptError,
};

/// A model that can be used to prompt completions from a completion model.
/// This is the simplest building block for creating an LLM powered application.
pub struct Model<M: CompletionModel> {
    /// Completion model (e.g.: OpenAI's gpt-3.5-turbo-1106, Cohere's command-r)
    model: M,
    /// Temperature of the model
    temperature: Option<f64>,
}

impl<M: CompletionModel> Completion<M> for Model<M> {
    async fn completion(
        &self,
        prompt: &str,
        chat_history: Vec<Message>,
    ) -> Result<CompletionRequestBuilder<M>, CompletionError> {
        Ok(self
            .model
            .completion_request(prompt)
            .messages(chat_history)
            .temperature_opt(self.temperature))
    }
}

impl<M: CompletionModel> Prompt for Model<M> {
    async fn chat(&self, prompt: &str, chat_history: Vec<Message>) -> Result<String, PromptError> {
        match self.completion(prompt, chat_history).await?.send().await? {
            CompletionResponse {
                choice: ModelChoice::Message(message),
                ..
            } => Ok(message),
            CompletionResponse {
                choice: ModelChoice::ToolCall(toolname, _),
                ..
            } => Err(PromptError::ToolError(
                crate::tool::ToolSetError::ToolNotFoundError(toolname),
            )),
        }
    }
}

pub struct ModelBuilder<M: CompletionModel> {
    model: M,
    pub temperature: Option<f64>,
}

impl<M: CompletionModel> ModelBuilder<M> {
    pub fn new(model: M) -> Self {
        Self {
            model,
            temperature: None,
        }
    }

    pub fn temperature(mut self, temperature: f64) -> Self {
        self.temperature = Some(temperature);
        self
    }

    pub fn build(self) -> Model<M> {
        Model {
            model: self.model,
            temperature: self.temperature,
        }
    }
}