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

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

302
README.md
View File

@@ -15,13 +15,16 @@ Rust client for the Mistral AI API.
- [As an environment variable](#as-an-environment-variable)
- [As a client argument](#as-a-client-argument)
- [Usage](#usage)
- [Chat without streaming](#chat-without-streaming)
- [Chat without streaming (async)](#chat-without-streaming-async)
- [Chat](#chat)
- [Chat (async)](#chat-async)
- [Chat with streaming (async)](#chat-with-streaming-async)
- [Chat with Function Calling](#chat-with-function-calling)
- [Chat with Function Calling (async)](#chat-with-function-calling-async)
- [Embeddings](#embeddings)
- [Embeddings (async)](#embeddings-async)
- [List models](#list-models)
- [List models (async)](#list-models-async)
- [Contributing](#contributing)
---
@@ -34,8 +37,8 @@ Rust client for the Mistral AI API.
- [x] Embedding (async)
- [x] List models
- [x] List models (async)
- [ ] Function Calling
- [ ] Function Calling (async)
- [x] Function Calling
- [x] Function Calling (async)
## Installation
@@ -53,6 +56,18 @@ You can get your Mistral API Key there: <https://docs.mistral.ai/#api-access>.
Just set the `MISTRAL_API_KEY` environment variable.
```rs
use mistralai_client::v1::client::Client;
fn main() {
let client = Client::new(None, None, None, None);
}
```
```sh
MISTRAL_API_KEY=your_api_key cargo run
```
#### As a client argument
```rs
@@ -67,11 +82,11 @@ fn main() {
## Usage
### Chat without streaming
### Chat
```rs
use mistralai_client::v1::{
chat_completion::{ChatCompletionParams, ChatMessage, ChatMessageRole},
chat::{ChatMessage, ChatMessageRole, ChatParams},
client::Client,
constants::Model,
};
@@ -82,10 +97,11 @@ fn main() {
let model = Model::OpenMistral7b;
let messages = vec![ChatMessage {
role: ChatMessageRole::user,
role: ChatMessageRole::User,
content: "Just guess the next word: \"Eiffel ...\"?".to_string(),
tool_calls: None,
}];
let options = ChatCompletionRequestOptions {
let options = ChatParams {
temperature: Some(0.0),
random_seed: Some(42),
..Default::default()
@@ -93,15 +109,15 @@ fn main() {
let result = client.chat(model, messages, Some(options)).unwrap();
println!("Assistant: {}", result.choices[0].message.content);
// => "Assistant: Tower. [...]"
// => "Assistant: Tower. The Eiffel Tower is a famous landmark in Paris, France."
}
```
### Chat without streaming (async)
### Chat (async)
```rs
use mistralai_client::v1::{
chat_completion::{ChatCompletionParams, ChatMessage, ChatMessageRole},
chat::{ChatMessage, ChatMessageRole, ChatParams},
client::Client,
constants::Model,
};
@@ -113,18 +129,25 @@ async fn main() {
let model = Model::OpenMistral7b;
let messages = vec![ChatMessage {
role: ChatMessageRole::user,
role: ChatMessageRole::User,
content: "Just guess the next word: \"Eiffel ...\"?".to_string(),
tool_calls: None,
}];
let options = ChatCompletionRequestOptions {
let options = ChatParams {
temperature: Some(0.0),
random_seed: Some(42),
..Default::default()
};
let result = client.chat_async(model, messages, Some(options)).await.unwrap();
println!("Assistant: {}", result.choices[0].message.content);
// => "Assistant: Tower. [...]"
let result = client
.chat_async(model, messages, Some(options))
.await
.unwrap();
println!(
"{:?}: {}",
result.choices[0].message.role, result.choices[0].message.content
);
// => "Assistant: Tower. The Eiffel Tower is a famous landmark in Paris, France."
}
```
@@ -133,38 +156,206 @@ async fn main() {
```rs
use futures::stream::StreamExt;
use mistralai_client::v1::{
chat_completion::{ChatCompletionParams, ChatMessage, ChatMessageRole},
chat::{ChatMessage, ChatMessageRole, ChatParams},
client::Client,
constants::Model,
};
use std::io::{self, Write};
[#tokio::main]
#[tokio::main]
async fn main() {
// This example suppose you have set the `MISTRAL_API_KEY` environment variable.
let client = Client::new(None, None, None, None).unwrap();
let client = Client::new(None, None, None, None).unwrap();
let model = Model::OpenMistral7b;
let messages = vec![ChatMessage {
role: ChatMessageRole::user,
content: "Just guess the next word: \"Eiffel ...\"?".to_string(),
role: ChatMessageRole::User,
content: "Tell me a short happy story.".to_string(),
tool_calls: None,
}];
let options = ChatCompletionParams {
let options = ChatParams {
temperature: Some(0.0),
random_seed: Some(42),
..Default::default()
};
let stream_result = client.chat_stream(model, messages, Some(options)).await;
let mut stream = stream_result.expect("Failed to create stream.");
while let Some(chunk_result) = stream.next().await {
match chunk_result {
Ok(chunk) => {
println!("Assistant (message chunk): {}", chunk.choices[0].delta.content);
let stream_result = client
.chat_stream(model, messages, Some(options))
.await
.expect("Failed to create stream.");
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(e) => eprintln!("Error processing chunk: {:?}", e),
}
})
.await;
print!("\n") // To persist the last chunk output.
}
```
### Chat with Function Calling
```rs
use mistralai_client::v1::{
chat::{ChatMessage, ChatMessageRole, ChatParams},
client::Client,
constants::Model,
tool::{Function, Tool, ToolChoice, ToolFunctionParameter, ToolFunctionParameterType},
};
use serde::Deserialize;
use std::any::Any;
#[derive(Debug, Deserialize)]
struct GetCityTemperatureArguments {
city: String,
}
struct GetCityTemperatureFunction;
#[async_trait::async_trait]
impl Function for GetCityTemperatureFunction {
async fn execute(&self, arguments: String) -> Box<dyn Any + Send> {
// Deserialize arguments, perform the logic, and return the result
let GetCityTemperatureArguments { city } = serde_json::from_str(&arguments).unwrap();
let temperature = match city.as_str() {
"Paris" => "20°C",
_ => "Unknown city",
};
Box::new(temperature.to_string())
}
}
fn main() {
let tools = vec![Tool::new(
"get_city_temperature".to_string(),
"Get the current temperature in a city.".to_string(),
vec![ToolFunctionParameter::new(
"city".to_string(),
"The name of the city.".to_string(),
ToolFunctionParameterType::String,
)],
)];
// This example suppose you have set the `MISTRAL_API_KEY` environment variable.
let mut client = Client::new(None, None, None, None).unwrap();
client.register_function(
"get_city_temperature".to_string(),
Box::new(GetCityTemperatureFunction),
);
let model = Model::MistralSmallLatest;
let messages = vec![ChatMessage {
role: ChatMessageRole::User,
content: "What's the temperature in Paris?".to_string(),
tool_calls: None,
}];
let options = ChatParams {
temperature: Some(0.0),
random_seed: Some(42),
tool_choice: Some(ToolChoice::Auto),
tools: Some(tools),
..Default::default()
};
client.chat(model, messages, Some(options)).unwrap();
let temperature = client
.get_last_function_call_result()
.unwrap()
.downcast::<String>()
.unwrap();
println!("The temperature in Paris is: {}.", temperature);
// => "The temperature in Paris is: 20°C."
}
```
### Chat with Function Calling (async)
```rs
use mistralai_client::v1::{
chat::{ChatMessage, ChatMessageRole, ChatParams},
client::Client,
constants::Model,
tool::{Function, Tool, ToolChoice, ToolFunctionParameter, ToolFunctionParameterType},
};
use serde::Deserialize;
use std::any::Any;
#[derive(Debug, Deserialize)]
struct GetCityTemperatureArguments {
city: String,
}
struct GetCityTemperatureFunction;
#[async_trait::async_trait]
impl Function for GetCityTemperatureFunction {
async fn execute(&self, arguments: String) -> Box<dyn Any + Send> {
// Deserialize arguments, perform the logic, and return the result
let GetCityTemperatureArguments { city } = serde_json::from_str(&arguments).unwrap();
let temperature = match city.as_str() {
"Paris" => "20°C",
_ => "Unknown city",
};
Box::new(temperature.to_string())
}
}
#[tokio::main]
async fn main() {
let tools = vec![Tool::new(
"get_city_temperature".to_string(),
"Get the current temperature in a city.".to_string(),
vec![ToolFunctionParameter::new(
"city".to_string(),
"The name of the city.".to_string(),
ToolFunctionParameterType::String,
)],
)];
// This example suppose you have set the `MISTRAL_API_KEY` environment variable.
let mut client = Client::new(None, None, None, None).unwrap();
client.register_function(
"get_city_temperature".to_string(),
Box::new(GetCityTemperatureFunction),
);
let model = Model::MistralSmallLatest;
let messages = vec![ChatMessage {
role: ChatMessageRole::User,
content: "What's the temperature in Paris?".to_string(),
tool_calls: None,
}];
let options = ChatParams {
temperature: Some(0.0),
random_seed: Some(42),
tool_choice: Some(ToolChoice::Auto),
tools: Some(tools),
..Default::default()
};
client
.chat_async(model, messages, Some(options))
.await
.unwrap();
let temperature = client
.get_last_function_call_result()
.unwrap()
.downcast::<String>()
.unwrap();
println!("The temperature in Paris is: {}.", temperature);
// => "The temperature in Paris is: 20°C."
}
```
### Embeddings
@@ -174,18 +365,18 @@ use mistralai_client::v1::{client::Client, constants::EmbedModel};
fn main() {
// This example suppose you have set the `MISTRAL_API_KEY` environment variable.
let client: Client = Client::new(None, None, None, None).unwrap();
let client: Client = Client::new(None, None, None, None).unwrap();
let model = EmbedModel::MistralEmbed;
let input = vec!["Embed this sentence.", "As well as this one."]
.iter()
.map(|s| s.to_string())
.collect();
let options = None;
let model = EmbedModel::MistralEmbed;
let input = vec!["Embed this sentence.", "As well as this one."]
.iter()
.map(|s| s.to_string())
.collect();
let options = None;
let response = client.embeddings(model, input, options).unwrap();
println!("Embeddings: {:?}", response.data);
// => "Embeddings: [{...}, {...}]"
let response = client.embeddings(model, input, options).unwrap();
println!("First Embedding: {:?}", response.data[0]);
// => "First Embedding: {...}"
}
```
@@ -197,18 +388,21 @@ use mistralai_client::v1::{client::Client, constants::EmbedModel};
#[tokio::main]
async fn main() {
// This example suppose you have set the `MISTRAL_API_KEY` environment variable.
let client: Client = Client::new(None, None, None, None).unwrap();
let client: Client = Client::new(None, None, None, None).unwrap();
let model = EmbedModel::MistralEmbed;
let input = vec!["Embed this sentence.", "As well as this one."]
.iter()
.map(|s| s.to_string())
.collect();
let options = None;
let model = EmbedModel::MistralEmbed;
let input = vec!["Embed this sentence.", "As well as this one."]
.iter()
.map(|s| s.to_string())
.collect();
let options = None;
let response = client.embeddings_async(model, input, options).await.unwrap();
println!("Embeddings: {:?}", response.data);
// => "Embeddings: [{...}, {...}]"
let response = client
.embeddings_async(model, input, options)
.await
.unwrap();
println!("First Embedding: {:?}", response.data[0]);
// => "First Embedding: {...}"
}
```
@@ -235,10 +429,14 @@ use mistralai_client::v1::client::Client;
#[tokio::main]
async fn main() {
// This example suppose you have set the `MISTRAL_API_KEY` environment variable.
let client = Client::new(None, None, None, None).await.unwrap();
let client = Client::new(None, None, None, None).unwrap();
let result = client.list_models_async().unwrap();
let result = client.list_models_async().await.unwrap();
println!("First Model ID: {:?}", result.data[0].id);
// => "First Model ID: open-mistral-7b"
}
```
## Contributing
Please read [CONTRIBUTING.md](./CONTRIBUTING.md) for details on how to contribute to this library.