diff --git a/atlas/src/api/claude.rs b/atlas/src/api/claude.rs new file mode 100644 index 000000000..bbd951950 --- /dev/null +++ b/atlas/src/api/claude.rs @@ -0,0 +1,171 @@ +//! Anthropic's Claude API client, see [Anthropic Claude API](https://docs.anthropic.com/en/api/messages). + +use serde::{Deserialize, Serialize}; + +use crate::{AskModel, Model}; + +/// Refer to the [Claude Models](https://docs.anthropic.com/en/docs/about-claude/models) for more information. +#[derive(Debug, Clone)] +pub enum ClaudeModels { + Claude3_5Sonnet, + Claude3Opus, + Claude3Sonnet, + Claude3Haiku, +} + +impl Model for ClaudeModels { + fn as_str(&self) -> &str { + match self { + ClaudeModels::Claude3_5Sonnet => "claude-3-5-sonnet-20240620", + ClaudeModels::Claude3Opus => "claude-3-opus-20240229", + ClaudeModels::Claude3Sonnet => "claude-3-sonnet-20240229", + ClaudeModels::Claude3Haiku => "claude-3-haiku-20240307", + } + } +} + +impl ClaudeModels { + pub fn get_max_tokens(&self) -> usize { + match self { + ClaudeModels::Claude3_5Sonnet => 4096, + ClaudeModels::Claude3Opus => 4096, + ClaudeModels::Claude3Sonnet => 4096, + ClaudeModels::Claude3Haiku => 4096, + } + } +} + +const CLAUDE_MESSAGE_URL: &str = "https://api.anthropic.com/v1/messages"; +const ANTHROPIC_VERSION: &str = "2023-06-01"; + +#[derive(Debug, Clone)] +pub struct ClaudeClient { + api_key: String, + model: ClaudeModels, + http_client: reqwest::Client, +} + +impl ClaudeClient { + pub fn new(api_key: String, model: ClaudeModels) -> Self { + Self { + api_key, + model, + http_client: reqwest::Client::new(), + } + } +} + +impl AskModel for ClaudeClient { + async fn ask_model_with_context( + &self, + context: crate::ChatMessage, + ) -> Result> { + let mut request_body = MessageRequest { + model: self.model.as_str().to_string(), + max_tokens: self.model.get_max_tokens(), + messages: vec![], + }; + context.messages.iter().for_each(|(role, content)| { + request_body.messages.push(Message { + role: match role { + crate::ChatRole::User => MessageRole::User, + crate::ChatRole::Model => MessageRole::Assistant, + }, + content: content.clone(), + }); + }); + let body = serde_json::to_string(&request_body)?; + let res = self + .http_client + .post(CLAUDE_MESSAGE_URL) + .body(body) + .header("Content-Type", "application/json; charset=utf-8") + .header("x-api-key", &self.api_key) + .header("anthropic-version", ANTHROPIC_VERSION) + .send() + .await?; + + let content = res.text().await?; + let response: MessageResponse = serde_json::from_str(&content) + .map_err(|e| format!("Failed to parse response from Claude: {}\n{}", content, e))?; + let content = &response.content[0]; + if response.content.len() != 1 || content.r#type != "text" { + return Err("Failed to parse response from Claude.".into()); + } + Ok(content.text.clone()) + } +} + +#[derive(Debug, Deserialize, Serialize)] +struct MessageRequest { + model: String, + max_tokens: usize, + messages: Vec, +} + +#[derive(Debug, Deserialize, Serialize)] +struct Message { + role: MessageRole, + content: String, +} + +#[derive(Debug, Deserialize, Serialize)] +#[serde(rename_all = "snake_case")] +enum MessageRole { + User, + Assistant, +} +/* { + "content": [ + { + "text": "Hi! My name is Claude.", + "type": "text" + } + ], + "id": "msg_013Zva2CMHLNnXjNJJKqJ2EF", + "model": "claude-3-5-sonnet-20240620", + "role": "assistant", + "stop_reason": "end_turn", + "stop_sequence": null, + "type": "message", + "usage": { + "input_tokens": 10, + "output_tokens": 25 + } +}*/ +#[derive(Debug, Deserialize)] +struct MessageResponse { + content: Vec, + // id: String, + // model: String, + // role: String, + // stop_reason: String, + // stop_sequence: Option, + // r#type: String, + // usage: Usage, +} + +#[derive(Debug, Deserialize)] +struct MessageContent { + text: String, + r#type: String, +} + +#[derive(Debug, Deserialize)] +#[allow(dead_code)] +struct Usage { + input_tokens: i32, + output_tokens: i32, +} + +#[cfg(test)] +mod test { + use super::*; + use crate::api::test::{get_claude_key, test_client_with_context}; + #[tokio::test] + async fn test_claude_client_with_context() { + let client = ClaudeClient::new(get_claude_key().unwrap(), ClaudeModels::Claude3_5Sonnet); + + test_client_with_context(client).await; + } +} diff --git a/atlas/src/api/gemini.rs b/atlas/src/api/gemini.rs index 0ea481a61..20b348662 100644 --- a/atlas/src/api/gemini.rs +++ b/atlas/src/api/gemini.rs @@ -1,7 +1,10 @@ +//! Google Gemini API client, see [Google Gemini API](https://ai.google.dev/api/caching) + use serde::{Deserialize, Serialize}; use crate::{AskModel, ChatMessage, ChatRole}; +#[derive(Debug, Clone)] pub enum GeminiModels { ChatBison001, TextBison001, @@ -48,6 +51,7 @@ impl GeminiModels { } } +#[derive(Debug, Clone)] pub struct GeminiClient { api_key: String, model: GeminiModels, @@ -149,49 +153,12 @@ struct UsageMegadata { #[cfg(test)] mod test { - use crate::{AskModel, ChatRole}; - - #[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!(); - } - } - } + use crate::api::test::test_client_with_context; #[tokio::test] async fn test_gemini_client_with_context() { let api_key = super::super::test::get_gemini_key().unwrap(); let client = super::GeminiClient::new(api_key, super::GeminiModels::Gemini15Flash); - let _context = crate::ChatMessage { - messages: vec![ - ( - ChatRole::User, - "Resposponse a '0' no matter what you receive".into(), - ), - ( - ChatRole::Model, - "Ok, I will response with a number 0.".into(), - ), - (ChatRole::User, "who are you".into()), - ], - }; - let res = client.ask_model_with_context(_context).await; - match res { - Ok(text) => { - println!("Google Gemini response with {}", text); - } - Err(e) => { - panic!("error: {}", e); - } - } + test_client_with_context(client).await; } } diff --git a/atlas/src/api/gitee.rs b/atlas/src/api/gitee.rs new file mode 100644 index 000000000..8187f835a --- /dev/null +++ b/atlas/src/api/gitee.rs @@ -0,0 +1,89 @@ +//! The GiteeAI Serverless API is similar to the OpenAI API, so we can use the same code with a different API base URL. +//! Note that GiteeAI does not guarantee to maintain the same API structure as OpenAI, so this may break in the future. +//! GiteeAI uses URL parameters to specify the model, so there's no need to set the model in the request body. + +use crate::{AskModel, Model}; + +use super::openai::OpenAIClient; + +/// [GiteeAI Serverless API](https://ai.gitee.com/serverless-api) +#[derive(Debug, Clone)] +pub enum GiteeServerlessModels { + Qwen2_7bInstruct, + Qwen2_72bInstruct, + Yi1_5_34bChat, +} +impl Model for GiteeServerlessModels { + fn as_str(&self) -> &str { + match self { + GiteeServerlessModels::Qwen2_7bInstruct => "YP3A1DT28TAJ", + GiteeServerlessModels::Qwen2_72bInstruct => "H87ZZLSFILML", + GiteeServerlessModels::Yi1_5_34bChat => "KIXIB7TOZA1U", + } + } +} + +const GITEE_SERVERLESS_API_BASE: &str = "https://ai.gitee.com/api/inference/serverless"; + +pub struct GiteeServerlessClient { + openai_client: OpenAIClient, +} + +impl GiteeServerlessClient { + pub fn new(api_key: String, model: GiteeServerlessModels) -> Self { + let config = async_openai::config::OpenAIConfig::new() + .with_api_key(&api_key) + .with_api_base(format!("{}/{}", GITEE_SERVERLESS_API_BASE, model.as_str())); + let client = async_openai::Client::with_config(config); + Self { + openai_client: OpenAIClient::from_client_and_model(client, Box::new(model)), + } + } +} + +impl AskModel for GiteeServerlessClient { + async fn ask_model_with_context( + &self, + context: crate::ChatMessage, + ) -> Result> { + self.openai_client.ask_model_with_context(context).await + } +} + +#[cfg(test)] +mod test { + use async_openai::types::{ + ChatCompletionRequestSystemMessageArgs, CreateChatCompletionRequestArgs, + }; + + use super::*; + use crate::api::test::{get_giteeai_key, test_client_with_context}; + #[tokio::test] + async fn test_openai_rs_with_gitee() { + let config = async_openai::config::OpenAIConfig::new() + .with_api_key(get_giteeai_key().unwrap()) + .with_api_base("https://ai.gitee.com/api/inference/serverless/H87ZZLSFILML"); + let client = async_openai::Client::with_config(config); + + let request = CreateChatCompletionRequestArgs::default() + .model("") + .messages([ChatCompletionRequestSystemMessageArgs::default() + .content("What is the meaning of life?") + .build() + .unwrap() + .into()]) + .build() + .unwrap(); + + let response = client.chat().create(request).await.unwrap(); + println!("{}", response.choices[0].message.content.clone().unwrap()); + } + + #[tokio::test] + async fn test_gitee_serverless_client_with_context() { + let api_key = get_giteeai_key().unwrap(); + let model = GiteeServerlessModels::Qwen2_7bInstruct; + let client = GiteeServerlessClient::new(api_key, model); + test_client_with_context(client).await; + } +} diff --git a/atlas/src/api/lingyiwanwu.rs b/atlas/src/api/lingyiwanwu.rs index eac2c361d..52db833c4 100644 --- a/atlas/src/api/lingyiwanwu.rs +++ b/atlas/src/api/lingyiwanwu.rs @@ -1,5 +1,5 @@ -//! 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. +//! According to the [Lingyiwanwu API Reference](https://platform.lingyiwanwu.com/docs/api-reference), the Lingyiwanwu API is identical to the OpenAI API. +//! Therefore, this is just a wrapper around the OpenAI API, with a different API base URL. use crate::AskModel; @@ -60,43 +60,15 @@ impl AskModel for LingyiwanwuClient { #[cfg(test)] mod tests { - use crate::{api::test::get_01_key, ChatRole}; + use crate::api::test::{get_01_key, test_client_with_context}; 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); - } - #[tokio::test] async fn test_lingyiwanwu_client_with_context() { let api_key = get_01_key().unwrap(); let model = LingyiwanwuModels::YiLarge; let client = LingyiwanwuClient::new(api_key, model); - let _context = crate::ChatMessage { - messages: vec![ - ( - ChatRole::User, - "Resposponse a '0' no matter what you receive".into(), - ), - ( - ChatRole::Model, - "Ok, I will response with a number 0.".into(), - ), - (ChatRole::User, "What is the meaning of life?".into()), - ], - }; - let response = client.ask_model_with_context(_context).await.unwrap(); - assert!(!response.is_empty()); - println!("Lingyiwanwu response: {}", response); + test_client_with_context(client).await; } } diff --git a/atlas/src/api/mod.rs b/atlas/src/api/mod.rs index 51e8cecad..b5efe6cdd 100644 --- a/atlas/src/api/mod.rs +++ b/atlas/src/api/mod.rs @@ -1,4 +1,6 @@ +pub mod claude; pub mod gemini; +pub mod gitee; pub mod lingyiwanwu; pub mod openai; #[cfg(test)] @@ -12,6 +14,43 @@ mod test { pub fn get_01_key() -> Option { // Some("".to_string()) - env::var("GOOGLE_GEMINI_KEY").unwrap().into() + env::var("LINGYI_KEY").unwrap().into() + } + + pub fn get_giteeai_key() -> Option { + // Some("".to_string()) + env::var("GITEEAI_KEY").unwrap().into() + } + + pub fn get_claude_key() -> Option { + // Some("".to_string()) + env::var("CLAUDE_KEY").unwrap().into() + } + + pub async fn test_client_with_context(client: impl crate::AskModel) { + let _context = crate::ChatMessage { + messages: vec![ + ( + crate::ChatRole::User, + "Resposponse a '0' no matter what you receive".into(), + ), + ( + crate::ChatRole::Model, + "Ok, I will response with a number 0.".into(), + ), + (crate::ChatRole::User, "who are you".into()), + ], + }; + let res = client.ask_model_with_context(_context).await; + match res { + Ok(text) => { + assert!(!text.is_empty()); + println!("model response with {}", text); + } + Err(e) => { + println!("{}", e); + panic!(); + } + } } } diff --git a/atlas/src/api/openai.rs b/atlas/src/api/openai.rs index f6b690471..399841974 100644 --- a/atlas/src/api/openai.rs +++ b/atlas/src/api/openai.rs @@ -78,6 +78,7 @@ impl AskModel for OpenAIClient { } } } + let request = CreateChatCompletionRequestArgs::default() .model(self.model.as_str()) .messages(messages)