Atlas add Gitee AI serverless, Claude API client implementation

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<String, Box<dyn std::error::Error>> {
+        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<Message>,
+}
+
+#[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<MessageContent>,
+    // id: String,
+    // model: String,
+    // role: String,
+    // stop_reason: String,
+    // stop_sequence: Option<String>,
+    // 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;
+    }
+}

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;
     }
 }

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<String, Box<dyn std::error::Error>> {
+        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;
+    }
+}

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;
     }
 }