From 51ee25ccdb40434017b351454fc59b679bb97a5e Mon Sep 17 00:00:00 2001 From: nathan Date: Thu, 13 Jun 2024 11:28:54 +0800 Subject: [PATCH] chore: ai tag api endpoint --- libs/appflowy-ai-client/src/client.rs | 16 ++++++- libs/appflowy-ai-client/src/dto.rs | 24 +++++++++++ libs/appflowy-ai-client/tests/row_test/mod.rs | 1 + .../tests/row_test/tag_test.rs | 25 +++++++++++ libs/client-api/src/http_ai.rs | 20 ++++++++- src/api/ai.rs | 25 ++++++++++- tests/ai_test/mod.rs | 2 +- tests/ai_test/summarize_row.rs | 22 ---------- tests/ai_test/tool_test.rs | 43 +++++++++++++++++++ 9 files changed, 150 insertions(+), 28 deletions(-) create mode 100644 libs/appflowy-ai-client/tests/row_test/tag_test.rs delete mode 100644 tests/ai_test/summarize_row.rs create mode 100644 tests/ai_test/tool_test.rs diff --git a/libs/appflowy-ai-client/src/client.rs b/libs/appflowy-ai-client/src/client.rs index 73aff13e3..69d58464e 100644 --- a/libs/appflowy-ai-client/src/client.rs +++ b/libs/appflowy-ai-client/src/client.rs @@ -1,7 +1,7 @@ use crate::dto::{ ChatAnswer, ChatQuestion, CompleteTextResponse, CompletionType, Document, MessageData, - RepeatedRelatedQuestion, SearchDocumentsRequest, SummarizeRowResponse, TranslateRowData, - TranslateRowResponse, + RepeatedRelatedQuestion, SearchDocumentsRequest, SummarizeRowResponse, TagRowData, + TagRowResponse, TranslateRowData, TranslateRowResponse, }; use crate::error::AIError; @@ -97,6 +97,18 @@ impl AppFlowyAIClient { .into_data() } + pub async fn tag_row(&self, data: TagRowData) -> Result { + let url = format!("{}/tag_row", self.url); + let resp = self + .http_client(Method::POST, &url)? + .json(&data) + .send() + .await?; + AIResponse::::from_response(resp) + .await? + .into_data() + } + pub async fn index_documents(&self, documents: &[Document]) -> Result<(), AIError> { let url = format!("{}/index_documents", self.url); let resp = self diff --git a/libs/appflowy-ai-client/src/dto.rs b/libs/appflowy-ai-client/src/dto.rs index 6d2fa482b..9b8903373 100644 --- a/libs/appflowy-ai-client/src/dto.rs +++ b/libs/appflowy-ai-client/src/dto.rs @@ -97,6 +97,12 @@ pub struct TranslateRowParams { pub data: TranslateRowData, } +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct TagRowParams { + pub workspace_id: String, + pub data: TagRowData, +} + /// Represents different types of content that can be used to summarize a database row. #[derive(Clone, Debug, Serialize, Deserialize)] pub struct TranslateRowData { @@ -114,3 +120,21 @@ pub struct TranslateItem { pub struct TranslateRowResponse { pub items: Vec>, } + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct TagRowData { + pub existing_tags: Vec, + pub items: Vec, + pub num_tags: i32, +} + +#[derive(Clone, Debug, Serialize, Deserialize)] +pub struct TagItem { + pub title: String, + pub content: String, +} + +#[derive(Clone, Debug, Default, Serialize, Deserialize)] +pub struct TagRowResponse { + pub tags: Vec, +} diff --git a/libs/appflowy-ai-client/tests/row_test/mod.rs b/libs/appflowy-ai-client/tests/row_test/mod.rs index 26c481a7a..411f534f7 100644 --- a/libs/appflowy-ai-client/tests/row_test/mod.rs +++ b/libs/appflowy-ai-client/tests/row_test/mod.rs @@ -1,2 +1,3 @@ mod summarize_test; +mod tag_test; mod translate_test; diff --git a/libs/appflowy-ai-client/tests/row_test/tag_test.rs b/libs/appflowy-ai-client/tests/row_test/tag_test.rs new file mode 100644 index 000000000..a05cd5eaf --- /dev/null +++ b/libs/appflowy-ai-client/tests/row_test/tag_test.rs @@ -0,0 +1,25 @@ +use crate::appflowy_ai_client; + +use appflowy_ai_client::dto::{TagItem, TagRowData}; + +#[tokio::test] +async fn translate_row_test() { + let client = appflowy_ai_client(); + + let mut items = Vec::new(); + for (key, value) in [("book name", "Atomic Habits"), ("author", "James Clear")].iter() { + items.push(TagItem { + title: key.to_string(), + content: value.to_string(), + }); + } + + let data = TagRowData { + existing_tags: vec![], + items, + num_tags: 3, + }; + + let result = client.tag_row(data).await.unwrap(); + assert_eq!(result.tags.len(), 3); +} diff --git a/libs/client-api/src/http_ai.rs b/libs/client-api/src/http_ai.rs index 44426c427..f3cad8b44 100644 --- a/libs/client-api/src/http_ai.rs +++ b/libs/client-api/src/http_ai.rs @@ -2,8 +2,8 @@ use crate::http::log_request_id; use crate::Client; use reqwest::Method; use shared_entity::dto::ai_dto::{ - CompleteTextParams, CompleteTextResponse, SummarizeRowParams, SummarizeRowResponse, - TranslateRowParams, TranslateRowResponse, + CompleteTextParams, CompleteTextResponse, SummarizeRowParams, SummarizeRowResponse, TagRowParams, + TagRowResponse, TranslateRowParams, TranslateRowResponse, }; use shared_entity::response::{AppResponse, AppResponseError}; use tracing::instrument; @@ -55,6 +55,22 @@ impl Client { .into_data() } + #[instrument(level = "info", skip_all)] + pub async fn tag_row(&self, params: TagRowParams) -> Result { + let url = format!("{}/api/ai/{}/tag_row", self.base_url, params.workspace_id); + + let resp = self + .http_client_with_auth(Method::POST, &url) + .await? + .json(¶ms) + .send() + .await?; + + log_request_id(&resp); + AppResponse::::from_response(resp) + .await? + .into_data() + } #[instrument(level = "info", skip_all)] pub async fn completion_text( &self, diff --git a/src/api/ai.rs b/src/api/ai.rs index ff42bb555..1de2bf3ae 100644 --- a/src/api/ai.rs +++ b/src/api/ai.rs @@ -2,7 +2,9 @@ use crate::state::AppState; use actix_web::web::{Data, Json}; use actix_web::{web, Scope}; use app_error::AppError; -use appflowy_ai_client::dto::{CompleteTextResponse, TranslateRowParams, TranslateRowResponse}; +use appflowy_ai_client::dto::{ + CompleteTextResponse, TagRowParams, TagRowResponse, TranslateRowParams, TranslateRowResponse, +}; use shared_entity::dto::ai_dto::{ CompleteTextParams, SummarizeRowData, SummarizeRowParams, SummarizeRowResponse, }; @@ -14,6 +16,7 @@ pub fn ai_completion_scope() -> Scope { .service(web::resource("/complete_text").route(web::post().to(complete_text_handler))) .service(web::resource("/summarize_row").route(web::post().to(summarize_row_handler))) .service(web::resource("/translate_row").route(web::post().to(translate_row_handler))) + .service(web::resource("/tag_row").route(web::post().to(tag_row_handler))) } async fn complete_text_handler( @@ -85,3 +88,23 @@ async fn translate_row_handler( }, } } + +#[instrument(level = "debug", skip(state, payload), err)] +async fn tag_row_handler( + state: web::Data, + payload: web::Json, +) -> actix_web::Result>> { + let params = payload.into_inner(); + + match state.ai_client.tag_row(params.data).await { + Ok(resp) => Ok(AppResponse::Ok().with_data(resp).into()), + Err(err) => { + error!("Failed to tag row: {:?}", err); + Ok( + AppResponse::Ok() + .with_data(TagRowResponse::default()) + .into(), + ) + }, + } +} diff --git a/tests/ai_test/mod.rs b/tests/ai_test/mod.rs index 70ff75b83..d2111e37d 100644 --- a/tests/ai_test/mod.rs +++ b/tests/ai_test/mod.rs @@ -1,3 +1,3 @@ mod chat_test; mod complete_text; -mod summarize_row; +mod tool_test; diff --git a/tests/ai_test/summarize_row.rs b/tests/ai_test/summarize_row.rs deleted file mode 100644 index c8e141cb0..000000000 --- a/tests/ai_test/summarize_row.rs +++ /dev/null @@ -1,22 +0,0 @@ -use client_api_test::TestClient; -use serde_json::json; -use shared_entity::dto::ai_dto::{SummarizeRowData, SummarizeRowParams}; - -#[tokio::test] -async fn summarize_row_test() { - let test_client = TestClient::new_user().await; - let workspace_id = test_client.workspace_id().await; - - let params = SummarizeRowParams { - workspace_id: workspace_id.clone(), - data: SummarizeRowData::Content( - json!({"name": "Jack", "age": 25, "city": "New York"}) - .as_object() - .unwrap() - .clone(), - ), - }; - - let resp = test_client.api_client.summarize_row(params).await.unwrap(); - assert!(!resp.text.is_empty()); -} diff --git a/tests/ai_test/tool_test.rs b/tests/ai_test/tool_test.rs new file mode 100644 index 000000000..05e9b656f --- /dev/null +++ b/tests/ai_test/tool_test.rs @@ -0,0 +1,43 @@ +use appflowy_ai_client::dto::{TagItem, TagRowData, TagRowParams}; +use client_api_test::TestClient; +use serde_json::json; +use shared_entity::dto::ai_dto::{SummarizeRowData, SummarizeRowParams}; + +#[tokio::test] +async fn summarize_row_test() { + let test_client = TestClient::new_user().await; + let workspace_id = test_client.workspace_id().await; + + let params = SummarizeRowParams { + workspace_id: workspace_id.clone(), + data: SummarizeRowData::Content( + json!({"name": "Jack", "age": 25, "city": "New York"}) + .as_object() + .unwrap() + .clone(), + ), + }; + + let resp = test_client.api_client.summarize_row(params).await.unwrap(); + assert!(!resp.text.is_empty()); +} +#[tokio::test] +async fn tag_row_test() { + let test_client = TestClient::new_user().await; + let workspace_id = test_client.workspace_id().await; + + let params = TagRowParams { + workspace_id: workspace_id.clone(), + data: TagRowData { + existing_tags: vec![], + items: vec![TagItem { + title: "Atomic habits".to_string(), + content: "Atomic Habits by James Clear discusses how small, consistent changes in habits can lead to significant improvements over time. The book provides strategies for building good habits and breaking bad ones, emphasizing the importance of making tiny adjustments that compound into remarkable results. Clear introduces the concepts of habit stacking, the two-minute rule, and the four laws of behavior change: making habits obvious, attractive, easy, and satisfying".to_string(), + }], + num_tags: 5, + }, + }; + + let resp = test_client.api_client.tag_row(params).await.unwrap(); + assert_eq!(resp.tags.len(), 5); +}