FEATURE: improve custom tool infra (#1463)

- Add support for `chain.streamCustomRaw(test)` that can be used to stream text from a JS tool direct to composer
- Add support for llm params in `llm.generate` which unlocks stuff like structured outputs
- Add discourse.createStagedUser, discourse.createTopic  and discourse.createPost - for content creation

lib/personas/tool_runner.rb

@@ -68,7 +68,7 @@ module DiscourseAi
 
         const llm = {
           truncate: _llm_truncate,
-          generate: _llm_generate,
+          generate: function(prompt, options) { return _llm_generate(prompt, options); },
         };
 
         const index = {
@@ -85,6 +85,7 @@ module DiscourseAi
 
         const chain = {
           setCustomRaw: _chain_set_custom_raw,
+          streamCustomRaw: _chain_stream_custom_raw,
         };
 
         const discourse = {
@@ -132,6 +133,27 @@ module DiscourseAi
             }
             return result;
           },
+          createStagedUser: function(params) {
+            const result = _discourse_create_staged_user(params);
+            if (result.error) {
+              throw new Error(result.error);
+            }
+            return result;
+          },
+          createTopic: function(params) {
+            const result = _discourse_create_topic(params);
+            if (result.error) {
+              throw new Error(result.error);
+            }
+            return result;
+          },
+          createPost: function(params) {
+            const result = _discourse_create_post(params);
+            if (result.error) {
+              throw new Error(result.error);
+            }
+            return result;
+          },
         };
 
         const context = #{JSON.generate(@context.to_json)};
@@ -182,11 +204,14 @@ module DiscourseAi
         t&.join
       end
 
-      def invoke
+      def invoke(progress_callback: nil)
+        @progress_callback = progress_callback
         mini_racer_context.eval(tool.script)
         eval_with_timeout("invoke(#{JSON.generate(parameters)})")
       rescue MiniRacer::ScriptTerminatedError
         { error: "Script terminated due to timeout" }
+      ensure
+        @progress_callback = nil
       end
 
       def has_custom_context?
@@ -258,12 +283,22 @@ module DiscourseAi
 
         mini_racer_context.attach(
           "_llm_generate",
-          ->(prompt) do
+          ->(prompt, options) do
             in_attached_function do
+              options ||= {}
+              response_format = options["response_format"]
+              if response_format && !response_format.is_a?(Hash)
+                raise Discourse::InvalidParameters.new("response_format must be a hash")
+              end
               @llm.generate(
                 convert_js_prompt_to_ruby(prompt),
                 user: llm_user,
                 feature_name: "custom_tool_#{tool.name}",
+                response_format: response_format,
+                temperature: options["temperature"],
+                top_p: options["top_p"],
+                max_tokens: options["max_tokens"],
+                stop_sequences: options["stop_sequences"],
               )
             end
           end,
@@ -316,6 +351,13 @@ module DiscourseAi
 
       def attach_chain(mini_racer_context)
         mini_racer_context.attach("_chain_set_custom_raw", ->(raw) { self.custom_raw = raw })
+        mini_racer_context.attach(
+          "_chain_stream_custom_raw",
+          ->(raw) do
+            self.custom_raw = raw
+            @progress_callback.call(raw) if @progress_callback
+          end,
+        )
       end
 
       # this is useful for polling apis
@@ -499,6 +541,166 @@ module DiscourseAi
           end,
         )
 
