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] 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) }