aither 0.2.0

Providing unified trait abstractions for AI models
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

use alloc::{string::String, vec::Vec};

use crate::{
    LanguageModel,
    llm::{Message, Tool, model::Parameters, tool::Tools, try_collect},
};

/// A struct representing an Assistant that interacts with a language model (LLM),
/// manages a collection of messages, and provides access to various tools.
///
/// # Type Parameters
/// - `LLM`: A type that implements the `LanguageModel` trait, representing the
///   underlying language model used by the Assistant.
///
/// # Fields
/// - `messages`: A vector of `Message` instances representing the conversation
///   history or context.
/// - `tools`: A collection of tools available to the Assistant for performing
///   various tasks.
/// - `llm`: The language model instance used by the Assistant for generating
///   responses or performing language-related tasks.
#[derive(Debug)]
pub struct Assistant<LLM: LanguageModel> {
    messages: Vec<Message>,
    tools: Tools,
    llm: LLM,
}

impl<LLM: LanguageModel> Assistant<LLM> {
    /// Creates a new `Assistant` instance with the provided language model.
    ///
    /// # Parameters
    /// - `llm`: The language model instance to be used by the assistant.
    ///
    /// # Returns
    /// Returns a new `Assistant` with empty messages and tools.
    #[must_use]
    pub const fn new(llm: LLM) -> Self {
        Self {
            messages: Vec::new(),
            tools: Tools::new(),
            llm,
        }
    }

    /// Adds a system message to the conversation history.
    ///
    /// # Parameters
    /// - `prompt`: The system prompt to add as a message.
    ///
    /// # Returns
    /// Returns the updated Assistant instance with the system message added.
    #[must_use]
    pub fn system(mut self, prompt: impl Into<String>) -> Self {
        self.messages.push(Message::system(prompt.into()));
        self
    }

    /// Registers a tool with the assistant, making it available for use in interactions.
    ///
    /// # Parameters
    /// - `tool`: The tool to register with the assistant.
    ///
    /// # Returns
    /// Returns the updated Assistant instance with the tool registered.
    #[must_use]
    pub fn tool(mut self, tool: impl Tool) -> Self {
        self.tools.register(tool);
        self
    }

    /// Sends a user message to the assistant, processes it with the language model, and appends the response to the conversation history.
    ///
    /// # Parameters
    /// - `message`: The user message to send to the assistant.
    ///
    /// # Errors
    /// Returns an error if the language model fails to generate a response or if message processing fails.
    pub async fn send(&mut self, message: impl Into<String>) -> anyhow::Result<()> {
        self.messages.push(Message::user(message));
        let binding = Parameters::default();
        let stream = self.llm.respond(&self.messages, &mut self.tools, &binding);

        let response = try_collect(stream).await?;
        self.messages.push(Message::assistant(response));
        Ok(())
    }

    /// Returns a slice of all messages in the conversation history.
    pub const fn messages(&self) -> &[Message] {
        self.messages.as_slice()
    }
}