feat: add client.list_models() method
This commit is contained in:
22
README.md
22
README.md
@@ -27,7 +27,7 @@ Rust client for the Mistral AI API.
|
|||||||
- [x] Chat without streaming
|
- [x] Chat without streaming
|
||||||
- [ ] Chat with streaming
|
- [ ] Chat with streaming
|
||||||
- [ ] Embedding
|
- [ ] Embedding
|
||||||
- [ ] List models
|
- [x] List models
|
||||||
- [ ] Function Calling
|
- [ ] Function Calling
|
||||||
|
|
||||||
## Installation
|
## Installation
|
||||||
@@ -63,11 +63,8 @@ fn main() {
|
|||||||
### Chat without streaming
|
### Chat without streaming
|
||||||
|
|
||||||
```rs
|
```rs
|
||||||
use mistralai::v1::{
|
use mistralai_client::v1::{
|
||||||
chat_completion::{
|
chat_completion::{ChatCompletionMessage, ChatCompletionMessageRole, ChatCompletionRequestOptions},
|
||||||
ChatCompletionMessage, ChatCompletionMessageRole, ChatCompletionRequest,
|
|
||||||
ChatCompletionRequestOptions,
|
|
||||||
},
|
|
||||||
client::Client,
|
client::Client,
|
||||||
constants::OPEN_MISTRAL_7B,
|
constants::OPEN_MISTRAL_7B,
|
||||||
};
|
};
|
||||||
@@ -103,4 +100,15 @@ _In progress._
|
|||||||
|
|
||||||
### List models
|
### List models
|
||||||
|
|
||||||
_In progress._
|
```rs
|
||||||
|
use mistralai_client::v1::client::Client;
|
||||||
|
|
||||||
|
fn main() {
|
||||||
|
// This example suppose you have set the `MISTRAL_API_KEY` environment variable.
|
||||||
|
let client = Client::new(None, None, None, None);
|
||||||
|
|
||||||
|
let result = client.list_models(model, messages, Some(options)).unwrap();
|
||||||
|
println!("First Model ID: {:?}", result.data[0].id);
|
||||||
|
// => "First Model ID: open-mistral-7b"
|
||||||
|
}
|
||||||
|
```
|
||||||
|
|||||||
@@ -44,9 +44,9 @@ pub struct ChatCompletionRequest {
|
|||||||
pub stream: Option<bool>,
|
pub stream: Option<bool>,
|
||||||
#[serde(skip_serializing_if = "Option::is_none")]
|
#[serde(skip_serializing_if = "Option::is_none")]
|
||||||
pub safe_prompt: Option<bool>,
|
pub safe_prompt: Option<bool>,
|
||||||
// TODO Check that prop (seen in official Python client but not in API doc).
|
// TODO Check this prop (seen in official Python client but not in API doc).
|
||||||
// pub tool_choice: Option<String>,
|
// pub tool_choice: Option<String>,
|
||||||
// TODO Check that prop (seen in official Python client but not in API doc).
|
// TODO Check this prop (seen in official Python client but not in API doc).
|
||||||
// pub response_format: Option<String>,
|
// pub response_format: Option<String>,
|
||||||
}
|
}
|
||||||
impl ChatCompletionRequest {
|
impl ChatCompletionRequest {
|
||||||
@@ -95,7 +95,7 @@ pub struct ChatCompletionChoice {
|
|||||||
pub index: u32,
|
pub index: u32,
|
||||||
pub message: ChatCompletionMessage,
|
pub message: ChatCompletionMessage,
|
||||||
pub finish_reason: String,
|
pub finish_reason: String,
|
||||||
// TODO Check that prop (seen in API responses but undocumented).
|
// TODO Check this prop (seen in API responses but undocumented).
|
||||||
// pub logprobs: ???
|
// pub logprobs: ???
|
||||||
}
|
}
|
||||||
|
|
||||||
|
|||||||
@@ -2,12 +2,13 @@ use crate::v1::error::APIError;
|
|||||||
use minreq::Response;
|
use minreq::Response;
|
||||||
|
|
||||||
use crate::v1::{
|
use crate::v1::{
|
||||||
chat_completion::{ChatCompletionRequest, ChatCompletionResponse},
|
chat_completion::{
|
||||||
|
ChatCompletionMessage, ChatCompletionParams, ChatCompletionRequest, ChatCompletionResponse,
|
||||||
|
},
|
||||||
constants::API_URL_BASE,
|
constants::API_URL_BASE,
|
||||||
|
list_models::ListModelsResponse,
|
||||||
};
|
};
|
||||||
|
|
||||||
use super::chat_completion::{ChatCompletionMessage, ChatCompletionParams};
|
|
||||||
|
|
||||||
pub struct Client {
|
pub struct Client {
|
||||||
pub api_key: String,
|
pub api_key: String,
|
||||||
pub endpoint: String,
|
pub endpoint: String,
|
||||||
@@ -38,7 +39,7 @@ impl Client {
|
|||||||
pub fn build_request(&self, request: minreq::Request) -> minreq::Request {
|
pub fn build_request(&self, request: minreq::Request) -> minreq::Request {
|
||||||
let authorization = format!("Bearer {}", self.api_key);
|
let authorization = format!("Bearer {}", self.api_key);
|
||||||
let user_agent = format!(
|
let user_agent = format!(
|
||||||
"ivangabriele/mistral-client-rs/{}",
|
"ivangabriele/mistralai-client-rs/{}",
|
||||||
env!("CARGO_PKG_VERSION")
|
env!("CARGO_PKG_VERSION")
|
||||||
);
|
);
|
||||||
|
|
||||||
@@ -53,20 +54,24 @@ impl Client {
|
|||||||
|
|
||||||
pub fn get(&self, path: &str) -> Result<Response, APIError> {
|
pub fn get(&self, path: &str) -> Result<Response, APIError> {
|
||||||
let url = format!("{}{}", self.endpoint, path);
|
let url = format!("{}{}", self.endpoint, path);
|
||||||
let request = self.build_request(minreq::post(url));
|
let request = self.build_request(minreq::get(url));
|
||||||
|
|
||||||
let result = request.send();
|
let result = request.send();
|
||||||
match result {
|
match result {
|
||||||
Ok(res) => {
|
Ok(response) => {
|
||||||
if (200..=299).contains(&res.status_code) {
|
if (200..=299).contains(&response.status_code) {
|
||||||
Ok(res)
|
Ok(response)
|
||||||
} else {
|
} else {
|
||||||
Err(APIError {
|
Err(APIError {
|
||||||
message: format!("{}: {}", res.status_code, res.as_str().unwrap()),
|
message: format!(
|
||||||
|
"{}: {}",
|
||||||
|
response.status_code,
|
||||||
|
response.as_str().unwrap()
|
||||||
|
),
|
||||||
})
|
})
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
Err(e) => Err(self.new_error(e)),
|
Err(error) => Err(self.new_error(error)),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
@@ -83,7 +88,7 @@ impl Client {
|
|||||||
let result = request.with_json(params).unwrap().send();
|
let result = request.with_json(params).unwrap().send();
|
||||||
match result {
|
match result {
|
||||||
Ok(response) => {
|
Ok(response) => {
|
||||||
print!("{:?}", response.as_str().unwrap());
|
// print!("{:?}", response.as_str().unwrap());
|
||||||
|
|
||||||
if (200..=299).contains(&response.status_code) {
|
if (200..=299).contains(&response.status_code) {
|
||||||
Ok(response)
|
Ok(response)
|
||||||
@@ -124,24 +129,6 @@ impl Client {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// pub fn completion(&self, req: CompletionRequest) -> Result<CompletionResponse, APIError> {
|
|
||||||
// let res = self.post("/completions", &req)?;
|
|
||||||
// let r = res.json::<CompletionResponse>();
|
|
||||||
// match r {
|
|
||||||
// Ok(r) => Ok(r),
|
|
||||||
// Err(e) => Err(self.new_error(e)),
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
|
|
||||||
// pub fn embedding(&self, req: EmbeddingRequest) -> Result<EmbeddingResponse, APIError> {
|
|
||||||
// let res = self.post("/embeddings", &req)?;
|
|
||||||
// let r = res.json::<EmbeddingResponse>();
|
|
||||||
// match r {
|
|
||||||
// Ok(r) => Ok(r),
|
|
||||||
// Err(e) => Err(self.new_error(e)),
|
|
||||||
// }
|
|
||||||
// }
|
|
||||||
|
|
||||||
pub fn chat(
|
pub fn chat(
|
||||||
&self,
|
&self,
|
||||||
model: String,
|
model: String,
|
||||||
@@ -158,35 +145,18 @@ impl Client {
|
|||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
|
pub fn list_models(&self) -> Result<ListModelsResponse, APIError> {
|
||||||
|
let response = self.get("/models")?;
|
||||||
|
let result = response.json::<ListModelsResponse>();
|
||||||
|
match result {
|
||||||
|
Ok(response) => Ok(response),
|
||||||
|
Err(error) => Err(self.new_error(error)),
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
fn new_error(&self, err: minreq::Error) -> APIError {
|
fn new_error(&self, err: minreq::Error) -> APIError {
|
||||||
APIError {
|
APIError {
|
||||||
message: err.to_string(),
|
message: err.to_string(),
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|
||||||
// fn query_params(
|
|
||||||
// limit: Option<i64>,
|
|
||||||
// order: Option<String>,
|
|
||||||
// after: Option<String>,
|
|
||||||
// before: Option<String>,
|
|
||||||
// mut url: String,
|
|
||||||
// ) -> String {
|
|
||||||
// let mut params = vec![];
|
|
||||||
// if let Some(limit) = limit {
|
|
||||||
// params.push(format!("limit={}", limit));
|
|
||||||
// }
|
|
||||||
// if let Some(order) = order {
|
|
||||||
// params.push(format!("order={}", order));
|
|
||||||
// }
|
|
||||||
// if let Some(after) = after {
|
|
||||||
// params.push(format!("after={}", after));
|
|
||||||
// }
|
|
||||||
// if let Some(before) = before {
|
|
||||||
// params.push(format!("before={}", before));
|
|
||||||
// }
|
|
||||||
// if !params.is_empty() {
|
|
||||||
// url = format!("{}?{}", url, params.join("&"));
|
|
||||||
// }
|
|
||||||
// url
|
|
||||||
// }
|
|
||||||
}
|
}
|
||||||
|
|||||||
39
src/v1/list_models.rs
Normal file
39
src/v1/list_models.rs
Normal file
@@ -0,0 +1,39 @@
|
|||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
|
||||||
|
#[derive(Clone, Debug, Deserialize, Serialize)]
|
||||||
|
pub struct ListModelsResponse {
|
||||||
|
pub object: String,
|
||||||
|
pub data: Vec<ListModelsModel>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone, Debug, Deserialize, Serialize)]
|
||||||
|
pub struct ListModelsModel {
|
||||||
|
pub id: String,
|
||||||
|
pub object: String,
|
||||||
|
/// Unix timestamp (in seconds).
|
||||||
|
pub created: u32,
|
||||||
|
pub owned_by: String,
|
||||||
|
pub permission: Vec<ListModelsModelPermission>,
|
||||||
|
// TODO Check this prop (seen in API responses but undocumented).
|
||||||
|
// pub root: ???,
|
||||||
|
// TODO Check this prop (seen in API responses but undocumented).
|
||||||
|
// pub parent: ???,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Clone, Debug, Deserialize, Serialize)]
|
||||||
|
pub struct ListModelsModelPermission {
|
||||||
|
pub id: String,
|
||||||
|
pub object: String,
|
||||||
|
/// Unix timestamp (in seconds).
|
||||||
|
pub created: u32,
|
||||||
|
pub allow_create_engine: bool,
|
||||||
|
pub allow_sampling: bool,
|
||||||
|
pub allow_logprobs: bool,
|
||||||
|
pub allow_search_indices: bool,
|
||||||
|
pub allow_view: bool,
|
||||||
|
pub allow_fine_tuning: bool,
|
||||||
|
pub organization: String,
|
||||||
|
pub is_blocking: bool,
|
||||||
|
// TODO Check this prop (seen in API responses but undocumented).
|
||||||
|
// pub group: ???,
|
||||||
|
}
|
||||||
@@ -3,3 +3,4 @@ pub mod client;
|
|||||||
pub mod common;
|
pub mod common;
|
||||||
pub mod constants;
|
pub mod constants;
|
||||||
pub mod error;
|
pub mod error;
|
||||||
|
pub mod list_models;
|
||||||
|
|||||||
17
tests/v1_list_models_test.rs
Normal file
17
tests/v1_list_models_test.rs
Normal file
@@ -0,0 +1,17 @@
|
|||||||
|
use jrest::expect;
|
||||||
|
use mistralai_client::v1::client::Client;
|
||||||
|
|
||||||
|
#[test]
|
||||||
|
fn test_list_models() {
|
||||||
|
extern crate dotenv;
|
||||||
|
|
||||||
|
use dotenv::dotenv;
|
||||||
|
dotenv().ok();
|
||||||
|
|
||||||
|
let client = Client::new(None, None, None, None);
|
||||||
|
|
||||||
|
let response = client.list_models().unwrap();
|
||||||
|
|
||||||
|
expect!(response.object).to_be("list".to_string());
|
||||||
|
expect!(response.data.len()).to_be_greater_than(0);
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user