feat: add chat completion without streaming

src/lib.rs

@@ -0,0 +1 @@
+pub mod v1;

src/v1/chat_completion.rs

@@ -0,0 +1,113 @@
+use serde::{Deserialize, Serialize};
+
+use crate::v1::common;
+
+#[derive(Debug)]
+pub struct ChatCompletionRequestOptions {
+    pub tools: Option<String>,
+    pub temperature: Option<f32>,
+    pub max_tokens: Option<u32>,
+    pub top_p: Option<f32>,
+    pub random_seed: Option<u32>,
+    pub stream: Option<bool>,
+    pub safe_prompt: Option<bool>,
+}
+impl Default for ChatCompletionRequestOptions {
+    fn default() -> Self {
+        Self {
+            tools: None,
+            temperature: None,
+            max_tokens: None,
+            top_p: None,
+            random_seed: None,
+            stream: None,
+            safe_prompt: None,
+        }
+    }
+}
+
+#[derive(Debug, Serialize, Deserialize)]
+pub struct ChatCompletionRequest {
+    pub messages: Vec<ChatCompletionMessage>,
+    pub model: String,
+    #[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>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub stream: Option<bool>,
+    #[serde(skip_serializing_if = "Option::is_none")]
+    pub safe_prompt: Option<bool>,
+    // TODO Check that prop (seen in official Python client but not in API doc).
+    // pub tool_choice: Option<String>,
+    // TODO Check that prop (seen in official Python client but not in API doc).
+    // pub response_format: Option<String>,
+}
+impl ChatCompletionRequest {
+    pub fn new(
+        model: String,
+        messages: Vec<ChatCompletionMessage>,
+        options: Option<ChatCompletionRequestOptions>,
+    ) -> Self {
+        let ChatCompletionRequestOptions {
+            tools,
+            temperature,
+            max_tokens,
+            top_p,
+            random_seed,
+            stream,
+            safe_prompt,
+        } = options.unwrap_or_default();
+
+        Self {
+            messages,
+            model,
+            tools,
+            temperature,
+            max_tokens,
+            top_p,
+            random_seed,
+            stream,
+            safe_prompt,
+        }
+    }
+}
+
+#[derive(Clone, Debug, Deserialize, Serialize)]
+pub struct ChatCompletionResponse {
+    pub id: String,
+    pub object: String,
+    /// Unix timestamp (in seconds).
+    pub created: u32,
+    pub model: String,
+    pub choices: Vec<ChatCompletionChoice>,
+    pub usage: common::ResponseUsage,
+}
+
+#[derive(Clone, Debug, Deserialize, Serialize)]
+pub struct ChatCompletionChoice {
+    pub index: u32,
+    pub message: ChatCompletionMessage,
+    pub finish_reason: String,
+    // TODO Check that prop (seen in API responses but undocumented).
+    // pub logprobs: ???
+}
+
+#[derive(Clone, Debug, Deserialize, Serialize)]
+pub struct ChatCompletionMessage {
+    pub role: ChatCompletionMessageRole,
+    pub content: String,
+}
+
+#[derive(Clone, Debug, Eq, PartialEq, Deserialize, Serialize)]
+#[allow(non_camel_case_types)]
+pub enum ChatCompletionMessageRole {
+    assistant,
+    user,
+}

src/v1/client.rs

