diff --git a/ee/lib/gitlab/duo/developments/setup.rb b/ee/lib/gitlab/duo/developments/setup.rb index 2a6557f25eac1c067451d0435fb5fce664cbdc27..f4931b67a2840d3818c3c5b988da791f3db418b6 100644 --- a/ee/lib/gitlab/duo/developments/setup.rb +++ b/ee/lib/gitlab/duo/developments/setup.rb @@ -250,6 +250,65 @@ def print_result Group.find_by_full_path(@namespace) end end + + class SeedWorkflowCatalog + def initialize + full_path = ENV[''] || 'gitlab-duo/test' + @project = Project.find_by(full_path: full_path) # rubocop:disable CodeReuse/ActiveRecord -- Development purpose + raise "#{full_path} project not found" unless project + + @group = project.root_ancestor + end + + def execute + ::Feature.enable(:global_ai_catalog) + seed_agents + end + + private + + attr_reader :project, :group + + def seed_agents + agents = [ + { name: 'Issue Labeler', description: 'Label new issues', public: true, definition: { + 'system_prompt' => '', + 'user_prompt' => '', + 'tools' => [] + } }, + { name: 'Merge Request Labeler', description: 'Label new merge requests', definition: { + 'system_prompt' => '', + 'user_prompt' => '', + 'tools' => [] + } }, + { name: 'Backlog Prioritizer', description: 'Prioritize issue backlog', definition: { + 'system_prompt' => '', + 'user_prompt' => '', + 'tools' => [] + } }, + { name: 'Spam Monitor', description: 'Identify potential spam comments', definition: { + 'system_prompt' => '', + 'user_prompt' => '', + 'tools' => [] + } } + ] + + agents.each do |agent| + Ai::Catalog::Item.create(agent.merge(item_type: :agent, project: project)) + end + end + + def seed_flows + Ai::Catalog::Item.last(3) do |agent| + Ai::Catalog::Item.create(item_type: :flow, project: project, definition: { 'steps' => [agent.id] }) + end + end + + def seed_consumers + Ai::Catalog::ItemConsumer.create(project: project, item: Ai::Catalog::Item.last, enabled: true) + Ai::Catalog::ItemConsumer.create(group: group, item: Ai::Catalog::Item.last, enabled: true) + end + end end end end diff --git a/ee/lib/tasks/gitlab/duo.rake b/ee/lib/tasks/gitlab/duo.rake index 12db701a64c505c72167cf4fecd0b73400ae482c..17afcc30a50de34d15af9a24707bc492b1771108 100644 --- a/ee/lib/tasks/gitlab/duo.rake +++ b/ee/lib/tasks/gitlab/duo.rake @@ -31,5 +31,10 @@ namespace :gitlab do task :verify_self_hosted_setup, [:username] => :gitlab_environment do |_, args| Gitlab::Duo::Administration::VerifySelfHostedSetup.new(args[:username]).execute end + + desc 'GitLab | Duo | Seed workflow catalog' + task seed_workflow_catalog: :gitlab_environment do + Gitlab::Duo::Developments::SeedWorkflowCatalog.new.execute + end end end diff --git a/ee/spec/factories/ai/catalog/item_consumers.rb b/ee/spec/factories/ai/catalog/item_consumers.rb index 694088ffd9ead7971b486474d42b622560710c3b..a01ab1f2e50a236b5656f4fc914942ad5a2aa044 100644 --- a/ee/spec/factories/ai/catalog/item_consumers.rb +++ b/ee/spec/factories/ai/catalog/item_consumers.rb @@ -2,7 +2,7 @@ FactoryBot.define do factory :ai_catalog_item_consumer, class: 'Ai::Catalog::ItemConsumer' do - item { association :ai_catalog_item } + item { association :ai_catalog_flow } enabled { true } locked { true } end diff --git a/ee/spec/factories/ai/catalog/item_versions.rb b/ee/spec/factories/ai/catalog/item_versions.rb index 8e8bbc6f6d724ccc5966015291bd6c266fbc05f9..58c58f883954fe2a75186a05e5f4897b30b9e7c8 100644 --- a/ee/spec/factories/ai/catalog/item_versions.rb +++ b/ee/spec/factories/ai/catalog/item_versions.rb @@ -4,18 +4,30 @@ factory :ai_catalog_item_version, class: 'Ai::Catalog::ItemVersion' do version { '1.0.0' } schema_version { 1 } + for_agent release_date { nil } - definition do - { - 'system_prompt' => 'Talk like a pirate!', - 'tools' => [1], - 'user_prompt' => 'What is a leap year?' - } - end - item { association :ai_catalog_item } + + factory :ai_catalog_agent_version, traits: [:for_agent] + factory :ai_catalog_flow_version, traits: [:for_flow] trait :released do release_date { Time.current } end + + trait :for_agent do + item { association :ai_catalog_agent } + definition do + { + 'system_prompt' => 'Talk like a pirate!', + 'tools' => [1], + 'user_prompt' => 'What is a leap year?' + } + end + end + + trait :for_flow do + item { association :ai_catalog_flow } + definition { { 'triggers' => [1] } } + end end end diff --git a/ee/spec/factories/ai/catalog/items.rb b/ee/spec/factories/ai/catalog/items.rb index f6b256472b2e38b4fc0518f1dcdc05c43336ea49..d16ccb07fe438f6ca9028a5dd7d734a12b5ba103 100644 --- a/ee/spec/factories/ai/catalog/items.rb +++ b/ee/spec/factories/ai/catalog/items.rb @@ -18,7 +18,11 @@ end trait :with_version do - versions { build_list(:ai_catalog_item_version, 1) } + # if item_type == :agent + # build_list(:ai_catalog_agent_version, 1) + # else + # build_list(:ai_catalog_flow_version, 1) + # end end after(:build) do |item, _|