aither_core/llm/
assistant.rs

1use alloc::{string::String, vec::Vec};
2
3use crate::{
4    LanguageModel,
5    llm::{
6        LLMRequest, Message, Tool, collect_text,
7        tool::{RegisterError, Tools},
8    },
9};
10
11/// A struct representing an Assistant that interacts with a language model (LLM),
12/// manages a collection of messages, and provides access to various tools.
13///
14/// # Type Parameters
15/// - `LLM`: A type that implements the `LanguageModel` trait, representing the
16///   underlying language model used by the Assistant.
17///
18/// # Fields
19/// - `messages`: A vector of `Message` instances representing the conversation
20///   history or context.
21/// - `tools`: A collection of tools available to the Assistant for performing
22///   various tasks.
23/// - `llm`: The language model instance used by the Assistant for generating
24///   responses or performing language-related tasks.
25#[derive(Debug)]
26pub struct Assistant<LLM: LanguageModel> {
27    messages: Vec<Message>,
28    tools: Tools,
29    llm: LLM,
30}
31
32impl<LLM: LanguageModel> Assistant<LLM> {
33    /// Creates a new `Assistant` instance with the provided language model.
34    ///
35    /// # Parameters
36    /// - `llm`: The language model instance to be used by the assistant.
37    ///
38    /// # Returns
39    /// Returns a new `Assistant` with empty messages and tools.
40    #[must_use]
41    pub const fn new(llm: LLM) -> Self {
42        Self {
43            messages: Vec::new(),
44            tools: Tools::new(),
45            llm,
46        }
47    }
48
49    /// Adds a system message to the conversation history.
50    ///
51    /// # Parameters
52    /// - `prompt`: The system prompt to add as a message.
53    ///
54    /// # Returns
55    /// Returns the updated Assistant instance with the system message added.
56    #[must_use]
57    pub fn system(mut self, prompt: impl Into<String>) -> Self {
58        self.messages.push(Message::system(prompt.into()));
59        self
60    }
61
62    /// Registers a tool with the assistant, making it available for use in interactions.
63    ///
64    /// # Parameters
65    /// - `tool`: The tool to register with the assistant.
66    ///
67    /// # Errors
68    /// Returns [`RegisterError`] if a tool of that name is already registered,
69    /// or the tool has no description for the model to read.
70    pub fn tool(mut self, tool: impl Tool + 'static) -> core::result::Result<Self, RegisterError> {
71        self.tools.register(tool)?;
72        Ok(self)
73    }
74
75    /// Sends a user message to the assistant, processes it with the language model, and appends the response to the conversation history.
76    ///
77    /// # Parameters
78    /// - `message`: The user message to send to the assistant.
79    ///
80    /// # Errors
81    /// Returns an error if the language model fails to generate a response or if message processing fails.
82    pub async fn send(&mut self, message: impl Into<String>) -> anyhow::Result<()> {
83        self.messages.push(Message::user(message));
84        let request = LLMRequest::new(self.messages.as_slice()).with_tools(&mut self.tools);
85        let stream = self.llm.respond_with_tools(request);
86
87        let response = collect_text(stream).await?;
88        self.messages.push(Message::assistant(response));
89        Ok(())
90    }
91
92    /// Returns a slice of all messages in the conversation history.
93    pub const fn messages(&self) -> &[Message] {
94        self.messages.as_slice()
95    }
96}
aither_core/llm/assistant.rs

aither_core/llm/
assistant.rs
