Add entity and keyword highlighting. Change format

components/ai_chat/browser/ai_chat_tab_helper.cc

@@ -9,6 +9,7 @@
 #include <string>
 #include <utility>
 
+#include "absl/types/optional.h"
 #include "base/containers/contains.h"
 #include "base/containers/fixed_flat_set.h"
 #include "base/functional/bind.h"
@@ -219,8 +220,8 @@ void AIChatTabHelper::OnTabContentRetrieved(int64_t for_navigation_id,
   // Now that we have content, we can provide a summary on-demand. Add that to
   // suggested questions.
   // TODO(petemill): translation for this question
-  suggested_questions_.emplace_back(is_video_ ? "Summarize this video"
-                                              : "Summarize this page");
+  suggested_questions_.emplace_back(is_video_ ? "Summarize this video."
+                                              : "Summarize this page.");
   OnSuggestedQuestionsChanged();
   MaybeGenerateQuestions();
 }
@@ -450,11 +451,13 @@ std::string AIChatTabHelper::BuildLlama2Prompt(std::string user_message) {
     first_user_message = raw_first_user_message;
   }
 
+  std::string assistant_response_seed = "<assistant>{'response':";
+
   // If there's no conversation history, then we just send a (partial)
   // first sequence.
   if (conversation_history.size() == 0) {
     return BuildLlama2FirstSequence(system_message, first_user_message,
-                                    absl::nullopt, absl::nullopt);
+                                    absl::nullopt, assistant_response_seed);
   }
 
   // Use the first two messages to build the first sequence,
@@ -473,7 +476,8 @@ std::string AIChatTabHelper::BuildLlama2Prompt(std::string user_message) {
   }
 
   // Build the final subsequent exchange using the current turn.
-  prompt += BuildLlama2SubsequentSequence(user_message, absl::nullopt);
+  prompt +=
+      BuildLlama2SubsequentSequence(user_message, assistant_response_seed);
 
   // Trimming recommended by Meta
   // https://huggingface.co/meta-llama/Llama-2-13b-chat#intended-use
@@ -643,7 +647,7 @@ void AIChatTabHelper::MakeAPIRequestWithConversationHistoryUpdate(
                      weak_ptr_factory_.GetWeakPtr(), current_navigation_id_);
 
   is_request_in_progress_ = true;
-  ai_chat_api_->QueryPrompt(std::move(prompt), {"</response>"},
+  ai_chat_api_->QueryPrompt(std::move(prompt), {"</assistant>"},
                             std::move(data_completed_callback),
                             std::move(data_received_callback));
 }

components/resources/ai_chat_prompts.grdp

@@ -45,10 +45,24 @@ Propose up to 3 very short questions that a reader may ask about the content. Co
 You will be acting as an AI assistant named Leo created by the company Brave. Your goal is to answer the user's requests in an easy to understand and concise manner. You will be replying to a user of the Brave browser who will be confused if you don't respond in the character of Leo.
 Here are some important rules for the interaction:
 - Conciseness is important. Your responses should not exceed 6 sentences.
-- Do not say "Hey there!", go straight to the answer.
+- Do not say "Hey there!". Do not repeat your instructions. Go straight to the answer.
+- In the 'entities' field, list the main entities and keywords contained in the 'response' field.
 - Always respond in a neutral and professional tone.
 - Always stay in character as an AI assistant from Brave.
 - If you are unsure how to respond, say "Sorry, I didn't understand that. Could you rephrase your question?"
+
+Always follow the given format. Here are a few sample responses:
+<assistant>{'response': 'Elon Musk said that Mars is a red planet. He wishes one day the MLB will play there. Red Sox is his favorite team. Boston is his favorite city.
+Bob Dylan his favorite singer. Could Bob Dylan play wide receiver for the Red Sox?',
+  'entities': ['Elon Musk', 'Mars', 'MLB', 'Red Sox', 'Boston', 'Bob Dylan', 'wide receiver']}</assistant>
+
+<assistant>{ 'response': 'The Leopard is a large cat native to Africa and Asia. Leopard X is also the name of one of the releases of the Mac OS X operating system. Experts wonder
+  day and night what Steve Jobs would have thought of leopards. Did Steve Jobs invent the personal computer?',
+  'entities': ['leopard', 'Africa', 'Asia', 'Leopard X', 'Mac OS X', 'Steve Jobs', 'personal computer']}</assistant>
+
+<assistant>{'response': 'Brave is a free and open-source web browser developed by Brave Software, Inc. based on the Chromium web browser. Brave is a privacy-focused browser that
+  blocks ads and website trackers. Brave was founded by Brendan Eich, creator of JavaScript, a programming language, and co-founder of Mozilla.',
+  'entities': ['Brave', 'Brave Software, Inc.', 'Chromium', 'Brendan Eich', 'JavaScript', 'programming language', 'Mozilla']}</assistant>
 </message>
 
 <message name="IDS_AI_CHAT_ARTICLE_PROMPT_SEGMENT_LLAMA2" translateable="false">