+        mini_racer_context.attach(
+          "_discourse_create_staged_user",
+          ->(params) do
+            in_attached_function do
+              params = params.symbolize_keys
+              email = params[:email]
+              username = params[:username]
+              name = params[:name]
+
+              # Validate parameters
+              return { error: "Missing required parameter: email" } if email.blank?
+              return { error: "Missing required parameter: username" } if username.blank?
+
+              # Check if user already exists
+              existing_user = User.find_by_email(email) || User.find_by_username(username)
+              return { error: "User already exists", user_id: existing_user.id } if existing_user
+
+              begin
+                user =
+                  User.create!(
+                    email: email,
+                    username: username,
+                    name: name || username,
+                    staged: true,
+                    approved: true,
+                    trust_level: TrustLevel[0],
+                  )
+
+                { success: true, user_id: user.id, username: user.username, email: user.email }
+              rescue => e
+                { error: "Failed to create staged user: #{e.message}" }
+              end
+            end
+          end,
+        )
+
+        mini_racer_context.attach(
+          "_discourse_create_topic",
+          ->(params) do
+            in_attached_function do
+              params = params.symbolize_keys
+              category_name = params[:category_name]
+              category_id = params[:category_id]
+              title = params[:title]
+              raw = params[:raw]
+              username = params[:username]
+              tags = params[:tags]
+
+              if category_id.blank? && category_name.blank?
+                return { error: "Missing required parameter: category_id or category_name" }
+              end
+              return { error: "Missing required parameter: title" } if title.blank?
+              return { error: "Missing required parameter: raw" } if raw.blank?
+
+              user =
+                if username.present?
+                  User.find_by(username: username)
+                else
+                  Discourse.system_user
+                end
+              return { error: "User not found: #{username}" } if user.nil?
+
+              category =
+                if category_id.present?
+                  Category.find_by(id: category_id)
+                else
+                  Category.find_by(name: category_name) || Category.find_by(slug: category_name)
+                end
+
+              return { error: "Category not found" } if category.nil?
+
+              begin
+                post_creator =
+                  PostCreator.new(
+                    user,
+                    title: title,
+                    raw: raw,
+                    category: category.id,
+                    tags: tags,
+                    skip_validations: true,
+                    guardian: Guardian.new(Discourse.system_user),
+                  )
+
+                post = post_creator.create
+
+                if post_creator.errors.present?
+                  return { error: post_creator.errors.full_messages.join(", ") }
+                end
+
+                {
+                  success: true,
+                  topic_id: post.topic_id,
+                  post_id: post.id,
+                  topic_slug: post.topic.slug,
+                  topic_url: post.topic.url,
+                }
+              rescue => e
+                { error: "Failed to create topic: #{e.message}" }
+              end
+            end
+          end,
+        )
+
+        mini_racer_context.attach(
+          "_discourse_create_post",
+          ->(params) do
+            in_attached_function do
+              params = params.symbolize_keys
+              topic_id = params[:topic_id]
+              raw = params[:raw]
+              username = params[:username]
+              reply_to_post_number = params[:reply_to_post_number]
+
+              # Validate parameters
+              return { error: "Missing required parameter: topic_id" } if topic_id.blank?
+              return { error: "Missing required parameter: raw" } if raw.blank?
+
+              # Find the user
+              user =
+                if username.present?
+                  User.find_by(username: username)
+                else
+                  Discourse.system_user
+                end
+              return { error: "User not found: #{username}" } if user.nil?
+
+              # Verify topic exists
+              topic = Topic.find_by(id: topic_id)
+              return { error: "Topic not found" } if topic.nil?
+
+              begin
+                post_creator =
+                  PostCreator.new(
+                    user,
+                    raw: raw,
+                    topic_id: topic_id,
+                    reply_to_post_number: reply_to_post_number,
+                    skip_validations: true,
+                    guardian: Guardian.new(Discourse.system_user),
+                  )
+
+                post = post_creator.create
+
+                if post_creator.errors.present?
+                  return { error: post_creator.errors.full_messages.join(", ") }
+                end
+
+                {
+                  success: true,
+                  post_id: post.id,
+                  post_number: post.post_number,
+                  cooked: post.cooked,
+                }
+              rescue => e
+                { error: "Failed to create post: #{e.message}" }
+              end
+            end
+          end,
+        )
+
         mini_racer_context.attach(
           "_discourse_search",
           ->(params) do

lib/personas/tools/custom.rb

@@ -66,8 +66,16 @@ module DiscourseAi
           super(*args, **kwargs)
         end
 
-        def invoke
-          result = runner.invoke
+        def invoke(&blk)
+          callback =
+            proc do |raw|
+              if blk
+                self.custom_raw = raw
+                @chain_next_response = false
+                blk.call(raw, true)
+              end
+            end
+          result = runner.invoke(progress_callback: callback)
           if runner.custom_raw
             self.custom_raw = runner.custom_raw
             @chain_next_response = false

spec/models/ai_tool_spec.rb

@@ -879,4 +879,301 @@ RSpec.describe AiTool do
       expect(result).to be_nil
     end
   end
+
+  context "when creating staged users" do
+    it "can create a staged user" do
+      script = <<~JS
+        function invoke(params) {
+          return discourse.createStagedUser({
+            email: params.email,
+            username: params.username,
+            name: params.name
+          });
+        }
+      JS
+
+      tool = create_tool(script: script)
+      runner =
+        tool.runner(
+          { email: "testuser@example.com", username: "testuser123", name: "Test User" },
+          llm: nil,
+          bot_user: nil,
+        )
+
+      result = runner.invoke
+
+      expect(result["success"]).to eq(true)
+      expect(result["username"]).to eq("testuser123")
+      expect(result["email"]).to eq("testuser@example.com")
+
+      user = User.find_by(id: result["user_id"])
+      expect(user).not_to be_nil
+      expect(user.staged).to eq(true)
+      expect(user.username).to eq("testuser123")
+      expect(user.email).to eq("testuser@example.com")
+      expect(user.name).to eq("Test User")
+    end
+
+    it "returns an error if user already exists" do
+      existing_user = Fabricate(:user, email: "existing@example.com", username: "existinguser")
+
+      script = <<~JS
+        function invoke(params) {
+          try {
+          return discourse.createStagedUser({
+            email: params.email,
+            username: params.username
+          });
+          } catch (e) {
+            return { error: e.message };
+          }
+        }
+      JS
+
+      tool = create_tool(script: script)
+      runner =
+        tool.runner(
+          { email: existing_user.email, username: "newusername" },
+          llm: nil,
+          bot_user: nil,
+        )
+
+      result = runner.invoke
+
+      expect(result["error"]).to eq("User already exists")
+    end
+  end
+
+  context "when creating topics" do
+    fab!(:category)
+    fab!(:user) { Fabricate(:admin) }
+
+    it "can create a topic" do
+      script = <<~JS
+        function invoke(params) {
+          return discourse.createTopic({
+            category_id: params.category_id,
+            title: params.title,
+            raw: params.raw,
+            username: params.username,
+            tags: params.tags
+          });
+        }
+      JS
+
+      tool = create_tool(script: script)
+      runner =
+        tool.runner(
+          {
+            category_id: category.id,
+            title: "Test Topic Title",
+            raw: "This is the content of the test topic",
+            username: user.username,
+            tags: %w[test example],
+          },
+          llm: nil,
+          bot_user: nil,
+        )
+
+      result = runner.invoke
+
+      expect(result["success"]).to eq(true)
+      expect(result["topic_id"]).to be_present
+      expect(result["post_id"]).to be_present
+
+      topic = Topic.find_by(id: result["topic_id"])
+      expect(topic).not_to be_nil
+      expect(topic.title).to eq("Test Topic Title")
+      expect(topic.category_id).to eq(category.id)
+      expect(topic.user_id).to eq(user.id)
+      expect(topic.archetype).to eq("regular")
+      expect(topic.tags.pluck(:name)).to contain_exactly("test", "example")
+
+      post = Post.find_by(id: result["post_id"])
+      expect(post).not_to be_nil
+      expect(post.raw).to eq("This is the content of the test topic")
+    end
+
+    it "can create a topic without username (uses system user)" do
+      script = <<~JS
+        function invoke(params) {
+          return discourse.createTopic({
+            category_id: params.category_id,
+            title: params.title,
+            raw: params.raw
+          });
+        }
+      JS
+
+      tool = create_tool(script: script)
+      runner =
+        tool.runner(
+          { category_id: category.id, title: "System User Topic", raw: "Created by system" },
+          llm: nil,
+          bot_user: nil,
+        )
+
+      result = runner.invoke
+
+      expect(result["success"]).to eq(true)
+
+      topic = Topic.find_by(id: result["topic_id"])
+      expect(topic.user_id).to eq(Discourse.system_user.id)
+    end
+
+    it "returns an error for invalid category" do
+      script = <<~JS
+        function invoke(params) {
+          return discourse.createTopic({
+            category_id: 99999,
+            title: "Test",
+            raw: "Test"
+          });
+        }
+      JS
+
+      tool = create_tool(script: script)
+      runner = tool.runner({}, llm: nil, bot_user: nil)
+
+      expect { runner.invoke }.to raise_error(MiniRacer::RuntimeError, /Category not found/)
+    end
+  end
+
+  context "when creating posts" do
+    fab!(:topic) { Fabricate(:post).topic }
+    fab!(:user)
+
+    it "can create a post in a topic" do
+      script = <<~JS
+        function invoke(params) {
+          return discourse.createPost({
+            topic_id: params.topic_id,
+            raw: params.raw,
+            username: params.username
+          });
+        }
+      JS
+
+      tool = create_tool(script: script)
+      runner =
+        tool.runner(
+          { topic_id: topic.id, raw: "This is a reply to the topic", username: user.username },
+          llm: nil,
+          bot_user: nil,
+        )
+
+      result = runner.invoke
+
+      expect(result["success"]).to eq(true)
+      expect(result["post_id"]).to be_present
+      expect(result["post_number"]).to be > 1
+
+      post = Post.find_by(id: result["post_id"])
+      expect(post).not_to be_nil
+      expect(post.raw).to eq("This is a reply to the topic")
+      expect(post.topic_id).to eq(topic.id)
+      expect(post.user_id).to eq(user.id)
+    end
+
+    it "can create a reply to a specific post" do
+      _original_post = Fabricate(:post, topic: topic, post_number: 2)
+
+      script = <<~JS
+        function invoke(params) {
+          return discourse.createPost({
+            topic_id: params.topic_id,
+            raw: params.raw,
+            reply_to_post_number: params.reply_to_post_number
+          });
+        }
+      JS
+
+      tool = create_tool(script: script)
+      runner =
+        tool.runner(
+          { topic_id: topic.id, raw: "This is a reply to post #2", reply_to_post_number: 2 },
+          llm: nil,
+          bot_user: nil,
+        )
+
+      result = runner.invoke
+
+      expect(result["success"]).to eq(true)
+
+      post = Post.find_by(id: result["post_id"])
+      expect(post.reply_to_post_number).to eq(2)
+    end
+
+    it "returns an error for invalid topic" do
+      script = <<~JS
+        function invoke(params) {
+          return discourse.createPost({
+            topic_id: 99999,
+            raw: "Test"
+          });
+        }
+      JS
+
+      tool = create_tool(script: script)
+      runner = tool.runner({}, llm: nil, bot_user: nil)
+
+      expect { runner.invoke }.to raise_error(MiniRacer::RuntimeError, /Topic not found/)
+    end
+  end
+
+  context "when seeding a category with topics" do
+    fab!(:category)
+
+    it "can seed a category with a topic and post" do
+      script = <<~JS
+        function invoke(params) {
+          // Create a staged user
+          const user = discourse.createStagedUser({
+            email: 'testuser@example.com',
+            username: 'testuser',
+            name: 'Test User'
+          });
+
+          // Create a topic
+          const topic = discourse.createTopic({
+            category_name: params.category_name,
+            title: 'Test Topic 123 123 123',
+            raw: 'This is the initial post content.',
+            username: user.username
+          });
+
+          // Add an extra post to the topic
+          const post = discourse.createPost({
+            topic_id: topic.topic_id,
+            raw: 'This is a reply to the topic.',
+            username: user.username
+          });
+
+          return {
+            success: true,
+            user: user,
+            topic: topic,
+            post: post
+          };
+        }
+      JS
+
+      tool = create_tool(script: script)
+      runner = tool.runner({ category_name: category.name }, llm: nil, bot_user: nil)
+
+      result = runner.invoke
+
+      expect(result["success"]).to eq(true)
+
+      user = User.find_by(username: "testuser")
+      expect(user).not_to be_nil
+      expect(user.staged).to eq(true)
+
+      topic = Topic.find_by(id: result["topic"]["topic_id"])
+      expect(topic).not_to be_nil
+      expect(topic.category_id).to eq(category.id)
+
+      expect(topic.posts.count).to eq(2)
+    end
+  end
 end