From e5432ff779641473dbd8e4466ac8307aae87de8d Mon Sep 17 00:00:00 2001
From: sigoden <sigoden@gmail.com>
Date: Thu, 25 Apr 2024 20:41:25 +0800
Subject: refactor: simplify openai compatible module (#440)

---
 src/client/mistral.rs | 69 +++++++++++----------------------------------------
 1 file changed, 14 insertions(+), 55 deletions(-)

(limited to 'src/client/mistral.rs')

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),
+    ]
+);
-- 
cgit v1.2.3