discourse-ai/spec/models/ai_tool_spec.rb

# frozen_string_literal: true

RSpec.describe AiTool do
  fab!(:llm_model) { Fabricate(:llm_model, name: "claude-2") }
  let(:llm) { DiscourseAi::Completions::Llm.proxy("custom:#{llm_model.id}") }
  fab!(:topic)
  fab!(:post) { Fabricate(:post, topic: topic, raw: "bananas are a tasty fruit") }
  fab!(:bot_user) { Discourse.system_user }

  def create_tool(
    parameters: nil,
    script: nil,
    rag_chunk_tokens: nil,
    rag_chunk_overlap_tokens: nil
  )
    AiTool.create!(
      name: "test #{SecureRandom.uuid}",
      tool_name: "test_#{SecureRandom.uuid.underscore}",
      description: "test",
      parameters:
        parameters || [{ name: "query", type: "string", description: "perform a search" }],
      script: script || "function invoke(params) { return params; }",
      created_by_id: 1,
      summary: "Test tool summary",
      rag_chunk_tokens: rag_chunk_tokens || 374,
      rag_chunk_overlap_tokens: rag_chunk_overlap_tokens || 10,
    )
  end

  it "it can run a basic tool" do
    tool = create_tool

    expect(tool.signature).to eq(
      {
        name: tool.tool_name,
        description: "test",
        parameters: [{ name: "query", type: "string", description: "perform a search" }],
      },
    )

    runner = tool.runner({ "query" => "test" }, llm: nil, bot_user: nil)

    expect(runner.invoke).to eq("query" => "test")
  end

  it "can perform HTTP requests with various verbs" do
    %i[post put delete patch].each do |verb|
      script = <<~JS
      function invoke(params) {
        result = http.#{verb}("https://example.com/api",
          {
            headers: { TestHeader: "TestValue" },
            body: JSON.stringify({ data: params.data })
          }
        );

        return result.body;
      }
    JS

      tool = create_tool(script: script)
      runner = tool.runner({ "data" => "test data" }, llm: nil, bot_user: nil)

      stub_request(verb, "https://example.com/api").with(
        body: "{\"data\":\"test data\"}",
        headers: {
          "Accept" => "*/*",
          "Testheader" => "TestValue",
          "User-Agent" => "Discourse AI Bot 1.0 (https://www.discourse.org)",
        },
      ).to_return(status: 200, body: "Success", headers: {})

      result = runner.invoke

      expect(result).to eq("Success")
    end
  end

  it "can perform GET HTTP requests, with 1 param" do
    script = <<~JS
      function invoke(params) {
        result = http.get("https://example.com/" + params.query);
        return result.body;
      }
    JS

    tool = create_tool(script: script)
    runner = tool.runner({ "query" => "test" }, llm: nil, bot_user: nil)

    stub_request(:get, "https://example.com/test").with(
      headers: {
        "Accept" => "*/*",
        "User-Agent" => "Discourse AI Bot 1.0 (https://www.discourse.org)",
      },
    ).to_return(status: 200, body: "Hello World", headers: {})

    result = runner.invoke

    expect(result).to eq("Hello World")
  end

  it "is limited to MAX http requests" do
    script = <<~JS
      function invoke(params) {
        let i = 0;
        while (i < 21) {
          http.get("https://example.com/");
          i += 1;
        }
        return "will not happen";
      }
      JS

    tool = create_tool(script: script)
    runner = tool.runner({}, llm: nil, bot_user: nil)

    stub_request(:get, "https://example.com/").to_return(
      status: 200,
      body: "Hello World",
      headers: {
      },
    )

    expect { runner.invoke }.to raise_error(DiscourseAi::Personas::ToolRunner::TooManyRequestsError)
  end

  it "can perform GET HTTP requests" do
    script = <<~JS
      function invoke(params) {
        result = http.get("https://example.com/" + params.query,
          { headers: { TestHeader: "TestValue" } }
        );

        return result.body;
      }
    JS

    tool = create_tool(script: script)
    runner = tool.runner({ "query" => "test" }, llm: nil, bot_user: nil)

    stub_request(:get, "https://example.com/test").with(
      headers: {
        "Accept" => "*/*",
        "Testheader" => "TestValue",
        "User-Agent" => "Discourse AI Bot 1.0 (https://www.discourse.org)",
      },
    ).to_return(status: 200, body: "Hello World", headers: {})

    result = runner.invoke

    expect(result).to eq("Hello World")
  end

  it "will not timeout on slow HTTP reqs" do
    script = <<~JS
      function invoke(params) {
        result = http.get("https://example.com/" + params.query,
          { headers: { TestHeader: "TestValue" } }
        );

        return result.body;
      }
    JS

    tool = create_tool(script: script)
    runner = tool.runner({ "query" => "test" }, llm: nil, bot_user: nil)

    stub_request(:get, "https://example.com/test").to_return do
      sleep 0.01
      { status: 200, body: "Hello World", headers: {} }
    end

    tool = create_tool(script: script)
    runner = tool.runner({ "query" => "test" }, llm: nil, bot_user: nil)

    runner.timeout = 10

    result = runner.invoke

    expect(result).to eq("Hello World")
  end

  it "has access to llm truncation tools" do
    script = <<~JS
      function invoke(params) {
        return llm.truncate("Hello World", 1);
      }
    JS

    tool = create_tool(script: script)

    runner = tool.runner({}, llm: llm, bot_user: nil)
    result = runner.invoke

    expect(result).to eq("Hello")
  end

  it "is able to run llm completions" do
    script = <<~JS
      function invoke(params) {
        return llm.generate("question two") + llm.generate(
          { messages: [
            { type: "system", content: "system message" },
            { type: "user", content: "user message" }
          ]}
        );
      }
    JS

    tool = create_tool(script: script)

    result = nil
    prompts = nil
    responses = ["Hello ", "World"]

    DiscourseAi::Completions::Llm.with_prepared_responses(responses) do |_, _, _prompts|
      runner = tool.runner({}, llm: llm, bot_user: nil)
      result = runner.invoke
      prompts = _prompts
    end

    prompt =
      DiscourseAi::Completions::Prompt.new(
        "system message",
        messages: [{ type: :user, content: "user message" }],
      )
    expect(result).to eq("Hello World")
    expect(prompts[0]).to eq("question two")
    expect(prompts[1]).to eq(prompt)
  end

  it "can timeout slow JS" do
    script = <<~JS
      function invoke(params) {
        while (true) {}
      }
    JS

    tool = create_tool(script: script)
    runner = tool.runner({ "query" => "test" }, llm: nil, bot_user: nil)

    runner.timeout = 5

    result = runner.invoke
    expect(result[:error]).to eq("Script terminated due to timeout")
  end

  context "when defining RAG fragments" do
    fab!(:cloudflare_embedding_def)

    before do
      SiteSetting.authorized_extensions = "txt"
      SiteSetting.ai_embeddings_selected_model = cloudflare_embedding_def.id
      SiteSetting.ai_embeddings_enabled = true
      Jobs.run_immediately!
    end

    def create_upload(content, filename)
      upload = nil
      Tempfile.create(filename) do |file|
        file.write(content)
        file.rewind

        upload = UploadCreator.new(file, filename).create_for(Discourse.system_user.id)
      end
      upload
    end

    def stub_embeddings
      # this is a trick, we get ever increasing embeddings, this gives us in turn
      # 100% consistent search results
      @counter = 0
      stub_request(:post, cloudflare_embedding_def.url).to_return(
        status: 200,
        body: lambda { |req| { result: { data: [([@counter += 1] * 1024)] } }.to_json },
        headers: {
        },
      )
    end

    it "allows search within uploads" do
      stub_embeddings

      upload1 = create_upload(<<~TXT, "test.txt")
        1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30
      TXT

      upload2 = create_upload(<<~TXT, "test.txt")
        30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50
      TXT

      tool = create_tool(rag_chunk_tokens: 10, rag_chunk_overlap_tokens: 4, script: <<~JS)
        function invoke(params) {
          let result1 = index.search("testing a search", { limit: 1 });
          let result2 = index.search("testing another search", { limit: 3, filenames: ["test.txt"] });

          return [result1, result2];
        }
      JS

      RagDocumentFragment.link_target_and_uploads(tool, [upload1.id, upload2.id])

      result = tool.runner({}, llm: nil, bot_user: nil).invoke

      expected = [
        [{ "fragment" => "44 45 46 47 48 49 50", "metadata" => nil }],
        [
          { "fragment" => "44 45 46 47 48 49 50", "metadata" => nil },
          { "fragment" => "36 37 38 39 40 41 42 43 44 45", "metadata" => nil },
          { "fragment" => "30 31 32 33 34 35 36 37", "metadata" => nil },
        ],
      ]

      expect(result).to eq(expected)

      # will force a reindex
      tool.rag_chunk_tokens = 5
      tool.rag_chunk_overlap_tokens = 2
      tool.save!

      # this part of the API is a bit awkward, maybe we should do it
      # automatically
      RagDocumentFragment.update_target_uploads(tool, [upload1.id, upload2.id])
      result = tool.runner({}, llm: nil, bot_user: nil).invoke

      expected = [
        [{ "fragment" => "48 49 50", "metadata" => nil }],
        [
          { "fragment" => "48 49 50", "metadata" => nil },
          { "fragment" => "45 46 47", "metadata" => nil },
          { "fragment" => "42 43 44", "metadata" => nil },
        ],
      ]

      expect(result).to eq(expected)
    end
  end

  context "when using the topic API" do
    it "can fetch topic details" do
      script = <<~JS
        function invoke(params) {
          return discourse.getTopic(params.topic_id);
        }
      JS

      tool = create_tool(script: script)
      runner = tool.runner({ "topic_id" => topic.id }, llm: nil, bot_user: nil)

      result = runner.invoke

      expect(result["id"]).to eq(topic.id)
      expect(result["title"]).to eq(topic.title)
      expect(result["archetype"]).to eq("regular")
      expect(result["posts_count"]).to eq(1)
    end
  end

  context "when using the post API" do
    it "can fetch post details" do
      script = <<~JS
        function invoke(params) {
          const post = discourse.getPost(params.post_id);
          return {
            post: post,
            topic: post.topic
          }
        }
      JS

      tool = create_tool(script: script)
      runner = tool.runner({ "post_id" => post.id }, llm: nil, bot_user: nil)

      result = runner.invoke
      post_hash = result["post"]
      topic_hash = result["topic"]

      expect(post_hash["id"]).to eq(post.id)
      expect(post_hash["topic_id"]).to eq(topic.id)
      expect(post_hash["raw"]).to eq(post.raw)

      expect(topic_hash["id"]).to eq(topic.id)
    end
  end

  context "when using the search API" do
    before { SearchIndexer.enable }
    after { SearchIndexer.disable }

    it "can perform a discourse search" do
      SearchIndexer.index(topic, force: true)
      SearchIndexer.index(post, force: true)

      script = <<~JS
        function invoke(params) {
          return discourse.search({ search_query: params.query });
        }
      JS

      tool = create_tool(script: script)
      runner = tool.runner({ "query" => "banana" }, llm: nil, bot_user: nil)

      result = runner.invoke

      expect(result["rows"].length).to be > 0
      expect(result["rows"].first["title"]).to eq(topic.title)
    end
  end

  context "when using the chat API" do
    before(:each) do
      skip "Chat plugin tests skipped because Chat module is not defined." unless defined?(Chat)
      SiteSetting.chat_enabled = true
    end

    fab!(:chat_user) { Fabricate(:user) }
    fab!(:chat_channel) do
      Fabricate(:chat_channel).tap do |channel|
        Fabricate(
          :user_chat_channel_membership,
          user: chat_user,
          chat_channel: channel,
          following: true,
        )
      end
    end

    it "can create a chat message" do
      script = <<~JS
        function invoke(params) {
          return discourse.createChatMessage({
            channel_name: params.channel_name,
            username: params.username,
            message: params.message
          });
        }
      JS

      tool = create_tool(script: script)
      runner =
        tool.runner(
          {
            "channel_name" => chat_channel.name,
            "username" => chat_user.username,
            "message" => "Hello from the tool!",
          },
          llm: nil,
          bot_user: bot_user, # The user *running* the tool doesn't affect sender
        )

      initial_message_count = Chat::Message.count
      result = runner.invoke

      expect(result["success"]).to eq(true), "Tool invocation failed: #{result["error"]}"
      expect(result["message"]).to eq("Hello from the tool!")
      expect(result["created_at"]).to be_present
      expect(result).not_to have_key("error")

      # Verify message was actually created in the database
      expect(Chat::Message.count).to eq(initial_message_count + 1)
      created_message = Chat::Message.find_by(id: result["message_id"])

      expect(created_message).not_to be_nil
      expect(created_message.message).to eq("Hello from the tool!")
      expect(created_message.user_id).to eq(chat_user.id) # Message is sent AS the specified user
      expect(created_message.chat_channel_id).to eq(chat_channel.id)
    end

    it "can create a chat message using channel slug" do
      chat_channel.update!(name: "My Test Channel", slug: "my-test-channel")
      expect(chat_channel.slug).to eq("my-test-channel")

      script = <<~JS
        function invoke(params) {
          return discourse.createChatMessage({
            channel_name: params.channel_slug, // Using slug here
            username: params.username,
            message: params.message
          });
        }
      JS

      tool = create_tool(script: script)
      runner =
        tool.runner(
          {
            "channel_slug" => chat_channel.slug,
            "username" => chat_user.username,
            "message" => "Hello via slug!",
          },
          llm: nil,
          bot_user: bot_user,
        )

      result = runner.invoke

      expect(result["success"]).to eq(true), "Tool invocation failed: #{result["error"]}"
      # see: https://github.com/rubyjs/mini_racer/issues/348
      # expect(result["message_id"]).to be_a(Integer)

      created_message = Chat::Message.find_by(id: result["message_id"])
      expect(created_message).not_to be_nil
      expect(created_message.message).to eq("Hello via slug!")
      expect(created_message.chat_channel_id).to eq(chat_channel.id)
    end

    it "returns an error if the channel is not found" do
      script = <<~JS
        function invoke(params) {
          return discourse.createChatMessage({
            channel_name: "non_existent_channel",
            username: params.username,
            message: params.message
          });
        }
      JS

      tool = create_tool(script: script)
      runner =
        tool.runner(
          { "username" => chat_user.username, "message" => "Test" },
          llm: nil,
          bot_user: bot_user,
        )

      initial_message_count = Chat::Message.count
      expect { runner.invoke }.to raise_error(
        MiniRacer::RuntimeError,
        /Channel not found: non_existent_channel/,
      )

      expect(Chat::Message.count).to eq(initial_message_count) # Verify no message created
    end

    it "returns an error if the user is not found" do
      script = <<~JS
        function invoke(params) {
          return discourse.createChatMessage({
            channel_name: params.channel_name,
            username: "non_existent_user",
            message: params.message
          });
        }
      JS

      tool = create_tool(script: script)
      runner =
        tool.runner(
          { "channel_name" => chat_channel.name, "message" => "Test" },
          llm: nil,
          bot_user: bot_user,
        )

      initial_message_count = Chat::Message.count
      expect { runner.invoke }.to raise_error(
        MiniRacer::RuntimeError,
        /User not found: non_existent_user/,
      )

      expect(Chat::Message.count).to eq(initial_message_count) # Verify no message created
    end
  end

  context "when updating personas" do
    fab!(:ai_persona) do
      Fabricate(:ai_persona, name: "TestPersona", system_prompt: "Original prompt")
    end

    it "can update a persona with proper permissions" do
      script = <<~JS
        function invoke(params) {
          return discourse.updatePersona(params.persona_name, {
            system_prompt: params.new_prompt,
            temperature: 0.7,
            top_p: 0.9
          });
        }
      JS

      tool = create_tool(script: script)
      runner =
        tool.runner(
          { persona_name: "TestPersona", new_prompt: "Updated system prompt" },
          llm: nil,
          bot_user: bot_user,
        )

      result = runner.invoke
      expect(result["success"]).to eq(true)
      expect(result["persona"]["system_prompt"]).to eq("Updated system prompt")
      expect(result["persona"]["temperature"]).to eq(0.7)

      ai_persona.reload
      expect(ai_persona.system_prompt).to eq("Updated system prompt")
      expect(ai_persona.temperature).to eq(0.7)
      expect(ai_persona.top_p).to eq(0.9)
    end
  end

  context "when fetching persona information" do
    fab!(:ai_persona) do
      Fabricate(
        :ai_persona,
        name: "TestPersona",
        description: "Test description",
        system_prompt: "Test system prompt",
        temperature: 0.8,
        top_p: 0.9,
        vision_enabled: true,
        tools: ["Search", ["WebSearch", { param: "value" }, true]],
      )
    end

    it "can fetch a persona by name" do
      script = <<~JS
        function invoke(params) {
          const persona = discourse.getPersona(params.persona_name);
          return persona;
        }
      JS

      tool = create_tool(script: script)
      runner = tool.runner({ persona_name: "TestPersona" }, llm: nil, bot_user: bot_user)

      result = runner.invoke

      expect(result["id"]).to eq(ai_persona.id)
      expect(result["name"]).to eq("TestPersona")
      expect(result["description"]).to eq("Test description")
      expect(result["system_prompt"]).to eq("Test system prompt")
      expect(result["temperature"]).to eq(0.8)
      expect(result["top_p"]).to eq(0.9)
      expect(result["vision_enabled"]).to eq(true)
      expect(result["tools"]).to include("Search")
      expect(result["tools"][1]).to be_a(Array)
    end

    it "raises an error when the persona doesn't exist" do
      script = <<~JS
        function invoke(params) {
          return discourse.getPersona("NonExistentPersona");
        }
      JS

      tool = create_tool(script: script)
      runner = tool.runner({}, llm: nil, bot_user: bot_user)

      expect { runner.invoke }.to raise_error(MiniRacer::RuntimeError, /Persona not found/)
    end

    it "can update a persona after fetching it" do
      script = <<~JS
        function invoke(params) {
          const persona = discourse.getPersona("TestPersona");
          return persona.update({
            system_prompt: "Updated through getPersona().update()",
            temperature: 0.5
          });
        }
      JS

      tool = create_tool(script: script)
      runner = tool.runner({}, llm: nil, bot_user: bot_user)

      result = runner.invoke
      expect(result["success"]).to eq(true)

      ai_persona.reload
      expect(ai_persona.system_prompt).to eq("Updated through getPersona().update()")
      expect(ai_persona.temperature).to eq(0.5)
    end
  end
end