unai 0.0.1

Common AI abstractions
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

use std::sync::Arc;

use async_trait::async_trait;

use super::completion::Completion;
use super::llm::{GenerateCompletionError, GenerateCompletionInput, GenerateCompletionResult, LLM};
use super::settings::Settings;

pub type MockLLMResult = Result<Completion, String>;

/// A mock LLM for testing purposes
pub struct MockLLM {
    result: MockLLMResult,
}

impl MockLLM {
    pub fn new(result: MockLLMResult) -> Self {
        Self { result }
    }
}

#[async_trait]
impl LLM for MockLLM {
    async fn generate_completion(
        &self,
        _input: GenerateCompletionInput,
    ) -> GenerateCompletionResult {
        match &self.result {
            Ok(completion) => Ok(completion.clone()),
            Err(err_msg) => Err(GenerateCompletionError::Unknown(err_msg.clone().into())),
        }
    }
}

pub struct WithDefaultSettings {
    llm: Box<dyn LLM>,
    default_settings: Settings,
}

impl WithDefaultSettings {
    pub fn new(llm: Box<dyn LLM>, default_settings: Settings) -> Self {
        Self {
            llm,
            default_settings,
        }
    }
}

#[async_trait]
impl LLM for WithDefaultSettings {
    async fn generate_completion(
        &self,
        mut input: GenerateCompletionInput,
    ) -> GenerateCompletionResult {
        let input_settings = input.settings.unwrap_or_default();

        input.settings = Some(input_settings.merge(self.default_settings.clone()));

        self.llm.generate_completion(input).await
    }
}