chk
/
AppFlowy-Cloud
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

test: add test (#1077)

This commit is contained in:
Nathan.fooo 2024-12-16 16:19:41 +08:00 committed by GitHub
parent ec9c38254b
commit 91c2a925bc
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 317 additions and 128 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
libs/database-entity/src/dto.rs
View File

@ -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]

124
tests/ai_test/chat_with_doc_test.rs
View File

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

298
tests/ai_test/chat_with_selected_doc_test.rs Normal file
View File

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

2
tests/ai_test/mod.rs
View File

@ -4,4 +4,4 @@ mod complete_text;
mod summarize_row;
mod util;
mod chat_with_doc_test;
mod chat_with_selected_doc_test;

15
tests/collab/util.rs
View File

@ -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.",