From a8143f6cb2e358c23affdb33bc2babded0931011 Mon Sep 17 00:00:00 2001 From: Pete Miller Date: Fri, 30 Aug 2024 16:23:23 -0700 Subject: [PATCH] refactor search query fetching --- .../content/browser/ai_chat_tab_helper.cc | 4 + .../core/browser/associated_content_driver.cc | 123 +++++++++++++++++- .../core/browser/associated_content_driver.h | 24 +++- .../core/browser/conversation_handler.cc | 83 +++++++++++- .../core/browser/conversation_handler.h | 16 +++ 5 files changed, 239 insertions(+), 11 deletions(-) diff --git a/components/ai_chat/content/browser/ai_chat_tab_helper.cc b/components/ai_chat/content/browser/ai_chat_tab_helper.cc index cb045bcd40dc..b784c2ad2e5f 100644 --- a/components/ai_chat/content/browser/ai_chat_tab_helper.cc +++ b/components/ai_chat/content/browser/ai_chat_tab_helper.cc @@ -17,6 +17,7 @@ #include "base/strings/string_util.h" #include "brave/components/ai_chat/content/browser/page_content_fetcher.h" #include "brave/components/ai_chat/content/browser/pdf_utils.h" +#include "brave/components/ai_chat/core/browser/associated_content_driver.h" #include "brave/components/ai_chat/core/browser/constants.h" #include "brave/components/ai_chat/core/browser/utils.h" #include "brave/components/ai_chat/core/common/mojom/page_content_extractor.mojom.h" @@ -88,6 +89,9 @@ AIChatTabHelper::AIChatTabHelper( ExtractPrintPreviewContentFunction extract_print_preview_content_function) : content::WebContentsObserver(web_contents), content::WebContentsUserData(*web_contents), + AssociatedContentDriver(web_contents->GetBrowserContext() + ->GetDefaultStoragePartition() + ->GetURLLoaderFactoryForBrowserProcess()), extract_print_preview_content_function_( extract_print_preview_content_function) { CHECK(extract_print_preview_content_function); diff --git a/components/ai_chat/core/browser/associated_content_driver.cc b/components/ai_chat/core/browser/associated_content_driver.cc index d2eb6636377c..faec5e53ff98 100644 --- a/components/ai_chat/core/browser/associated_content_driver.cc +++ b/components/ai_chat/core/browser/associated_content_driver.cc @@ -17,13 +17,46 @@ #include "base/one_shot_event.h" #include "base/ranges/algorithm.h" #include "base/strings/string_util.h" +#include "brave/brave_domains/service_domains.h" +#include "brave/components/ai_chat/core/browser/brave_search_responses.h" #include "brave/components/ai_chat/core/browser/constants.h" #include "brave/components/ai_chat/core/browser/conversation_handler.h" - +#include "brave/components/ai_chat/core/browser/utils.h" +#include "net/base/url_util.h" +#include "net/traffic_annotation/network_traffic_annotation.h" namespace ai_chat { -AssociatedContentDriver::AssociatedContentDriver() - : on_page_text_fetch_complete_(std::make_unique()) {} +namespace { + +net::NetworkTrafficAnnotationTag +GetSearchQuerySummaryNetworkTrafficAnnotationTag() { + return net::DefineNetworkTrafficAnnotation( + "ai_chat_associated_content_driver", + R"( + semantics { + sender: "Brave Leo AI Chat" + description: + "This sender is used to get search query summary from Brave search." + trigger: + "Triggered by uses of Brave Leo AI Chat on Brave Search SERP." + data: + "User's search query and the corresponding summary." + destination: WEBSITE + } + policy { + cookies_allowed: NO + policy_exception_justification: + "Not implemented." + } + )"); +} + +} // namespace + +AssociatedContentDriver::AssociatedContentDriver( + scoped_refptr url_loader_factory) + : url_loader_factory_(url_loader_factory), + on_page_text_fetch_complete_(std::make_unique()) {} AssociatedContentDriver::~AssociatedContentDriver() { for (auto& conversation : associated_conversations_) { @@ -139,6 +172,89 @@ bool AssociatedContentDriver::GetCachedIsVideo() { return is_video_; } +void AssociatedContentDriver::GetStagedEntriesFromContent( + ConversationHandler::GetStagedEntriesCallback callback) { + // At the moment we only know about staged entries from: + // - Brave Search results page + if (!IsBraveSearchSERP(GetPageURL())) { + std::move(callback).Run(std::nullopt); + return; + } + GetSearchSummarizerKey( + base::BindOnce(&AssociatedContentDriver::OnSearchSummarizerKeyFetched, + weak_ptr_factory_.GetWeakPtr(), std::move(callback), + current_navigation_id_)); +} + +void AssociatedContentDriver::OnSearchSummarizerKeyFetched( + ConversationHandler::GetStagedEntriesCallback callback, + int64_t navigation_id, + const std::optional& key) { + if (!key || key->empty() || navigation_id != current_navigation_id_) { + std::move(callback).Run(std::nullopt); + return; + } + + if (!api_request_helper_) { + api_request_helper_ = + std::make_unique( + GetSearchQuerySummaryNetworkTrafficAnnotationTag(), + url_loader_factory_); + } + + // https://search.brave.com/api/chatllm/raw_data?key= + GURL base_url( + base::StrCat({url::kHttpsScheme, url::kStandardSchemeSeparator, + brave_domains::GetServicesDomain(kBraveSearchURLPrefix), + "/api/chatllm/raw_data"})); + CHECK(base_url.is_valid()); + GURL url = net::AppendQueryParameter(base_url, "key", *key); + + api_request_helper_->Request( + "GET", url, "", "application/json", + base::BindOnce(&AssociatedContentDriver::OnSearchQuerySummaryFetched, + weak_ptr_factory_.GetWeakPtr(), std::move(callback), + navigation_id), + {}, {}); +} + +void AssociatedContentDriver::OnSearchQuerySummaryFetched( + ConversationHandler::GetStagedEntriesCallback callback, + int64_t navigation_id, + api_request_helper::APIRequestResult result) { + if (!result.Is2XXResponseCode() || navigation_id != current_navigation_id_) { + std::move(callback).Run(std::nullopt); + return; + } + + auto search_query_summary = + ParseSearchQuerySummaryResponse(result.value_body()); + if (!search_query_summary) { + std::move(callback).Run(std::nullopt); + return; + } + + std::move(callback).Run(search_query_summary); +} + +// static +std::optional +AssociatedContentDriver::ParseSearchQuerySummaryResponse( + const base::Value& value) { + auto search_query_response = + brave_search_responses::QuerySummaryResponse::FromValue(value); + if (!search_query_response || search_query_response->conversation.empty()) { + return std::nullopt; + } + + const auto& query_summary = search_query_response->conversation[0]; + if (query_summary.answer.empty()) { + return std::nullopt; + } + + return SearchQuerySummary(query_summary.query, query_summary.answer[0].text); +} + void AssociatedContentDriver::OnFaviconImageDataChanged() { for (auto& conversation : associated_conversations_) { conversation->OnFaviconImageDataChanged(); @@ -161,6 +277,7 @@ void AssociatedContentDriver::OnNewPage(int64_t navigation_id) { cached_text_content_.clear(); content_invalidation_token_.clear(); is_video_ = false; + api_request_helper_.reset(); ConversationHandler::AssociatedContentDelegate::OnNewPage(navigation_id); } diff --git a/components/ai_chat/core/browser/associated_content_driver.h b/components/ai_chat/core/browser/associated_content_driver.h index d9a247930c04..c0a07dbb28ee 100644 --- a/components/ai_chat/core/browser/associated_content_driver.h +++ b/components/ai_chat/core/browser/associated_content_driver.h @@ -18,6 +18,8 @@ #include "brave/components/ai_chat/core/browser/ai_chat_service.h" #include "brave/components/ai_chat/core/browser/conversation_handler.h" #include "brave/components/ai_chat/core/browser/model_service.h" +#include "brave/components/ai_chat/core/common/mojom/page_content_extractor.mojom.h" +#include "brave/components/api_request_helper/api_request_helper.h" class PrefService; @@ -35,7 +37,8 @@ class AssociatedContentDriver virtual void OnAssociatedContentNavigated(int new_navigation_id) {} }; - AssociatedContentDriver(); + explicit AssociatedContentDriver( + scoped_refptr url_loader_factory); ~AssociatedContentDriver() override; AssociatedContentDriver(const AssociatedContentDriver&) = delete; @@ -55,6 +58,8 @@ class AssociatedContentDriver ConversationHandler::GetPageContentCallback callback) override; std::string_view GetCachedTextContent() override; bool GetCachedIsVideo() override; + void GetStagedEntriesFromContent( + ConversationHandler::GetStagedEntriesCallback callback) override; // // Implementer should use alternative method of page content fetching // void PrintPreviewFallback(ConversationHandler::GetPageContentCallback // callback) override; @@ -66,6 +71,9 @@ class AssociatedContentDriver protected: virtual GURL GetPageURL() const = 0; virtual std::u16string GetPageTitle() const = 0; + // Get summarizer-key meta tag content from Brave Search SERP if exists. + virtual void GetSearchSummarizerKey( + mojom::PageContentExtractor::GetSearchSummarizerKeyCallback callback) = 0; // Implementer should fetch content from the "page" associated with this // conversation. @@ -120,12 +128,26 @@ class AssociatedContentDriver ConversationHandler::GetPageContentCallback callback, int64_t navigation_id); + void OnSearchSummarizerKeyFetched( + ConversationHandler::GetStagedEntriesCallback callback, + int64_t navigation_id, + const std::optional& key); + void OnSearchQuerySummaryFetched( + ConversationHandler::GetStagedEntriesCallback callback, + int64_t navigation_id, + api_request_helper::APIRequestResult result); + static std::optional ParseSearchQuerySummaryResponse( + const base::Value& value); + raw_ptr pref_service_; raw_ptr ai_chat_metrics_; std::unique_ptr credential_manager_; scoped_refptr url_loader_factory_; + // Used for fetching search query summary. + std::unique_ptr api_request_helper_; + base::ObserverList observers_; std::unique_ptr on_page_text_fetch_complete_; bool is_page_text_fetch_in_progress_ = false; diff --git a/components/ai_chat/core/browser/conversation_handler.cc b/components/ai_chat/core/browser/conversation_handler.cc index e73992e4dd7e..aae91d97a040 100644 --- a/components/ai_chat/core/browser/conversation_handler.cc +++ b/components/ai_chat/core/browser/conversation_handler.cc @@ -26,6 +26,7 @@ #include "brave/components/ai_chat/core/browser/constants.h" #include "brave/components/ai_chat/core/browser/leo_local_models_updater.h" #include "brave/components/ai_chat/core/browser/model_service.h" +#include "brave/components/ai_chat/core/browser/types.h" #include "brave/components/ai_chat/core/common/features.h" #include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom-forward.h" #include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom.h" @@ -116,6 +117,11 @@ void ConversationHandler::AssociatedContentDelegate::OnNewPage( } } +void ConversationHandler::AssociatedContentDelegate:: + GetStagedEntriesFromContent(GetStagedEntriesCallback callback) { + std::move(callback).Run(std::nullopt); +} + void ConversationHandler::AssociatedContentDelegate:: GetTopSimilarityWithPromptTilContextLimit( const std::string& prompt, @@ -230,6 +236,7 @@ void ConversationHandler::Bind( if (!pending_conversation_entry_.is_null()) { OnHistoryUpdate(); } + MaybeFetchOrClearContentStagedConversation(); } bool ConversationHandler::IsAnyClientConnected() { @@ -343,6 +350,7 @@ void ConversationHandler::SetAssociatedContentDelegate( should_send_page_contents_ = true; MaybeSeedOrClearSuggestions(); + MaybeFetchOrClearContentStagedConversation(); OnAssociatedContentInfoChanged(); } @@ -472,7 +480,7 @@ void ConversationHandler::SubmitHumanConversationEntry( mojom::ConversationTurnPtr turn = mojom::ConversationTurn::New( CharacterType::HUMAN, mojom::ActionType::UNSPECIFIED, mojom::ConversationTurnVisibility::VISIBLE, input, std::nullopt, - std::nullopt, base::Time::Now(), std::nullopt); + std::nullopt, base::Time::Now(), std::nullopt, false); SubmitHumanConversationEntry(std::move(turn)); } @@ -619,7 +627,7 @@ void ConversationHandler::ModifyConversation(uint32_t turn_index, auto edited_turn = mojom::ConversationTurn::New( turn->character_type, turn->action_type, turn->visibility, trimmed_input, std::nullopt /* selected_text */, std::move(events), - base::Time::Now(), std::nullopt /* edits */); + base::Time::Now(), std::nullopt /* edits */, false); edited_turn->events->at(*completion_event_index) ->get_completion_event() ->completion = trimmed_input; @@ -650,7 +658,8 @@ void ConversationHandler::ModifyConversation(uint32_t turn_index, auto edited_turn = mojom::ConversationTurn::New( turn->character_type, turn->action_type, turn->visibility, sanitized_input, std::nullopt /* selected_text */, - std::nullopt /* events */, base::Time::Now(), std::nullopt /* edits */); + std::nullopt /* events */, base::Time::Now(), std::nullopt /* edits */, + false); if (!turn->edits) { turn->edits.emplace(); } @@ -673,7 +682,7 @@ void ConversationHandler::SubmitSummarizationRequest() { CharacterType::HUMAN, mojom::ActionType::SUMMARIZE_PAGE, mojom::ConversationTurnVisibility::VISIBLE, l10n_util::GetStringUTF8(IDS_CHAT_UI_SUMMARIZE_PAGE), std::nullopt, - std::nullopt, base::Time::Now(), std::nullopt); + std::nullopt, base::Time::Now(), std::nullopt, false); SubmitHumanConversationEntry(std::move(turn)); } @@ -751,6 +760,7 @@ void ConversationHandler::SetShouldSendPageContents(bool should_send) { OnAssociatedContentInfoChanged(); MaybeSeedOrClearSuggestions(); + MaybeFetchOrClearContentStagedConversation(); } void ConversationHandler::RetryAPIRequest() { @@ -844,7 +854,7 @@ void ConversationHandler::SubmitSelectedTextWithQuestion( mojom::ConversationTurnPtr turn = mojom::ConversationTurn::New( CharacterType::HUMAN, action_type, mojom::ConversationTurnVisibility::VISIBLE, question, selected_text, - std::nullopt, base::Time::Now(), std::nullopt); + std::nullopt, base::Time::Now(), std::nullopt, false); SubmitHumanConversationEntry(std::move(turn)); } else { @@ -894,7 +904,7 @@ void ConversationHandler::AddSubmitSelectedTextError( mojom::ConversationTurnPtr turn = mojom::ConversationTurn::New( CharacterType::HUMAN, action_type, mojom::ConversationTurnVisibility::VISIBLE, question, selected_text, - std::nullopt, base::Time::Now(), std::nullopt); + std::nullopt, base::Time::Now(), std::nullopt, false); AddToConversationHistory(std::move(turn)); SetAPIError(error); } @@ -972,7 +982,7 @@ void ConversationHandler::UpdateOrCreateLastAssistantEntry( CharacterType::ASSISTANT, mojom::ActionType::RESPONSE, mojom::ConversationTurnVisibility::VISIBLE, "", std::nullopt, std::vector{}, base::Time::Now(), - std::nullopt); + std::nullopt, false); chat_history_.push_back(std::move(entry)); } @@ -1056,6 +1066,65 @@ void ConversationHandler::MaybeSeedOrClearSuggestions() { } } +void ConversationHandler::MaybeFetchOrClearContentStagedConversation() { + const bool can_check_for_staged_conversation = + ai_chat_service_->HasUserOptedIn() && IsContentAssociationPossible() && + should_send_page_contents_; + if (!can_check_for_staged_conversation) { + // Clear any staged conversation entries since user might have unassociated + // content with this conversation + // For now, we assume all staged conversations are 2 entries (question and + // answer). + if (chat_history_.size() != 2) { + return; + } + + const auto& last_turn = chat_history_.back(); + if (last_turn->from_brave_search_SERP) { + chat_history_.clear(); // Clear the staged query and summary. + OnHistoryUpdate(); + } + } + + // Try later when we get a connected client + if (!IsAnyClientConnected()) { + return; + } + + // Currently only have search query summary at the start of a conversation. + if (!chat_history_.empty()) { + return; + } + + associated_content_delegate_->GetStagedEntriesFromContent( + base::BindOnce(&ConversationHandler::OnGetStagedEntriesFromContent, + weak_ptr_factory_.GetWeakPtr())); +} + +void ConversationHandler::OnGetStagedEntriesFromContent( + const std::optional& search_query_summary) { + // Check if all requirements are still met. + if (!search_query_summary || !chat_history_.empty() || + !IsContentAssociationPossible() || !should_send_page_contents_) { + return; + } + + // Add the query & summary to the conversation history and call + // OnHistoryUpdate to update UI. + chat_history_.push_back(mojom::ConversationTurn::New( + CharacterType::HUMAN, mojom::ActionType::QUERY, + mojom::ConversationTurnVisibility::VISIBLE, search_query_summary->query, + std::nullopt, std::nullopt, base::Time::Now(), std::nullopt, true)); + std::vector events; + events.push_back(mojom::ConversationEntryEvent::NewCompletionEvent( + mojom::CompletionEvent::New(search_query_summary->summary))); + chat_history_.push_back(mojom::ConversationTurn::New( + CharacterType::ASSISTANT, mojom::ActionType::RESPONSE, + mojom::ConversationTurnVisibility::VISIBLE, search_query_summary->summary, + std::nullopt, std::move(events), base::Time::Now(), std::nullopt, true)); + OnHistoryUpdate(); +} + void ConversationHandler::GeneratePageContent(GetPageContentCallback callback) { VLOG(1) << __func__; DCHECK(should_send_page_contents_); diff --git a/components/ai_chat/core/browser/conversation_handler.h b/components/ai_chat/core/browser/conversation_handler.h index 2fbfc48ebf60..6c7301ee2f4f 100644 --- a/components/ai_chat/core/browser/conversation_handler.h +++ b/components/ai_chat/core/browser/conversation_handler.h @@ -19,6 +19,7 @@ #include "brave/components/ai_chat/core/browser/engine/engine_consumer.h" #include "brave/components/ai_chat/core/browser/model_service.h" #include "brave/components/ai_chat/core/browser/text_embedder.h" +#include "brave/components/ai_chat/core/browser/types.h" #include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom-forward.h" #include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom.h" #include "mojo/public/cpp/bindings/pending_receiver.h" @@ -53,6 +54,11 @@ class ConversationHandler : public mojom::ConversationHandler, using GeneratedTextCallback = base::RepeatingCallback; + // TODO(petemill): consider making SearchQuerySummary generic (StagedEntries) + // or a list of ConversationTurn objects. + using GetStagedEntriesCallback = base::OnceCallback& search_query_summary)>; + // Supplements a conversation with associated page content class AssociatedContentDelegate { public: @@ -77,6 +83,10 @@ class ConversationHandler : public mojom::ConversationHandler, // fetch for the content. virtual std::string_view GetCachedTextContent() = 0; virtual bool GetCachedIsVideo() = 0; + // Get summarizer-key meta tag content from Brave Search SERP if exists and + // use it to fetch search query and summary from Brave search chatllm + // endpoint. + virtual void GetStagedEntriesFromContent(GetStagedEntriesCallback callback); void GetTopSimilarityWithPromptTilContextLimit( const std::string& prompt, @@ -247,6 +257,12 @@ class ConversationHandler : public mojom::ConversationHandler, void UpdateOrCreateLastAssistantEntry(mojom::ConversationEntryEventPtr text); void MaybeSeedOrClearSuggestions(); + // Some associated content may provide some conversation that the user wants + // to continue, e.g. Brave Search. + void MaybeFetchOrClearContentStagedConversation(); + void OnGetStagedEntriesFromContent( + const std::optional& search_query_summary); + void GeneratePageContent(GetPageContentCallback callback); void SetPageContent(std::string contents_text, bool is_video,