@@ -0,0 +1,175 @@
+use crate::v1::error::APIError;
+use minreq::Response;
+
+use crate::v1::{
+    chat_completion::{ChatCompletionRequest, ChatCompletionResponse},
+    constants::API_URL_BASE,
+};
+
+pub struct Client {
+    pub api_key: String,
+    pub endpoint: String,
+    pub max_retries: u32,
+    pub timeout: u32,
+}
+
+impl Client {
+    pub fn new(
+        api_key: Option<String>,
+        endpoint: Option<String>,
+        max_retries: Option<u32>,
+        timeout: Option<u32>,
+    ) -> Self {
+        let api_key = api_key.unwrap_or(std::env::var("MISTRAL_API_KEY").unwrap());
+        let endpoint = endpoint.unwrap_or(API_URL_BASE.to_string());
+        let max_retries = max_retries.unwrap_or(5);
+        let timeout = timeout.unwrap_or(120);
+
+        Self {
+            api_key,
+            endpoint,
+            max_retries,
+            timeout,
+        }
+    }
+
+    pub fn build_request(&self, request: minreq::Request) -> minreq::Request {
+        let authorization = format!("Bearer {}", self.api_key);
+        let user_agent = format!(
+            "ivangabriele/mistral-client-rs/{}",
+            env!("CARGO_PKG_VERSION")
+        );
+
+        let request = request
+            .with_header("Authorization", authorization)
+            .with_header("Accept", "application/json")
+            .with_header("Content-Type", "application/json")
+            .with_header("User-Agent", user_agent);
+
+        request
+    }
+
+    pub fn get(&self, path: &str) -> Result<Response, APIError> {
+        let url = format!("{}{}", self.endpoint, path);
+        let request = self.build_request(minreq::post(url));
+
+        let result = request.send();
+        match result {
+            Ok(res) => {
+                if (200..=299).contains(&res.status_code) {
+                    Ok(res)
+                } else {
+                    Err(APIError {
+                        message: format!("{}: {}", res.status_code, res.as_str().unwrap()),
+                    })
+                }
+            }
+            Err(e) => Err(self.new_error(e)),
+        }
+    }
+
+    pub fn post<T: serde::ser::Serialize + std::fmt::Debug>(
+        &self,
+        path: &str,
+        params: &T,
+    ) -> Result<Response, APIError> {
+        // print!("{:?}", params);
+
+        let url = format!("{}{}", self.endpoint, path);
+        let request = self.build_request(minreq::post(url));
+
+        let result = request.with_json(params).unwrap().send();
+        match result {
+            Ok(res) => {
+                print!("{:?}", res.as_str().unwrap());
+
+                if (200..=299).contains(&res.status_code) {
+                    Ok(res)
+                } else {
+                    Err(APIError {
+                        message: format!("{}: {}", res.status_code, res.as_str().unwrap()),
+                    })
+                }
+            }
+            Err(e) => Err(self.new_error(e)),
+        }
+    }
+
+    pub fn delete(&self, path: &str) -> Result<Response, APIError> {
+        let url = format!("{}{}", self.endpoint, path);
+        let request = self.build_request(minreq::post(url));
+
+        let result = request.send();
+        match result {
+            Ok(res) => {
+                if (200..=299).contains(&res.status_code) {
+                    Ok(res)
+                } else {
+                    Err(APIError {
+                        message: format!("{}: {}", res.status_code, res.as_str().unwrap()),
+                    })
+                }
+            }
+            Err(e) => Err(self.new_error(e)),
+        }
+    }
+
+    // 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(&self, request: ChatCompletionRequest) -> Result<ChatCompletionResponse, APIError> {
+        let response = self.post("/chat/completions", &request)?;
+        let result = response.json::<ChatCompletionResponse>();
+        match result {
+            Ok(r) => Ok(r),
+            Err(e) => Err(self.new_error(e)),
+        }
+    }
+
+    fn new_error(&self, err: minreq::Error) -> APIError {
+        APIError {
+            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
+    // }
+}

src/v1/common.rs

@@ -0,0 +1,8 @@
+use serde::{Deserialize, Serialize};
+
+#[derive(Clone, Debug, Deserialize, Serialize)]
+pub struct ResponseUsage {
+    pub prompt_tokens: u32,
+    pub completion_tokens: u32,
+    pub total_tokens: u32,
+}

src/v1/constants.rs

@@ -0,0 +1,7 @@
+pub const API_URL_BASE: &str = "https://api.mistral.ai/v1";
+
+pub const OPEN_MISTRAL_7B: &str = "open-mistral-7b";
+pub const OPEN_MISTRAL_8X7B: &str = "open-mixtral-8x7b";
+pub const MISTRAL_SMALL_LATEST: &str = "mistral-small-latest";
+pub const MISTRAL_MEDIUM_LATEST: &str = "mistral-medium-latest";
+pub const MISTRAL_LARGE_LATEST: &str = "mistral-large-latest";

src/v1/error.rs

@@ -0,0 +1,15 @@
+use std::error::Error;
+use std::fmt;
+
+#[derive(Debug)]
+pub struct APIError {
+    pub message: String,
+}
+
+impl fmt::Display for APIError {
+    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        write!(f, "APIError: {}", self.message)
+    }
+}
+
+impl Error for APIError {}

src/v1/mod.rs

@@ -0,0 +1,5 @@
+pub mod chat_completion;
+pub mod client;
+pub mod common;
+pub mod constants;
+pub mod error;