diff --git a/app/models/ai/list_collections_tool.rb b/app/models/ai/list_collections_tool.rb index 3e85eefd3..d7d73f1c0 100644 --- a/app/models/ai/list_collections_tool.rb +++ b/app/models/ai/list_collections_tool.rb @@ -13,6 +13,10 @@ class Ai::ListCollectionsTool < Ai::Tool type: :string, desc: "Which page to return. Leave blank to get the first page", required: false + param :ids, + type: :string, + desc: "If provided, will return only the collections with the given IDs (comma-separated)", + required: false attr_reader :user @@ -21,7 +25,7 @@ class Ai::ListCollectionsTool < Ai::Tool end def execute(**params) - collections = user.collections + collections = Filter.new(scope: user.collections, filters: params).filter # TODO: The serialization here is temporary until we add an API, # then we can re-use the jbuilder views and caching from that diff --git a/app/models/ai/list_collections_tool/filter.rb b/app/models/ai/list_collections_tool/filter.rb new file mode 100644 index 000000000..0d843762c --- /dev/null +++ b/app/models/ai/list_collections_tool/filter.rb @@ -0,0 +1,8 @@ +class Ai::ListCollectionsTool::Filter < Ai::Tool::Filter + register_filter :ids, :apply_ids_filter + + private + def apply_ids_filter(scope) + scope.where(id: filters[:ids].split(",").map(&:strip)) + end +end diff --git a/test/models/ai/list_collections_tool_test.rb b/test/models/ai/list_collections_tool_test.rb new file mode 100644 index 000000000..bef20ecd1 --- /dev/null +++ b/test/models/ai/list_collections_tool_test.rb @@ -0,0 +1,28 @@ +require "test_helper" + +class Ai::ListCollectionsToolTest < ActiveSupport::TestCase + include McpHelper + + setup do + @tool = Ai::ListCollectionsTool.new(user: users(:kevin)) + end + + test "execute" do + response = @tool.execute + page = parse_paginated_response(response) + + assert page[:records].is_a?(Array) + end + + test "execute when filtering by ids" do + collections = collections(:writebook, :private) + collection_ids = collections.pluck(:id) + + response = @tool.execute(ids: collection_ids.join(", ")) + page = parse_paginated_response(response) + record_ids = page[:records].map { |collection| collection["id"].to_i } + + assert_equal 2, record_ids.count + assert_equal collection_ids.sort, record_ids.sort + end +end