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

131
src/v1/chat_completion.rs → src/v1/chat.rs
View File

@@ -1,6 +1,6 @@
use serde::{Deserialize, Serialize};
use crate::v1::{common, constants};
use crate::v1::{common, constants, tool};
// -----------------------------------------------------------------------------
// Definitions
@@ -9,88 +9,115 @@ use crate::v1::{common, constants};
pub struct ChatMessage {
pub role: ChatMessageRole,
pub content: String,
pub tool_calls: Option<Vec<tool::ToolCall>>,
}
impl ChatMessage {
pub fn new_assistant_message(content: &str, tool_calls: Option<Vec<tool::ToolCall>>) -> Self {
Self {
role: ChatMessageRole::Assistant,
content: content.to_string(),
tool_calls,
}
}
pub fn new_user_message(content: &str) -> Self {
Self {
role: ChatMessageRole::User,
content: content.to_string(),
tool_calls: None,
}
}
}
#[derive(Clone, Debug, strum_macros::Display, Eq, PartialEq, Deserialize, Serialize)]
#[allow(non_camel_case_types)]
#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)]
pub enum ChatMessageRole {
assistant,
user,
#[serde(rename = "assistant")]
Assistant,
#[serde(rename = "user")]
User,
}
// -----------------------------------------------------------------------------
// Request
#[derive(Debug)]
pub struct ChatCompletionParams {
pub tools: Option<String>,
pub temperature: Option<f32>,
pub struct ChatParams {
pub max_tokens: Option<u32>,
pub top_p: Option<f32>,
pub random_seed: Option<u32>,
pub safe_prompt: Option<bool>,
pub temperature: Option<f32>,
pub tool_choice: Option<tool::ToolChoice>,
pub tools: Option<Vec<tool::Tool>>,
pub top_p: Option<f32>,
}
impl Default for ChatCompletionParams {
impl Default for ChatParams {
fn default() -> Self {
Self {
tools: None,
temperature: None,
max_tokens: None,
top_p: None,
random_seed: None,
safe_prompt: None,
temperature: None,
tool_choice: None,
tools: None,
top_p: None,
}
}
}
#[derive(Debug, Serialize, Deserialize)]
pub struct ChatCompletionRequest {
pub struct ChatRequest {
pub messages: Vec<ChatMessage>,
pub model: constants::Model,
#[serde(skip_serializing_if = "Option::is_none")]
pub tools: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub temperature: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub max_tokens: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub top_p: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub random_seed: Option<u32>,
pub stream: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub safe_prompt: Option<bool>,
pub stream: bool,
#[serde(skip_serializing_if = "Option::is_none")]
pub temperature: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tool_choice: Option<tool::ToolChoice>,
#[serde(skip_serializing_if = "Option::is_none")]
pub tools: Option<Vec<tool::Tool>>,
#[serde(skip_serializing_if = "Option::is_none")]
pub top_p: Option<f32>,
// TODO Check this prop (seen in official Python client but not in API doc).
// pub tool_choice: Option<String>,
// TODO Check this prop (seen in official Python client but not in API doc).
// pub response_format: Option<String>,
}
impl ChatCompletionRequest {
impl ChatRequest {
pub fn new(
model: constants::Model,
messages: Vec<ChatMessage>,
stream: bool,
options: Option<ChatCompletionParams>,
options: Option<ChatParams>,
) -> Self {
let ChatCompletionParams {
tools,
temperature,
let ChatParams {
max_tokens,
top_p,
random_seed,
safe_prompt,
temperature,
tool_choice,
tools,
top_p,
} = options.unwrap_or_default();
Self {
messages,
model,
tools,
temperature,
max_tokens,
top_p,
random_seed,
stream,
safe_prompt,
stream,
temperature,
tool_choice,
tools,
top_p,
}
}
}
@@ -99,51 +126,29 @@ impl ChatCompletionRequest {
// Response
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct ChatCompletionResponse {
pub struct ChatResponse {
pub id: String,
pub object: String,
/// Unix timestamp (in seconds).
pub created: u32,
pub model: constants::Model,
pub choices: Vec<ChatCompletionResponseChoice>,
pub choices: Vec<ChatResponseChoice>,
pub usage: common::ResponseUsage,
}
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct ChatCompletionResponseChoice {
pub struct ChatResponseChoice {
pub index: u32,
pub message: ChatMessage,
pub finish_reason: String,
pub finish_reason: ChatResponseChoiceFinishReason,
// TODO Check this prop (seen in API responses but undocumented).
// pub logprobs: ???
}
// -----------------------------------------------------------------------------
// Stream
#[derive(Debug, Deserialize)]
pub struct ChatCompletionStreamChunk {
pub id: String,
pub object: String,
/// Unix timestamp (in seconds).
pub created: u32,
pub model: constants::Model,
pub choices: Vec<ChatCompletionStreamChunkChoice>,
// TODO Check this prop (seen in API responses but undocumented).
// pub usage: ???,
}
#[derive(Debug, Deserialize)]
pub struct ChatCompletionStreamChunkChoice {
pub index: u32,
pub delta: ChatCompletionStreamChunkChoiceDelta,
pub finish_reason: Option<String>,
// TODO Check this prop (seen in API responses but undocumented).
// pub logprobs: ???,
}
#[derive(Debug, Deserialize)]
pub struct ChatCompletionStreamChunkChoiceDelta {
pub role: Option<ChatMessageRole>,
pub content: String,
#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)]
pub enum ChatResponseChoiceFinishReason {
#[serde(rename = "stop")]
Stop,
#[serde(rename = "tool_calls")]
ToolCalls,
}

57
src/v1/chat_stream.rs Normal file
View File

@@ -0,0 +1,57 @@
use serde::{Deserialize, Serialize};
use serde_json::from_str;
use crate::v1::{chat, common, constants, error};
// -----------------------------------------------------------------------------
// Response
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct ChatStreamChunk {
pub id: String,
pub object: String,
/// Unix timestamp (in seconds).
pub created: u32,
pub model: constants::Model,
pub choices: Vec<ChatStreamChunkChoice>,
pub usage: Option<common::ResponseUsage>,
// TODO Check this prop (seen in API responses but undocumented).
// pub logprobs: ???,
}
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct ChatStreamChunkChoice {
pub index: u32,
pub delta: ChatStreamChunkChoiceDelta,
pub finish_reason: Option<String>,
// TODO Check this prop (seen in API responses but undocumented).
// pub logprobs: ???,
}
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct ChatStreamChunkChoiceDelta {
pub role: Option<chat::ChatMessageRole>,
pub content: String,
}
/// Extracts serialized chunks from a stream message.
pub fn get_chunk_from_stream_message_line(
line: &str,
) -> Result<Option<Vec<ChatStreamChunk>>, error::ApiError> {
if line.trim() == "data: [DONE]" {
return Ok(None);
}
let chunk_as_json = line.trim_start_matches("data: ").trim();
if chunk_as_json.is_empty() {
return Ok(Some(vec![]));
}
// Attempt to deserialize the JSON string into ChatStreamChunk
match from_str::<ChatStreamChunk>(chunk_as_json) {
Ok(chunk) => Ok(Some(vec![chunk])),
Err(e) => Err(error::ApiError {
message: e.to_string(),
}),
}
}

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(),
}
}

6
src/v1/embedding.rs
View File

@@ -2,6 +2,9 @@ use serde::{Deserialize, Serialize};
use crate::v1::{common, constants};
// -----------------------------------------------------------------------------
// Request
#[derive(Debug)]
pub struct EmbeddingRequestOptions {
pub encoding_format: Option<EmbeddingRequestEncodingFormat>,
@@ -43,6 +46,9 @@ pub enum EmbeddingRequestEncodingFormat {
float,
}
// -----------------------------------------------------------------------------
// Response
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct EmbeddingResponse {
pub id: String,

4
src/v1/error.rs
View File

@@ -14,7 +14,9 @@ impl Error for ApiError {}
#[derive(Debug, PartialEq, thiserror::Error)]
pub enum ClientError {
#[error("You must either set the `MISTRAL_API_KEY` environment variable or specify it in `Client::new(api_key, ...).")]
#[error(
"You must either set the `MISTRAL_API_KEY` environment variable or specify it in `Client::new(api_key, ...)."
)]
MissingApiKey,
#[error("Failed to read the response text.")]
UnreadableResponseText,

5
src/v1/mod.rs
View File

@@ -1,7 +1,10 @@
pub mod chat_completion;
pub mod chat;
pub mod chat_stream;
pub mod client;
pub mod common;
pub mod constants;
pub mod embedding;
pub mod error;
pub mod model_list;
pub mod tool;
pub mod utils;

3
src/v1/model_list.rs
View File

@@ -1,5 +1,8 @@
use serde::{Deserialize, Serialize};
// -----------------------------------------------------------------------------
// Response
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct ModelListResponse {
pub object: String,

134
src/v1/tool.rs Normal file
View File

@@ -0,0 +1,134 @@
use async_trait::async_trait;
use serde::{Deserialize, Serialize};
use std::{any::Any, collections::HashMap};
// -----------------------------------------------------------------------------
// Definitions
#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)]
pub struct ToolCall {
pub function: ToolCallFunction,
}
#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)]
pub struct ToolCallFunction {
pub name: String,
pub arguments: String,
}
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct Tool {
pub r#type: ToolType,
pub function: ToolFunction,
}
impl Tool {
pub fn new(
function_name: String,
function_description: String,
function_parameters: Vec<ToolFunctionParameter>,
) -> Self {
let properties: HashMap<String, ToolFunctionParameterProperty> = function_parameters
.into_iter()
.map(|param| {
(
param.name,
ToolFunctionParameterProperty {
r#type: param.r#type,
description: param.description,
},
)
})
.collect();
let property_names = properties.keys().cloned().collect();
let parameters = ToolFunctionParameters {
r#type: ToolFunctionParametersType::Object,
properties,
required: property_names,
};
Self {
r#type: ToolType::Function,
function: ToolFunction {
name: function_name,
description: function_description,
parameters,
},
}
}
}
// -----------------------------------------------------------------------------
// Request
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct ToolFunction {
name: String,
description: String,
parameters: ToolFunctionParameters,
}
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct ToolFunctionParameter {
name: String,
description: String,
r#type: ToolFunctionParameterType,
}
impl ToolFunctionParameter {
pub fn new(name: String, description: String, r#type: ToolFunctionParameterType) -> Self {
Self {
name,
r#type,
description,
}
}
}
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct ToolFunctionParameters {
r#type: ToolFunctionParametersType,
properties: HashMap<String, ToolFunctionParameterProperty>,
required: Vec<String>,
}
#[derive(Clone, Debug, Deserialize, Serialize)]
pub struct ToolFunctionParameterProperty {
r#type: ToolFunctionParameterType,
description: String,
}
#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)]
pub enum ToolFunctionParametersType {
#[serde(rename = "object")]
Object,
}
#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)]
pub enum ToolFunctionParameterType {
#[serde(rename = "string")]
String,
}
#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)]
pub enum ToolType {
#[serde(rename = "function")]
Function,
}
#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)]
pub enum ToolChoice {
#[serde(rename = "any")]
Any,
#[serde(rename = "auto")]
Auto,
#[serde(rename = "none")]
None,
}
// -----------------------------------------------------------------------------
// Custom
#[async_trait]
pub trait Function {
async fn execute(&self, arguments: String) -> Box<dyn Any + Send>;
}

32
src/v1/utils.rs Normal file
View File

@@ -0,0 +1,32 @@
use std::fmt::Debug;
use log::debug;
use serde::Serialize;
pub fn prettify_json_string(json: &String) -> String {
match serde_json::from_str::<serde_json::Value>(&json) {
Ok(json_value) => {
serde_json::to_string_pretty(&json_value).unwrap_or_else(|_| json.to_owned())
}
Err(_) => json.to_owned(),
}
}
pub fn prettify_json_struct<T: Debug + Serialize>(value: T) -> String {
match serde_json::to_string_pretty(&value) {
Ok(pretty_json) => pretty_json,
Err(_) => format!("{:?}", value),
}
}
pub fn debug_pretty_json_from_string(label: &str, json: &String) -> () {
let pretty_json = prettify_json_string(json);
debug!("{label}: {}", pretty_json);
}
pub fn debug_pretty_json_from_struct<T: Debug + Serialize>(label: &str, value: &T) -> () {
let pretty_json = prettify_json_struct(value);
debug!("{label}: {}", pretty_json);
}