feat: defining tools at the resource level

README.md

@@ -199,6 +199,25 @@ end
 Expose these actions as tools. Use `AshAi.build_tools_and_registry/1` to get ReqLLM tools and callbacks,
 or use `AshAi.ToolLoop.run/2` / `AshAi.ToolLoop.stream/2` to execute the full model + tool loop.
 
+You can also define tools directly on a resource using the 2-argument shorthand:
+
+```elixir
+defmodule MyApp.Blog.Post do
+  use Ash.Resource, extensions: [AshAi]
+
+  tools do
+    tool :read_posts, :read
+    tool :create_post, :create
+  end
+end
+```
+
+This shorthand expands to the same tool definition as `tool :name, MyApp.Blog.Post, :action`.
+
+When discovering tools with `AshAi.exposed_tools/1` and `actions: [{Resource, ...}]`, AshAi includes:
+- tools defined on the domain for that resource
+- tools defined directly on the resource
+
 For migration guidance around `extra_tools`, `req_llm_opts`, and legacy adapter mapping, see [LangChain to ReqLLM Migration Guide](/documentation/topics/langchain-to-reqllm-migration.md).
 
 ## Expose content as MCP resources

lib/ash_ai.ex

@@ -16,7 +16,11 @@ defmodule AshAi do
   use Spark.Dsl.Extension,
     sections: AshAi.Dsl.sections(),
     imports: [AshAi.Actions],
-    transformers: [AshAi.Transformers.Vectorize, AshAi.Transformers.McpApps],
+    transformers: [
+      AshAi.Transformers.Vectorize,
+      AshAi.Transformers.ResourceTools,
+      AshAi.Transformers.McpApps
+    ],
     verifiers: [AshAi.Verifiers.McpResourceActionsReturnString]
 
   defmodule Tool do
@@ -690,32 +694,19 @@ defmodule AshAi do
       if opts.actions do
         Enum.flat_map(opts.actions, fn
           {resource, actions} ->
-            domain = Ash.Resource.Info.domain(resource)
-
-            if !domain do
-              raise "Cannot use an ash resource that does not have a domain"
-            end
-
-            tools = AshAi.Info.tools(domain)
+            tools = tools_for_resource(resource)
 
             if !Enum.any?(tools, fn tool ->
-                 tool.resource == resource && (actions == :* || tool.action in actions)
+                 actions == :* || tool.action.name in actions
                end) do
               raise "Cannot use an action that is not exposed as a tool"
             end
 
             if actions == :* do
               tools
-              |> Enum.filter(&(&1.resource == resource))
-              |> Enum.map(fn tool ->
-                %{tool | domain: domain, action: Ash.Resource.Info.action(resource, tool.action)}
-              end)
             else
               tools
-              |> Enum.filter(&(&1.resource == resource && &1.action in actions))
-              |> Enum.map(fn tool ->
-                %{tool | domain: domain, action: Ash.Resource.Info.action(resource, tool.action)}
-              end)
+              |> Enum.filter(&(&1.action.name in actions))
             end
         end)
       else
@@ -724,8 +715,8 @@ defmodule AshAi do
         end
 
         for domain <- Application.get_env(opts.otp_app, :ash_domains) || [],
-            tool <- AshAi.Info.tools(domain) do
-          %{tool | domain: domain, action: Ash.Resource.Info.action(tool.resource, tool.action)}
+            tool <- tools_for_domain(domain) do
+          tool
         end
       end
     end
@@ -757,6 +748,39 @@ defmodule AshAi do
     )
   end
 
+  defp tools_for_domain(domain) do
+    domain_tools = attach_tool_runtime_details(AshAi.Info.tools(domain), domain)
+
+    resource_tools =
+      domain
+      |> Ash.Domain.Info.resources()
+      |> Enum.flat_map(fn resource ->
+        resource
+        |> AshAi.Info.tools()
+        |> attach_tool_runtime_details(domain)
+      end)
+
+    Enum.uniq(domain_tools ++ resource_tools)
+  end
+
+  defp tools_for_resource(resource) do
+    domain = Ash.Resource.Info.domain(resource)
+
+    if !domain do
+      raise "Cannot use an ash resource that does not have a domain"
+    end
+
+    domain
+    |> tools_for_domain()
+    |> Enum.filter(&(&1.resource == resource))
+  end
+
+  defp attach_tool_runtime_details(tools, domain) do
+    Enum.map(tools, fn tool ->
+      %{tool | domain: domain, action: Ash.Resource.Info.action(tool.resource, tool.action)}
+    end)
+  end
+
   def has_vectorize_change?(%Ash.Changeset{} = changeset) do
     full_text_attrs =
       AshAi.Info.vectorize(changeset.resource) |> Enum.flat_map(& &1.used_attributes)

lib/ash_ai/dsl.ex

