From 043219ab0c248262772c34173eb866f00078716f Mon Sep 17 00:00:00 2001
From: Lee Tickett <ltickett@gitlab.com>
Date: Fri, 18 Jul 2025 17:48:44 +0100
Subject: [PATCH] Experiment with AI catalog item type version abstraction

---
 ee/app/models/ai/catalog/item_version.rb  | 14 +++++++++--
 ee/lib/ai/catalog/agent_version_schema.rb | 18 ++++++++++++++
 ee/lib/ai/catalog/flow_version_schema.rb  | 16 ++++++++++++
 ee/lib/ai/catalog/item_version_loader.rb  | 30 +++++++++++++++++++++++
 ee/spec/models/ai/catalog/item_spec.rb    | 13 ++++++++++
 5 files changed, 89 insertions(+), 2 deletions(-)
 create mode 100644 ee/lib/ai/catalog/agent_version_schema.rb
 create mode 100644 ee/lib/ai/catalog/flow_version_schema.rb
 create mode 100644 ee/lib/ai/catalog/item_version_loader.rb

diff --git a/ee/app/models/ai/catalog/item_version.rb b/ee/app/models/ai/catalog/item_version.rb
index 34943378090c55..a0db320b1be359 100644
--- a/ee/app/models/ai/catalog/item_version.rb
+++ b/ee/app/models/ai/catalog/item_version.rb
@@ -4,8 +4,7 @@ module Ai
   module Catalog
     class ItemVersion < ApplicationRecord
       include SafelyChangeColumnDefault
-
-      AGENT_SCHEMA_VERSION = 1
+      include ItemVersionLoader
 
       self.table_name = "ai_catalog_item_versions"
 
@@ -18,6 +17,8 @@ class ItemVersion < ApplicationRecord
 
       validates :definition, json_schema: { filename: 'ai_catalog_item_version_definition', size_limit: 64.kilobytes }
 
+      validate :validate_type_schema
+
       belongs_to :item, class_name: 'Ai::Catalog::Item',
         foreign_key: :ai_catalog_item_id, inverse_of: :versions, optional: false, autosave: true
       belongs_to :organization, class_name: 'Organizations::Organization'
@@ -41,6 +42,15 @@ def draft?
         !released?
       end
 
+      def method_missing(method_name, *args, &block)
+        return super if @schema_methods_defined
+
+        ItemVersionLoader.define_schema_methods(self)
+        @schema_methods_defined = true
+
+        send(method_name, *args, &block)
+      end
+
       private
 
       def populate_organization
diff --git a/ee/lib/ai/catalog/agent_version_schema.rb b/ee/lib/ai/catalog/agent_version_schema.rb
new file mode 100644
index 00000000000000..cc7368d7a01b5b
--- /dev/null
+++ b/ee/lib/ai/catalog/agent_version_schema.rb
@@ -0,0 +1,18 @@
+module Ai
+  module Catalog
+    class AgentVersionSchema
+      SCHEMA_VERSION = 1
+      SCHEMA = {
+        system_prompt: { type: String, required: true, max_length: 2048 },
+        user_prompt: { type: String, required: false, max_length: 1024 },
+        tools: { type: Array, required: false }
+      }.freeze
+
+      def self.define_validation_method(instance)
+        instance.define_singleton_method(:validate_type_schema) do
+          puts 'AOK'
+        end
+      end
+    end
+  end
+end
diff --git a/ee/lib/ai/catalog/flow_version_schema.rb b/ee/lib/ai/catalog/flow_version_schema.rb
new file mode 100644
index 00000000000000..49fdfda208611d
--- /dev/null
+++ b/ee/lib/ai/catalog/flow_version_schema.rb
@@ -0,0 +1,16 @@
+module Ai
+  module Catalog
+    class FlowVersionSchema
+      SCHEMA_VERSION = 1
+      SCHEMA = {
+        agents: { type: Array, required: false }
+      }.freeze
+
+      def self.define_validation_method(instance)
+        instance.define_singleton_method(:validate_type_schema) do
+          errors.add(:base, 'oh oh')
+        end
+      end
+    end
+  end
+end
diff --git a/ee/lib/ai/catalog/item_version_loader.rb b/ee/lib/ai/catalog/item_version_loader.rb
new file mode 100644
index 00000000000000..4f7500584ca982
--- /dev/null
+++ b/ee/lib/ai/catalog/item_version_loader.rb
@@ -0,0 +1,30 @@
+# frozen_string_literal: true
+
+# Can we load
+
+module Ai
+  module Catalog
+    module ItemVersionLoader
+      ITEM_TYPE_DEFINITIONS = {
+        Ai::Catalog::Item::AGENT_TYPE => Ai::Catalog::AgentVersionSchema,
+        Ai::Catalog::Item::FLOW_TYPE => Ai::Catalog::FlowVersionSchema
+      }
+
+      def self.define_schema_methods(instance)
+        schema_class = ITEM_TYPE_DEFINITIONS[instance.item.item_type.to_sym]
+        schema_class::SCHEMA.each_key do |attr_name|
+          instance.define_singleton_method(attr_name) do
+            definition&.dig(attr_name.to_s)
+          end
+
+          instance.define_singleton_method("#{attr_name}=") do |value|
+            self.definition ||= {}
+            self.definition[attr_name.to_s] = value
+          end
+        end
+
+        schema_class.define_validation_method(instance)
+      end
+    end
+  end
+end
diff --git a/ee/spec/models/ai/catalog/item_spec.rb b/ee/spec/models/ai/catalog/item_spec.rb
index 74989a97730e40..1d4088cccb0d9e 100644
--- a/ee/spec/models/ai/catalog/item_spec.rb
+++ b/ee/spec/models/ai/catalog/item_spec.rb
@@ -113,4 +113,17 @@
       expect { item.soft_delete }.to change { item.deleted_at }.from(nil)
     end
   end
+
+  describe 'magic' do
+    it 'is magic' do
+      agent_version = create(:ai_catalog_item, :with_version).latest_version
+
+      expect(agent_version).to be_valid
+
+      flow_version = create(:ai_catalog_item, :with_version, item_type: 2).latest_version
+
+      expect(flow_version).not_to be_valid
+      expect(flow_version.errors.full_messages).to match_array(["oh oh"])
+    end
+  end
 end
-- 
GitLab