From ce7504bc98080bac4d6a90dca3e975531915aed6 Mon Sep 17 00:00:00 2001 From: Hou Xiaoxuan Date: Thu, 1 Aug 2024 20:00:40 +0800 Subject: [PATCH 1/2] Add `atlas` and impl gemini client Signed-off-by: Hou Xiaoxuan --- Cargo.toml | 3 +- atlas/Cargo.toml | 10 +++ atlas/src/api/gemini.rs | 156 ++++++++++++++++++++++++++++++++++++++++ atlas/src/api/mod.rs | 10 +++ atlas/src/lib.rs | 8 +++ 5 files changed, 186 insertions(+), 1 deletion(-) create mode 100644 atlas/Cargo.toml create mode 100644 atlas/src/api/gemini.rs create mode 100644 atlas/src/api/mod.rs create mode 100644 atlas/src/lib.rs diff --git a/Cargo.toml b/Cargo.toml index 7ae589ac9..acab7c5ac 100644 --- a/Cargo.toml +++ b/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" diff --git a/atlas/Cargo.toml b/atlas/Cargo.toml new file mode 100644 index 000000000..47715dfcc --- /dev/null +++ b/atlas/Cargo.toml @@ -0,0 +1,10 @@ +[package] +name = "atlas" +version = "0.1.0" +edition = "2021" + +[dependencies] +reqwest.workspace = true +serde.workspace = true +serde_json.workspace = true +tokio = { workspace = true, features = ["macros"] } diff --git a/atlas/src/api/gemini.rs b/atlas/src/api/gemini.rs new file mode 100644 index 000000000..c44e5301c --- /dev/null +++ b/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> { + 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, + // prompt_feedback: String, +} + +#[derive(Debug, Deserialize)] +#[serde(rename_all = "camelCase")] +struct Candidate { + content: Content, + // finish_reason: String, + // index: i32, + // safetry_ratings: Vec, +} + +#[derive(Debug, Deserialize)] +struct Content { + parts: Vec, + 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!(); + } + } + } +} diff --git a/atlas/src/api/mod.rs b/atlas/src/api/mod.rs new file mode 100644 index 000000000..48ba0e1fe --- /dev/null +++ b/atlas/src/api/mod.rs @@ -0,0 +1,10 @@ +pub mod gemini; + +#[cfg(test)] +mod test { + use std::env; + + pub fn get_gemini_key() -> Option { + env::var("GOOGLE_GEMINI_KEY").unwrap().into() + } +} diff --git a/atlas/src/lib.rs b/atlas/src/lib.rs new file mode 100644 index 000000000..ef5c86543 --- /dev/null +++ b/atlas/src/lib.rs @@ -0,0 +1,8 @@ +pub mod api; + +pub trait AskModel { + fn ask_model( + &self, + question: &str, + ) -> impl std::future::Future>> + Send; +} From 2e4adf783ce9d2866a4a5fde9198c88610546c19 Mon Sep 17 00:00:00 2001 From: Hou Xiaoxuan Date: Fri, 2 Aug 2024 18:25:12 +0800 Subject: [PATCH 2/2] =?UTF-8?q?impl=20openai=20client=E3=80=81impl=20lingy?= =?UTF-8?q?iwanwu=20client?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Signed-off-by: Hou Xiaoxuan --- atlas/Cargo.toml | 1 + atlas/src/api/lingyiwanwu.rs | 76 ++++++++++++++++++++++++++++++++++++ atlas/src/api/mod.rs | 9 ++++- atlas/src/api/openai.rs | 75 +++++++++++++++++++++++++++++++++++ atlas/src/lib.rs | 4 ++ 5 files changed, 164 insertions(+), 1 deletion(-) create mode 100644 atlas/src/api/lingyiwanwu.rs create mode 100644 atlas/src/api/openai.rs diff --git a/atlas/Cargo.toml b/atlas/Cargo.toml index 47715dfcc..0faccf1b1 100644 --- a/atlas/Cargo.toml +++ b/atlas/Cargo.toml @@ -4,6 +4,7 @@ version = "0.1.0" edition = "2021" [dependencies] +async-openai = "0.23.4" reqwest.workspace = true serde.workspace = true serde_json.workspace = true diff --git a/atlas/src/api/lingyiwanwu.rs b/atlas/src/api/lingyiwanwu.rs new file mode 100644 index 000000000..3618f4dc0 --- /dev/null +++ b/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> { + 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); + } +} diff --git a/atlas/src/api/mod.rs b/atlas/src/api/mod.rs index 48ba0e1fe..51e8cecad 100644 --- a/atlas/src/api/mod.rs +++ b/atlas/src/api/mod.rs @@ -1,10 +1,17 @@ pub mod gemini; - +pub mod lingyiwanwu; +pub mod openai; #[cfg(test)] mod test { use std::env; pub fn get_gemini_key() -> Option { + // Some("".to_string()) + env::var("GOOGLE_GEMINI_KEY").unwrap().into() + } + + pub fn get_01_key() -> Option { + // Some("".to_string()) env::var("GOOGLE_GEMINI_KEY").unwrap().into() } } diff --git a/atlas/src/api/openai.rs b/atlas/src/api/openai.rs new file mode 100644 index 000000000..a87992f41 --- /dev/null +++ b/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, + client: Client, +} + +/// 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, + model: Box, + ) -> Self { + Self { model, client } + } +} + +impl AskModel for OpenAIClient { + async fn ask_model(&self, question: &str) -> Result> { + 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()) + } +} diff --git a/atlas/src/lib.rs b/atlas/src/lib.rs index ef5c86543..9804ee240 100644 --- a/atlas/src/lib.rs +++ b/atlas/src/lib.rs @@ -1,5 +1,9 @@ pub mod api; +pub trait Model: Sync { + fn as_str(&self) -> &str; +} + pub trait AskModel { fn ask_model( &self,