atlas: add method ask_model_with_context

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<String, Box<dyn std::error::Error>> {
+    async fn ask_model_with_context(
+        &self,
+        _context: ChatMessage,
+    ) -> 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 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<Content>,
+}
+
 #[derive(Debug, Deserialize)]
 #[serde(rename_all = "camelCase")]
 struct GeminiResponse {
@@ -113,13 +126,13 @@ struct Candidate {
     // safetry_ratings: Vec<SafetyRating>,
 }
 
-#[derive(Debug, Deserialize)]
+#[derive(Debug, Deserialize, Serialize)]
 struct Content {
     parts: Vec<Part>,
     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);
+            }
+        }
+    }
 }

atlas/src/api/lingyiwanwu.rs

@@ -50,14 +50,17 @@ impl LingyiwanwuClient {
 }
 
 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
+    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 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);
+    }
 }

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<String, Box<dyn std::error::Error>> {
+    async fn ask_model_with_context(
+        &self,
+        _context: ChatMessage,
+    ) -> Result<String, Box<dyn std::error::Error>> {
+        let mut messages: Vec<ChatCompletionRequestMessage> = 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()

atlas/src/lib.rs

@@ -8,5 +8,25 @@ pub trait AskModel {
     fn ask_model(
         &self,
         question: &str,
+    ) -> impl std::future::Future<Output = Result<String, Box<dyn std::error::Error>>> + 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<Output = Result<String, Box<dyn std::error::Error>>> + 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)>,
+}