Discharge7443/brave-core
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

[AI Chat]: Suggested Models in Model Selector (#31165)

Browse Source 
* [AI Chat]: Suggested Models in Model Selector

* fix iOS build with updated ai chat delegate method.

add a new parameter in ai_chat_delegate method onMethodChanged from mojom update.

---------

Co-authored-by: Nuo Xu <nuoxu.nx@gmail.com>
This commit is contained in:
Douglas Daniel
2025-09-17 09:52:39 -05:00
committed by GitHub
co-authored by Nuo Xu
parent c0dd287af3
commit 505024dbfd
17 changed files with 194 additions and 40 deletions
Download Patch File Download Diff File Expand all files Collapse all files
browser/ai_chat/ai_chat_render_view_context_menu_browsertest.cc
+1
View File
@@ -87,6 +87,7 @@ class MockConversationHandlerClient : public mojom::ConversationUI {
MOCK_METHOD(void,
OnModelDataChanged,
(const std::string& conversation_model_key,
const std::string& default_model_key,
std::vector<mojom::ModelPtr> all_models),
(override));
browser/resources/settings/brave_leo_assistant_page/model_config_ui.ts
+2 -1
View File
@@ -161,7 +161,8 @@ export class ModelConfigUI extends ModelConfigUIBase {
key: modelKey,
displayName: this.label,
visionSupport: this.hasVisionSupport,
supportsTools: false
supportsTools: false,
isSuggestedModel: false,
}
this.fire('save', { modelConfig })
components/ai_chat/core/browser/ai_chat_service_unittest.cc
+3 -1
View File
@@ -139,7 +139,9 @@ class MockConversationHandlerClient : public mojom::ConversationUI {
MOCK_METHOD(void,
OnModelDataChanged,
(const std::string&, std::vector<mojom::ModelPtr>),
(const std::string& conversation_model_key,
const std::string& default_model_key,
std::vector<mojom::ModelPtr> all_models),
(override));
MOCK_METHOD(void,
components/ai_chat/core/browser/conversation_handler.cc
+12 -2
View File
@@ -344,6 +344,7 @@ void ConversationHandler::GetState(GetStateCallback callback) {
std::transform(models.cbegin(), models.cend(), models_copy.begin(),
[](auto& model) { return model.Clone(); });
auto model_key = GetCurrentModel().key;
auto default_model_key = model_service_->GetDefaultModelKey();
std::vector<std::string> suggestions;
std::ranges::transform(suggestions_, std::back_inserter(suggestions),
@@ -351,7 +352,8 @@ void ConversationHandler::GetState(GetStateCallback callback) {
mojom::ConversationStatePtr state = mojom::ConversationState::New(
metadata_->uuid, is_request_in_progress_, std::move(models_copy),
model_key, std::move(suggestions), suggestion_generation_status_,
model_key, default_model_key, std::move(suggestions),
suggestion_generation_status_,
associated_content_manager_->GetAssociatedContent(), current_error_,
metadata_->temporary);
@@ -1509,6 +1511,12 @@ void ConversationHandler::OnDefaultModelChanged(const std::string& old_key,
DVLOG(1) << "Default model changed from " << old_key << " to " << new_key;
if (model_key_ == old_key) {
ChangeModel(new_key);
} else {
// If this conversation is not using the old default model, we still need to
// notify the UI about the default model change without changing the current
// model. This ensures the UI gets updated with the new default model
// information even when the conversation is using a different model.
OnModelDataChanged();
}
}
@@ -1526,13 +1534,15 @@ void ConversationHandler::OnModelRemoved(const std::string& removed_key) {
void ConversationHandler::OnModelDataChanged() {
const std::vector<mojom::ModelPtr>& models = model_service_->GetModels();
auto default_model_key = model_service_->GetDefaultModelKey();
for (const mojo::Remote<mojom::ConversationUI>& client :
conversation_ui_handlers_) {
std::vector<mojom::ModelPtr> models_copy(models.size());
std::transform(models.cbegin(), models.cend(), models_copy.begin(),
[](auto& model) { return model.Clone(); });
client->OnModelDataChanged(model_key_, std::move(models_copy));
client->OnModelDataChanged(model_key_, default_model_key,
std::move(models_copy));
}
OnStateForConversationEntriesChanged();
}
components/ai_chat/core/browser/conversation_handler_unittest.cc
+1
View File
@@ -122,6 +122,7 @@ class MockConversationHandlerClient : public mojom::ConversationUI {
MOCK_METHOD(void,
OnModelDataChanged,
(const std::string& conversation_model_key,
const std::string& default_model_key,
std::vector<mojom::ModelPtr> all_models),
(override));
components/ai_chat/core/browser/model_service.cc
+7
View File
@@ -114,6 +114,7 @@ const std::vector<mojom::ModelPtr>& GetLeoModels() {
model->display_name = "Automatic";
model->vision_support = true;
model->supports_tools = false;
model->is_suggested_model = true;
model->options =
mojom::ModelOptions::NewLeoModelOptions(std::move(options));
models.push_back(std::move(model));
@@ -134,6 +135,7 @@ const std::vector<mojom::ModelPtr>& GetLeoModels() {
model->display_name = "DeepSeek R1";
model->vision_support = false;
model->supports_tools = false;
model->is_suggested_model = false;
model->options =
mojom::ModelOptions::NewLeoModelOptions(std::move(options));
@@ -154,6 +156,7 @@ const std::vector<mojom::ModelPtr>& GetLeoModels() {
model->display_name = "Claude Haiku";
model->vision_support = true;
model->supports_tools = true;
model->is_suggested_model = false;
model->options =
mojom::ModelOptions::NewLeoModelOptions(std::move(options));
@@ -174,6 +177,7 @@ const std::vector<mojom::ModelPtr>& GetLeoModels() {
model->display_name = "Claude Sonnet";
model->vision_support = true;
model->supports_tools = true;
model->is_suggested_model = true;
model->options =
mojom::ModelOptions::NewLeoModelOptions(std::move(options));
@@ -196,6 +200,7 @@ const std::vector<mojom::ModelPtr>& GetLeoModels() {
model->display_name = "Llama 3.1 8B";
model->vision_support = false;
model->supports_tools = false;
model->is_suggested_model = true;
model->options =
mojom::ModelOptions::NewLeoModelOptions(std::move(options));
@@ -218,6 +223,7 @@ const std::vector<mojom::ModelPtr>& GetLeoModels() {
model->display_name = "Qwen 14B";
model->vision_support = false;
model->supports_tools = false;
model->is_suggested_model = false;
model->options =
mojom::ModelOptions::NewLeoModelOptions(std::move(options));
@@ -240,6 +246,7 @@ const std::vector<mojom::ModelPtr>& GetLeoModels() {
model->display_name = "Gemma 12B";
model->vision_support = true;
model->supports_tools = false;
model->is_suggested_model = false;
model->options =
mojom::ModelOptions::NewLeoModelOptions(std::move(options));
components/ai_chat/core/common/mojom/ai_chat.mojom
+6 -1
View File
@@ -389,6 +389,8 @@ struct Model {
bool vision_support;
// Model supports tool use
bool supports_tools = false;
// Is a suggested model
bool is_suggested_model;
};
struct ModelWithSubtitle {
@@ -545,6 +547,7 @@ struct ConversationState {
bool is_request_in_progress;
array<Model> all_models;
string current_model_key;
string default_model_key;
array<string> suggested_questions;
SuggestionGenerationStatus suggestion_status;
array<AssociatedContent> associated_content;
@@ -684,7 +687,9 @@ interface ConversationUI {
// it could be changed from the browser.
// When the browser changes the model, it will also provide a new list
// of models.
OnModelDataChanged(string conversation_model_key, array<Model> all_models);
OnModelDataChanged(string conversation_model_key,
string default_model_key,
array<Model> all_models);
// Provides the full list of questions for the active tab in the page's
// Browser window.
OnSuggestedQuestionsChanged(
components/ai_chat/ios/browser/ai_chat_delegate.h
+1
View File
@@ -28,6 +28,7 @@ OBJC_EXPORT
- (void)onAPIRequestInProgress:(bool)inProgress;
- (void)onAPIResponseError:(AiChatAPIError)error;
- (void)onModelChanged:(NSString*)modelKey
defaultModelKey:(NSString*)defaultModelKey
modelList:(NSArray<AiChatModel*>*)modelList;
- (void)onSuggestedQuestionsChanged:(NSArray<NSString*>*)questions
status:(AiChatSuggestionGenerationStatus)status;
components/ai_chat/ios/browser/conversation_client.h
+1
View File
@@ -39,6 +39,7 @@ class ConversationClient : public mojom::ConversationUI,
void OnAPIResponseError(mojom::APIError error) override;
void OnModelDataChanged(
const std::string& model_key,
const std::string& default_model_key,
std::vector<ai_chat::mojom::ModelPtr> model_list) override;
void OnSuggestedQuestionsChanged(
const std::vector<std::string>& questions,
components/ai_chat/ios/browser/conversation_client.mm
+4 -1
View File
@@ -49,6 +49,7 @@ void ConversationClient::OnAPIResponseError(mojom::APIError error) {
void ConversationClient::OnModelDataChanged(
const std::string& model_key,
const std::string& default_model_key,
std::vector<mojom::ModelPtr> model_list) {
NSMutableArray* models =
[[NSMutableArray alloc] initWithCapacity:model_list.size()];
@@ -57,7 +58,9 @@ void ConversationClient::OnModelDataChanged(
[models addObject:[[AiChatModel alloc] initWithModelPtr:model->Clone()]];
}
[bridge_ onModelChanged:base::SysUTF8ToNSString(model_key) modelList:models];
[bridge_ onModelChanged:base::SysUTF8ToNSString(model_key)
defaultModelKey:base::SysUTF8ToNSString(default_model_key)
modelList:models];
}
void ConversationClient::OnSuggestedQuestionsChanged(
components/ai_chat/resources/page/components/model_selector/index.test.tsx
+49 -11
View File
@@ -27,9 +27,30 @@ describe('ModelSelector', () => {
>
const mockModels = [
{
key: 'chat-automatic',
displayName: 'Automatic',
isSuggestedModel: true,
options: {
leoModelOptions: {
access: Mojom.ModelAccess.BASIC_AND_PREMIUM,
},
},
},
{
key: 'chat-basic',
displayName: 'Basic Model',
isSuggestedModel: true,
options: {
leoModelOptions: {
access: Mojom.ModelAccess.BASIC_AND_PREMIUM,
},
},
},
{
key: 'another-chat-basic',
displayName: 'Another Basic Model',
isSuggestedModel: false,
options: {
leoModelOptions: {
access: Mojom.ModelAccess.BASIC_AND_PREMIUM,
@@ -39,12 +60,24 @@ describe('ModelSelector', () => {
{
key: 'chat-premium',
displayName: 'Premium Model',
isSuggestedModel: true,
options: {
leoModelOptions: {
access: Mojom.ModelAccess.PREMIUM,
},
},
},
{
key: 'another-chat-premium',
displayName: 'Another Premium Model',
isSuggestedModel: false,
options: {
leoModelOptions: {
access: Mojom.ModelAccess.PREMIUM,
},
},
},
{
key: 'chat-custom',
displayName: 'Custom Model',
@@ -62,7 +95,7 @@ describe('ModelSelector', () => {
const defaultConversationContext = {
allModels: mockModels,
currentModel: mockModels[0],
currentModel: mockModels[1],
setCurrentModel: jest.fn(),
}
@@ -112,8 +145,10 @@ describe('ModelSelector', () => {
// Make sure the default menu items are visible
const menuItems = document.querySelectorAll<HTMLElement>('leo-menu-item')
expect(menuItems).toHaveLength(2)
expect(menuItems[0]).toHaveTextContent('Basic Model')
expect(menuItems).toHaveLength(4)
expect(menuItems[0]).toHaveTextContent('Automatic')
expect(menuItems[1]).toHaveTextContent('Basic Model')
expect(menuItems[2]).toHaveTextContent('Premium Model')
})
it('shows all models if Show all models button is clicked', async () => {
@@ -143,10 +178,13 @@ describe('ModelSelector', () => {
// Check that all model items are visible
const allMenuItems = document.querySelectorAll<HTMLElement>('leo-menu-item')
expect(allMenuItems).toHaveLength(4)
expect(allMenuItems[0]).toHaveTextContent('Basic Model')
expect(allMenuItems[1]).toHaveTextContent('Premium Model')
expect(allMenuItems[2]).toHaveTextContent('Custom Model')
expect(allMenuItems).toHaveLength(7)
expect(allMenuItems[0]).toHaveTextContent('Automatic')
expect(allMenuItems[1]).toHaveTextContent('Basic Model')
expect(allMenuItems[2]).toHaveTextContent('Another Basic Model')
expect(allMenuItems[3]).toHaveTextContent('Premium Model')
expect(allMenuItems[4]).toHaveTextContent('Another Premium Model')
expect(allMenuItems[5]).toHaveTextContent('Custom Model')
const labels = document.querySelectorAll<HTMLElement>('leo-label')
@@ -158,9 +196,9 @@ describe('ModelSelector', () => {
)
// Check that local label is visible
expect(labels[1]).toBeInTheDocument()
expect(labels[1]).toBeVisible()
expect(labels[1]).toHaveTextContent('CHAT_UI_MODEL_LOCAL_LABEL')
expect(labels[2]).toBeInTheDocument()
expect(labels[2]).toBeVisible()
expect(labels[2]).toHaveTextContent('CHAT_UI_MODEL_LOCAL_LABEL')
})
it('should call setCurrentModel when a model is clicked', async () => {
@@ -191,7 +229,7 @@ describe('ModelSelector', () => {
// Select another model
const allMenuItems = document.querySelectorAll<HTMLElement>('leo-menu-item')
expect(allMenuItems).toHaveLength(4)
expect(allMenuItems).toHaveLength(7)
await act(async () => {
allMenuItems[1].click()
})
components/ai_chat/resources/page/components/model_selector/index.tsx
+75 -13
View File
@@ -16,6 +16,8 @@ import { useConversation } from '../../state/conversation_context'
import { getModelIcon } from '../../../common/constants'
import styles from './style.module.scss'
const AUTO_MODEL_KEY = 'chat-automatic'
export function ModelSelector() {
const aiChatContext = useAIChat()
const conversationContext = useConversation()
@@ -24,16 +26,64 @@ export function ModelSelector() {
const [isOpen, setIsOpen] = React.useState(false)
const [showAllModels, setShowAllModels] = React.useState(false)
// Computed
const leoModels = React.useMemo(
() =>
conversationContext.allModels.filter(
// Memos
const suggestedModels = React.useMemo(() => {
return conversationContext.allModels.filter(
(model) => model.isSuggestedModel,
)
}, [conversationContext.allModels])
const models = React.useMemo(() => {
// Show all BASIC_AND_PREMIUM models if showAllModels is true
if (showAllModels) {
return conversationContext.allModels.filter(
(model) =>
model.options.leoModelOptions?.access
=== Mojom.ModelAccess.BASIC_AND_PREMIUM,
),
[conversationContext.allModels],
)
)
}
// Find the Auto model (chat-automatic)
const autoModel = conversationContext.allModels.find(
(model) => model.key === AUTO_MODEL_KEY,
)
const defaultModel = conversationContext.userDefaultModel
const currentModel = conversationContext.currentModel
const recommendedList: Mojom.Model[] = []
// Keep Auto model first in list
if (autoModel) {
recommendedList.push(autoModel)
}
// Add defaultModel if it exists and is not Auto
if (defaultModel && defaultModel.key !== AUTO_MODEL_KEY) {
recommendedList.push(defaultModel)
}
// Add currentModel if it exists and is not Auto or defaultModel
if (
currentModel
&& currentModel.key !== AUTO_MODEL_KEY
&& currentModel.key !== defaultModel?.key
) {
recommendedList.push(currentModel)
}
// Add suggestedModels that are not already in the list
const existingKeys = new Set(recommendedList.map((model) => model.key))
const filteredSuggestedModels = suggestedModels.filter(
(model) => !existingKeys.has(model.key),
)
recommendedList.push(...filteredSuggestedModels)
return recommendedList
}, [
showAllModels,
conversationContext.allModels,
conversationContext.userDefaultModel,
suggestedModels,
])
const premiumModels = React.useMemo(
() =>
@@ -86,7 +136,8 @@ export function ModelSelector() {
</div>
</Button>
{leoModels.map((model) => {
{models.map((model) => {
const isCustomModel = model.options.customModelOptions
return (
<leo-menu-item
key={model.key}
@@ -104,11 +155,13 @@ export function ModelSelector() {
<div className={styles.menuText}>
<div>{model.displayName}</div>
<p className={styles.modelSubtitle}>
{getLocale(
`CHAT_UI_${model.key
.toUpperCase()
.replaceAll('-', '_')}_SUBTITLE`,
)}
{isCustomModel
? model.options.customModelOptions?.modelRequestName
: getLocale(
`CHAT_UI_${model.key
.toUpperCase()
.replaceAll('-', '_')}_SUBTITLE`,
)}
</p>
</div>
{model.options.leoModelOptions?.access === Mojom.ModelAccess.PREMIUM
@@ -121,6 +174,15 @@ export function ModelSelector() {
{getLocale(S.CHAT_UI_MODEL_PREMIUM_LABEL_NON_PREMIUM)}
</Label>
)}
{isCustomModel && (
<Label
className={styles.modelLabel}
mode='default'
color='blue'
>
{getLocale(S.CHAT_UI_MODEL_LOCAL_LABEL)}
</Label>
)}
</leo-menu-item>
)
})}
components/ai_chat/resources/page/state/conversation_context.tsx
+14 -3
View File
@@ -43,6 +43,7 @@ export type ConversationContext = SendFeedbackState
associatedContentInfo: Mojom.AssociatedContent[]
allModels: Mojom.Model[]
currentModel?: Mojom.Model
userDefaultModel?: Mojom.Model
suggestedQuestions: string[]
isGenerating: boolean
suggestionStatus: Mojom.SuggestionGenerationStatus
@@ -226,11 +227,13 @@ export function ConversationContextProvider(props: React.PropsWithChildren) {
const getModelContext = (
currentModelKey: string,
defaultModelKey: string,
allModels: Mojom.Model[],
): Partial<ConversationContext> => {
return {
allModels,
currentModel: allModels.find((m) => m.key === currentModelKey),
userDefaultModel: allModels.find((m) => m.key === defaultModelKey),
}
}
@@ -265,6 +268,7 @@ export function ConversationContextProvider(props: React.PropsWithChildren) {
isRequestInProgress: isGenerating,
allModels: models,
currentModelKey,
defaultModelKey,
suggestedQuestions,
suggestionStatus,
associatedContent,
@@ -275,7 +279,7 @@ export function ConversationContextProvider(props: React.PropsWithChildren) {
setPartialContext({
conversationUuid,
isGenerating,
...getModelContext(currentModelKey, models),
...getModelContext(currentModelKey, defaultModelKey, models),
suggestedQuestions,
suggestionStatus,
associatedContentInfo: associatedContent,
@@ -312,8 +316,15 @@ export function ConversationContextProvider(props: React.PropsWithChildren) {
listenerIds.push(id)
id = callbackRouter.onModelDataChanged.addListener(
(conversationModelKey: string, allModels: Mojom.Model[]) =>
setPartialContext(getModelContext(conversationModelKey, allModels)),
(
conversationModelKey: string,
defaultModelKey: string,
allModels: Mojom.Model[],
) => {
setPartialContext(
getModelContext(conversationModelKey, defaultModelKey, allModels),
)
},
)
listenerIds.push(id)
components/ai_chat/resources/page/stories/components_panel.tsx
+4
View File
@@ -744,6 +744,7 @@ const MODELS: Mojom.Model[] = [
displayName: 'Model One',
visionSupport: false,
supportsTools: false,
isSuggestedModel: true,
options: {
leoModelOptions: {
name: 'model-one',
@@ -761,6 +762,7 @@ const MODELS: Mojom.Model[] = [
displayName: 'Model Two',
visionSupport: true,
supportsTools: true,
isSuggestedModel: true,
options: {
leoModelOptions: {
name: 'model-two-premium',
@@ -778,6 +780,7 @@ const MODELS: Mojom.Model[] = [
displayName: 'Model Three',
visionSupport: false,
supportsTools: false,
isSuggestedModel: false,
options: {
leoModelOptions: {
name: 'model-three-freemium',
@@ -795,6 +798,7 @@ const MODELS: Mojom.Model[] = [
displayName: 'Microsoft Phi-3',
visionSupport: false,
supportsTools: true,
isSuggestedModel: false,
options: {
leoModelOptions: undefined,
customModelOptions: {
ios/brave-ios/Sources/AIChat/Components/Messages/AIChatIntroMessageView.swift
+2 -1
View File
@@ -113,7 +113,8 @@ struct AIChatIntroMessageView_Previews: PreviewProvider {
key: "mixtral_8x7b",
displayName: "Mixtral 8x7b",
visionSupport: false,
supportsTools: false
supportsTools: false,
isSuggestedModel: false
)
)
.previewLayout(.sizeThatFits)
ios/brave-ios/Sources/AIChat/Components/Settings/AIChatMenuView.swift
+10 -5
View File
@@ -294,7 +294,8 @@ struct AIChatMenuView_Preview: PreviewProvider {
key: "mixtral_8x7b",
displayName: "Mixtral 8x7b",
visionSupport: false,
supportsTools: false
supportsTools: false,
isSuggestedModel: false
),
modelOptions: [
.init(
@@ -311,7 +312,8 @@ struct AIChatMenuView_Preview: PreviewProvider {
key: "mixtral_8x7b",
displayName: "Mixtral 8x7b",
visionSupport: false,
supportsTools: false
supportsTools: false,
isSuggestedModel: false
),
.init(
options: .init(
@@ -327,7 +329,8 @@ struct AIChatMenuView_Preview: PreviewProvider {
key: "claude_instant",
displayName: "Claude-Instant",
visionSupport: false,
supportsTools: false
supportsTools: false,
isSuggestedModel: false
),
.init(
options: .init(
@@ -343,7 +346,8 @@ struct AIChatMenuView_Preview: PreviewProvider {
key: "llama_2x13b",
displayName: "Llama-2 13b",
visionSupport: false,
supportsTools: false
supportsTools: false,
isSuggestedModel: false
),
.init(
options: .init(
@@ -359,7 +363,8 @@ struct AIChatMenuView_Preview: PreviewProvider {
key: "llama_2x70b",
displayName: "Llama-2 70b",
visionSupport: false,
supportsTools: false
supportsTools: false,
isSuggestedModel: false
),
],
onModelChanged: {
ios/brave-ios/Sources/AIChat/ModelView/AIChatViewModel.swift
+2 -1
View File
@@ -334,7 +334,8 @@ extension AIChatViewModel: AIChatDelegate {
self.suggestionsStatus = status
}
public func onModelChanged(_ modelKey: String, modelList: [AiChat.Model]) {
public func onModelChanged(_ modelKey: String, defaultModelKey: String, modelList: [AiChat.Model])
{
self.currentModel = self.models.first(where: { $0.key == modelKey })
self.models = modelList
}