studio/mistralai-client-rs
2
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat!: add function calling support to client.chat() & client.chat_async()

Browse Source 
BREAKING CHANGE: Too many to count in this version. Check the README examples.
This commit is contained in:
Ivan Gabriele
2024-03-09 11:28:50 +01:00
parent 9430d42382
commit 74bf8a96ee
30 changed files with 1510 additions and 322 deletions
Download Patch File Download Diff File Expand all files Collapse all files

417
src/v1/client.rs
View File

@@ -1,27 +1,23 @@
use futures::stream::StreamExt;
use futures::Stream;
use log::debug;
use reqwest::Error as ReqwestError;
use serde_json::from_str;
use crate::v1::error::ApiError;
use crate::v1::{
chat_completion::{
ChatCompletionParams, ChatCompletionRequest, ChatCompletionResponse, ChatMessage,
},
constants::{EmbedModel, Model, API_URL_BASE},
embedding::{EmbeddingRequest, EmbeddingRequestOptions, EmbeddingResponse},
error::ClientError,
model_list::ModelListResponse,
use std::{
any::Any,
collections::HashMap,
sync::{Arc, Mutex},
};
use super::chat_completion::ChatCompletionStreamChunk;
use crate::v1::{chat, chat_stream, constants, embedding, error, model_list, tool, utils};
pub struct Client {
pub api_key: String,
pub endpoint: String,
pub max_retries: u32,
pub timeout: u32,
functions: Arc<Mutex<HashMap<String, Box<dyn tool::Function>>>>,
last_function_call_result: Arc<Mutex<Option<Box<dyn Any + Send>>>>,
}
impl Client {
@@ -53,20 +49,28 @@ impl Client {
endpoint: Option<String>,
max_retries: Option<u32>,
timeout: Option<u32>,
) -> Result<Self, ClientError> {
) -> Result<Self, error::ClientError> {
let api_key = match api_key {
Some(api_key_from_param) => api_key_from_param,
None => std::env::var("MISTRAL_API_KEY").map_err(|_| ClientError::MissingApiKey)?,
None => {
std::env::var("MISTRAL_API_KEY").map_err(|_| error::ClientError::MissingApiKey)?
}
};
let endpoint = endpoint.unwrap_or(API_URL_BASE.to_string());
let endpoint = endpoint.unwrap_or(constants::API_URL_BASE.to_string());
let max_retries = max_retries.unwrap_or(5);
let timeout = timeout.unwrap_or(120);
let functions: Arc<_> = Arc::new(Mutex::new(HashMap::new()));
let last_function_call_result = Arc::new(Mutex::new(None));
Ok(Self {
api_key,
endpoint,
max_retries,
timeout,
functions,
last_function_call_result,
})
}
@@ -76,42 +80,49 @@ impl Client {
///
/// * `model` - The [Model] to use for the chat completion.
/// * `messages` - A vector of [ChatMessage] to send as part of the chat.
/// * `options` - Optional [ChatCompletionParams] to customize the request.
/// * `options` - Optional [ChatParams] to customize the request.
///
/// # Returns
///
/// Returns a [Result] containing the `ChatCompletionResponse` if the request is successful,
/// Returns a [Result] containing the `ChatResponse` if the request is successful,
/// or an [ApiError] if there is an error.
///
/// # Examples
///
/// ```
/// use mistralai_client::v1::{
/// chat_completion::{ChatMessage, ChatMessageRole},
/// chat::{ChatMessage, ChatMessageRole},
/// client::Client,
/// constants::Model,
/// };
///
/// let client = Client::new(None, None, None, None).unwrap();
/// let messages = vec![ChatMessage {
/// role: ChatMessageRole::user,
/// role: ChatMessageRole::User,
/// content: "Hello, world!".to_string(),
/// tool_calls: None,
/// }];
/// let response = client.chat(Model::OpenMistral7b, messages, None).unwrap();
/// println!("{}: {}", response.choices[0].message.role, response.choices[0].message.content);
/// println!("{:?}: {}", response.choices[0].message.role, response.choices[0].message.content);
/// ```
pub fn chat(
&self,
model: Model,
messages: Vec<ChatMessage>,
options: Option<ChatCompletionParams>,
) -> Result<ChatCompletionResponse, ApiError> {
let request = ChatCompletionRequest::new(model, messages, false, options);
model: constants::Model,
messages: Vec<chat::ChatMessage>,
options: Option<chat::ChatParams>,
) -> Result<chat::ChatResponse, error::ApiError> {
let request = chat::ChatRequest::new(model, messages, false, options);
let response = self.post_sync("/chat/completions", &request)?;
let result = response.json::<ChatCompletionResponse>();
let result = response.json::<chat::ChatResponse>();
match result {
Ok(response) => Ok(response),
Ok(data) => {
utils::debug_pretty_json_from_struct("Response Data", &data);
self.call_function_if_any(data.clone());
Ok(data)
}
Err(error) => Err(self.to_api_error(error)),
}
}
@@ -122,18 +133,18 @@ impl Client {
///
/// * `model` - The [Model] to use for the chat completion.
/// * `messages` - A vector of [ChatMessage] to send as part of the chat.
/// * `options` - Optional [ChatCompletionParams] to customize the request.
/// * `options` - Optional [ChatParams] to customize the request.
///
/// # Returns
///
/// Returns a [Result] containing a `Stream` of `ChatCompletionStreamChunk` if the request is successful,
/// Returns a [Result] containing a `Stream` of `ChatStreamChunk` if the request is successful,
/// or an [ApiError] if there is an error.
///
/// # Examples
///
/// ```
/// use mistralai_client::v1::{
/// chat_completion::{ChatMessage, ChatMessageRole},
/// chat::{ChatMessage, ChatMessageRole},
/// client::Client,
/// constants::Model,
/// };
@@ -142,25 +153,32 @@ impl Client {
/// async fn main() {
/// let client = Client::new(None, None, None, None).unwrap();
/// let messages = vec![ChatMessage {
/// role: ChatMessageRole::user,
/// role: ChatMessageRole::User,
/// content: "Hello, world!".to_string(),
/// tool_calls: None,
/// }];
/// let response = client.chat_async(Model::OpenMistral7b, messages, None).await.unwrap();
/// println!("{}: {}", response.choices[0].message.role, response.choices[0].message.content);
/// println!("{:?}: {}", response.choices[0].message.role, response.choices[0].message.content);
/// }
/// ```
pub async fn chat_async(
&self,
model: Model,
messages: Vec<ChatMessage>,
options: Option<ChatCompletionParams>,
) -> Result<ChatCompletionResponse, ApiError> {
let request = ChatCompletionRequest::new(model, messages, false, options);
model: constants::Model,
messages: Vec<chat::ChatMessage>,
options: Option<chat::ChatParams>,
) -> Result<chat::ChatResponse, error::ApiError> {
let request = chat::ChatRequest::new(model, messages, false, options);
let response = self.post_async("/chat/completions", &request).await?;
let result = response.json::<ChatCompletionResponse>().await;
let result = response.json::<chat::ChatResponse>().await;
match result {
Ok(response) => Ok(response),
Ok(data) => {
utils::debug_pretty_json_from_struct("Response Data", &data);
self.call_function_if_any_async(data.clone()).await;
Ok(data)
}
Err(error) => Err(self.to_api_error(error)),
}
}
@@ -171,11 +189,11 @@ impl Client {
///
/// * `model` - The [Model] to use for the chat completion.
/// * `messages` - A vector of [ChatMessage] to send as part of the chat.
/// * `options` - Optional [ChatCompletionParams] to customize the request.
/// * `options` - Optional [ChatParams] to customize the request.
///
/// # Returns
///
/// Returns a [Result] containing a `Stream` of `ChatCompletionStreamChunk` if the request is successful,
/// Returns a [Result] containing a `Stream` of `ChatStreamChunk` if the request is successful,
/// or an [ApiError] if there is an error.
///
/// # Examples
@@ -183,124 +201,176 @@ impl Client {
/// ```
/// use futures::stream::StreamExt;
/// use mistralai_client::v1::{
/// chat_completion::{ChatMessage, ChatMessageRole},
/// chat::{ChatMessage, ChatMessageRole},
/// client::Client,
/// constants::Model,
/// };
/// use std::io::{self, Write};
///
/// #[tokio::main]
/// async fn main() {
/// let client = Client::new(None, None, None, None).unwrap();
/// let messages = vec![ChatMessage {
/// role: ChatMessageRole::user,
/// role: ChatMessageRole::User,
/// content: "Hello, world!".to_string(),
/// tool_calls: None,
/// }];
/// let mut stream = client.chat_stream(Model::OpenMistral7b, messages, None).await.unwrap();
/// while let Some(chunk_result) = stream.next().await {
/// match chunk_result {
/// Ok(chunk) => {
/// print!("{}", chunk.choices[0].delta.content);
///
/// let stream_result = client
/// .chat_stream(Model::OpenMistral7b,messages, None)
/// .await
/// .unwrap();
/// stream_result
/// .for_each(|chunk_result| async {
/// match chunk_result {
/// Ok(chunks) => chunks.iter().for_each(|chunk| {
/// print!("{}", chunk.choices[0].delta.content);
/// io::stdout().flush().unwrap();
/// // => "Once upon a time, [...]"
/// }),
/// Err(error) => {
/// eprintln!("Error processing chunk: {:?}", error)
/// }
/// }
/// Err(error) => {
/// println!("Error: {}", error.message);
/// }
/// }
/// }
/// })
/// .await;
/// print!("\n") // To persist the last chunk output.
/// }
pub async fn chat_stream(
&self,
model: Model,
messages: Vec<ChatMessage>,
options: Option<ChatCompletionParams>,
) -> Result<impl Stream<Item = Result<ChatCompletionStreamChunk, ApiError>>, ApiError> {
let request = ChatCompletionRequest::new(model, messages, true, options);
model: constants::Model,
messages: Vec<chat::ChatMessage>,
options: Option<chat::ChatParams>,
) -> Result<
impl Stream<Item = Result<Vec<chat_stream::ChatStreamChunk>, error::ApiError>>,
error::ApiError,
> {
let request = chat::ChatRequest::new(model, messages, true, options);
let response = self
.post_stream("/chat/completions", &request)
.await
.map_err(|e| ApiError {
.map_err(|e| error::ApiError {
message: e.to_string(),
})?;
if !response.status().is_success() {
let status = response.status();
let text = response.text().await.unwrap_or_default();
return Err(ApiError {
return Err(error::ApiError {
message: format!("{}: {}", status, text),
});
}
let deserialized_stream =
response
.bytes_stream()
.map(|item| -> Result<ChatCompletionStreamChunk, ApiError> {
match item {
Ok(bytes) => {
let text = String::from_utf8(bytes.to_vec()).map_err(|e| ApiError {
message: e.to_string(),
})?;
let text_trimmed = text.trim_start_matches("data: ");
from_str(&text_trimmed).map_err(|e| ApiError {
message: e.to_string(),
})
}
Err(e) => Err(ApiError {
message: e.to_string(),
}),
let deserialized_stream = response.bytes_stream().then(|bytes_result| async move {
match bytes_result {
Ok(bytes) => match String::from_utf8(bytes.to_vec()) {
Ok(message) => {
let chunks = message
.lines()
.filter_map(
|line| match chat_stream::get_chunk_from_stream_message_line(line) {
Ok(Some(chunks)) => Some(chunks),
Ok(None) => None,
Err(_error) => None,
},
)
.flatten()
.collect();
Ok(chunks)
}
});
Err(e) => Err(error::ApiError {
message: e.to_string(),
}),
},
Err(e) => Err(error::ApiError {
message: e.to_string(),
}),
}
});
Ok(deserialized_stream)
}
pub fn embeddings(
&self,
model: EmbedModel,
model: constants::EmbedModel,
input: Vec<String>,
options: Option<EmbeddingRequestOptions>,
) -> Result<EmbeddingResponse, ApiError> {
let request = EmbeddingRequest::new(model, input, options);
options: Option<embedding::EmbeddingRequestOptions>,
) -> Result<embedding::EmbeddingResponse, error::ApiError> {
let request = embedding::EmbeddingRequest::new(model, input, options);
let response = self.post_sync("/embeddings", &request)?;
let result = response.json::<EmbeddingResponse>();
let result = response.json::<embedding::EmbeddingResponse>();
match result {
Ok(response) => Ok(response),
Ok(data) => {
utils::debug_pretty_json_from_struct("Response Data", &data);
Ok(data)
}
Err(error) => Err(self.to_api_error(error)),
}
}
pub async fn embeddings_async(
&self,
model: EmbedModel,
model: constants::EmbedModel,
input: Vec<String>,
options: Option<EmbeddingRequestOptions>,
) -> Result<EmbeddingResponse, ApiError> {
let request = EmbeddingRequest::new(model, input, options);
options: Option<embedding::EmbeddingRequestOptions>,
) -> Result<embedding::EmbeddingResponse, error::ApiError> {
let request = embedding::EmbeddingRequest::new(model, input, options);
let response = self.post_async("/embeddings", &request).await?;
let result = response.json::<EmbeddingResponse>().await;
let result = response.json::<embedding::EmbeddingResponse>().await;
match result {
Ok(response) => Ok(response),
Ok(data) => {
utils::debug_pretty_json_from_struct("Response Data", &data);
Ok(data)
}
Err(error) => Err(self.to_api_error(error)),
}
}
pub fn list_models(&self) -> Result<ModelListResponse, ApiError> {
pub fn get_last_function_call_result(&self) -> Option<Box<dyn Any + Send>> {
let mut result_lock = self.last_function_call_result.lock().unwrap();
result_lock.take()
}
pub fn list_models(&self) -> Result<model_list::ModelListResponse, error::ApiError> {
let response = self.get_sync("/models")?;
let result = response.json::<ModelListResponse>();
let result = response.json::<model_list::ModelListResponse>();
match result {
Ok(response) => Ok(response),
Ok(data) => {
utils::debug_pretty_json_from_struct("Response Data", &data);
Ok(data)
}
Err(error) => Err(self.to_api_error(error)),
}
}
pub async fn list_models_async(&self) -> Result<ModelListResponse, ApiError> {
pub async fn list_models_async(
&self,
) -> Result<model_list::ModelListResponse, error::ApiError> {
let response = self.get_async("/models").await?;
let result = response.json::<ModelListResponse>().await;
let result = response.json::<model_list::ModelListResponse>().await;
match result {
Ok(response) => Ok(response),
Ok(data) => {
utils::debug_pretty_json_from_struct("Response Data", &data);
Ok(data)
}
Err(error) => Err(self.to_api_error(error)),
}
}
pub fn register_function(&mut self, name: String, function: Box<dyn tool::Function>) {
let mut functions = self.functions.lock().unwrap();
functions.insert(name, function);
}
fn build_request_sync(
&self,
request: reqwest::blocking::RequestBuilder,
@@ -349,9 +419,70 @@ impl Client {
request_builder
}
fn get_sync(&self, path: &str) -> Result<reqwest::blocking::Response, ApiError> {
fn call_function_if_any(&self, response: chat::ChatResponse) -> () {
let next_result = match response.choices.get(0) {
Some(first_choice) => match first_choice.message.tool_calls.to_owned() {
Some(tool_calls) => match tool_calls.get(0) {
Some(first_tool_call) => {
let functions = self.functions.lock().unwrap();
match functions.get(&first_tool_call.function.name) {
Some(function) => {
let runtime = tokio::runtime::Runtime::new().unwrap();
let result = runtime.block_on(async {
function
.execute(first_tool_call.function.arguments.to_owned())
.await
});
Some(result)
}
None => None,
}
}
None => None,
},
None => None,
},
None => None,
};
let mut last_result_lock = self.last_function_call_result.lock().unwrap();
*last_result_lock = next_result;
}
async fn call_function_if_any_async(&self, response: chat::ChatResponse) -> () {
let next_result = match response.choices.get(0) {
Some(first_choice) => match first_choice.message.tool_calls.to_owned() {
Some(tool_calls) => match tool_calls.get(0) {
Some(first_tool_call) => {
let functions = self.functions.lock().unwrap();
match functions.get(&first_tool_call.function.name) {
Some(function) => {
let result = function
.execute(first_tool_call.function.arguments.to_owned())
.await;
Some(result)
}
None => None,
}
}
None => None,
},
None => None,
},
None => None,
};
let mut last_result_lock = self.last_function_call_result.lock().unwrap();
*last_result_lock = next_result;
}
fn get_sync(&self, path: &str) -> Result<reqwest::blocking::Response, error::ApiError> {
let reqwest_client = reqwest::blocking::Client::new();
let url = format!("{}{}", self.endpoint, path);
debug!("Request URL: {}", url);
let request = self.build_request_sync(reqwest_client.get(url));
let result = request.send();
@@ -360,22 +491,27 @@ impl Client {
if response.status().is_success() {
Ok(response)
} else {
let status = response.status();
let text = response.text().unwrap();
Err(ApiError {
message: format!("{}: {}", status, text),
let response_status = response.status();
let response_body = response.text().unwrap_or_default();
debug!("Response Status: {}", &response_status);
utils::debug_pretty_json_from_string("Response Data", &response_body);
Err(error::ApiError {
message: format!("{}: {}", response_status, response_body),
})
}
}
Err(error) => Err(ApiError {
Err(error) => Err(error::ApiError {
message: error.to_string(),
}),
}
}
async fn get_async(&self, path: &str) -> Result<reqwest::Response, ApiError> {
async fn get_async(&self, path: &str) -> Result<reqwest::Response, error::ApiError> {
let reqwest_client = reqwest::Client::new();
let url = format!("{}{}", self.endpoint, path);
debug!("Request URL: {}", url);
let request_builder = reqwest_client.get(url);
let request = self.build_request_async(request_builder);
@@ -385,26 +521,32 @@ impl Client {
if response.status().is_success() {
Ok(response)
} else {
let status = response.status();
let text = response.text().await.unwrap_or_default();
Err(ApiError {
message: format!("{}: {}", status, text),
let response_status = response.status();
let response_body = response.text().await.unwrap_or_default();
debug!("Response Status: {}", &response_status);
utils::debug_pretty_json_from_string("Response Data", &response_body);
Err(error::ApiError {
message: format!("{}: {}", response_status, response_body),
})
}
}
Err(error) => Err(ApiError {
Err(error) => Err(error::ApiError {
message: error.to_string(),
}),
}
}
fn post_sync<T: serde::ser::Serialize + std::fmt::Debug>(
fn post_sync<T: std::fmt::Debug + serde::ser::Serialize>(
&self,
path: &str,
params: &T,
) -> Result<reqwest::blocking::Response, ApiError> {
) -> Result<reqwest::blocking::Response, error::ApiError> {
let reqwest_client = reqwest::blocking::Client::new();
let url = format!("{}{}", self.endpoint, path);
debug!("Request URL: {}", url);
utils::debug_pretty_json_from_struct("Request Body", params);
let request_builder = reqwest_client.post(url).json(params);
let request = self.build_request_sync(request_builder);
@@ -414,14 +556,17 @@ impl Client {
if response.status().is_success() {
Ok(response)
} else {
let status = response.status();
let text = response.text().unwrap_or_default();
Err(ApiError {
message: format!("{}: {}", status, text),
let response_status = response.status();
let response_body = response.text().unwrap_or_default();
debug!("Response Status: {}", &response_status);
utils::debug_pretty_json_from_string("Response Data", &response_body);
Err(error::ApiError {
message: format!("{}: {}", response_body, response_status),
})
}
}
Err(error) => Err(ApiError {
Err(error) => Err(error::ApiError {
message: error.to_string(),
}),
}
@@ -431,9 +576,12 @@ impl Client {
&self,
path: &str,
params: &T,
) -> Result<reqwest::Response, ApiError> {
) -> Result<reqwest::Response, error::ApiError> {
let reqwest_client = reqwest::Client::new();
let url = format!("{}{}", self.endpoint, path);
debug!("Request URL: {}", url);
utils::debug_pretty_json_from_struct("Request Body", params);
let request_builder = reqwest_client.post(url).json(params);
let request = self.build_request_async(request_builder);
@@ -443,14 +591,17 @@ impl Client {
if response.status().is_success() {
Ok(response)
} else {
let status = response.status();
let text = response.text().await.unwrap_or_default();
Err(ApiError {
message: format!("{}: {}", status, text),
let response_status = response.status();
let response_body = response.text().await.unwrap_or_default();
debug!("Response Status: {}", &response_status);
utils::debug_pretty_json_from_string("Response Data", &response_body);
Err(error::ApiError {
message: format!("{}: {}", response_status, response_body),
})
}
}
Err(error) => Err(ApiError {
Err(error) => Err(error::ApiError {
message: error.to_string(),
}),
}
@@ -460,9 +611,12 @@ impl Client {
&self,
path: &str,
params: &T,
) -> Result<reqwest::Response, ApiError> {
) -> Result<reqwest::Response, error::ApiError> {
let reqwest_client = reqwest::Client::new();
let url = format!("{}{}", self.endpoint, path);
debug!("Request URL: {}", url);
utils::debug_pretty_json_from_struct("Request Body", params);
let request_builder = reqwest_client.post(url).json(params);
let request = self.build_request_stream(request_builder);
@@ -472,21 +626,24 @@ impl Client {
if response.status().is_success() {
Ok(response)
} else {
let status = response.status();
let text = response.text().await.unwrap_or_default();
Err(ApiError {
message: format!("{}: {}", status, text),
let response_status = response.status();
let response_body = response.text().await.unwrap_or_default();
debug!("Response Status: {}", &response_status);
utils::debug_pretty_json_from_string("Response Data", &response_body);
Err(error::ApiError {
message: format!("{}: {}", response_status, response_body),
})
}
}
Err(error) => Err(ApiError {
Err(error) => Err(error::ApiError {
message: error.to_string(),
}),
}
}
fn to_api_error(&self, err: ReqwestError) -> ApiError {
ApiError {
fn to_api_error(&self, err: ReqwestError) -> error::ApiError {
error::ApiError {
message: err.to_string(),
}
}