From 91c2a925bc26ac8026dce16d4145c8ab3e923d9e Mon Sep 17 00:00:00 2001
From: "Nathan.fooo" <86001920+appflowy@users.noreply.github.com>
Date: Mon, 16 Dec 2024 16:19:41 +0800
Subject: [PATCH] test: add test (#1077)

---
 libs/database-entity/src/dto.rs              |   6 +-
 tests/ai_test/chat_with_doc_test.rs          | 124 --------
 tests/ai_test/chat_with_selected_doc_test.rs | 298 +++++++++++++++++++
 tests/ai_test/mod.rs                         |   2 +-
 tests/collab/util.rs                         |  15 +
 5 files changed, 317 insertions(+), 128 deletions(-)
 delete mode 100644 tests/ai_test/chat_with_doc_test.rs
 create mode 100644 tests/ai_test/chat_with_selected_doc_test.rs

diff --git a/libs/database-entity/src/dto.rs b/libs/database-entity/src/dto.rs
index 98ebe9bbb..253b68b3e 100644
--- a/libs/database-entity/src/dto.rs
+++ b/libs/database-entity/src/dto.rs
@@ -1224,10 +1224,10 @@ pub struct WorkspaceNamespace {
 #[cfg(test)]
 mod test {
   use crate::dto::{CollabParams, CollabParamsV0};
-  use crate::error::EntityError;
+
   use bytes::Bytes;
-  use collab_entity::{proto, CollabType};
-  use prost::Message;
+  use collab_entity::CollabType;
+
   use uuid::Uuid;
 
   #[test]
diff --git a/tests/ai_test/chat_with_doc_test.rs b/tests/ai_test/chat_with_doc_test.rs
deleted file mode 100644
index 38e0acccc..000000000
--- a/tests/ai_test/chat_with_doc_test.rs
+++ /dev/null
@@ -1,124 +0,0 @@
-use crate::collab::util::{alex_banker_story, alex_software_engineer_story, empty_document_editor};
-use client_api_test::{ai_test_enabled, collect_answer, TestClient};
-use collab_entity::CollabType;
-use database_entity::dto::CreateCollabParams;
-use shared_entity::dto::chat_dto::{CreateChatMessageParams, CreateChatParams};
-use uuid::Uuid;
-
-#[tokio::test]
-async fn chat_with_embedded_document() {
-  if !ai_test_enabled() {
-    return;
-  }
-  let object_id = Uuid::new_v4().to_string();
-  let mut editor = empty_document_editor(&object_id);
-  let contents = alex_software_engineer_story();
-  editor.insert_paragraphs(contents.into_iter().map(|s| s.to_string()).collect());
-  let encode_collab = editor.encode_collab();
-
-  let test_client = TestClient::new_user().await;
-  let workspace_id = test_client.workspace_id().await;
-  let params = CreateCollabParams {
-    workspace_id: workspace_id.clone(),
-    object_id: object_id.clone(),
-    encoded_collab_v1: encode_collab.encode_to_bytes().unwrap(),
-    collab_type: CollabType::Document,
-  };
-  test_client.api_client.create_collab(params).await.unwrap();
-  test_client
-    .wait_until_get_embedding(&workspace_id, &object_id)
-    .await;
-
-  // chat with document
-  let chat_id = uuid::Uuid::new_v4().to_string();
-  let params = CreateChatParams {
-    chat_id: chat_id.clone(),
-    name: "my first chat".to_string(),
-    rag_ids: vec![object_id.clone()],
-  };
-
-  // create a chat
-  test_client
-    .api_client
-    .create_chat(&workspace_id, params)
-    .await
-    .unwrap();
-
-  // ask question to check the chat is using document embedding or not
-  let params = CreateChatMessageParams::new_user(
-    "What are some of the sports Alex enjoys, and what are his experiences with them",
-  );
-  let question = test_client
-    .api_client
-    .create_question(&workspace_id, &chat_id, params)
-    .await
-    .unwrap();
-  let answer_stream = test_client
-    .api_client
-    .stream_answer_v2(&workspace_id, &chat_id, question.message_id)
-    .await
-    .unwrap();
-  let answer = collect_answer(answer_stream).await;
-  let expected = r#"
-  Alex enjoys a variety of sports that keep him active and engaged:
-	1.	Tennis: Learned in Singapore, he plays on weekends with friends.
-	2.	Basketball: Enjoys casual play, though specific details aren’t provided.
-	3.	Cycling: Brought his bike to Singapore and looks forward to exploring parks.
-	4.	Badminton: Enjoys it, though details aren’t given.
-	5.	Snowboarding: Had an unforgettable experience on challenging slopes in Lake Tahoe.
-Overall, Alex balances his work as a software programmer with his passion for sports, finding excitement and freedom in each activity.
-  "#;
-  test_client
-    .assert_similarity(&workspace_id, &answer, expected, 0.8)
-    .await;
-
-  // remove all content for given document
-  editor.clear();
-
-  // Simulate insert new content
-  let contents = alex_banker_story();
-  editor.insert_paragraphs(contents.into_iter().map(|s| s.to_string()).collect());
-  let text = editor.document.to_plain_text(false, false).unwrap();
-  let expected = alex_banker_story().join("");
-  assert_eq!(text, expected);
-
-  // full sync
-  let encode_collab = editor.encode_collab();
-  test_client
-    .api_client
-    .collab_full_sync(
-      &workspace_id,
-      &object_id,
-      CollabType::Document,
-      encode_collab.doc_state.to_vec(),
-      encode_collab.state_vector.to_vec(),
-    )
-    .await
-    .unwrap();
-
-  // after full sync, chat with the same question. After update the document content, the chat
-  // should not reply with previous context.
-  let params = CreateChatMessageParams::new_user(
-    "What are some of the sports Alex enjoys, and what are his experiences with them",
-  );
-  let question = test_client
-    .api_client
-    .create_question(&workspace_id, &chat_id, params)
-    .await
-    .unwrap();
-  let answer_stream = test_client
-    .api_client
-    .stream_answer_v2(&workspace_id, &chat_id, question.message_id)
-    .await
-    .unwrap();
-  let answer = collect_answer(answer_stream).await;
-  let expected = r#"
- Alex does not enjoy sports or physical activities. Instead, he prefers to relax and finds joy in
- exploring delicious food and trying new restaurants. For Alex, food is a form of relaxation and self-care,
- making it his favorite way to unwind rather than engaging in sports. While he may not have experiences with sports,
-  he certainly has many experiences in the culinary world, where he enjoys savoring flavors and discovering new dishes
-  "#;
-  test_client
-    .assert_similarity(&workspace_id, &answer, expected, 0.8)
-    .await;
-}
diff --git a/tests/ai_test/chat_with_selected_doc_test.rs b/tests/ai_test/chat_with_selected_doc_test.rs
new file mode 100644
index 000000000..333a33755
--- /dev/null
+++ b/tests/ai_test/chat_with_selected_doc_test.rs
@@ -0,0 +1,298 @@
+use crate::collab::util::{
+  alex_banker_story, alex_software_engineer_story, empty_document_editor,
+  snowboarding_in_japan_plan, TestDocumentEditor,
+};
+use client_api_test::{ai_test_enabled, collect_answer, TestClient};
+use collab_entity::CollabType;
+use database_entity::dto::CreateCollabParams;
+use futures_util::future::join_all;
+use shared_entity::dto::chat_dto::{CreateChatMessageParams, CreateChatParams, UpdateChatParams};
+use std::sync::Arc;
+use uuid::Uuid;
+
+struct TestDoc {
+  object_id: String,
+  editor: TestDocumentEditor,
+}
+
+impl TestDoc {
+  fn new(contents: Vec<&'static str>) -> Self {
+    let object_id = Uuid::new_v4().to_string();
+    let mut editor = empty_document_editor(&object_id);
+    editor.insert_paragraphs(contents.into_iter().map(|s| s.to_string()).collect());
+
+    Self { object_id, editor }
+  }
+}
+
+#[tokio::test]
+async fn chat_with_multiple_selected_source_test() {
+  if !ai_test_enabled() {
+    return;
+  }
+
+  let docs = vec![
+    TestDoc::new(alex_software_engineer_story()),
+    TestDoc::new(snowboarding_in_japan_plan()),
+  ];
+
+  let test_client = Arc::new(TestClient::new_user().await);
+  let workspace_id = test_client.workspace_id().await;
+
+  // Use futures' join_all to run async tasks concurrently
+  let tasks: Vec<_> = docs
+    .iter()
+    .map(|doc| {
+      let params = CreateCollabParams {
+        workspace_id: workspace_id.clone(),
+        object_id: doc.object_id.clone(),
+        encoded_collab_v1: doc.editor.encode_collab().encode_to_bytes().unwrap(),
+        collab_type: CollabType::Document,
+      };
+
+      let object_id = doc.object_id.clone();
+      let cloned_workspace_id = workspace_id.clone();
+      let cloned_test_client = Arc::clone(&test_client);
+      async move {
+        // Create collaboration and wait for embedding in parallel
+        cloned_test_client
+          .api_client
+          .create_collab(params)
+          .await
+          .unwrap();
+        cloned_test_client
+          .wait_until_get_embedding(&cloned_workspace_id, &object_id)
+          .await;
+      }
+    })
+    .collect();
+
+  // Run all tasks concurrently
+  join_all(tasks).await;
+
+  // create chat
+  let chat_id = uuid::Uuid::new_v4().to_string();
+  let params = CreateChatParams {
+    chat_id: chat_id.clone(),
+    name: "my first chat".to_string(),
+    rag_ids: vec![],
+  };
+  test_client
+    .api_client
+    .create_chat(&workspace_id, params)
+    .await
+    .unwrap();
+
+  // use alex_software_engineer_story as chat context
+  let params = UpdateChatParams {
+    name: None,
+    metadata: None,
+    rag_ids: Some(vec![docs[0].object_id.clone()]),
+  };
+  test_client
+    .api_client
+    .update_chat_settings(&workspace_id, &chat_id, params)
+    .await
+    .unwrap();
+
+  // ask question that relate to the plan to Japan. The chat doesn't know any plan to Japan because
+  // I have added the snowboarding_in_japan_plan as a chat context.
+  let answer = ask_question(
+    &test_client,
+    &workspace_id,
+    &chat_id,
+    "When do we take off to Japan? Just tell me the date, and if you’re not sure, please let me know you don’t know",
+  )
+  .await;
+  let expected_unknown_japan_answer = r#"
+  I'm sorry, but I don't know the date for your trip to Japan.
+  "#;
+  test_client
+    .assert_similarity(&workspace_id, &answer, expected_unknown_japan_answer, 0.8)
+    .await;
+
+  // update chat context to snowboarding_in_japan_plan
+  let params = UpdateChatParams {
+    name: None,
+    metadata: None,
+    rag_ids: Some(vec![docs[0].object_id.clone(), docs[1].object_id.clone()]),
+  };
+  test_client
+    .api_client
+    .update_chat_settings(&workspace_id, &chat_id, params)
+    .await
+    .unwrap();
+  let answer = ask_question(
+    &test_client,
+    &workspace_id,
+    &chat_id,
+    "when do we take off to Japan? Just tell me the date",
+  )
+  .await;
+  let expected = r#"
+  You take off to Japan on **January 7th**
+  "#;
+  test_client
+    .assert_similarity(&workspace_id, &answer, expected, 0.8)
+    .await;
+
+  // Ask question for alex to make sure two documents are treated as chat context
+  let answer = ask_question(
+    &test_client,
+    &workspace_id,
+    &chat_id,
+    "Can you list the sports Alex enjoys? Please provide just the names, separated by commas",
+  )
+  .await;
+  let expected = r#"Tennis, basketball, cycling, badminton, snowboarding."#;
+  test_client
+    .assert_similarity(&workspace_id, &answer, expected, 0.8)
+    .await;
+
+  // remove the Japan plan and check the response. After remove the Japan plan, the chat should not
+  // know about the plan to Japan.
+  let params = UpdateChatParams {
+    name: None,
+    metadata: None,
+    rag_ids: Some(vec![]),
+  };
+  test_client
+    .api_client
+    .update_chat_settings(&workspace_id, &chat_id, params)
+    .await
+    .unwrap();
+  let answer = ask_question(
+    &test_client,
+    &workspace_id,
+    &chat_id,
+    "When do we take off to Japan? Just tell me the date, and if you’re not sure, please let me know you don’t know",
+  )
+  .await;
+  test_client
+    .assert_similarity(&workspace_id, &answer, expected_unknown_japan_answer, 0.8)
+    .await;
+}
+
+#[tokio::test]
+async fn chat_with_selected_source_override_test() {
+  if !ai_test_enabled() {
+    return;
+  }
+  let object_id = Uuid::new_v4().to_string();
+  let mut editor = empty_document_editor(&object_id);
+  let contents = alex_software_engineer_story();
+  editor.insert_paragraphs(contents.into_iter().map(|s| s.to_string()).collect());
+  let encode_collab = editor.encode_collab();
+
+  let test_client = TestClient::new_user().await;
+  let workspace_id = test_client.workspace_id().await;
+  let params = CreateCollabParams {
+    workspace_id: workspace_id.clone(),
+    object_id: object_id.clone(),
+    encoded_collab_v1: encode_collab.encode_to_bytes().unwrap(),
+    collab_type: CollabType::Document,
+  };
+  test_client.api_client.create_collab(params).await.unwrap();
+  test_client
+    .wait_until_get_embedding(&workspace_id, &object_id)
+    .await;
+
+  // chat with document
+  let chat_id = uuid::Uuid::new_v4().to_string();
+  let params = CreateChatParams {
+    chat_id: chat_id.clone(),
+    name: "my first chat".to_string(),
+    rag_ids: vec![object_id.clone()],
+  };
+
+  // create a chat
+  test_client
+    .api_client
+    .create_chat(&workspace_id, params)
+    .await
+    .unwrap();
+
+  // ask question to check the chat is using document embedding or not
+  let answer = ask_question(
+    &test_client,
+    &workspace_id,
+    &chat_id,
+    "What are some of the sports Alex enjoys, and what are his experiences with them",
+  )
+  .await;
+  let expected = r#"
+  Alex enjoys a variety of sports that keep him active and engaged:
+	1.	Tennis: Learned in Singapore, he plays on weekends with friends.
+	2.	Basketball: Enjoys casual play, though specific details aren’t provided.
+	3.	Cycling: Brought his bike to Singapore and looks forward to exploring parks.
+	4.	Badminton: Enjoys it, though details aren’t given.
+	5.	Snowboarding: Had an unforgettable experience on challenging slopes in Lake Tahoe.
+Overall, Alex balances his work as a software programmer with his passion for sports, finding excitement and freedom in each activity.
+  "#;
+  test_client
+    .assert_similarity(&workspace_id, &answer, expected, 0.8)
+    .await;
+
+  // remove all content for given document
+  editor.clear();
+
+  // Simulate insert new content
+  let contents = alex_banker_story();
+  editor.insert_paragraphs(contents.into_iter().map(|s| s.to_string()).collect());
+  let text = editor.document.to_plain_text(false, false).unwrap();
+  let expected = alex_banker_story().join("");
+  assert_eq!(text, expected);
+
+  // full sync
+  let encode_collab = editor.encode_collab();
+  test_client
+    .api_client
+    .collab_full_sync(
+      &workspace_id,
+      &object_id,
+      CollabType::Document,
+      encode_collab.doc_state.to_vec(),
+      encode_collab.state_vector.to_vec(),
+    )
+    .await
+    .unwrap();
+
+  // after full sync, chat with the same question. After update the document content, the chat
+  // should not reply with previous context.
+  let answer = ask_question(
+    &test_client,
+    &workspace_id,
+    &chat_id,
+    "What are some of the sports Alex enjoys, and what are his experiences with them",
+  )
+  .await;
+  let expected = r#"
+ Alex does not enjoy sports or physical activities. Instead, he prefers to relax and finds joy in
+ exploring delicious food and trying new restaurants. For Alex, food is a form of relaxation and self-care,
+ making it his favorite way to unwind rather than engaging in sports. While he may not have experiences with sports,
+  he certainly has many experiences in the culinary world, where he enjoys savoring flavors and discovering new dishes
+  "#;
+  test_client
+    .assert_similarity(&workspace_id, &answer, expected, 0.8)
+    .await;
+}
+
+async fn ask_question(
+  test_client: &TestClient,
+  workspace_id: &str,
+  chat_id: &str,
+  question: &str,
+) -> String {
+  let params = CreateChatMessageParams::new_user(question);
+  let question = test_client
+    .api_client
+    .create_question(workspace_id, chat_id, params)
+    .await
+    .unwrap();
+  let answer_stream = test_client
+    .api_client
+    .stream_answer_v2(workspace_id, chat_id, question.message_id)
+    .await
+    .unwrap();
+  collect_answer(answer_stream).await
+}
diff --git a/tests/ai_test/mod.rs b/tests/ai_test/mod.rs
index 58987b2ef..bbcc4d7dc 100644
--- a/tests/ai_test/mod.rs
+++ b/tests/ai_test/mod.rs
@@ -4,4 +4,4 @@ mod complete_text;
 mod summarize_row;
 mod util;
 
-mod chat_with_doc_test;
+mod chat_with_selected_doc_test;
diff --git a/tests/collab/util.rs b/tests/collab/util.rs
index 14dbf9db7..e6528e0c2 100644
--- a/tests/collab/util.rs
+++ b/tests/collab/util.rs
@@ -121,6 +121,21 @@ pub fn alex_software_engineer_story() -> Vec<&'static str> {
   ]
 }
 
+pub fn snowboarding_in_japan_plan() -> Vec<&'static str> {
+  vec![
+    "Our trip begins with a flight from American to Tokyo on January 7th.",
+    "In Tokyo, we’ll spend three days, from February 7th to 10th, exploring the city’s tech scene and snowboarding gear shops.",
+    "We’ll visit popular spots like Shibuya, Shinjuku, and Odaiba before heading to our next destination.",
+    "From Tokyo, we fly to Sendai and then travel to Zao Onsen for a 3-day stay from February 10th to 14th.",
+    "Zao Onsen is famous for its beautiful snow and the iconic ice trees, which will make for a unique snowboarding experience.",
+    "After Zao Onsen, we fly from Sendai to Chitose, then head to Sapporo for a 2-day visit, exploring the city’s vibrant atmosphere and winter attractions.",
+    "On the next day, we’ll spend time at Sapporo Tein, a ski resort that offers great runs and stunning views of the city and the sea.",
+    "Then we head to Rusutsu for 5 days, one of the top ski resorts in Japan, known for its deep powder snow and extensive runs.",
+    "Finally, we’ll fly back to Singapore after experiencing some of the best snowboarding Japan has to offer.",
+    "Ski resorts to visit include Niseko (二世谷), Rusutsu (留寿都), Sapporo Tein (札幌和海景), and Zao Onsen Ski Resort (冰树).",
+  ]
+}
+
 pub fn alex_banker_story() -> Vec<&'static str> {
   vec![
     "Alex is a banker who spends most of their time working with numbers.",