Implement engine consumer and conversation api client for tab organization feature

- GetSuggestedTopics, DedupeTopics, GetFocusTabs are implemented to get
  topics and classified tabs result from the server.

Author: yrliou

components/ai_chat/core/browser/engine/conversation_api_client.cc

@@ -200,6 +200,10 @@ base::Value::List ConversationEventsToList(
            {ConversationEventType::RequestSuggestedActions,
             "requestSuggestedActions"},
            {ConversationEventType::SuggestedActions, "suggestedActions"},
+           {ConversationEventType::GetSuggestedTopicsForFocusTabs,
+            "suggestFocusTopics"},
+           {ConversationEventType::DedupeTopics, "dedupeFocusTopics"},
+           {ConversationEventType::GetFocusTabsForTopic, "classifyTabs"},
            {ConversationEventType::UploadImage, "uploadImage"}});
 
   base::Value::List events;
@@ -217,6 +221,11 @@ base::Value::List ConversationEventsToList(
     event_dict.Set("type", type_it->second);
 
     event_dict.Set("content", event.content);
+
+    if (event.type == ConversationEventType::GetFocusTabsForTopic) {
+      event_dict.Set("topic", event.topic);
+    }
+
     events.Append(std::move(event_dict));
   }
   return events;

components/ai_chat/core/browser/engine/conversation_api_client.h

