summaryrefslogtreecommitdiffstats
path: root/src/client/mistral.rs
diff options
context:
space:
mode:
authorsigoden <sigoden@gmail.com>2024-04-25 20:41:25 +0800
committerGitHub <noreply@github.com>2024-04-25 20:41:25 +0800
commite5432ff779641473dbd8e4466ac8307aae87de8d (patch)
treee836d436ccde84cad74ac1d8439f4b8306e5d2b5 /src/client/mistral.rs
parent1a56e38fe22336e6b6587ff7d99218295e28235f (diff)
downloadaichat-e5432ff779641473dbd8e4466ac8307aae87de8d.tar.gz
refactor: simplify openai compatible module (#440)
Diffstat (limited to 'src/client/mistral.rs')
-rw-r--r--src/client/mistral.rs69
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),
+ ]
+);