diff --git a/ee/app/graphql/resolvers/ai/chat/available_models_resolver.rb b/ee/app/graphql/resolvers/ai/chat/available_models_resolver.rb index 508654de6620a68f31bfa729fd6c5c043dafcea5..f930828ee295e63622e80e5f1427376e0ab98a44 100644 --- a/ee/app/graphql/resolvers/ai/chat/available_models_resolver.rb +++ b/ee/app/graphql/resolvers/ai/chat/available_models_resolver.rb @@ -20,29 +20,11 @@ def resolve(root_namespace_id:) result = ::Ai::ModelSelection::FetchModelDefinitionsService .new(current_user, model_selection_scope: namespace) - .execute + .execute(feature_setting: "duo_agent_platform") + return { default_model: nil, selectable_models: [] } unless result&.success? - feature_settings = result["unit_primitives"].find do |setting| - setting["feature_setting"] == "duo_agent_platform" - end - - return { default_model: nil, selectable_models: [] } unless feature_settings - - models = result["models"] - identifiers = feature_settings["selectable_models"] - duo_chat_models = models - .select { |model| identifiers.include?(model["identifier"]) } - .map { |model| { name: model["name"], ref: model["identifier"] } } - - default_model = duo_chat_models.find do |model| - model[:ref] == feature_settings["default_model"] - end - - { - default_model: default_model, - selectable_models: duo_chat_models - } + result.payload end end end diff --git a/ee/app/services/ai/model_selection/fetch_model_definitions_service.rb b/ee/app/services/ai/model_selection/fetch_model_definitions_service.rb index febb7b5af9aee109bfcc2fdc34dae2279bd44a1f..009ac0bb11a601fb31d6ef6ae93d6367a38fec1b 100644 --- a/ee/app/services/ai/model_selection/fetch_model_definitions_service.rb +++ b/ee/app/services/ai/model_selection/fetch_model_definitions_service.rb @@ -14,13 +14,19 @@ def initialize(user, model_selection_scope:) @model_selection_scope = model_selection_scope end - def execute(force_api_call: false) + def execute(force_api_call: false, feature_setting: nil) return unless model_selection_enabled? return ServiceResponse.success(payload: nil) if ::License.current&.offline_cloud_license? - return cached_response if Rails.cache.exist?(RESPONSE_CACHE_NAME) && !force_api_call + result = if Rails.cache.exist?(RESPONSE_CACHE_NAME) && !force_api_call + cached_response + else + fetch_model_definitions + end - fetch_model_definitions + return result unless result&.success? && feature_setting.present? + + parse_feature_setting_result(result.payload, feature_setting) end private @@ -83,6 +89,32 @@ def endpoint "#{base_url}/v1/#{endpoint_route}" end + + def parse_feature_setting_result(payload, feature_setting) + return ServiceResponse.success(payload: { default_model: nil, selectable_models: [] }) unless payload + + feature_settings = payload["unit_primitives"]&.find do |setting| + setting["feature_setting"] == feature_setting.to_s + end + + return ServiceResponse.success(payload: { default_model: nil, selectable_models: [] }) unless feature_settings + + models = payload["models"] + identifiers = feature_settings["selectable_models"] + + selectable_models = models + .select { |model| identifiers.include?(model["identifier"]) } + .map { |model| { name: model["name"], ref: model["identifier"] } } + + default_model = selectable_models.find do |model| + model[:ref] == feature_settings["default_model"] + end + + ServiceResponse.success(payload: { + default_model: default_model, + selectable_models: selectable_models + }) + end end end end diff --git a/ee/spec/graphql/resolvers/ai/chat/available_models_resolver_spec.rb b/ee/spec/graphql/resolvers/ai/chat/available_models_resolver_spec.rb index f51f6950bbcf2b8805a23b7d0239c1f88176d70d..687f6eb790b519c4f4b0d0f650ce19831d85bbc2 100644 --- a/ee/spec/graphql/resolvers/ai/chat/available_models_resolver_spec.rb +++ b/ee/spec/graphql/resolvers/ai/chat/available_models_resolver_spec.rb @@ -24,32 +24,17 @@ context "when service returns successful result" do let(:service_result) do ServiceResponse.success(payload: { - "models" => [ - { "name" => "Claude Sonnet 4.0 - Anthropic", "identifier" => "claude_sonnet_4_20250514" }, - { "name" => "Claude Sonnet 4.0 - Vertex", "identifier" => "claude_sonnet_4_20250514_vertex" }, - { "name" => "Claude Sonnet 3.7 - Anthropic", "identifier" => "claude_sonnet_3_7_20250219" }, - { "name" => "Claude Sonnet 3.7 - Vertex", "identifier" => "claude_sonnet_3_7_20250219_vertex" } - ], - "unit_primitives" => [ - { - "feature_setting" => "duo_chat", - "default_model" => "claude_sonnet_4_20250514_vertex", - "selectable_models" => %w[claude_sonnet_4_20250514 claude_sonnet_4_20250514_vertex], - "beta_models" => [] - }, - { - "feature_setting" => "duo_agent_platform", - "default_model" => "claude_sonnet_4_20250514", - "selectable_models" => %w[claude_sonnet_4_20250514 claude_sonnet_3_7_20250219], - "beta_models" => [] - } + default_model: { name: "Claude Sonnet 4.0 - Anthropic", ref: "claude_sonnet_4_20250514" }, + selectable_models: [ + { name: "Claude Sonnet 4.0 - Anthropic", ref: "claude_sonnet_4_20250514" }, + { name: "Claude Sonnet 3.7 - Anthropic", ref: "claude_sonnet_3_7_20250219" } ] }) end before do allow_next_instance_of(::Ai::ModelSelection::FetchModelDefinitionsService) do |service| - allow(service).to receive(:execute).and_return(service_result) + allow(service).to receive(:execute).with(feature_setting: "duo_agent_platform").and_return(service_result) end end @@ -63,20 +48,11 @@ }) end - context "when duo_chat feature setting is not found" do + context "when duo_agent_platform feature setting is not found" do let(:service_result) do ServiceResponse.success(payload: { - "models" => [ - { "name" => "Claude Sonnet", "identifier" => "claude-sonnet" } - ], - "unit_primitives" => [ - { - "feature_setting" => "code_suggestions", - "default_model" => "claude-sonnet", - "selectable_models" => ["claude-sonnet"], - "beta_models" => [] - } - ] + default_model: nil, + selectable_models: [] }) end diff --git a/ee/spec/services/ai/model_selection/fetch_model_definitions_service_spec.rb b/ee/spec/services/ai/model_selection/fetch_model_definitions_service_spec.rb index bab6daf8f16505c5f5ab08a0802c6029c6fe5f51..f7bdc5cdb1b812e5058cc943ef327b7eaa23174a 100644 --- a/ee/spec/services/ai/model_selection/fetch_model_definitions_service_spec.rb +++ b/ee/spec/services/ai/model_selection/fetch_model_definitions_service_spec.rb @@ -80,6 +80,68 @@ end end + context 'when feature_setting parameter is provided' do + let(:feature_setting) { 'duo_agent_platform' } + let(:model_definitions_with_duo_agent_platform) do + { + 'models' => [ + { 'name' => 'Claude Sonnet 4.0', 'identifier' => 'claude_sonnet_4' }, + { 'name' => 'Claude Sonnet 3.7', 'identifier' => 'claude_sonnet_3_7' } + ], + 'unit_primitives' => [ + { + 'feature_setting' => 'duo_chat', + 'default_model' => 'claude-sonnet', + 'selectable_models' => %w[claude-sonnet], + 'beta_models' => [] + }, + { + 'feature_setting' => 'duo_agent_platform', + 'default_model' => 'claude_sonnet_4', + 'selectable_models' => %w[claude_sonnet_4 claude_sonnet_3_7], + 'beta_models' => [] + } + ] + } + end + + subject(:service_with_feature_setting) { initialized_class.execute(feature_setting: feature_setting) } + + before do + allow(initialized_class).to receive(:model_selection_enabled?).and_return(true) + allow(Rails.cache).to receive(:exist?).with(cache_key).and_return(false) + stub_request(:get, endpoint_url) + .to_return( + status: 200, + body: model_definitions_with_duo_agent_platform.to_json, + headers: { 'Content-Type' => 'application/json' } + ) + end + + it 'returns parsed feature setting result' do + expect(service_with_feature_setting).to be_success + expect(service_with_feature_setting.payload).to eq({ + default_model: { name: 'Claude Sonnet 4.0', ref: 'claude_sonnet_4' }, + selectable_models: [ + { name: 'Claude Sonnet 4.0', ref: 'claude_sonnet_4' }, + { name: 'Claude Sonnet 3.7', ref: 'claude_sonnet_3_7' } + ] + }) + end + + context 'when feature_setting is not found' do + let(:feature_setting) { 'non_existent_feature' } + + it 'returns empty result' do + expect(service_with_feature_setting).to be_success + expect(service_with_feature_setting.payload).to eq({ + default_model: nil, + selectable_models: [] + }) + end + end + end + context 'when model switching is enabled' do before do allow(initialized_class).to receive(:model_selection_enabled?).and_return(true)