@@ -64,6 +64,9 @@ class ConversationAPIClient {
     RequestRewrite,
     SuggestedActions,
     UploadImage,
+    GetSuggestedTopicsForFocusTabs,
+    DedupeTopics,
+    GetFocusTabsForTopic,
     // TODO(petemill):
     // - Search in-progress?
     // - Sources?
@@ -76,6 +79,7 @@ class ConversationAPIClient {
     mojom::CharacterType role;
     ConversationEventType type;
     std::string content;
+    std::string topic;  // Used in GetFocusTabsForTopic event.
   };
 
   ConversationAPIClient(

components/ai_chat/core/browser/engine/conversation_api_client_unittest.cc

@@ -64,6 +64,44 @@ using Ticket = api_request_helper::APIRequestHelper::Ticket;
 
 namespace ai_chat {
 
+namespace {
+
+const std::pair<std::vector<ConversationAPIClient::ConversationEvent>,
+                std::string>&
+GetMockEventsAndExpectedEventsBody() {
+  static base::NoDestructor<std::pair<
+      std::vector<ConversationAPIClient::ConversationEvent>, std::string>>
+      mock_events_and_expected_events_body{
+          std::vector<ConversationAPIClient::ConversationEvent>{
+              {mojom::CharacterType::HUMAN, ConversationAPIClient::PageText,
+               "This is a page about The Mandalorian."},
+              {mojom::CharacterType::HUMAN, ConversationAPIClient::PageExcerpt,
+               "The Mandalorian"},
+              {mojom::CharacterType::HUMAN, ConversationAPIClient::ChatMessage,
+               "Est-ce lié à une série plus large?"},
+              {mojom::CharacterType::HUMAN,
+               ConversationAPIClient::GetSuggestedTopicsForFocusTabs,
+               "GetSuggestedTopicsForFocusTabs"},
+              {mojom::CharacterType::HUMAN, ConversationAPIClient::DedupeTopics,
+               "DedupeTopics"},
+              {mojom::CharacterType::HUMAN,
+               ConversationAPIClient::GetFocusTabsForTopic,
+               "GetFocusTabsForTopics", "C++"},
+          },
+          R"([
+      {"role": "user", "type": "pageText", "content": "This is a page about The Mandalorian."},
+      {"role": "user", "type": "pageExcerpt", "content": "The Mandalorian"},
+      {"role": "user", "type": "chatMessage", "content": "Est-ce lié à une série plus large?"},
+      {"role": "user", "type": "suggestFocusTopics", "content": "GetSuggestedTopicsForFocusTabs"},
+      {"role": "user", "type": "dedupeFocusTopics", "content": "DedupeTopics"},
+      {"role": "user", "type": "classifyTabs", "content": "GetFocusTabsForTopics", "topic": "C++"}
+    ])"};
+
+  return *mock_events_and_expected_events_body;
+}
+
+}  // namespace
+
 using ConversationEvent = ConversationAPIClient::ConversationEvent;
 
 class MockCallbacks {
@@ -198,18 +236,10 @@ TEST_F(ConversationAPIUnitTest, PerformRequest_PremiumHeaders) {
   //  - ConversationEvent is correctly formatted into JSON
   //  - completion response is parsed and passed through to the callbacks
   std::string expected_crediential = "unit_test_credential";
-  std::vector<ConversationAPIClient::ConversationEvent> events = {
-      {mojom::CharacterType::HUMAN, ConversationAPIClient::PageText,
-       "This is a page about The Mandalorian."},
-      {mojom::CharacterType::HUMAN, ConversationAPIClient::PageExcerpt,
-       "The Mandalorian"},
-      {mojom::CharacterType::HUMAN, ConversationAPIClient::ChatMessage,
-       "Est-ce lié à une série plus large?"}};
-  std::string expected_events_body = R"([
-    {"role": "user", "type": "pageText", "content": "This is a page about The Mandalorian."},
-    {"role": "user", "type": "pageExcerpt", "content": "The Mandalorian"},
-    {"role": "user", "type": "chatMessage", "content": "Est-ce lié à une série plus large?"}
-  ])";
+  const std::vector<ConversationAPIClient::ConversationEvent>& events =
+      GetMockEventsAndExpectedEventsBody().first;
+  const std::string& expected_events_body =
+      GetMockEventsAndExpectedEventsBody().second;
   std::string expected_system_language = "en_KY";
   const brave_l10n::test::ScopedDefaultLocale scoped_default_locale(
       expected_system_language);
@@ -406,18 +436,10 @@ TEST_F(ConversationAPIUnitTest, PerformRequest_NonPremium) {
   //  - ConversationEvent is correctly formatted into JSON
   //  - completion response is parsed and passed through to the callbacks
   std::string expected_crediential = "unit_test_credential";
-  std::vector<ConversationAPIClient::ConversationEvent> events = {
-      {mojom::CharacterType::HUMAN, ConversationAPIClient::PageText,
-       "This is a page about The Mandalorian."},
-      {mojom::CharacterType::HUMAN, ConversationAPIClient::PageExcerpt,
-       "The Mandalorian"},
-      {mojom::CharacterType::HUMAN, ConversationAPIClient::ChatMessage,
-       "Est-ce lié à une série plus large?"}};
-  std::string expected_events_body = R"([
-    {"role": "user", "type": "pageText", "content": "This is a page about The Mandalorian."},
-    {"role": "user", "type": "pageExcerpt", "content": "The Mandalorian"},
-    {"role": "user", "type": "chatMessage", "content": "Est-ce lié à une série plus large?"}
-  ])";
+  const std::vector<ConversationAPIClient::ConversationEvent>& events =
+      GetMockEventsAndExpectedEventsBody().first;
+  const std::string& expected_events_body =
+      GetMockEventsAndExpectedEventsBody().second;
   std::string expected_system_language = "en_KY";
   const brave_l10n::test::ScopedDefaultLocale scoped_default_locale(
       expected_system_language);

components/ai_chat/core/browser/engine/engine_consumer.cc

@@ -49,4 +49,17 @@ bool EngineConsumer::CanPerformCompletionRequest(
   return true;
 }
 
+void EngineConsumer::GetSuggestedTopics(const std::vector<Tab>& tabs,
+                                        GetSuggestedTopicsCallback callback) {
+  NOTIMPLEMENTED();
+  std::move(callback).Run(base::unexpected(mojom::APIError::InternalError));
+}
+
+void EngineConsumer::GetFocusTabs(const std::vector<Tab>& tabs,
+                                  const std::string& topic,
+                                  GetFocusTabsCallback callback) {
+  NOTIMPLEMENTED();
+  std::move(callback).Run(base::unexpected(mojom::APIError::InternalError));
+}
+
 }  // namespace ai_chat

components/ai_chat/core/browser/engine/engine_consumer.h

@@ -16,6 +16,7 @@
 #include "base/functional/callback_forward.h"
 #include "base/types/expected.h"
 #include "brave/components/ai_chat/core/browser/engine/remote_completion_client.h"
