125 lines
4.5 KiB
Rust
125 lines
4.5 KiB
Rust
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;
|
||
}
|