diff options
| author | sigoden <sigoden@gmail.com> | 2024-04-25 20:41:25 +0800 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2024-04-25 20:41:25 +0800 |
| commit | e5432ff779641473dbd8e4466ac8307aae87de8d (patch) | |
| tree | e836d436ccde84cad74ac1d8439f4b8306e5d2b5 /src/client/mistral.rs | |
| parent | 1a56e38fe22336e6b6587ff7d99218295e28235f (diff) | |
| download | aichat-e5432ff779641473dbd8e4466ac8307aae87de8d.tar.gz | |
refactor: simplify openai compatible module (#440)
Diffstat (limited to 'src/client/mistral.rs')
| -rw-r--r-- | src/client/mistral.rs | 69 |
1 files changed, 14 insertions, 55 deletions
diff --git a/src/client/mistral.rs b/src/client/mistral.rs index 23bf1f2..c20abe7 100644 --- a/src/client/mistral.rs +++ b/src/client/mistral.rs @@ -1,55 +1,14 @@ -use super::openai::openai_build_body; -use super::{ExtraConfig, MistralClient, Model, ModelConfig, PromptType, SendData}; - -use crate::utils::PromptKind; - -use anyhow::Result; -use async_trait::async_trait; -use reqwest::{Client as ReqwestClient, RequestBuilder}; -use serde::Deserialize; - -const API_URL: &str = "https://api.mistral.ai/v1/chat/completions"; - -#[derive(Debug, Clone, Deserialize)] -pub struct MistralConfig { - pub name: Option<String>, - pub api_key: Option<String>, - #[serde(default)] - pub models: Vec<ModelConfig>, - pub extra: Option<ExtraConfig>, -} - -openai_compatible_client!(MistralClient); - -impl MistralClient { - list_models_fn!( - MistralConfig, - [ - // https://docs.mistral.ai/platform/endpoints/ - ("open-mixtral-8x22b", "text", 64000), - ("mistral-small-latest", "text", 32000), - ("mistral-large-latest", "text", 32000), - ] - ); - config_get_fn!(api_key, get_api_key); - - pub const PROMPTS: [PromptType<'static>; 1] = - [("api_key", "API Key:", false, PromptKind::String)]; - - fn request_builder(&self, client: &ReqwestClient, data: SendData) -> Result<RequestBuilder> { - let api_key = self.get_api_key().ok(); - - let body = openai_build_body(data, &self.model); - - let url = API_URL; - - debug!("Mistral Request: {url} {body}"); - - let mut builder = client.post(url).json(&body); - if let Some(api_key) = api_key { - builder = builder.bearer_auth(api_key); - } - - Ok(builder) - } -} +openai_compatible_module!( + MistralConfig, + MistralClient, + "https://api.mistral.ai/v1", + [ + // https://docs.mistral.ai/platform/endpoints/ + ("open-mistral-7b", "text", 32000), + ("open-mixtral-8x7b", "text", 32000), + ("open-mixtral-8x22b", "text", 64000), + ("mistral-small-latest", "text", 32000), + ("mistral-medium-latest", "text", 32000), + ("mistral-large-latest", "text", 32000), + ] +); |