+#include "brave/components/ai_chat/core/browser/types.h"
 #include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom-forward.h"
 
 namespace ai_chat {
@@ -43,6 +44,11 @@ class EngineConsumer {
 
   using ConversationHistory = std::vector<mojom::ConversationTurnPtr>;
 
+  using GetSuggestedTopicsCallback = base::OnceCallback<void(
+      base::expected<std::vector<std::string>, mojom::APIError>)>;
+  using GetFocusTabsCallback = base::OnceCallback<void(
+      base::expected<std::vector<std::string>, mojom::APIError>)>;
+
   static std::string GetPromptForEntry(const mojom::ConversationTurnPtr& entry);
 
   EngineConsumer();
@@ -86,6 +92,15 @@ class EngineConsumer {
 
   virtual void UpdateModelOptions(const mojom::ModelOptions& options) = 0;
 
+  // Given a list of tabs, return a list of suggested topics from the server.
+  virtual void GetSuggestedTopics(const std::vector<Tab>& tabs,
+                                  GetSuggestedTopicsCallback callback);
+  // Given a list of tabs and a specific topic, return a list of tabs to be
+  // focused on from the server.
+  virtual void GetFocusTabs(const std::vector<Tab>& tabs,
+                            const std::string& topic,
+                            GetFocusTabsCallback callback);
+
   void SetMaxAssociatedContentLengthForTesting(
       uint32_t max_associated_content_length) {
     max_associated_content_length_ = max_associated_content_length;

components/ai_chat/core/browser/engine/engine_consumer_conversation_api.cc

@@ -5,24 +5,31 @@
 
 #include "brave/components/ai_chat/core/browser/engine/engine_consumer_conversation_api.h"
 
+#include <algorithm>
 #include <optional>
 #include <string>
 #include <string_view>
 #include <type_traits>
 #include <vector>
 
+#include "base/barrier_callback.h"
 #include "base/check.h"
 #include "base/functional/bind.h"
 #include "base/functional/callback.h"
 #include "base/functional/callback_helpers.h"
+#include "base/json/json_reader.h"
+#include "base/json/json_writer.h"
 #include "base/memory/scoped_refptr.h"
 #include "base/memory/weak_ptr.h"
 #include "base/numerics/clamped_math.h"
 #include "base/strings/string_split.h"
+#include "base/strings/string_util.h"
 #include "base/time/time.h"
 #include "base/types/expected.h"
+#include "base/values.h"
 #include "brave/components/ai_chat/core/common/mojom/ai_chat.mojom.h"
 #include "services/network/public/cpp/shared_url_loader_factory.h"
+#include "third_party/re2/src/re2/re2.h"
 
 namespace ai_chat {
 
@@ -31,6 +38,9 @@ namespace {
 using ConversationEvent = ConversationAPIClient::ConversationEvent;
 using ConversationEventType = ConversationAPIClient::ConversationEventType;
 
+constexpr size_t kChunkSize = 75;
+constexpr char kStrArrPattern[] = R"((\[.+?\]))";
+
 }  // namespace
 
 EngineConsumerConversationAPI::EngineConsumerConversationAPI(
@@ -49,6 +59,66 @@ void EngineConsumerConversationAPI::ClearAllQueries() {
   api_->ClearAllQueries();
 }
 
+// static
+base::expected<std::vector<std::string>, mojom::APIError>
+EngineConsumerConversationAPI::GetStrArrFromTabOrganizationResponses(
+    std::vector<EngineConsumer::GenerationResult>& results) {
+  // Rust implementation of JSON reader is required to parse the response
+  // safely. This function currently is only called on Desktop which uses a
+  // rust JSON reader. Chromium is in progress of making all platforms use the
+  // rust JSON reader and updating the rule of 2 documentation to explicitly
+  // point out that JSON parser in base is considered safe.
+  if (!base::JSONReader::UsingRust()) {
+    return base::unexpected(mojom::APIError::InternalError);
+  }
+
+  std::vector<std::string> str_arr;
+  mojom::APIError error = mojom::APIError::None;
+  for (auto& result : results) {
+    // Fail the operation if server returns an error, such as rate limiting.
+    // On the other hand, ignore the result which cannot be parsed as expected.
+    if (!result.has_value()) {
+      error = result.error();
+      break;
+    }
+
+    // Skip empty results.
+    if (result->empty()) {
+      continue;
+    }
+
+    // Remove newline characters from the result.
+    base::ReplaceChars(*result, "\n", "", &result.value());
+    std::string strArr = "";
+    if (!RE2::PartialMatch(*result, kStrArrPattern, &strArr)) {
+      continue;
+    }
+    auto value = base::JSONReader::Read(strArr, base::JSON_PARSE_RFC);
+    if (!value) {
+      continue;
+    }
+
+    auto* list = value->GetIfList();
+    if (!list) {
+      continue;
+    }
+
+    for (const auto& item : *list) {
+      auto* str = item.GetIfString();
+      if (!str || str->empty()) {
+        continue;
+      }
+      str_arr.push_back(*str);
+    }
+  }
+
+  if (error != mojom::APIError::None) {
+    return base::unexpected(error);
+  }
+
+  return str_arr;
+}
+
 void EngineConsumerConversationAPI::GenerateRewriteSuggestion(
     std::string text,
     const std::string& question,
@@ -183,4 +253,104 @@ EngineConsumerConversationAPI::GetAssociatedContentConversationEvent(
   return event;
 }
 
+void EngineConsumerConversationAPI::DedupeTopics(
+    base::expected<std::vector<std::string>, mojom::APIError> topics_result,
+    GetSuggestedTopicsCallback callback) {
+  if (!topics_result.has_value() || topics_result->empty()) {
+    std::move(callback).Run(topics_result);
+    return;
+  }
+
+  base::Value::List topic_list;
+  for (const auto& topic : *topics_result) {
+    topic_list.Append(topic);
+  }
+  std::vector<ConversationEvent> conversation;
+  conversation.push_back({mojom::CharacterType::HUMAN,
+                          ConversationEventType::DedupeTopics,
+                          base::WriteJson(topic_list).value_or(std::string())});
+  api_->PerformRequest(
+      std::move(conversation), "" /* selected_language */,
+      base::NullCallback() /* data_received_callback */,
+      base::BindOnce(
+          [](GetSuggestedTopicsCallback callback,
+             EngineConsumer::GenerationResult result) {
+            // Return deduped topics from the response.
+            std::vector<EngineConsumer::GenerationResult> results = {result};
+            std::move(callback).Run(
+                EngineConsumerConversationAPI::
+                    GetStrArrFromTabOrganizationResponses(results));
+          },
+          std::move(callback)));
+}
+
+void EngineConsumerConversationAPI::ProcessTabChunks(
+    const std::vector<Tab>& tabs,
+    ConversationEventType event_type,
+    base::OnceCallback<void(std::vector<GenerationResult>)> merge_callback,
+    const std::string& topic) {
+  CHECK(event_type == ConversationEventType::GetSuggestedTopicsForFocusTabs ||
+        event_type == ConversationEventType::GetFocusTabsForTopic);
+
+  // Split tab into chunks of 75
+  size_t num_chunks = (tabs.size() + kChunkSize - 1) / kChunkSize;
+  const auto barrier_callback = base::BarrierCallback<GenerationResult>(
+      num_chunks, std::move(merge_callback));
+
+  for (size_t chunk = 0; chunk < num_chunks; ++chunk) {
+    base::Value::List tab_value_list;
+    for (size_t i = chunk * kChunkSize;
+         i < std::min((chunk + 1) * kChunkSize, tabs.size()); ++i) {
+      tab_value_list.Append(base::Value::Dict()
+                                .Set("id", tabs[i].id)
+                                .Set("title", tabs[i].title)
+                                .Set("url", tabs[i].origin.Serialize()));
+    }
+
+    std::vector<ConversationEvent> conversation;
+    conversation.push_back(
+        {mojom::CharacterType::HUMAN, event_type,
+         base::WriteJson(tab_value_list).value_or(std::string()), topic});
+
+    api_->PerformRequest(std::move(conversation), "" /* selected_language */,
+                         base::NullCallback() /* data_received_callback */,
+                         barrier_callback /* data_completed_callback */);
+  }
+}
+
+void EngineConsumerConversationAPI::MergeSuggestTopicsResults(
+    GetSuggestedTopicsCallback callback,
+    std::vector<GenerationResult> results) {
+  // Merge the result and send another request to dedupe topics.
+  DedupeTopics(GetStrArrFromTabOrganizationResponses(results),
+               std::move(callback));
+}
+
+void EngineConsumerConversationAPI::GetSuggestedTopics(
+    const std::vector<Tab>& tabs,
+    GetSuggestedTopicsCallback callback) {
+  ProcessTabChunks(
+      tabs, ConversationEventType::GetSuggestedTopicsForFocusTabs,
+      base::BindOnce(&EngineConsumerConversationAPI::MergeSuggestTopicsResults,
+                     weak_ptr_factory_.GetWeakPtr(), std::move(callback)),
+      "" /* topic */);
+}
+
+void EngineConsumerConversationAPI::GetFocusTabs(
+    const std::vector<Tab>& tabs,
+    const std::string& topic,
+    EngineConsumer::GetFocusTabsCallback callback) {
+  ProcessTabChunks(tabs, ConversationEventType::GetFocusTabsForTopic,
+                   base::BindOnce(
+                       [&](EngineConsumer::GetFocusTabsCallback callback,
+                           std::vector<GenerationResult> results) {
+                         // Merge the results and call callback with tab IDs or
+                         // error.
+                         std::move(callback).Run(
+                             GetStrArrFromTabOrganizationResponses(results));
+                       },
+                       std::move(callback)),
+                   topic);
+}
+
 }  // namespace ai_chat