use alloc::{string::String, vec::Vec};
use crate::{
LanguageModel,
llm::{
LLMRequest, Message, Tool, collect_text,
tool::{RegisterError, Tools},
},
};
#[derive(Debug)]
pub struct Assistant<LLM: LanguageModel> {
messages: Vec<Message>,
tools: Tools,
llm: LLM,
}
impl<LLM: LanguageModel> Assistant<LLM> {
#[must_use]
pub const fn new(llm: LLM) -> Self {
Self {
messages: Vec::new(),
tools: Tools::new(),
llm,
}
}
#[must_use]
pub fn system(mut self, prompt: impl Into<String>) -> Self {
self.messages.push(Message::system(prompt.into()));
self
}
pub fn tool(mut self, tool: impl Tool + 'static) -> core::result::Result<Self, RegisterError> {
self.tools.register(tool)?;
Ok(self)
}
pub async fn send(&mut self, message: impl Into<String>) -> anyhow::Result<()> {
self.messages.push(Message::user(message));
let request = LLMRequest::new(self.messages.as_slice()).with_tools(&mut self.tools);
let stream = self.llm.respond_with_tools(request);
let response = collect_text(stream).await?;
self.messages.push(Message::assistant(response));
Ok(())
}
pub const fn messages(&self) -> &[Message] {
self.messages.as_slice()
}
}