@@ -45,7 +45,7 @@ defmodule AshAi.Dsl do
 
   @tool_schema [
     name: [type: :atom, required: true],
-    resource: [type: {:spark, Ash.Resource}, required: true],
+    resource: [type: {:spark, Ash.Resource}, required: false],
     action: [type: :atom, required: true],
     action_parameters: [
       type: {:list, :atom},
@@ -208,7 +208,7 @@ defmodule AshAi.Dsl do
     ],
     target: AshAi.Tool,
     schema: @tool_schema,
-    args: [:name, :resource, :action],
+    args: [:name, {:optional, :resource}, :action],
     entities: [
       arguments: [@tool_argument]
     ]

lib/ash_ai/transformers/resource_tools.ex

@@ -0,0 +1,53 @@
+# SPDX-FileCopyrightText: 2024 ash_ai contributors <https://github.com/ash-project/ash_ai/graphs/contributors>
+#
+# SPDX-License-Identifier: MIT
+
+defmodule AshAi.Transformers.ResourceTools do
+  @moduledoc false
+  use Spark.Dsl.Transformer
+
+  alias Spark.Dsl.Transformer
+
+  def transform(dsl_state) do
+    module = Transformer.get_persisted(dsl_state, :module)
+    resource_dsl? = not is_nil(Ash.Resource.Info.data_layer(dsl_state))
+
+    dsl_state
+    |> Transformer.get_entities([:tools])
+    |> Enum.reduce(dsl_state, fn tool, dsl ->
+      cond do
+        resource_dsl? and is_nil(tool.resource) ->
+          Transformer.replace_entity(
+            dsl,
+            [:tools],
+            %{tool | resource: module},
+            &(&1.name == tool.name)
+          )
+
+        resource_dsl? ->
+          raise Spark.Error.DslError,
+            module: module,
+            path: [:tools, tool.name, :resource],
+            message: """
+            Resource-level tools cannot set `resource`.
+
+            Inside an Ash.Resource, define tools as `tool :name, :action`.
+            """
+
+        is_nil(tool.resource) ->
+          raise Spark.Error.DslError,
+            module: module,
+            path: [:tools, tool.name, :resource],
+            message: """
+            Tool `#{tool.name}` is missing a resource.
+
+            On domains, define tools as `tool :name, Resource, :action`.
+            """
+
+        true ->
+          dsl
+      end
+    end)
+    |> then(&{:ok, &1})
+  end
+end

test/ash_ai/resource_tools_test.exs

@@ -0,0 +1,127 @@
+# SPDX-FileCopyrightText: 2024 ash_ai contributors <https://github.com/ash-project/ash_ai/graphs/contributors>
+#
+# SPDX-License-Identifier: MIT
+
+defmodule AshAi.ResourceToolsTest do
+  use ExUnit.Case, async: true
+
+  alias __MODULE__.{ResourceToolDomain, ResourceToolResource}
+
+  defmodule ResourceToolResource do
+    use Ash.Resource,
+      domain: ResourceToolDomain,
+      extensions: [AshAi],
+      data_layer: Ash.DataLayer.Ets,
+      validate_domain_inclusion?: false
+
+    attributes do
+      uuid_v7_primary_key(:id, writable?: true)
+      attribute :name, :string, public?: true
+    end
+
+    actions do
+      default_accept([:id, :name])
+      defaults([:read, :create])
+    end
+
+    tools do
+      tool(:resource_read, :read)
+      tool(:resource_create, :create)
+    end
+  end
+
+  defmodule ResourceToolDomain do
+    use Ash.Domain, extensions: [AshAi], validate_config_inclusion?: false
+
+    resources do
+      resource ResourceToolResource
+    end
+
+    tools do
+      tool :domain_create_alias, ResourceToolResource, :create
+    end
+  end
+
+  describe "resource-level tools shorthand" do
+    test "tool :name, :action populates resource and action metadata" do
+      tool =
+        ResourceToolResource
+        |> AshAi.Info.tools()
+        |> Enum.find(&(&1.name == :resource_read))
+
+      assert tool.resource == ResourceToolResource
+      assert tool.action == :read
+    end
+
+    test "resource-level tools reject explicit resource argument" do
+      module =
+        Module.concat(__MODULE__, :"InvalidResourceTool#{System.unique_integer([:positive])}")
+
+      assert_raise Spark.Error.DslError, ~r/Resource-level tools cannot set `resource`/, fn ->
+        Module.create(
+          module,
+          quote do
+            use Ash.Resource,
+              domain: ResourceToolDomain,
+              extensions: [AshAi],
+              data_layer: Ash.DataLayer.Ets,
+              validate_domain_inclusion?: false
+
+            attributes do
+              uuid_v7_primary_key(:id, writable?: true)
+            end
+
+            actions do
+              defaults([:read])
+            end
+
+            tools do
+              tool :resource_read, ResourceToolResource, :read
+            end
+          end,
+          Macro.Env.location(__ENV__)
+        )
+      end
+    end
+  end
+
+  describe "domain-level tools requirements" do
+    test "domain-level tools require resource argument" do
+      module =
+        Module.concat(__MODULE__, :"InvalidDomainTool#{System.unique_integer([:positive])}")
+
+      assert_raise Spark.Error.DslError, ~r/is missing a resource/, fn ->
+        Module.create(
+          module,
+          quote do
+            use Ash.Domain, extensions: [AshAi], validate_config_inclusion?: false
+
+            tools do
+              tool(:missing_resource, :read)
+            end
+          end,
+          Macro.Env.location(__ENV__)
+        )
+      end
+    end
+  end
+
+  describe "AshAi.exposed_tools/1 discovery" do
+    test "actions filter includes both domain-level and resource-level tools" do
+      tools = AshAi.exposed_tools(actions: [{ResourceToolResource, :*}])
+
+      assert tools
+             |> Enum.map(& &1.name)
+             |> MapSet.new() ==
+               MapSet.new([:resource_read, :resource_create, :domain_create_alias])
+    end
+
+    test "actions filter by specific action keeps matching tools from both levels" do
+      tools = AshAi.exposed_tools(actions: [{ResourceToolResource, [:create]}])
+
+      assert tools
+             |> Enum.map(& &1.name)
+             |> MapSet.new() == MapSet.new([:resource_create, :domain_create_alias])
+    end
+  end
+end