From 960f9b9a933905f7d251885f8d2a6a17115681a7 Mon Sep 17 00:00:00 2001 From: nathan Date: Thu, 26 Dec 2024 13:04:22 +0800 Subject: [PATCH 1/3] chore: support set rag only --- .../ai_chat/application/chat_bloc.dart | 21 ++++++-- frontend/rust-lib/flowy-ai/src/ai_manager.rs | 54 +++++++++++-------- frontend/rust-lib/flowy-ai/src/entities.rs | 11 +++- .../rust-lib/flowy-ai/src/event_handler.rs | 6 ++- 4 files changed, 64 insertions(+), 28 deletions(-) diff --git a/frontend/appflowy_flutter/lib/plugins/ai_chat/application/chat_bloc.dart b/frontend/appflowy_flutter/lib/plugins/ai_chat/application/chat_bloc.dart index 58146c42f1f2f..08d3d7063a650 100644 --- a/frontend/appflowy_flutter/lib/plugins/ai_chat/application/chat_bloc.dart +++ b/frontend/appflowy_flutter/lib/plugins/ai_chat/application/chat_bloc.dart @@ -249,11 +249,20 @@ class ChatBloc extends Bloc { }, updateSelectedSources: (selectedSourcesIds) async { emit(state.copyWith(selectedSourceIds: selectedSourcesIds)); + final ragIds = RepeatedRagId(ragIds: selectedSourcesIds); + final payload = UpdateChatSettingsPB.create() + ..chatId = ChatId(value: chatId) + ..ragIds = ragIds; + + await AIEventUpdateChatSettings(payload) + .send() + .onFailure(Log.error); + }, + setRagOnly: (ragOnly) async { + final payload = UpdateChatSettingsPB.create() + ..chatId = ChatId(value: chatId) + ..ragOnly = ragOnly; - final payload = UpdateChatSettingsPB( - chatId: ChatId(value: chatId), - ragIds: selectedSourcesIds, - ); await AIEventUpdateChatSettings(payload) .send() .onFailure(Log.error); @@ -575,6 +584,10 @@ class ChatEvent with _$ChatEvent { required List selectedSourcesIds, }) = _UpdateSelectedSources; + const factory ChatEvent.setRagOnly({ + required bool ragOnly, + }) = _SetRagOnly; + // send message const factory ChatEvent.sendMessage({ required String message, diff --git a/frontend/rust-lib/flowy-ai/src/ai_manager.rs b/frontend/rust-lib/flowy-ai/src/ai_manager.rs index d2f82f6e0f445..b094342565ad0 100644 --- a/frontend/rust-lib/flowy-ai/src/ai_manager.rs +++ b/frontend/rust-lib/flowy-ai/src/ai_manager.rs @@ -19,6 +19,7 @@ use crate::notification::{chat_notification_builder, ChatNotification}; use flowy_storage_pub::storage::StorageService; use lib_infra::async_trait::async_trait; use lib_infra::util::timestamp; +use serde_json::json; use std::path::PathBuf; use std::sync::{Arc, Weak}; use tracing::{error, info, trace}; @@ -362,46 +363,55 @@ impl AIManager { Ok(settings.rag_ids) } - pub async fn update_rag_ids(&self, chat_id: &str, rag_ids: Vec) -> FlowyResult<()> { + pub async fn update_settings( + &self, + chat_id: &str, + rag_ids: Option>, + rag_only: Option, + ) -> FlowyResult<()> { info!("[Chat] update chat:{} rag ids: {:?}", chat_id, rag_ids); - let workspace_id = self.user_service.workspace_id()?; let update_setting = UpdateChatParams { name: None, - metadata: None, - rag_ids: Some(rag_ids.clone()), + metadata: rag_only.map(|rag_only| json!({"rag_only": rag_only})), + rag_ids: rag_ids.clone(), }; + self .cloud_service_wm .update_chat_settings(&workspace_id, chat_id, update_setting) .await?; let chat_setting_store_key = setting_store_key(chat_id); - - if let Some(settings) = self + if let Some(mut settings) = self .store_preferences .get_object::(&chat_setting_store_key) { - if let Err(err) = self.store_preferences.set_object( - &chat_setting_store_key, - &ChatSettings { - rag_ids: rag_ids.clone(), - ..settings - }, - ) { + if let Some(rag_only) = rag_only { + settings.metadata = json!({"rag_only": rag_only}); + } + + if let Some(rag_ids) = rag_ids { + settings.rag_ids = rag_ids.clone(); + let user_service = self.user_service.clone(); + let external_service = self.external_service.clone(); + tokio::spawn(async move { + if let Ok(workspace_id) = user_service.workspace_id() { + let _ = external_service + .sync_rag_documents(&workspace_id, rag_ids) + .await; + } + }); + } + + if let Err(err) = self + .store_preferences + .set_object(&chat_setting_store_key, &settings) + { error!("failed to set chat settings: {}", err); } } - let user_service = self.user_service.clone(); - let external_service = self.external_service.clone(); - tokio::spawn(async move { - if let Ok(workspace_id) = user_service.workspace_id() { - let _ = external_service - .sync_rag_documents(&workspace_id, rag_ids) - .await; - } - }); Ok(()) } } diff --git a/frontend/rust-lib/flowy-ai/src/entities.rs b/frontend/rust-lib/flowy-ai/src/entities.rs index 4eb80dd07140c..353323aed7d11 100644 --- a/frontend/rust-lib/flowy-ai/src/entities.rs +++ b/frontend/rust-lib/flowy-ai/src/entities.rs @@ -551,6 +551,15 @@ pub struct UpdateChatSettingsPB { #[validate(nested)] pub chat_id: ChatId, - #[pb(index = 2)] + #[pb(index = 2, one_of)] + pub rag_ids: Option, + + #[pb(index = 3, one_of)] + pub rag_only: Option, +} + +#[derive(Default, ProtoBuf, Clone, Debug)] +pub struct RepeatedRagId { + #[pb(index = 1)] pub rag_ids: Vec, } diff --git a/frontend/rust-lib/flowy-ai/src/event_handler.rs b/frontend/rust-lib/flowy-ai/src/event_handler.rs index 419ff5403fe4c..82ebf2939e967 100644 --- a/frontend/rust-lib/flowy-ai/src/event_handler.rs +++ b/frontend/rust-lib/flowy-ai/src/event_handler.rs @@ -459,7 +459,11 @@ pub(crate) async fn update_chat_settings_handler( let params = data.try_into_inner()?; let ai_manager = upgrade_ai_manager(ai_manager)?; ai_manager - .update_rag_ids(¶ms.chat_id.value, params.rag_ids) + .update_settings( + ¶ms.chat_id.value, + params.rag_ids.map(|v| v.rag_ids), + params.rag_only, + ) .await?; Ok(()) From 36d8a326707abd580805b2759f9d495401f97d84 Mon Sep 17 00:00:00 2001 From: Richard Shiue <71320345+richardshiue@users.noreply.github.com> Date: Thu, 26 Dec 2024 14:41:03 +0800 Subject: [PATCH 2/3] chore: implement ui --- .../ai_chat/application/chat_bloc.dart | 3 ++ .../chat_select_sources_cubit.dart | 8 +++++ .../lib/plugins/ai_chat/chat_page.dart | 22 +++++++++---- .../chat_input/desktop_ai_prompt_input.dart | 4 +-- .../chat_input/select_sources_menu.dart | 31 +++++++++++++++++-- frontend/resources/translations/en.json | 3 +- 6 files changed, 60 insertions(+), 11 deletions(-) diff --git a/frontend/appflowy_flutter/lib/plugins/ai_chat/application/chat_bloc.dart b/frontend/appflowy_flutter/lib/plugins/ai_chat/application/chat_bloc.dart index 08d3d7063a650..63e13aa4321e4 100644 --- a/frontend/appflowy_flutter/lib/plugins/ai_chat/application/chat_bloc.dart +++ b/frontend/appflowy_flutter/lib/plugins/ai_chat/application/chat_bloc.dart @@ -259,6 +259,7 @@ class ChatBloc extends Bloc { .onFailure(Log.error); }, setRagOnly: (ragOnly) async { + emit(state.copyWith(onlyUseSelectedSources: ragOnly)); final payload = UpdateChatSettingsPB.create() ..chatId = ChatId(value: chatId) ..ragOnly = ragOnly; @@ -624,12 +625,14 @@ class ChatEvent with _$ChatEvent { @freezed class ChatState with _$ChatState { const factory ChatState({ + required bool onlyUseSelectedSources, required List selectedSourceIds, required LoadChatMessageStatus loadingState, required PromptResponseState promptResponseState, }) = _ChatState; factory ChatState.initial() => const ChatState( + onlyUseSelectedSources: false, selectedSourceIds: [], loadingState: LoadChatMessageStatus.loading, promptResponseState: PromptResponseState.ready, diff --git a/frontend/appflowy_flutter/lib/plugins/ai_chat/application/chat_select_sources_cubit.dart b/frontend/appflowy_flutter/lib/plugins/ai_chat/application/chat_select_sources_cubit.dart index 06ead7d8d8134..25dd7f7e39213 100644 --- a/frontend/appflowy_flutter/lib/plugins/ai_chat/application/chat_select_sources_cubit.dart +++ b/frontend/appflowy_flutter/lib/plugins/ai_chat/application/chat_select_sources_cubit.dart @@ -107,6 +107,12 @@ class ChatSettingsCubit extends Cubit { List selectedSources = []; String filter = ''; + void updateOnlyUseSelectedSources(bool onlyUseSelectedSources) { + if (state.onlyUseSelectedSources != onlyUseSelectedSources) { + emit(state.copyWith(onlyUseSelectedSources: onlyUseSelectedSources)); + } + } + void updateSelectedSources(List newSelectedSourceIds) { selectedSourceIds = [...newSelectedSourceIds]; } @@ -380,11 +386,13 @@ class ChatSettingsCubit extends Cubit { @freezed class ChatSettingsState with _$ChatSettingsState { const factory ChatSettingsState({ + required bool onlyUseSelectedSources, required List visibleSources, required List selectedSources, }) = _ChatSettingsState; factory ChatSettingsState.initial() => const ChatSettingsState( + onlyUseSelectedSources: false, visibleSources: [], selectedSources: [], ); diff --git a/frontend/appflowy_flutter/lib/plugins/ai_chat/chat_page.dart b/frontend/appflowy_flutter/lib/plugins/ai_chat/chat_page.dart index dc6d9c649ce7a..626be04eb47ae 100644 --- a/frontend/appflowy_flutter/lib/plugins/ai_chat/chat_page.dart +++ b/frontend/appflowy_flutter/lib/plugins/ai_chat/chat_page.dart @@ -9,6 +9,7 @@ import 'package:appflowy_result/appflowy_result.dart'; import 'package:desktop_drop/desktop_drop.dart'; import 'package:easy_localization/easy_localization.dart'; import 'package:flowy_infra_ui/flowy_infra_ui.dart'; +import 'package:flutter/foundation.dart'; import 'package:flutter/material.dart'; import 'package:flutter_bloc/flutter_bloc.dart'; import 'package:flutter_chat_core/flutter_chat_core.dart'; @@ -296,12 +297,21 @@ class _ChatContentPage extends StatelessWidget { ), ); }, - onUpdateSelectedSources: (ids) { - chatBloc.add( - ChatEvent.updateSelectedSources( - selectedSourcesIds: ids, - ), - ); + onUpdateSelectedSources: (ragOnly, ids) { + if (ragOnly != chatBloc.state.onlyUseSelectedSources) { + chatBloc.add( + ChatEvent.setRagOnly( + ragOnly: ragOnly, + ), + ); + } + if (!listEquals(ids, chatBloc.state.selectedSourceIds)) { + chatBloc.add( + ChatEvent.updateSelectedSources( + selectedSourcesIds: ids, + ), + ); + } }, ) : MobileAIPromptInput( diff --git a/frontend/appflowy_flutter/lib/plugins/ai_chat/presentation/chat_input/desktop_ai_prompt_input.dart b/frontend/appflowy_flutter/lib/plugins/ai_chat/presentation/chat_input/desktop_ai_prompt_input.dart index a44f446433f99..e9146fceea9ea 100644 --- a/frontend/appflowy_flutter/lib/plugins/ai_chat/presentation/chat_input/desktop_ai_prompt_input.dart +++ b/frontend/appflowy_flutter/lib/plugins/ai_chat/presentation/chat_input/desktop_ai_prompt_input.dart @@ -32,7 +32,7 @@ class DesktopAIPromptInput extends StatefulWidget { final bool isStreaming; final void Function() onStopStreaming; final void Function(String, Map) onSubmitted; - final void Function(List) onUpdateSelectedSources; + final void Function(bool, List) onUpdateSelectedSources; @override State createState() => _DesktopAIPromptInputState(); @@ -497,7 +497,7 @@ class _PromptBottomActions extends StatelessWidget { final SendButtonState sendButtonState; final void Function() onSendPressed; final void Function() onStopStreaming; - final void Function(List) onUpdateSelectedSources; + final void Function(bool, List) onUpdateSelectedSources; @override Widget build(BuildContext context) { diff --git a/frontend/appflowy_flutter/lib/plugins/ai_chat/presentation/chat_input/select_sources_menu.dart b/frontend/appflowy_flutter/lib/plugins/ai_chat/presentation/chat_input/select_sources_menu.dart index 47cfec9b4ae14..a1d221ae2280d 100644 --- a/frontend/appflowy_flutter/lib/plugins/ai_chat/presentation/chat_input/select_sources_menu.dart +++ b/frontend/appflowy_flutter/lib/plugins/ai_chat/presentation/chat_input/select_sources_menu.dart @@ -27,7 +27,7 @@ class PromptInputDesktopSelectSourcesButton extends StatefulWidget { required this.onUpdateSelectedSources, }); - final void Function(List) onUpdateSelectedSources; + final void Function(bool, List) onUpdateSelectedSources; @override State createState() => @@ -80,6 +80,7 @@ class _PromptInputDesktopSelectSourcesButtonState return BlocListener( listener: (context, state) { cubit + ..updateOnlyUseSelectedSources(state.onlyUseSelectedSources) ..updateSelectedSources(state.selectedSourceIds) ..updateSelectedStatus(); }, @@ -95,7 +96,10 @@ class _PromptInputDesktopSelectSourcesButtonState } }, onClose: () { - widget.onUpdateSelectedSources(cubit.selectedSourceIds); + widget.onUpdateSelectedSources( + cubit.state.onlyUseSelectedSources, + cubit.selectedSourceIds, + ); if (spaceView != null) { context.read().refreshSources(spaceView); } @@ -188,6 +192,29 @@ class _PopoverContent extends StatelessWidget { context.read().updateFilter(value), ), ), + Container( + margin: const EdgeInsets.fromLTRB(8, 0, 8, 8), + height: 30, + child: FlowyButton( + text: FlowyText( + LocaleKeys.chat_onlyUseRags.tr(), + overflow: TextOverflow.ellipsis, + ), + onTap: () { + context + .read() + .updateOnlyUseSelectedSources( + !state.onlyUseSelectedSources, + ); + }, + rightIcon: state.onlyUseSelectedSources + ? FlowySvg( + FlowySvgs.check_s, + color: Theme.of(context).colorScheme.primary, + ) + : null, + ), + ), _buildDivider(), Flexible( child: ListView( diff --git a/frontend/resources/translations/en.json b/frontend/resources/translations/en.json index 6cc508ed4aa50..2eb651c99782f 100644 --- a/frontend/resources/translations/en.json +++ b/frontend/resources/translations/en.json @@ -212,7 +212,8 @@ "addToPageButton": "Add to page", "addToPageTitle": "Add message to...", "addToNewPage": "Add to a new page", - "addToNewPageName": "Messages extracted from \"{}\"" + "addToNewPageName": "Messages extracted from \"{}\"", + "onlyUseRags": "Only use selected sources to generate response" }, "trash": { "text": "Trash", From 8246af93c74c864d432fb36e3d73fa220e7a3b43 Mon Sep 17 00:00:00 2001 From: Richard Shiue <71320345+richardshiue@users.noreply.github.com> Date: Fri, 27 Dec 2024 10:10:24 +0800 Subject: [PATCH 3/3] chore: fetch rag only on open chat --- .../ai_chat/application/chat_bloc.dart | 5 +++- .../chat_input/select_sources_menu.dart | 7 +++--- frontend/resources/translations/en.json | 2 +- frontend/rust-lib/flowy-ai/src/ai_manager.rs | 10 +++----- frontend/rust-lib/flowy-ai/src/entities.rs | 25 +++++++++++++++++-- .../rust-lib/flowy-ai/src/event_handler.rs | 3 +-- 6 files changed, 37 insertions(+), 15 deletions(-) diff --git a/frontend/appflowy_flutter/lib/plugins/ai_chat/application/chat_bloc.dart b/frontend/appflowy_flutter/lib/plugins/ai_chat/application/chat_bloc.dart index 63e13aa4321e4..05846cef927fe 100644 --- a/frontend/appflowy_flutter/lib/plugins/ai_chat/application/chat_bloc.dart +++ b/frontend/appflowy_flutter/lib/plugins/ai_chat/application/chat_bloc.dart @@ -244,7 +244,10 @@ class ChatBloc extends Bloc { }, didReceiveChatSettings: (settings) { emit( - state.copyWith(selectedSourceIds: settings.ragIds), + state.copyWith( + selectedSourceIds: settings.ragIds.ragIds, + onlyUseSelectedSources: settings.ragOnly, + ), ); }, updateSelectedSources: (selectedSourcesIds) async { diff --git a/frontend/appflowy_flutter/lib/plugins/ai_chat/presentation/chat_input/select_sources_menu.dart b/frontend/appflowy_flutter/lib/plugins/ai_chat/presentation/chat_input/select_sources_menu.dart index a1d221ae2280d..7dae3f5fd8c97 100644 --- a/frontend/appflowy_flutter/lib/plugins/ai_chat/presentation/chat_input/select_sources_menu.dart +++ b/frontend/appflowy_flutter/lib/plugins/ai_chat/presentation/chat_input/select_sources_menu.dart @@ -43,9 +43,10 @@ class _PromptInputDesktopSelectSourcesButtonState void initState() { super.initState(); WidgetsBinding.instance.addPostFrameCallback((_) { - cubit.updateSelectedSources( - context.read().state.selectedSourceIds, - ); + final chatBlocState = context.read().state; + cubit + ..updateSelectedSources(chatBlocState.selectedSourceIds) + ..updateOnlyUseSelectedSources(chatBlocState.onlyUseSelectedSources); }); } diff --git a/frontend/resources/translations/en.json b/frontend/resources/translations/en.json index 2eb651c99782f..1562b8da8890c 100644 --- a/frontend/resources/translations/en.json +++ b/frontend/resources/translations/en.json @@ -213,7 +213,7 @@ "addToPageTitle": "Add message to...", "addToNewPage": "Add to a new page", "addToNewPageName": "Messages extracted from \"{}\"", - "onlyUseRags": "Only use selected sources to generate response" + "onlyUseRags": "Selected sources only" }, "trash": { "text": "Trash", diff --git a/frontend/rust-lib/flowy-ai/src/ai_manager.rs b/frontend/rust-lib/flowy-ai/src/ai_manager.rs index b094342565ad0..cc8452f3793b4 100644 --- a/frontend/rust-lib/flowy-ai/src/ai_manager.rs +++ b/frontend/rust-lib/flowy-ai/src/ai_manager.rs @@ -345,12 +345,12 @@ impl AIManager { Ok(()) } - pub async fn get_rag_ids(&self, chat_id: &str) -> FlowyResult> { + pub async fn get_chat_settings(&self, chat_id: &str) -> FlowyResult { if let Some(settings) = self .store_preferences .get_object::(&setting_store_key(chat_id)) { - return Ok(settings.rag_ids); + return Ok(settings.into()); } let settings = refresh_chat_setting( @@ -360,7 +360,7 @@ impl AIManager { chat_id, ) .await?; - Ok(settings.rag_ids) + Ok(settings.into()) } pub async fn update_settings( @@ -448,9 +448,7 @@ async fn refresh_chat_setting( } chat_notification_builder(chat_id, ChatNotification::DidUpdateChatSettings) - .payload(ChatSettingsPB { - rag_ids: settings.rag_ids.clone(), - }) + .payload(ChatSettingsPB::from(settings.clone())) .send(); Ok(settings) diff --git a/frontend/rust-lib/flowy-ai/src/entities.rs b/frontend/rust-lib/flowy-ai/src/entities.rs index 353323aed7d11..6cff39da5e7e9 100644 --- a/frontend/rust-lib/flowy-ai/src/entities.rs +++ b/frontend/rust-lib/flowy-ai/src/entities.rs @@ -4,7 +4,8 @@ use std::collections::HashMap; use crate::local_ai::local_llm_resource::PendingResource; use flowy_ai_pub::cloud::{ - ChatMessage, LLMModel, RelatedQuestion, RepeatedChatMessage, RepeatedRelatedQuestion, + ChatMessage, ChatSettings, LLMModel, RelatedQuestion, RepeatedChatMessage, + RepeatedRelatedQuestion, }; use flowy_derive::{ProtoBuf, ProtoBuf_Enum}; use lib_infra::validator_fn::required_not_empty_str; @@ -542,7 +543,27 @@ pub struct CreateChatContextPB { #[derive(Default, ProtoBuf, Clone, Debug)] pub struct ChatSettingsPB { #[pb(index = 1)] - pub rag_ids: Vec, + pub rag_ids: RepeatedRagId, + + #[pb(index = 2)] + pub rag_only: bool, +} + +impl From for ChatSettingsPB { + fn from(value: ChatSettings) -> Self { + let rag_ids = RepeatedRagId { + rag_ids: value.rag_ids.clone(), + }; + + let rag_only = value + .metadata + .as_object() + .and_then(|map| map.get("rag_only")) + .and_then(|value| value.as_bool()) + .unwrap_or_default(); + + Self { rag_ids, rag_only } + } } #[derive(Default, ProtoBuf, Clone, Debug, Validate)] diff --git a/frontend/rust-lib/flowy-ai/src/event_handler.rs b/frontend/rust-lib/flowy-ai/src/event_handler.rs index 82ebf2939e967..212157640342f 100644 --- a/frontend/rust-lib/flowy-ai/src/event_handler.rs +++ b/frontend/rust-lib/flowy-ai/src/event_handler.rs @@ -446,8 +446,7 @@ pub(crate) async fn get_chat_settings_handler( ) -> DataResult { let chat_id = data.try_into_inner()?.value; let ai_manager = upgrade_ai_manager(ai_manager)?; - let rag_ids = ai_manager.get_rag_ids(&chat_id).await?; - let pb = ChatSettingsPB { rag_ids }; + let pb = ai_manager.get_chat_settings(&chat_id).await?; data_result_ok(pb) }