diff --git a/atlas/src/api/gemini.rs b/atlas/src/api/gemini.rs index c44e5301c..0ea481a61 100644 --- a/atlas/src/api/gemini.rs +++ b/atlas/src/api/gemini.rs @@ -1,6 +1,6 @@ -use serde::Deserialize; +use serde::{Deserialize, Serialize}; -use crate::AskModel; +use crate::{AskModel, ChatMessage, ChatRole}; pub enum GeminiModels { ChatBison001, @@ -60,23 +60,30 @@ impl GeminiClient { } impl AskModel for GeminiClient { - async fn ask_model(&self, question: &str) -> Result> { + async fn ask_model_with_context( + &self, + _context: ChatMessage, + ) -> 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 mut contents = CachedContents { contents: vec![] }; + _context.messages.iter().for_each(|(role, content)| { + contents.contents.push(Content { + parts: vec![Part::Text { + text: content.clone(), + }], + role: match role { + ChatRole::User => "user".to_string(), + ChatRole::Model => "model".to_string(), + }, + }); + }); + + let body = serde_json::to_string(&contents).unwrap(); let res = client .post(&url) .body(body) @@ -97,6 +104,12 @@ impl AskModel for GeminiClient { } } +#[derive(Debug, Deserialize, Serialize)] +#[serde(rename_all = "camelCase")] +struct CachedContents { + contents: Vec, +} + #[derive(Debug, Deserialize)] #[serde(rename_all = "camelCase")] struct GeminiResponse { @@ -113,13 +126,13 @@ struct Candidate { // safetry_ratings: Vec, } -#[derive(Debug, Deserialize)] +#[derive(Debug, Deserialize, Serialize)] struct Content { parts: Vec, role: String, } -#[derive(Debug, Deserialize)] +#[derive(Debug, Deserialize, Serialize)] #[serde(untagged)] enum Part { Text { text: String }, @@ -136,7 +149,7 @@ struct UsageMegadata { #[cfg(test)] mod test { - use crate::AskModel; + use crate::{AskModel, ChatRole}; #[tokio::test] async fn test_gemini_client() { @@ -153,4 +166,32 @@ mod test { } } } + + #[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); + } + } + } } diff --git a/atlas/src/api/lingyiwanwu.rs b/atlas/src/api/lingyiwanwu.rs index 3618f4dc0..eac2c361d 100644 --- a/atlas/src/api/lingyiwanwu.rs +++ b/atlas/src/api/lingyiwanwu.rs @@ -50,14 +50,17 @@ impl LingyiwanwuClient { } impl AskModel for LingyiwanwuClient { - async fn ask_model(&self, question: &str) -> Result> { - self.openai_client.ask_model(question).await + async fn ask_model_with_context( + &self, + context: crate::ChatMessage, + ) -> Result> { + self.openai_client.ask_model_with_context(context).await } } #[cfg(test)] mod tests { - use crate::api::test::get_01_key; + use crate::{api::test::get_01_key, ChatRole}; use super::*; @@ -73,4 +76,27 @@ mod tests { 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); + } } diff --git a/atlas/src/api/openai.rs b/atlas/src/api/openai.rs index a87992f41..f6b690471 100644 --- a/atlas/src/api/openai.rs +++ b/atlas/src/api/openai.rs @@ -1,7 +1,10 @@ -use crate::AskModel; +use crate::{AskModel, ChatMessage, ChatRole}; use async_openai::{ config::OpenAIConfig, - types::{ChatCompletionRequestSystemMessageArgs, CreateChatCompletionRequestArgs}, + types::{ + ChatCompletionRequestMessage, ChatCompletionRequestSystemMessageArgs, + ChatCompletionRequestUserMessageArgs, CreateChatCompletionRequestArgs, + }, Client, }; @@ -50,19 +53,37 @@ impl OpenAIClient { } impl AskModel for OpenAIClient { - async fn ask_model(&self, question: &str) -> Result> { + async fn ask_model_with_context( + &self, + _context: ChatMessage, + ) -> Result> { + let mut messages: Vec = vec![]; + for (role, content) in _context.messages.iter() { + match role { + ChatRole::User => { + messages.push( + ChatCompletionRequestSystemMessageArgs::default() + .content(content) + .build()? + .into(), + ); + } + ChatRole::Model => { + messages.push( + ChatCompletionRequestUserMessageArgs::default() + .content(content.as_str()) + .build()? + .into(), + ); + } + } + } let request = CreateChatCompletionRequestArgs::default() .model(self.model.as_str()) - .messages([ChatCompletionRequestSystemMessageArgs::default() - .content(question) - .build()? - .into()]) + .messages(messages) .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() diff --git a/atlas/src/lib.rs b/atlas/src/lib.rs index 9804ee240..b05d86d92 100644 --- a/atlas/src/lib.rs +++ b/atlas/src/lib.rs @@ -8,5 +8,25 @@ pub trait AskModel { fn ask_model( &self, question: &str, + ) -> impl std::future::Future>> + Send { + self.ask_model_with_context(ChatMessage { + messages: vec![(ChatRole::User, question.to_string())], + }) + } + + /// ask model with context messages, the last message should be user's current message + /// see [openai docs](https://platform.openai.com/docs/api-reference/chat/create) for more details + fn ask_model_with_context( + &self, + context: ChatMessage, ) -> impl std::future::Future>> + Send; } +pub enum ChatRole { + User, + Model, +} + +/// ChatMessage is a vector of (role, content), and role must be `user` or `model`. +pub struct ChatMessage { + pub messages: Vec<(ChatRole, String)>, +}