Add the Atlas module to invoke the reach model

Cargo.toml

@@ -12,6 +12,7 @@ members = [
     "neptune",
     "saturn",
     "lunar/src-tauri",
+    "atlas",
 ]
 default-members = ["mega", "libra"]
 exclude = ["craft"]
@@ -31,7 +32,7 @@ neptune = { path = "neptune" }
 saturn = { path = "saturn" }
 mega = { path = "mega" }
 anyhow = "1.0.86"
-serde = "1.0.203"
+serde = {version = "1.0.203", features = ["derive"]}
 serde_json = "1.0.117"
 tracing = "0.1.40"
 tracing-subscriber = "0.3.18"

atlas/Cargo.toml

@@ -0,0 +1,11 @@
+[package]
+name = "atlas"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+async-openai = "0.23.4"
+reqwest.workspace = true
+serde.workspace = true
+serde_json.workspace = true
+tokio = { workspace = true, features = ["macros"] }

atlas/src/api/gemini.rs

@@ -0,0 +1,156 @@
+use serde::Deserialize;
+
+use crate::AskModel;
+
+pub enum GeminiModels {
+    ChatBison001,
+    TextBison001,
+    EmbeddingGecko001,
+    Gemini10ProLatest,
+    Gemini10Pro,
+    GeminiPro,
+    Gemini10Pro001,
+    Gemini10ProVisionLatest,
+    GeminiProVision,
+    Gemini15ProLatest,
+    Gemini15Pro001,
+    Gemini15Pro,
+    Gemini15FlashLatest,
+    Gemini15Flash001,
+    Gemini15Flash,
+    Embedding001,
+    TextEmbedding004,
+    AQA,
+}
+
+impl GeminiModels {
+    pub fn as_str(&self) -> &str {
+        match self {
+            GeminiModels::ChatBison001 => "models/chat-bison-001",
+            GeminiModels::TextBison001 => "models/text-bison-001",
+            GeminiModels::EmbeddingGecko001 => "models/embedding-gecko-001",
+            GeminiModels::Gemini10ProLatest => "models/gemini-1.0-pro-latest",
+            GeminiModels::Gemini10Pro => "models/gemini-1.0-pro",
+            GeminiModels::GeminiPro => "models/gemini-pro",
+            GeminiModels::Gemini10Pro001 => "models/gemini-1.0-pro-001",
+            GeminiModels::Gemini10ProVisionLatest => "models/gemini-1.0-pro-vision-latest",
+            GeminiModels::GeminiProVision => "models/gemini-pro-vision",
+            GeminiModels::Gemini15ProLatest => "models/gemini-1.5-pro-latest",
+            GeminiModels::Gemini15Pro001 => "models/gemini-1.5-pro-001",
+            GeminiModels::Gemini15Pro => "models/gemini-1.5-pro",
+            GeminiModels::Gemini15FlashLatest => "models/gemini-1.5-flash-latest",
+            GeminiModels::Gemini15Flash001 => "models/gemini-1.5-flash-001",
+            GeminiModels::Gemini15Flash => "models/gemini-1.5-flash",
+            GeminiModels::Embedding001 => "models/embedding-001",
+            GeminiModels::TextEmbedding004 => "models/text-embedding-004",
+            GeminiModels::AQA => "models/aqa",
+        }
+    }
+}
+
+pub struct GeminiClient {
+    api_key: String,
+    model: GeminiModels,
+}
+
+impl GeminiClient {
+    pub fn new(api_key: String, model: GeminiModels) -> Self {
+        Self { api_key, model }
+    }
+}
+
+impl AskModel for GeminiClient {
+    async fn ask_model(&self, question: &str) -> Result<String, Box<dyn std::error::Error>> {
+        let url = format!(
+            "https://generativelanguage.googleapis.com/v1beta/{}:generateContent?key={}",
+            self.model.as_str(),
+            self.api_key
+        );
+        let client = reqwest::Client::new();
+
+        let body = serde_json::json!({
+            "contents": [{
+                "parts": [{
+                    "text": question
+                }]
+            }]
+        })
+        .to_string();
+
+        let res = client
+            .post(&url)
+            .body(body)
+            .header("Content-Type", "application/json; charset=utf-8")
+            .send()
+            .await?;
+
+        let content = res.text().await?;
+        let response: GeminiResponse = serde_json::from_str(&content)
+            .map_err(|e| format!("Failed to parse response from Gemini: {}\n{}", content, e))?;
+        let content = &response.candidates[0].content;
+        if content.role != "model" && content.parts.len() != 1 {
+            return Err("Failed to parse response from Gemini".into());
+        }
+        Ok(match &content.parts[0] {
+            Part::Text { text } => text.clone(),
+        })
+    }
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+struct GeminiResponse {
+    candidates: Vec<Candidate>,
+    // prompt_feedback: String,
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+struct Candidate {
+    content: Content,
+    // finish_reason: String,
+    // index: i32,
+    // safetry_ratings: Vec<SafetyRating>,
+}
+
+#[derive(Debug, Deserialize)]
+struct Content {
+    parts: Vec<Part>,
+    role: String,
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(untagged)]
+enum Part {
+    Text { text: String },
+}
+
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase")]
+#[allow(dead_code)]
+struct UsageMegadata {
+    prompt_token_count: i32,
+    candidates_token_count: i32,
+    total_token_count: i32,
+}
+
+#[cfg(test)]
+mod test {
+    use crate::AskModel;
+
+    #[tokio::test]
+    async fn test_gemini_client() {
+        let api_key = super::super::test::get_gemini_key().unwrap();
+        let client = super::GeminiClient::new(api_key, super::GeminiModels::Gemini15Flash);
+        let res = client.ask_model("who are you").await;
+        match res {
+            Ok(text) => {
+                println!("{}", text);
+            }
+            Err(e) => {
+                println!("{}", e);
+                panic!();
+            }
+        }
+    }
+}

atlas/src/api/lingyiwanwu.rs

@@ -0,0 +1,76 @@
+//! According to the [Lingyiwanwu API Reference](https://platform.lingyiwanwu.com/docs/api-reference), the Lingyiwanwu API is the same as the OpenAI API.
+//! So this is just a wrapper around the OpenAI API, change api base.
+
+use crate::AskModel;
+
+use super::openai::OpenAIClient;
+/// yi-large, yi-medium, yi-vision, yi-medium-200k, yi-spark, vi-larqe-raq, yi-large-turbo, yi-large-fc
+pub enum LingyiwanwuModels {
+    YiLarge,
+    YiMedium,
+    YiVision,
+    YiMedium200k,
+    YiSpark,
+    ViLargeRaq,
+    YiLargeTurbo,
+    YiLargeFc,
+}
+
+impl crate::Model for LingyiwanwuModels {
+    fn as_str(&self) -> &str {
+        match self {
+            LingyiwanwuModels::YiLarge => "yi-large",
+            LingyiwanwuModels::YiMedium => "yi-medium",
+            LingyiwanwuModels::YiVision => "yi-vision",
+            LingyiwanwuModels::YiMedium200k => "yi-medium-200k",
+            LingyiwanwuModels::YiSpark => "yi-spark",
+            LingyiwanwuModels::ViLargeRaq => "vi-large-raq",
+            LingyiwanwuModels::YiLargeTurbo => "yi-large-turbo",
+            LingyiwanwuModels::YiLargeFc => "yi-large-fc",
+        }
+    }
+}
+
+const LING_YI_WAN_WU_API_BASE: &str = "https://api.lingyiwanwu.com/v1";
+
+pub struct LingyiwanwuClient {
+    openai_client: OpenAIClient,
+}
+
+impl LingyiwanwuClient {
+    pub fn new(api_key: String, model: LingyiwanwuModels) -> Self {
+        let config = async_openai::config::OpenAIConfig::new()
+            .with_api_key(&api_key)
+            .with_api_base(LING_YI_WAN_WU_API_BASE);
+        let client = async_openai::Client::with_config(config);
+        Self {
+            openai_client: OpenAIClient::from_client_and_model(client, Box::new(model)),
+        }
+    }
+}
+
+impl AskModel for LingyiwanwuClient {
+    async fn ask_model(&self, question: &str) -> Result<String, Box<dyn std::error::Error>> {
+        self.openai_client.ask_model(question).await
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use crate::api::test::get_01_key;
+
+    use super::*;
+
+    #[tokio::test]
+    async fn test_lingyiwanwu_client() {
+        let api_key = get_01_key().unwrap();
+        let model = LingyiwanwuModels::YiLarge;
+        let client = LingyiwanwuClient::new(api_key, model);
+        let response = client
+            .ask_model("What is the meaning of life?")
+            .await
+            .unwrap();
+        assert!(!response.is_empty());
+        println!("Lingyiwanwu response: {}", response);
+    }
+}

atlas/src/api/mod.rs

@@ -0,0 +1,17 @@
+pub mod gemini;
+pub mod lingyiwanwu;
+pub mod openai;
+#[cfg(test)]
+mod test {
+    use std::env;
+
+    pub fn get_gemini_key() -> Option<String> {
+        // Some("".to_string())
+        env::var("GOOGLE_GEMINI_KEY").unwrap().into()
+    }
+
+    pub fn get_01_key() -> Option<String> {
+        // Some("".to_string())
+        env::var("GOOGLE_GEMINI_KEY").unwrap().into()
+    }
+}

atlas/src/api/openai.rs

@@ -0,0 +1,75 @@
+use crate::AskModel;
+use async_openai::{
+    config::OpenAIConfig,
+    types::{ChatCompletionRequestSystemMessageArgs, CreateChatCompletionRequestArgs},
+    Client,
+};
+
+pub struct OpenAIClient {
+    model: Box<dyn crate::Model + 'static>,
+    client: Client<OpenAIConfig>,
+}
+
+/// gpt-4o, gpt-4o-mini, gpt-4-turbo, gpt-4, and gpt-3.5-turbo
+pub enum OpenAIModels {
+    GPT4O,
+    GPT4OMini,
+    GPT4Turbo,
+    GPT4,
+    GPT35Turbo,
+}
+
+impl crate::Model for OpenAIModels {
+    fn as_str(&self) -> &str {
+        match self {
+            OpenAIModels::GPT4O => "gpt-4o",
+            OpenAIModels::GPT4OMini => "gpt-4o-mini",
+            OpenAIModels::GPT4Turbo => "gpt-4-turbo",
+            OpenAIModels::GPT4 => "gpt-4",
+            OpenAIModels::GPT35Turbo => "gpt-3.5-turbo",
+        }
+    }
+}
+
+impl OpenAIClient {
+    pub fn new(api_key: String, model: OpenAIModels) -> Self {
+        let config = OpenAIConfig::new().with_api_key(&api_key);
+        let client = Client::with_config(config);
+        Self {
+            model: Box::new(model),
+            client,
+        }
+    }
+
+    pub fn from_client_and_model(
+        client: Client<OpenAIConfig>,
+        model: Box<dyn crate::Model + 'static>,
+    ) -> Self {
+        Self { model, client }
+    }
+}
+
+impl AskModel for OpenAIClient {
+    async fn ask_model(&self, question: &str) -> Result<String, Box<dyn std::error::Error>> {
+        let request = CreateChatCompletionRequestArgs::default()
+            .model(self.model.as_str())
+            .messages([ChatCompletionRequestSystemMessageArgs::default()
+                .content(question)
+                .build()?
+                .into()])
+            .build()
+            .unwrap();
+
+        // debug, make request to json
+        let json_str = serde_json::to_string(&request).unwrap();
+        println!("json_str: {}", json_str);
+        let response = self
+            .client
+            .chat()
+            .create(request)
+            .await
+            .map_err(|e| format!("Failed to get response : {}", e))?;
+
+        Ok(response.choices[0].message.content.clone().unwrap())
+    }
+}

atlas/src/lib.rs

@@ -0,0 +1,12 @@
+pub mod api;
+
+pub trait Model: Sync {
+    fn as_str(&self) -> &str;
+}
+
+pub trait AskModel {
+    fn ask_model(
+        &self,
+        question: &str,
+    ) -> impl std::future::Future<Output = Result<String, Box<dyn std::error::Error>>> + Send;
+}