better todos

This endpoint doesn't expect args of it's own in the body, but does accept the standard metabot conversation_id that
it echoes back.

.clj-kondo/config/modules/config.edn

@@ -1799,11 +1799,13 @@
            lib
            lib-be
            models
+           native-query-snippets
            permissions
            premium-features
            pulse
            query-processor
            request
+           search
            server
            settings
            sync
@@ -1812,7 +1814,8 @@
            util
            warehouse-schema
            warehouses
-           enterprise/documents}}
+           enterprise/documents
+           enterprise/transforms}}
 
   enterprise/permission-debug
   {:team "UX West"

enterprise/backend/src/metabase_enterprise/metabot_v3/api.clj

@@ -69,7 +69,6 @@
   "Send a chat message to the LLM via the AI Proxy."
   [_route-params
    _query-params
-
    body :- [:map
             [:profile_id {:optional true} :string]
             [:metabot_id {:optional true} :string]

enterprise/backend/src/metabase_enterprise/metabot_v3/api/metabot.clj

@@ -3,66 +3,19 @@
   (:require
    [medley.core :as m]
    [metabase-enterprise.metabot-v3.client :as metabot-v3.client]
+   [metabase-enterprise.metabot-v3.config :as metabot-v3.config]
    [metabase-enterprise.metabot-v3.dummy-tools :as metabot-v3.dummy-tools]
    [metabase-enterprise.metabot-v3.tools.util :as metabot-v3.tools.u]
    [metabase.api.common :as api]
    [metabase.api.macros :as api.macros]
    [metabase.api.routes.common :refer [+auth]]
    [metabase.app-db.core :as mdb]
-   [metabase.collections.models.collection :as collection]
    [metabase.lib-be.metadata.jvm :as lib.metadata.jvm]
+   [metabase.premium-features.core :as premium-features]
    [metabase.request.core :as request]
+   [metabase.util.i18n :refer [tru]]
    [toucan2.core :as t2]))
 
-;; TODO: Eventually this should be paged but since we are just going to hardcode two models for now
-;; lets not
-(api.macros/defendpoint :get "/"
-  "List configured metabot instances"
-  []
-  (api/check-superuser)
-  {:items (t2/select :model/Metabot {:order-by [[:name :asc]]})})
-
-(api.macros/defendpoint :get "/:id"
-  "Retrieve one metabot instance"
-  [{:keys [id]} :- [:map [:id pos-int?]]]
-  (api/check-superuser)
-  (api/check-404 (t2/select-one :model/Metabot :id id)))
-
-(def ^:private default-entities-page-size 200)
-
-(api.macros/defendpoint :get "/:id/entities"
-  "List the entities this metabot has access to"
-  [{:keys [id]} :- [:map [:id pos-int?]]]
-  (api/check-superuser)
-  (api/check-404 (t2/exists? :model/Metabot :id id))
-  (let [limit (or (request/limit) default-entities-page-size)
-        offset (or (request/offset) 0)
-        root-collection (collection/root-collection-with-ui-details nil)
-        entities (t2/select [:model/MetabotEntity
-                             :model_id
-                             :model
-                             [[:coalesce :card.name :collection.name] :name]
-                             :created_at
-                             [:parent_collection.id :collection_id]
-                             [:parent_collection.name :collection_name]]
-                            {:left-join [[:report_card :card] [:and [:in :model [[:inline "dataset"] [:inline "metric"]]]  [:= :model_id :card.id]]
-                                         [:collection :collection] [:and [:= :model [:inline "collection"]]  [:= :model_id :collection.id]]
-                                         ;; TODO: How to join parent collections from the location column?
-                                         [:collection :parent_collection] [:= :card.collection_id :parent_collection.id]]
-                             :where [:= :metabot_id id]
-                             :limit limit
-                             :offset offset
-                             :order-by [[:name :asc]]})
-        total (t2/count :model/MetabotEntity :metabot_id id)]
-    {:items (for [{:keys [collection_id model] :as entity} entities]
-              (cond-> entity
-                (and (nil? collection_id)
-                     (not= :collection model)) (assoc :collection_id (:id root-collection)
-                                                      :collection_name (:name root-collection))))
-     :total total
-     :limit limit
-     :offset offset}))
-
 (defn- column-input
   [answer-source-column]
   (some-> answer-source-column (select-keys [:name :type :description :table-reference])))
@@ -85,14 +38,14 @@
       (update :fields #(map column-input %))))
 
 (defn- generate-sample-prompts
-  [metabot-entity-ids]
+  [metabot-id]
   (lib.metadata.jvm/with-metadata-provider-cache
-    (let [scope (metabot-v3.tools.u/metabot-scope metabot-entity-ids)
+    (let [cards (metabot-v3.tools.u/get-metrics-and-models metabot-id)
           {metrics :metric, models :model}
-          (->> (for [[[card-type database-id] cards] (group-by (juxt :type :database_id) (keys scope))
+          (->> (for [[[card-type database-id] group-cards] (group-by (juxt :type :database_id) cards)
                      detail (map (fn [detail card] (assoc detail ::origin card))
-                                 (metabot-v3.dummy-tools/cards-details card-type database-id cards nil)
-                                 cards)]
+                                 (metabot-v3.dummy-tools/cards-details card-type database-id group-cards nil)
+                                 group-cards)]
                  detail)
                (group-by :type))
           metric-inputs (map metric-input metrics)
@@ -100,9 +53,9 @@
           {:keys [table_questions metric_questions]}
           (metabot-v3.client/generate-example-questions {:metrics metric-inputs, :tables model-inputs})
           ->prompt (fn [{:keys [questions]} {::keys [origin]}]
-                     (let [base {:metabot_entity_id (scope origin)
-                                 :model             (:type origin)
-                                 :card_id           (:id origin)}]
+                     (let [base {:metabot_id metabot-id
+                                 :model      (:type origin)
+                                 :card_id    (:id origin)}]
                        (map #(assoc base :prompt %) questions)))
           metric-prompts (mapcat ->prompt metric_questions metrics)
           model-prompts (mapcat ->prompt table_questions models)]
@@ -111,52 +64,54 @@
       (when (seq model-prompts)
         (t2/insert! :model/MetabotPrompt model-prompts)))))
 
-(api.macros/defendpoint :put "/:id/entities"
-  "Update the entities this metabot has access to"
-  [{:keys [id]} :- [:map [:id pos-int?]]
-   _query-params
-   {:keys [items]} :- [:map
-                       [:items [:sequential [:map
-                                             [:id pos-int?]
-                                             [:model [:enum "dataset" "metric" "collection"]]]]]]]
-  (api/check-superuser)
-  (t2/with-transaction [_conn]
-    (api/check-404 (t2/exists? :model/Metabot :id id))
-    (let [new-entity-ids (into []
-                               (keep (fn [{model-id :id model :model}]
-                                       (when-not (t2/exists? :model/MetabotEntity
-                                                             :metabot_id id
-                                                             :model model
-                                                             :model_id model-id)
-                                         (t2/insert-returning-pk! :model/MetabotEntity
-                                                                  {:metabot_id id
-                                                                   :model model
-                                                                   :model_id model-id}))))
-                               items)]
-      (generate-sample-prompts new-entity-ids)))
-  api/generic-204-no-content)
-
-(api.macros/defendpoint :delete ["/:id/entities/:model/:model-id" :model #"dataset|metric|collection"]
-  "Remove an entity from this metabot's access list"
-  [{:keys [id model model-id]} :- [:map
-                                   [:id pos-int?]
-                                   [:model [:enum "dataset" "metric" "collection"]]
-                                   [:model-id pos-int?]]]
-  (api/check-superuser)
-  (api/check-404 (t2/exists? :model/Metabot :id id))
-  (t2/delete! :model/MetabotEntity
-              :metabot_id id
-              :model model
-              :model_id model-id)
-  api/generic-204-no-content)
-
 (defn- delete-all-metabot-prompts
   [metabot-id]
-  (t2/delete! :model/MetabotPrompt {:where [:exists {:select [:*]
-                                                     :from   [[:metabot_entity :mbe]]
-                                                     :where  [:and
-                                                              [:= :mbe.id :metabot_prompt.metabot_entity_id]
-                                                              [:= :mbe.metabot_id metabot-id]]}]}))
+  (t2/delete! :model/MetabotPrompt {:where [:= :metabot_id metabot-id]}))
+
+;; TODO: Eventually this should be paged but since we are just going to hardcode two models for now
+;; lets not
+(api.macros/defendpoint :get "/"
+  "List configured metabot instances"
+  []
+  (api/check-superuser)
+  {:items (t2/select :model/Metabot {:order-by [[:name :asc]]})})
+
+(api.macros/defendpoint :get "/:id"
+  "Retrieve one metabot instance"
+  [{:keys [id]} :- [:map [:id pos-int?]]]
+  (api/check-superuser)
+  (api/check-404 (t2/select-one :model/Metabot :id id)))
+
+(api.macros/defendpoint :put "/:id"
+  "Update a metabot instance"
+  [{:keys [id]} :- [:map [:id pos-int?]]
+   _query-params
+   metabot-updates :- [:map
+                       [:use_verified_content {:optional true} :boolean]
+                       [:collection_id {:optional true} [:maybe pos-int?]]]]
+  (api/check-superuser)
+  (api/check-404 (t2/exists? :model/Metabot :id id))
+  (let [old-metabot (t2/select-one :model/Metabot :id id)]
+    ;; Prevent updating collection_id on the primary metabot instance
+    (when (and (contains? metabot-updates :collection_id)
+               (= (:entity_id old-metabot)
+                  (get-in metabot-v3.config/metabot-config [metabot-v3.config/internal-metabot-id :entity-id])))
+      (api/check-400 false "Cannot update collection_id for the primary metabot instance."))
+    ;; Prevent enabling verified content without the premium feature
+    (when (and (contains? metabot-updates :use_verified_content)
+               (:use_verified_content metabot-updates))
+      (premium-features/assert-has-feature :content-verification (tru "Content verification")))
+    (let [verified-content-changed? (and (contains? metabot-updates :use_verified_content)
+                                         (not= (:use_verified_content old-metabot)
+                                               (:use_verified_content metabot-updates)))
+          collection-changed? (and (contains? metabot-updates :collection_id)
+                                   (not= (:collection_id old-metabot)
+                                         (:collection_id metabot-updates)))]
+      (t2/update! :model/Metabot id metabot-updates)
+      (when (or verified-content-changed? collection-changed?)
+        (delete-all-metabot-prompts id)
+        (generate-sample-prompts id))
+      (t2/select-one :model/Metabot :id id))))
 
 (api.macros/defendpoint :post "/:id/prompt-suggestions/regenerate"
   "Remove any existing prompt suggestions for the Metabot instance with `id` and generate new ones."
@@ -164,9 +119,8 @@
   (api/check-superuser)
   (t2/with-transaction [_conn]
     (api/check-404 (t2/exists? :model/Metabot :id id))
-    (when-let [entity-ids (not-empty (t2/select-pks-vec :model/MetabotEntity :metabot_id id))]
-      (delete-all-metabot-prompts id)
-      (generate-sample-prompts entity-ids)))
+    (delete-all-metabot-prompts id)
+    (generate-sample-prompts id))
   api/generic-204-no-content)
 
 (api.macros/defendpoint :get "/:id/prompt-suggestions"
@@ -180,15 +134,14 @@
         rand-fn (case (mdb/db-type)
                   :postgres :random
                   :rand)
-        base-query (cond-> {:join  [[{:select [:id :name :type :metabot_entity_id]
-                                      :from   [[(metabot-v3.tools.u/metabot-scope-query
-                                                 [:= :mbe.metabot_id id])
+        base-query (cond-> {:join  [[{:select [:id :name :type]
+                                      :from   [[(metabot-v3.tools.u/metabot-metrics-and-models-query id)
                                                 :scope]]}
                                      :card]
                                     [:and
-                                     [:= :card.id                :metabot_prompt.card_id]
-                                     [:= :card.metabot_entity_id :metabot_prompt.metabot_entity_id]]]
-                            :where [:and]}
+                                     [:= :card.id :metabot_prompt.card_id]]]
+                            :where [:and
+                                    [:= :metabot_prompt.metabot_id id]]}
                      model    (update :where conj [:= :card.type model])
                      model_id (update :where conj [:= :card.id model_id]))
         total (t2/count :model/MetabotPrompt base-query)
@@ -228,11 +181,7 @@
   (api/check-superuser)
   (t2/delete! :model/MetabotPrompt {:where [:and
                                             [:= :id prompt-id]
-                                            [:exists {:select [:*]
-                                                      :from   [[:metabot_entity :mbe]]
-                                                      :where  [:and
-                                                               [:= :mbe.id :metabot_prompt.metabot_entity_id]
-                                                               [:= :mbe.metabot_id id]]}]]})
+                                            [:= :metabot_id id]]})
   api/generic-204-no-content)
 
 (def ^{:arglists '([request respond raise])} routes

enterprise/backend/src/metabase_enterprise/metabot_v3/config.clj

@@ -14,9 +14,9 @@
 
 (def metabot-config
   "The name of the collection exposed by the answer-sources tool."
-  {internal-metabot-id {:profile-id "experimental"
+  {internal-metabot-id {:profile-id "internal"
                         :entity-id "metabotmetabotmetabot"}
-   embedded-metabot-id {:profile-id "default"
+   embedded-metabot-id {:profile-id "embedding"
                         :entity-id "embeddedmetabotmetabo"}})
 
 (defn metabot-id->profile-id

enterprise/backend/src/metabase_enterprise/metabot_v3/context.clj

@@ -53,42 +53,50 @@
                                           :api/endpoints)]
     (str "backend:/api/ee/metabot-tools" url)))
 
-(def ^:private max-database-tables
-  "If the number of tables in the database doesn't exceed this number, we send them all to the agent."
-  100)
+(defn- query-for-sql-parsing
+  "Given an item in context, return the query if it is a native query or SQL transform that can have table usage parsed
+  from it, otherwise nil."
+  [item]
+  (when-let [query (case (:type item)
+                     "transform" (-> item :source :query)
+                     "adhoc" (-> item :query)
+                     nil)]
+    (when (and (#{:native "native"} (:type query))
+               (:database query))
+      query)))
 
 (defn- database-tables-for-context
-  "Get database tables formatted for metabot context, prioritizing tables used in the query (if provided), then filling up to the limit with most viewed tables."
-  ([database-id] (database-tables-for-context database-id nil))
-  ([database-id {:keys [all-tables-limit query] :or {all-tables-limit max-database-tables}}]
-   (when database-id
-     (try
-       (let [used-tables (if query
-                           (table-utils/used-tables query)
-                           [])
-             used-table-ids (set (map :id used-tables))]
-         (table-utils/database-tables database-id
-                                      {:all-tables-limit all-tables-limit
-                                       :priority-tables used-tables
-                                       :exclude-table-ids used-table-ids}))
-       (catch Exception e
-         (log/error e "Error getting database tables for context")
-         ;; If we can't get table info, just return empty - don't break the context
-         [])))))
+  "Get database tables formatted for metabot context. Only includes tables used in the query, formatted for API output.
+   Removes duplicate tables by id while preserving first occurrence order."
+  [{:keys [query]}]
+  (try
+    (if query
+      (let [used-tables (table-utils/used-tables query)
+            tables (table-utils/enhanced-database-tables (:database query)
+                                                         {:priority-tables used-tables
+                                                          :all-tables-limit (count used-tables)})]
+        ;; Ensure no duplicate tables (by id) are returned, preserve first occurrence order
+        (reduce (fn [acc t]
+                  (if (some #(= (:id %) (:id t)) acc)
+                    acc
+                    (conj acc t)))
+                []
+                tables))
+      [])
+    (catch Exception e
+      (log/error e "Error getting database tables for context")
+      [])))
 
 (defn- enhance-context-with-schema
-  "Enhance context by adding table schema information for native queries"
+  "Enhance context by adding table schema information for native queries & SQL transforms"
   [context]
   (if-let [user-viewing (get context :user_is_viewing)]
     (let [enhanced-viewing
           (mapv (fn [item]
-                  (if (and (#{:native "native"} (get-in item [:query :type]))
-                           (get-in item [:query :database]))
-                    (let [database-id (get-in item [:query :database])
-                          tables (database-tables-for-context database-id {:query (:query item)})]
-                      (if (seq tables)
-                        (assoc item :database_schema tables)
-                        item))
+                  (if-let [query (query-for-sql-parsing item)]
+                    (if-let [tables (seq (database-tables-for-context {:query query}))]
+                      (assoc item :used_tables tables)
+                      item)
                     item))
                 user-viewing)]
       (assoc context :user_is_viewing enhanced-viewing))

enterprise/backend/src/metabase_enterprise/metabot_v3/dummy_tools.clj

@@ -9,7 +9,6 @@
    [metabase.lib.core :as lib]
    [metabase.lib.metadata :as lib.metadata]
    [metabase.lib.types.isa :as lib.types.isa]
-   [metabase.lib.util :as lib.util]
    [metabase.util :as u]
    [metabase.util.humanization :as u.humanization]
    [metabase.warehouse-schema.models.field-values :as field-values]
@@ -17,6 +16,16 @@
 
 (set! *warn-on-reflection* true)
 
+(defn- verified-review?
+  "Return true if the most recent ModerationReview for the given item id/type is verified."
+  [id item-type]
+  (let [review (t2/select-one [:model/ModerationReview :status]
+                              :moderated_item_id id
+                              :moderated_item_type item-type
+                              :most_recent true
+                              {:order-by [[:id :desc]]})]
+    (= (:status review) "verified")))
+
 (defn get-current-user
   "Get information about the current user."
   []
@@ -34,7 +43,11 @@
   [{:keys [dashboard-id]}]
   (if-let [dashboard (t2/select-one [:model/Dashboard :id :description :name :collection_id] dashboard-id)]
     (do (api/read-check dashboard)
-        {:structured-output (-> dashboard (dissoc :collection_id) (assoc :type :dashboard))})
+        {:structured-output
+         (-> dashboard
+             (dissoc :collection_id)
+             (assoc :type :dashboard
+                    :verified (verified-review? dashboard-id "dashboard")))})
     {:output "dashboard not found"}))
 
 (defn- get-field-values [id->values id]
@@ -50,14 +63,6 @@
       (map #(m/assoc-some % :field-values (some->> % :id (get-field-values id->values))) cols))
     cols))
 
-(defn- add-table-reference
-  [query col]
-  (cond-> col
-    (and (:fk-field-id col)
-         (:table-id col))
-    (assoc :table-reference (-> (lib/display-name query (lib.metadata/field query (:fk-field-id col)))
-                                lib.util/strip-id))))
-
 (defn metric-details
   "Get metric details as returned by tools."
   ([id] (metric-details id nil))
@@ -78,7 +83,7 @@
                       (lib/remove-all-breakouts metric-query))
          visible-cols (when query-needed?
                         (->> (lib/visible-columns base-query)
-                             (map #(add-table-reference base-query %))))
+                             (map #(metabot-v3.tools.u/add-table-reference base-query %))))
          col->index (when query-needed?
                       (into {} (map-indexed (fn [i col] [col i])) visible-cols))
          col-index (when query-needed?
@@ -98,10 +103,11 @@
                                                       default-temporal-breakout
                                                       (col-index default-temporal-breakout)
                                                       field-id-prefix)
-                                                     :field-id))}
+                                                     :field-id))
+              :verified (verified-review? id "card")}
        with-queryable-dimensions?
        (assoc :queryable-dimensions (into []
-                                          (comp (map #(add-table-reference base-query %))
+                                          (comp (map #(metabot-v3.tools.u/add-table-reference base-query %))
                                                 (map #(metabot-v3.tools.u/->result-column
                                                        metric-query % (col-index %) field-id-prefix)))
                                           (->> (lib/filterable-columns base-query)
@@ -129,25 +135,29 @@
         :as   options}]
    (when-let [base (if metadata-provider
                      (lib.metadata/table metadata-provider id)
-                     (metabot-v3.tools.u/get-table id :db_id :description :name))]
+                     (metabot-v3.tools.u/get-table id :db_id :description :name :schema))]
      (let [query-needed? (or with-fields? with-metrics?)
+           db-id (if metadata-provider (:db-id base) (:db_id base))
            mp (when query-needed?
                 (or metadata-provider
-                    (lib.metadata.jvm/application-database-metadata-provider (:db_id base))))
+                    (lib.metadata.jvm/application-database-metadata-provider db-id)))
            table-query (when query-needed?
                          (lib/query mp (lib.metadata/table mp id)))
            cols (when with-fields?
                   (->> (lib/visible-columns table-query)
                        field-values-fn
-                       (map #(add-table-reference table-query %))))
+                       (map #(metabot-v3.tools.u/add-table-reference table-query %))))
            field-id-prefix (when with-fields?
                              (metabot-v3.tools.u/table-field-id-prefix id))]
        (-> {:id id
             :type :table
             :fields (into [] (map-indexed #(metabot-v3.tools.u/->result-column table-query %2 %1 field-id-prefix)) cols)
-            ;; :name should be (lib/display-name table-query), but we want to avoid creating the query if possible
-            :name (some->> (:name base)
-                           (u.humanization/name->human-readable-name :simple))}
+            :name (:name base)
+            ;; :display_name should be (lib/display-name table-query), but we want to avoid creating the query if possible
+            :display_name (some->> (:name base)
+                                   (u.humanization/name->human-readable-name :simple))
+            :database_id db-id
+            :database_schema (:schema base)}
            (m/assoc-some :description (:description base)
                          :metrics (when with-metrics?
                                     (not-empty (mapv #(convert-metric % mp options)
@@ -181,13 +191,18 @@
          cols (when with-fields?
                 (->> (lib/visible-columns card-query)
                      field-values-fn
-                     (map #(add-table-reference card-query %))))
+                     (map #(metabot-v3.tools.u/add-table-reference card-query %))))
          field-id-prefix (metabot-v3.tools.u/card-field-id-prefix id)]
      (-> {:id id
           :type card-type
           :fields (into [] (map-indexed #(metabot-v3.tools.u/->result-column card-query %2 %1 field-id-prefix)) cols)
           :name (:name base)
-          :queryable-foreign-key-tables []}
+          :display_name (some->> (:name base)
+                                 (u.humanization/name->human-readable-name :simple))
+          :database_id (:database_id base)
+          :queryable-foreign-key-tables []
+          :verified (verified-review? id "card")}
+
          (m/assoc-some :description (:description base)
                        :metrics (when with-metrics?
                                   (not-empty (mapv #(convert-metric % metadata-provider options)
@@ -293,7 +308,7 @@
           details (if (int? report-id)
                     (let [details (card-details report-id options)]
                       (some-> details
-                              (select-keys [:id :type :description :name])
+                              (select-keys [:id :type :description :name :verified])
                               (assoc :result-columns (:fields details))))
                     "invalid report_id")]
       (if (map? details)

enterprise/backend/src/metabase_enterprise/metabot_v3/models/metabot.clj

@@ -18,13 +18,13 @@
 
 ;;; --------------------------------------------------- Hydration ----------------------------------------------------
 
-(methodical/defmethod t2/batched-hydrate [:model/Metabot :entities]
-  "Hydrate the list of entities for a collection of metabots."
+(methodical/defmethod t2/batched-hydrate [:model/Metabot :prompts]
+  "Hydrate the list of prompts for a collection of metabots."
   [_model k metabots]
   (mi/instances-with-hydrated-data
    metabots k
    #(group-by :metabot_id
-              (t2/select :model/MetabotEntity {:where [:in :metabot_id (map :id metabots)]}))
+              (t2/select :model/MetabotPrompt {:where [:in :metabot_id (map :id metabots)]}))
    :id
    {:default []}))
 
@@ -35,14 +35,15 @@
   [:name])
 
 (defmethod serdes/dependencies "Metabot"
-  [{:keys [entities]}]
-  (into #{} (mapcat serdes/dependencies) entities))
+  [{:keys [prompts]}]
+  (set (mapcat serdes/dependencies prompts)))
 
 (defmethod serdes/generate-path "Metabot" [_ metabot]
   [(serdes/infer-self-path "Metabot" metabot)])
 
 (defmethod serdes/make-spec "Metabot" [_model-name opts]
-  {:copy      [:name :description :entity_id]
-   :transform {:created_at (serdes/date)
-               :updated_at (serdes/date)
-               :entities   (serdes/nested :model/MetabotEntity :metabot_id opts)}})
+  {:copy      [:name :description :entity_id :use_verified_content]
+   :transform {:created_at    (serdes/date)
+               :updated_at    (serdes/date)
+               :collection_id (serdes/fk :model/Collection)
+               :prompts       (serdes/nested :model/MetabotPrompt :metabot_id opts)}})

enterprise/backend/src/metabase_enterprise/metabot_v3/models/metabot_entity.clj

@@ -1,54 +0,0 @@
-(ns metabase-enterprise.metabot-v3.models.metabot-entity
-  (:require
-   [metabase.models.interface :as mi]
-   [metabase.models.serialization :as serdes]
-   [methodical.core :as methodical]
-   [toucan2.core :as t2]))
-
-;;; --------------------------------------------------- Lifecycle ----------------------------------------------------
-
-(methodical/defmethod t2/table-name :model/MetabotEntity [_model] :metabot_entity)
-
-(doto :model/MetabotEntity
-  (derive :metabase/model)
-  (derive ::mi/read-policy.superuser)
-  (derive ::mi/write-policy.superuser)
-  (derive :hook/entity-id)
-  (derive :hook/created-at-timestamped?))
-
-(t2/deftransforms :model/MetabotEntity
-  {:model mi/transform-keyword})
-
-;;; ------------------------------------------------- Serialization -------------------------------------------------
-
-(defmethod serdes/hash-fields :model/MetabotEntity
-  [_table]
-  [:metabot_id :model :model_id])
-
-(defmethod serdes/dependencies "MetabotEntity"
-  [{:keys [model model_id prompts]}]
-  (into #{[{:model (case model
-                     "collection" "Collection"
-                     ("dataset" "metric") "Card")
-            :id model_id}]}
-        (mapcat serdes/dependencies prompts)))
-
-(defmethod serdes/generate-path "MetabotEntity" [_ entity]
-  [(serdes/infer-self-path "Metabot" (t2/select-one :model/Metabot :id (:metabot_id entity)))
-   (serdes/infer-self-path "MetabotEntity" entity)])
-
-(defmethod serdes/make-spec "MetabotEntity" [_model-name opts]
-  {:copy      [:entity_id]
-   :transform {:created_at (serdes/date)
-               :model      (serdes/kw)
-               :model_id   {::fk true
-                            :export-with-context (fn [{:keys [model model_id]} _ _]
-                                                   (serdes/*export-fk* model_id (case model
-                                                                                  :collection        :model/Collection
-                                                                                  (:dataset :metric) :model/Card)))
-                            :import-with-context (fn [{:keys [model model_id]} _ _]
-                                                   (serdes/*import-fk* model_id (case model
-                                                                                  "collection"         :model/Collection
-                                                                                  ("dataset" "metric") :model/Card)))}
-               :metabot_id (serdes/parent-ref)
-               :prompts    (serdes/nested :model/MetabotPrompt :metabot_entity_id opts)}})

enterprise/backend/src/metabase_enterprise/metabot_v3/models/metabot_prompt.clj

@@ -22,10 +22,10 @@
 
 (defmethod serdes/hash-fields :model/MetabotPrompt
   [_table]
-  [:metabot_entity_id :model :card_id :prompt])
+  [:metabot_id :model :card_id :prompt])
 
 (defmethod serdes/generate-path "MetabotPrompt" [_ entity]
-  (conj (serdes/generate-path "MetabotEntity" (t2/select-one :model/MetabotEntity (:metabot_entity_id entity)))
+  (conj (serdes/generate-path "Metabot" (t2/select-one :model/Metabot (:metabot_id entity)))
         (serdes/infer-self-path "MetabotPrompt" entity)))
 
 (defmethod serdes/dependencies "MetabotPrompt" [prompt]
@@ -33,8 +33,8 @@
 
 (defmethod serdes/make-spec "MetabotPrompt" [_model-name _opts]
   {:copy      [:entity_id :prompt]
-   :transform {:created_at        (serdes/date)
-               :updated_at        (serdes/date)
-               :model             (serdes/kw)
-               :card_id           (serdes/fk :model/Card)
-               :metabot_entity_id (serdes/parent-ref)}})
+   :transform {:created_at (serdes/date)
+               :updated_at (serdes/date)
+               :model      (serdes/kw)
+               :card_id    (serdes/fk :model/Card)
+               :metabot_id (serdes/parent-ref)}})

enterprise/backend/src/metabase_enterprise/metabot_v3/table_utils.clj

@@ -3,9 +3,14 @@
   (:require
    [clojure.set :as set]
    [metabase-enterprise.metabot-v3.query-analyzer :as query-analyzer]
+   [metabase-enterprise.metabot-v3.tools.util :as metabot-v3.tools.u]
    [metabase.api.common :as api]
+   [metabase.lib-be.metadata.jvm :as lib.metadata.jvm]
+   [metabase.lib.core :as lib]
+   [metabase.lib.metadata :as lib.metadata]
    [metabase.models.interface :as mi]
    [metabase.util :as u]
+   [metabase.util.humanization :as u.humanization]
    [toucan2.core :as t2]
    [toucan2.realize :as t2.realize])
   (:import
@@ -59,13 +64,62 @@
          all-tables (concat priority-tables fill-tables)
          all-tables (take all-tables-limit all-tables)]
      (mapv (fn [{:keys [fields] :as table}]
-             (merge (select-keys table [:name :schema :description])
+             (merge (select-keys table [:id :name :schema :description])
                     {:columns (mapv (fn [{:keys [database_type] :as field}]
-                                      (merge (select-keys field [:name :description])
+                                      (merge (select-keys field [:id :name :description])
                                              {:data_type database_type}))
                                     fields)}))
            all-tables))))
 
+(defn enhanced-database-tables
+  "Get database tables formatted with the new metabot tools schema format.
+
+  Returns tables with :type, :display_name, :database_id, :database_schema, :fields (with field-id), :metrics.
+  This format is used by metabot context and other modern tools."
+  ([database-id]
+   (enhanced-database-tables database-id nil))
+  ([database-id {:keys [all-tables-limit priority-tables exclude-table-ids]
+                 :or {all-tables-limit max-database-tables
+                      priority-tables []
+                      exclude-table-ids #{}}}]
+   (let [priority-table-ids (set (map :id priority-tables))
+         ;; Fetch most viewed tables, excluding priority tables and excluded tables
+         fill-tables (t2/select [:model/Table :id :db_id :name :schema :description]
+                                :db_id database-id
+                                :active true
+                                :visibility_type nil
+                                {:where    (mi/visible-filter-clause :model/Table
+                                                                     :id
+                                                                     {:user-id       api/*current-user-id*
+                                                                      :is-superuser? api/*is-superuser?*}
+                                                                     {:perms/view-data      :unrestricted
+                                                                      :perms/create-queries :query-builder-and-native})
+                                 :order-by [[:view_count :desc]]
+                                 :limit    all-tables-limit})
+         fill-tables (remove #(or (priority-table-ids (:id %))
+                                  (exclude-table-ids (:id %))) fill-tables)
+         all-tables (concat priority-tables fill-tables)
+         all-tables (take all-tables-limit all-tables)]
+     (lib.metadata.jvm/with-metadata-provider-cache
+       (let [mp (lib.metadata.jvm/application-database-metadata-provider database-id)
+             table-ids (map :id all-tables)
+             _ (lib.metadata/bulk-metadata mp :metadata/table table-ids)]
+         (mapv (fn [{:keys [id name schema description]}]
+                 (let [table-query (lib/query mp (lib.metadata/table mp id))
+                       cols (->> (lib/visible-columns table-query)
+                                 (map #(metabot-v3.tools.u/add-table-reference table-query %)))
+                       field-id-prefix (metabot-v3.tools.u/table-field-id-prefix id)]
+                   {:id id
+                    :type :table
+                    :name name
+                    :display_name (u.humanization/name->human-readable-name :simple name)
+                    :database_id database-id
+                    :database_schema schema
+                    :description description
+                    :fields (into [] (map-indexed #(metabot-v3.tools.u/->result-column table-query %2 %1 field-id-prefix) cols))
+                    :metrics []}))
+               all-tables))))))
+
 (defn similar?
   "Check if two strings are similar using Levenshtein distance with a max distance of 4."
   [left right]

enterprise/backend/src/metabase_enterprise/metabot_v3/tools/api.clj

@@ -16,9 +16,12 @@
     :as metabot-v3.tools.create-dashboard-subscription]
    [metabase-enterprise.metabot-v3.tools.field-stats :as metabot-v3.tools.field-stats]
    [metabase-enterprise.metabot-v3.tools.filters :as metabot-v3.tools.filters]
-   [metabase-enterprise.metabot-v3.tools.find-metric :as metabot-v3.tools.find-metric]
    [metabase-enterprise.metabot-v3.tools.find-outliers :as metabot-v3.tools.find-outliers]
    [metabase-enterprise.metabot-v3.tools.generate-insights :as metabot-v3.tools.generate-insights]
+   [metabase-enterprise.metabot-v3.tools.search :as metabot-v3.tools.search]
+   [metabase-enterprise.metabot-v3.tools.search-data-sources :as metabot-v3.tools.search-data-sources]
+   [metabase-enterprise.metabot-v3.tools.snippets :as metabot-v3.tools.snippets]
+   [metabase-enterprise.metabot-v3.tools.transforms :as metabot-v3.tools.transforms]
    [metabase-enterprise.metabot-v3.util :as metabot-v3.u]
    [metabase.api.macros :as api.macros]
    [metabase.api.response :as api.response]
@@ -218,6 +221,23 @@
     [:limit {:optional true} [:maybe :int]]]
    [:map {:encode/tool-api-request #(update-keys % metabot-v3.u/safe->kebab-case-en)}]])
 
+(mr/def ::query-datasource-arguments
+  [:and
+   [:map
+    [:table_id {:optional true} :int]
+    [:model_id {:optional true} :int]
+    [:fields {:optional true} [:maybe [:sequential ::field]]]
+    [:filters {:optional true} [:maybe [:sequential ::filter]]]
+    [:aggregations {:optional true} [:maybe [:sequential ::aggregation]]]
+    [:group_by {:optional true} [:maybe [:sequential ::group-by]]]
+    [:order_by {:optional true} [:maybe [:sequential [:map
+                                                      [:field ::field]
+                                                      [:direction [:enum {:encode/tool-api-request keyword} "asc" "desc"]]]]]]
+    [:limit {:optional true} [:maybe :int]]]
+   [:fn {:error/message "Exactly one of table_id and model_id required"}
+    #(= (count (select-keys % [:table_id :model_id])) 1)]
+   [:map {:encode/tool-api-request #(update-keys % metabot-v3.u/safe->kebab-case-en)}]])
+
 (mr/def ::count
   [:and
    :int
@@ -379,7 +399,8 @@
   [:merge
    ::basic-metric
    [:map {:decode/tool-api-response #(update-keys % metabot-v3.u/safe->snake_case_en)}
-    [:queryable_dimensions {:optional true} ::columns]]])
+    [:queryable_dimensions {:optional true} ::columns]
+    [:verified {:optional true} :boolean]]])
 
 (mr/def ::find-metric-result
   [:or
@@ -443,7 +464,6 @@
                          [:name :string]
                          [:email_address :string]]]]
    [:map [:output :string]]])
-
 (mr/def ::get-dashboard-details-result
   [:or
    [:map
@@ -452,7 +472,8 @@
                          [:id :int]
                          [:type [:= :dashboard]]
                          [:name :string]
-                         [:description {:optional true} :string]]]]
+                         [:description {:optional true} :string]
+                         [:verified {:optional true} :boolean]]]]
    [:map [:output :string]]])
 
 (mr/def ::get-metric-details-arguments
@@ -507,7 +528,8 @@
                          [:type [:= :question]]
                          [:name :string]
                          [:description {:optional true} [:maybe :string]]
-                         [:result_columns ::columns]]]]
+                         [:result_columns ::columns]
+                         [:verified {:optional true} :boolean]]]]
    [:map [:output :string]]])
 
 (mr/def ::get-document-details-arguments
@@ -581,6 +603,9 @@
    [:id :int]
    [:type [:enum :model :table]]
    [:name :string]
+   [:display_name :string]
+   [:database_id :int]
+   [:database_schema {:optional true} [:maybe :string]] ; Schema name, if applicable
    [:fields ::columns]
    [:description {:optional true} [:maybe :string]]
    [:metrics {:optional true} [:sequential ::basic-metric]]])
@@ -597,6 +622,43 @@
     [:structured_output ::full-table]]
    [:map [:output :string]]])
 
+(mr/def ::basic-transform
+  [:map
+   {:decode/tool-api-response #(update-keys % metabot-v3.u/safe->snake_case_en)}
+   [:id :int]
+   [:name :string]
+   [:description {:optional true} [:maybe :string]]
+   [:entity_id {:optional true} [:maybe :string]]
+   [:source ::metabot-v3.tools.transforms/transform-source]])
+
+(mr/def ::full-transform
+  [:merge
+   ::basic-transform
+   [:map {:decode/tool-api-response #(update-keys % metabot-v3.u/safe->snake_case_en)}
+    [:created_at ms/TemporalString]
+    [:updated_at ms/TemporalString]
+    [:target ::metabot-v3.tools.transforms/transform-target]]])
+
+(mr/def ::get-transforms-result
+  [:or
+   [:map
+    {:decode/tool-api-response #(update-keys % metabot-v3.u/safe->snake_case_en)}
+    [:structured_output [:sequential ::basic-transform]]]
+   [:map [:output :string]]])
+
+(mr/def ::get-transform-details-arguments
+  [:and
+   [:map
+    [:transform_id :int]]
+   [:map {:encode/tool-api-request
+          #(set/rename-keys % {:transform_id :transform-id})}]])
+
+(mr/def ::get-transform-details-result
+  [:or
+   [:map {:decode/tool-api-response #(update-keys % metabot-v3.u/safe->snake_case_en)}
+    [:structured_output [:sequential ::full-transform]]]
+   [:map [:output :string]]])
+
 (mr/def ::answer-sources-result
   [:or
    [:map
@@ -606,6 +668,109 @@
                          [:models  [:sequential ::full-table]]]]]
    [:map [:output :string]]])
 
+(mr/def ::search-data-sources-arguments
+  [:and
+   [:map
+    [:keywords {:optional true} [:maybe [:sequential :string]]]
+    [:description {:optional true} [:maybe :string]]
+    [:database_id {:optional true} [:maybe :int]]
+    [:entity_types {:optional true} [:maybe [:sequential [:enum "table" "model" "question" "dashboard" "metric"]]]]
+    [:limit {:optional true, :default 50} [:and :int [:fn #(<= 1 % 100)]]]]
+   [:map {:encode/tool-api-request
+          #(set/rename-keys % {:database_id :database-id
+                               :entity_types :entity-types})}]])
+
+(mr/def ::search-data-sources-table-result
+  "Schema for table/model search results"
+  [:map {:decode/tool-api-response #(update-keys % metabot-v3.u/safe->snake_case_en)}
+   [:id :int]
+   [:type [:enum :table :model]]
+   [:name :string]                                   ; Table name (technical for tables, display for models)
+   [:display_name {:optional true} [:maybe :string]] ; Display name
+   [:description {:optional true} [:maybe :string]]
+   [:database_id {:optional true} [:maybe :int]]
+   [:database_schema {:optional true} [:maybe :string]]])
+
+(mr/def ::search-data-sources-dashboard-result
+  "Schema for dashboard search results"
+  [:map {:decode/tool-api-response #(update-keys % metabot-v3.u/safe->snake_case_en)}
+   [:id :int]
+   [:type [:= :dashboard]]
+   [:name :string]                                    ; Display name
+   [:description {:optional true} [:maybe :string]]
+   [:verified {:optional true} :boolean]])
+
+(mr/def ::search-data-sources-question-result
+  "Schema for question search results"
+  [:map {:decode/tool-api-response #(update-keys % metabot-v3.u/safe->snake_case_en)}
+   [:id :int]
+   [:type [:= :question]]
+   [:name :string]                                    ; Display name
+   [:description {:optional true} [:maybe :string]]
+   [:verified {:optional true} :boolean]])
+
+(mr/def ::search-data-sources-metric-result
+  "Schema for metric search results"
+  [:map {:decode/tool-api-response #(update-keys % metabot-v3.u/safe->snake_case_en)}
+   [:id :int]
+   [:type [:= :metric]]
+   [:name :string]                                    ; Display name
+   [:description {:optional true} [:maybe :string]]
+   [:verified {:optional true} :boolean]])
+
+(mr/def ::search-data-sources-result-item
+  "Union of all search result types.
+
+   We use dedicated schemas per entity type because:
+   - Tables need technical names and database context
+   - Models need display names and database context
+   - Other entities use display names and don't need database fields
+   - Type-specific validation ensures correct field usage"
+  [:or
+   ::search-data-sources-table-result
+   ::search-data-sources-dashboard-result
+   ::search-data-sources-question-result
+   ::search-data-sources-metric-result])
+
+(mr/def ::search-data-sources-result
+  [:or
+   [:map
+    {:decode/tool-api-response #(update-keys % metabot-v3.u/safe->snake_case_en)}
+    [:structured_output [:map
+                         [:data [:sequential ::search-data-sources-result-item]]
+                         [:total_count :int]]]]
+   [:map [:output :string]]])
+
+(mr/def ::basic-snippet
+  [:map {:decode/tool-api-response #(update-keys % metabot-v3.u/safe->snake_case_en)}
+   [:id :int]
+   [:name :string]
+   [:description {:optional true} [:maybe :string]]])
+
+(mr/def ::full-snippet
+  [:merge
+   ::basic-snippet
+   [:map {:decode/tool-api-response #(update-keys % metabot-v3.u/safe->snake_case_en)}
+    [:content :string]]])
+
+(mr/def ::get-snippets-result
+  "Schema for SQL snippet list results"
+  [:map
+   {:decode/tool-api-response #(update-keys % metabot-v3.u/safe->snake_case_en)}
+   [:structured_output [:sequential ::basic-snippet]]])
+
+(mr/def ::get-snippet-details-arguments
+  [:and
+   [:map
+    [:snippet_id :int]]
+   [:map {:encode/tool-api-request
+          #(set/rename-keys % {:snippet_id :snippet-id})}]])
+
+(mr/def ::get-snippet-details-result
+  "Schema for SQL snippet detail results"
+  [:map {:decode/tool-api-response #(update-keys % metabot-v3.u/safe->snake_case_en)}
+   [:structured_output ::full-snippet]])
+
 (api.macros/defendpoint :post "/answer-sources" :- [:merge ::answer-sources-result ::tool-request]
   "Return top level meta information about available information sources."
   [_route-params
@@ -674,20 +839,6 @@
               (assoc :conversation_id conversation_id))
       (metabot-v3.context/log :llm.log/be->llm))))
 
-(api.macros/defendpoint :post "/find-metric" :- [:merge ::find-metric-result ::tool-request]
-  "Find a metric matching a description."
-  [_route-params
-   _query-params
-   {:keys [arguments conversation_id] :as body} :- [:merge
-                                                    [:map [:arguments [:map [:message :string]]]]
-                                                    ::tool-request]]
-  (metabot-v3.context/log (assoc body :api :find-metric) :llm.log/llm->be)
-  (doto (-> (mc/decode ::find-metric-result
-                       (metabot-v3.tools.find-metric/find-metric arguments)
-                       (mtx/transformer {:name :tool-api-response}))
-            (assoc :conversation_id conversation_id))
-    (metabot-v3.context/log :llm.log/be->llm)))
-
 (api.macros/defendpoint :post "/find-outliers" :- [:merge ::find-outliers-result ::tool-request]
   "Find outliers in the values provided by a data source for a given column."
   [_route-params
@@ -841,6 +992,34 @@
               (assoc :conversation_id conversation_id))
       (metabot-v3.context/log :llm.log/be->llm))))
 
+(api.macros/defendpoint :post "/get-transforms" :- [:merge ::get-transforms-result ::tool-request]
+  "Get a list of all known transforms."
+  [_route-params
+   _query-params
+   {:keys [conversation_id] :as body} :- ::tool-request]
+  (metabot-v3.context/log (assoc body :api :get-transforms) :llm.log/llm->be)
+  (doto (-> (mc/decode ::get-transforms-result
+                       (metabot-v3.tools.transforms/get-transforms)
+                       (mtx/transformer {:name :tool-api-response}))
+            (assoc :conversation_id conversation_id))
+    (metabot-v3.context/log :llm.log/be->llm)))
+
+(api.macros/defendpoint :post "/get-transform-details" :- [:merge ::get-transform-details-result ::tool-request]
+  "Get information about a transform."
+  [_route-params
+   _query-params
+   {:keys [arguments conversation_id] :as body} :- [:merge
+                                                    [:map [:arguments ::get-transform-details-arguments]]
+                                                    ::tool-request]]
+  (metabot-v3.context/log (assoc body :api :get-transform-details) :llm.log/llm->be)
+  (let [arguments (mc/encode ::get-transform-details-arguments arguments (mtx/transformer {:name :tool-api-request}))
+        transform-id (:transform-id arguments)]
+    (doto (-> (mc/decode ::get-transform-details-result
+                         (metabot-v3.tools.transforms/get-transform-details transform-id)
+                         (mtx/transformer {:name :tool-api-response}))
+              (assoc :conversation_id conversation_id))
+      (metabot-v3.context/log :llm.log/be->llm))))
+
 (api.macros/defendpoint :post "/query-metric" :- [:merge ::filtering-result ::tool-request]
   "Construct a query from a metric."
   [_route-params
@@ -873,6 +1052,147 @@
               (assoc :conversation_id conversation_id))
       (metabot-v3.context/log :llm.log/be->llm))))
 
+;; TODO tsplude - drop the `/query-model` endpoint and filter logic in favor of this
+(api.macros/defendpoint :post "/query-datasource" :- [:merge ::filtering-result ::tool-request]
+  "Construct a query from a model or table data source."
+  [_route-params
+   _query-params
+   {:keys [arguments conversation_id] :as body} :- [:merge
+                                                    [:map [:arguments ::query-datasource-arguments]]
+                                                    ::tool-request]]
+  (metabot-v3.context/log (assoc body :api :query-datasource) :llm.log/llm->be)
+  (let [arguments (mc/encode ::query-datasource-arguments
+                             arguments (mtx/transformer {:name :tool-api-request}))]
+    (doto (-> (mc/decode ::filtering-result
+                         (metabot-v3.tools.filters/query-datasource arguments)
+                         (mtx/transformer {:name :tool-api-response}))
+              (assoc :conversation_id conversation_id))
+      (metabot-v3.context/log :llm.log/be->llm))))
+
+(api.macros/defendpoint :post "/search-data-sources" :- [:merge ::search-data-sources-result ::tool-request]
+  "Search for data sources (tables, models, cards, dashboards, metrics) in Metabase."
+  [_route-params
+   _query-params
+   {:keys [arguments conversation_id] :as body} :- [:merge
+                                                    [:map [:arguments {:optional true} ::search-data-sources-arguments]]
+                                                    ::tool-request]
+   request]
+  (metabot-v3.context/log (assoc body :api :search-data-sources) :llm.log/llm->be)
+  (try
+    (let [options (mc/encode ::search-data-sources-arguments
+                             arguments (mtx/transformer {:name :tool-api-request}))
+          metabot-id (:metabot-v3/metabot-id request)
+          results (metabot-v3.tools.search-data-sources/search-data-sources
+                   (assoc options :metabot-id metabot-id))
+          response-data {:data results
+                         :total_count (count results)}]
+      (doto (-> (mc/decode ::search-data-sources-result
+                           {:structured_output response-data}
+                           (mtx/transformer {:name :tool-api-response}))
+                (assoc :conversation_id conversation_id))
+        (metabot-v3.context/log :llm.log/be->llm)))
+    (catch Exception e
+      (log/error e "Error in search-data-sources")
+      (doto (-> {:output (str "Search failed: " (or (ex-message e) "Unknown error"))}
+                (assoc :conversation_id conversation_id))
+        (metabot-v3.context/log :llm.log/be->llm)))))
+
+(mr/def ::search-arguments
+  [:and
+   [:map
+    [:term_queries     {:optional true} [:maybe [:sequential :string]]]
+    [:semantic_queries {:optional true} [:maybe [:sequential :string]]]
+    [:entity_types     {:optional true} [:maybe [:sequential [:enum "table" "model" "question" "dashboard" "metric" "database"]]]]
+    [:database_id      {:optional true} [:maybe :int]]
+    [:created_at       {:optional true} [:maybe ms/NonBlankString]]
+    [:last_edited_at   {:optional true} [:maybe ms/NonBlankString]]
+    [:limit            {:optional true, :default 50} [:and :int [:fn #(<= 1 % 100)]]]]
+   [:map {:encode/tool-api-request
+          #(update-keys % (comp keyword u/->kebab-case-en))}]])
+
+(mr/def ::search-result-item
+  "Unified schema for search result items."
+  [:map {:decode/tool-api-response #(update-keys % metabot-v3.u/safe->snake_case_en)}
+   [:id :int]
+   [:type [:enum :table :model :dashboard :question :metric :database]]
+   [:name :string]
+   [:display_name {:optional true} [:maybe :string]]
+   [:description {:optional true} [:maybe :string]]
+   [:database_id {:optional true} [:maybe :int]]
+   [:database_schema {:optional true} [:maybe :string]]
+   [:verified {:optional true} [:maybe :boolean]]
+   [:updated_at {:optional true} [:maybe :string]]
+   [:created_at {:optional true} [:maybe :string]]
+   [:collection {:optional true} [:maybe [:map
+                                          [:name {:optional true} [:maybe :string]]
+                                          [:authority_level {:optional true} [:maybe :string]]]]]])
+
+(mr/def ::search-result
+  [:or
+   [:map
+    {:decode/tool-api-response #(update-keys % metabot-v3.u/safe->snake_case_en)}
+    [:structured_output [:map
+                         [:data [:sequential ::search-result-item]]
+                         [:total_count :int]]]]
+   [:map [:output :string]]])
+
+(api.macros/defendpoint :post "/search" :- [:merge ::search-result ::tool-request]
+  "Enhanced search with term and semantic queries using Reciprocal Rank Fusion."
+  [_route-params
+   _query-params
+   {:keys [arguments conversation_id] :as body} :- [:merge
+                                                    [:map [:arguments {:optional true} ::search-arguments]]
+                                                    ::tool-request]
+   request]
+  (metabot-v3.context/log (assoc body :api :search) :llm.log/llm->be)
+  (try
+    (let [options (mc/encode ::search-arguments
+                             arguments (mtx/transformer {:name :tool-api-request}))
+          metabot-id (:metabot-v3/metabot-id request)
+          results (metabot-v3.tools.search/search
+                   (assoc options :metabot-id metabot-id))
+          response-data {:data results
+                         :total_count (count results)}]
+      (doto (-> (mc/decode ::search-result
+                           {:structured_output response-data}
+                           (mtx/transformer {:name :tool-api-response}))
+                (assoc :conversation_id conversation_id))
+        (metabot-v3.context/log :llm.log/be->llm)))
+    (catch Exception e
+      (log/error e "Error in search")
+      (doto (-> {:output (str "Search failed: " (or (ex-message e) "Unknown error"))}
+                (assoc :conversation_id conversation_id))
+        (metabot-v3.context/log :llm.log/be->llm)))))
+
+(api.macros/defendpoint :post "/get-snippets" :- [:merge ::get-snippets-result ::tool-request]
+  "Get a list of all known SQL snippets."
+  [_route-params
+   _query-params
+   {:keys [conversation_id] :as body} :- ::tool-request]
+  (metabot-v3.context/log (assoc body :api :get-snippets) :llm.log/llm->be)
+  (doto (-> (mc/decode ::get-snippets-result
+                       (metabot-v3.tools.snippets/get-snippets)
+                       (mtx/transformer {:name :tool-api-response}))
+            (assoc :conversation_id conversation_id))
+    (metabot-v3.context/log :llm.log/be->llm)))
+
+(api.macros/defendpoint :post "/get-snippet-details" :- [:merge ::get-snippet-details-result ::tool-request]
+  "Get the content of a single SQL snippet."
+  [_route-params
+   _query-params
+   {:keys [arguments conversation_id] :as body} :- [:merge
+                                                    [:map [:arguments ::get-snippet-details-arguments]]
+                                                    ::tool-request]]
+
+  (metabot-v3.context/log (assoc body :api :get-snippet-details) :llm.log/llm->be)
+  (let [arguments (mc/encode ::get-snippet-details-arguments arguments (mtx/transformer {:name :tool-api-request}))
+        snippet-id (:snippet-id arguments)]
+    (doto (-> (mc/decode ::get-snippet-details-result
+                         (metabot-v3.tools.snippets/get-snippet-details snippet-id)
+                         (mtx/transformer {:name :tool-api-response}))
+              (assoc :conversation_id conversation_id))
+      (metabot-v3.context/log :llm.log/be->llm))))
+
 (defn- enforce-authentication
   "Middleware that returns a 401 response if no `ai-session` can be found for  `request`."
   [handler]

enterprise/backend/src/metabase_enterprise/metabot_v3/tools/filters.clj

@@ -278,6 +278,73 @@
         {:output (str "No model found with model_id " model-id)}
         (metabot-v3.tools.u/handle-agent-error e)))))
 
+(defn- resolve-datasource
+  "Resolve datasource parameters to [field-id-prefix base-query] tuple.
+   Accepts either {:table-id id} or {:model-id id}."
+  [{:keys [table-id model-id]}]
+  (cond
+    model-id
+    [(metabot-v3.tools.u/card-field-id-prefix model-id) (metabot-v3.tools.u/card-query model-id)]
+
+    table-id
+    [(metabot-v3.tools.u/table-field-id-prefix table-id) (metabot-v3.tools.u/table-query table-id)]
+
+    :else
+    (throw (ex-info "Either table-id or model-id must be provided" {:agent-error? true}))))
+
+(defn- query-datasource*
+  [{:keys [fields filters aggregations group-by order-by limit] :as arguments}]
+  (let [[filter-field-id-prefix base-query] (resolve-datasource arguments)
+        visible-cols (lib/visible-columns base-query)
+        resolve-visible-column  #(metabot-v3.tools.u/resolve-column % filter-field-id-prefix visible-cols)
+        resolve-order-by-column (fn [{:keys [field direction]}] {:field (resolve-visible-column field) :direction direction})
+        projection (map (comp (juxt filter-bucketed-column (fn [{:keys [column bucket]}]
+                                                             (let [column (cond-> column
+                                                                            bucket (assoc :unit bucket))]
+                                                               (lib/display-name base-query -1 column :long))))
+                              resolve-visible-column)
+                        fields)
+        reduce-query (fn [query f coll] (reduce f query coll))
+        query (-> base-query
+                  (reduce-query (fn [query [expr-or-column expr-name]]
+                                  (lib/expression query expr-name expr-or-column))
+                                (filter (comp expression? first) projection))
+                  (add-fields projection)
+                  (reduce-query add-filter (map resolve-visible-column filters))
+                  (reduce-query add-aggregation (map resolve-visible-column aggregations))
+                  (reduce-query add-breakout (map resolve-visible-column group-by))
+                  (reduce-query add-order-by (map resolve-order-by-column order-by))
+                  (add-limit limit))
+        query-id (u/generate-nano-id)
+        query-field-id-prefix (metabot-v3.tools.u/query-field-id-prefix query-id)
+        returned-cols (lib/returned-columns query)]
+    {:type :query
+     :query-id query-id
+     ;; existing usage, don't do this going forward -- use Lib instead and persist MBQL 5 to the app DB
+     :query #_{:clj-kondo/ignore [:discouraged-var]} (lib/->legacy-MBQL query)
+     :result-columns (into []
+                           (map-indexed #(metabot-v3.tools.u/->result-column query %2 %1 query-field-id-prefix))
+                           returned-cols)}))
+
+(defn query-datasource
+  "Create a query based on a datasource (table or model)."
+  [{:keys [table-id model-id] :as arguments}]
+  (try
+    (cond
+      (and table-id model-id) (throw (ex-info "Cannot provide both table_id and model_id" {:agent-error? true}))
+      (int? model-id) {:structured-output (query-datasource* arguments)}
+      (int? table-id) {:structured-output (query-datasource* arguments)}
+      model-id        {:output (str "Invalid model_id " model-id)}
+      table-id        {:output (str "Invalid table_id " table-id)}
+      :else           {:output "Either table_id or model_id must be provided"})
+    (catch Exception e
+      (if (= (:status-code (ex-data e)) 404)
+        {:output (cond
+                   table-id (str "No table found with table_id " table-id)
+                   model-id (str "No model found with model_id " model-id)
+                   :else "Resource not found")}
+        (metabot-v3.tools.u/handle-agent-error e)))))
+
 (defn- base-query
   [data-source]
   (let [{:keys [table-id query query-id report-id]} data-source

enterprise/backend/src/metabase_enterprise/metabot_v3/tools/find_metric.clj

@@ -1,24 +0,0 @@
-(ns metabase-enterprise.metabot-v3.tools.find-metric
-  (:require
-   [metabase-enterprise.metabot-v3.client :as client]
-   [metabase-enterprise.metabot-v3.dummy-tools :as dummy-tools]
-   [metabase.lib-be.metadata.jvm :as lib.metadata.jvm]
-   [metabase.models.interface :as mi]
-   [toucan2.core :as t2]))
-
-(defn find-metric
-  "Find a metric described by `message`.
-
-  See [[metabase-enterprise.metabot-v3.dummy-tools/metric-details]] for the output if a metric can be found."
-  [{:keys [message]}]
-  (lib.metadata.jvm/with-metadata-provider-cache
-    (let [{:keys [id]} (client/select-metric-request
-                        (->> (t2/select [:model/Card :id :name :description]
-                                        :type [:= "metric"]
-                                        :archived [:= false])
-                             (filter mi/can-read?))
-                        message)]
-      (if-let [result (when id
-                        (dummy-tools/metric-details id))]
-        {:structured-output result}
-        {:output "Metric not found."}))))

enterprise/backend/src/metabase_enterprise/metabot_v3/tools/search.clj

@@ -0,0 +1,135 @@
+(ns metabase-enterprise.metabot-v3.tools.search
+  (:require
+   [clojure.set :as set]
+   [metabase-enterprise.metabot-v3.config :as metabot-v3.config]
+   [metabase-enterprise.metabot-v3.reactions]
+   [metabase.api.common :as api]
+   [metabase.permissions.core :as perms]
+   [metabase.search.core :as search]
+   [toucan2.core :as t2]))
+
+(def ^:private metabot-search-models
+  #{"table" "dataset" "card" "dashboard" "metric" "database"})
+
+(def ^:private search-model-mappings
+  "Maps metabot entity types to search engine model types"
+  {"model"    "dataset"
+   "question" "card"})
+
+(defn- entity-type->search-model
+  [entity-type]
+  (get search-model-mappings entity-type entity-type))
+
+(defn- search-model->result-type
+  [search-model]
+  (get (set/map-invert search-model-mappings) search-model search-model))
+
+(defn- transform-search-result
+  "Transform a single search result to match the appropriate entity-specific schema."
+  [{:keys [verified moderated_status collection] :as result}]
+  (let [model (:model result)
+        verified? (or (boolean verified) (= moderated_status "verified"))
+        collection-info (select-keys collection [:name :authority_level])
+        common-fields {:id           (:id result)
+                       :type         (search-model->result-type model)
+                       :name         (:name result)
+                       :description  (:description result)
+                       :updated_at   (:updated_at result)
+                       :created_at   (:created_at result)}]
+    (case model
+      "database"
+      common-fields
+
+      "table"
+      (merge common-fields
+             {:name            (:table_name result)
+              :display_name    (:name result)
+              :database_id     (:database_id result)
+              :database_schema (:table_schema result)})
+
+      "dashboard"
+      (merge common-fields
+             {:verified        verified?
+              :collection      collection-info})
+
+      ;; Questions, metrics, and datasets
+      (merge common-fields
+             {:database_id     (:database_id result)
+              :verified        verified?
+              :collection      collection-info}))))
+
+(defn- search-result-id
+  "Generate a unique identifier for a search result based on its id and model."
+  [search-result]
+  (juxt :id :model) search-result)
+
+(defn- reciprocal-rank-fusion
+  "Combine multiple ranked search result lists using Reciprocal Rank Fusion (RRF).
+
+  Takes a list of search result lists and combines them by:
+  1. Calculating RRF scores for each item based on its rank in each list
+  2. Summing scores for items that appear in multiple lists
+  3. Returning items sorted by total RRF score (descending)
+
+  The RRF score is calculated as: 1 / (k + r) where k defaults to 60 (typical RRF constant)"
+  ([result-lists]
+   (reciprocal-rank-fusion result-lists 60))
+  ([result-lists k]
+   (let [rrf-results (reduce
+                      (fn [acc-map result-list]
+                        (reduce-kv
+                         (fn [acc rank search-result]
+                           (let [id (search-result-id search-result)
+                                 rrf-score (/ 1.0 (+ k (inc rank)))]
+                             (if (contains? acc id)
+                               (update-in acc [id :rrf] + rrf-score)
+                               (assoc acc id {:search-result search-result
+                                              :rrf rrf-score}))))
+                         acc-map
+                         (vec result-list)))
+                      {}
+                      result-lists)]
+     (->> rrf-results
+          vals
+          (sort-by :rrf >)
+          (map :search-result)))))
+
+(defn search
+  "Search for data sources (tables, models, cards, dashboards, metrics) in Metabase.
+  Abstracted from the API endpoint logic."
+  [{:keys [term-queries semantic-queries database-id created-at last-edited-at
+           entity-types limit metabot-id]}]
+  (let [search-models (if (seq entity-types)
+                        (set (distinct (keep entity-type->search-model entity-types)))
+                        metabot-search-models)
+        all-queries   (distinct (concat (or term-queries []) (or semantic-queries [])))
+        metabot (t2/select-one :model/Metabot :entity_id (get-in metabot-v3.config/metabot-config [metabot-id :entity-id] metabot-id))
+        use-verified-content? (if metabot-id
+                                (:use_verified_content metabot)
+                                false)
+        search-fn (fn [query]
+                    (let [search-context (search/search-context
+                                          (merge
+                                           {:search-string query
+                                            :models search-models
+                                            :table-db-id database-id
+                                            :created-at created-at
+                                            :last-edited-at last-edited-at
+                                            :current-user-id api/*current-user-id*
+                                            :is-impersonated-user? (perms/impersonated-user?)
+                                            :is-sandboxed-user? (perms/sandboxed-user?)
+                                            :is-superuser? api/*is-superuser?*
+                                            :current-user-perms @api/*current-user-permissions-set*
+                                            :context :metabot
+                                            :archived false
+                                            :limit (or limit 50)
+                                            :offset 0}
+                                           (when use-verified-content?
+                                             {:verified true})))
+                          search-results (search/search search-context)]
+                      (:data search-results)))
+        ;; Create futures for parallel execution
+        futures (mapv #(future (search-fn %)) all-queries)
+        result-lists (mapv deref futures)
+        fused-results (reciprocal-rank-fusion result-lists)]
+    (map transform-search-result fused-results)))

enterprise/backend/src/metabase_enterprise/metabot_v3/tools/search_data_sources.clj

@@ -0,0 +1,130 @@
+(ns metabase-enterprise.metabot-v3.tools.search-data-sources
+  "Tools for searching data sources in the Metabot v3 context."
+  (:require
+   [clojure.string :as str]
+   [medley.core :as m]
+   [metabase-enterprise.metabot-v3.config :as metabot-v3.config]
+   [metabase-enterprise.metabot-v3.reactions]
+   [metabase.api.common :as api]
+   [metabase.permissions.core :as perms]
+   [metabase.search.core :as search]
+   [toucan2.core :as t2]))
+
+(def ^:private entity-type->search-model
+  "Maps API entity types to search engine model types"
+  {"table" "table"
+   "model" "dataset"
+   "question" "card"
+   "dashboard" "dashboard"
+   "metric" "metric"})
+
+(def ^:private search-model->result-type
+  "Maps search engine model types to result types"
+  {"dashboard" :dashboard
+   "table" :table
+   "dataset" :model
+   "card" :question
+   "metric" :metric})
+
+(defn transform-search-result
+  "Transform a single search result to match the appropriate entity-specific schema.
+
+   Field mapping:
+   - Tables: name = technical table name (table_name), include database_id and database_schema
+   - Models: name = display name, include database_id and database_schema
+   - Other entities: name = display name, no database fields"
+  [result]
+  (let [model (:model result)
+        result-type (search-model->result-type model)]
+    (case model
+      "table"
+      {:id (:id result)
+       :type result-type
+       :name (:table_name result)           ; Technical table name for tables
+       :display_name (:name result)
+       :description (:description result)
+       :database_id (:database_id result)
+       :database_schema (:table_schema result)}
+
+      "dataset"  ; Models
+      {:id (:id result)
+       :type result-type
+       :name (:name result)                 ; Use the display name here as well since there is no technical name
+       :display_name (:name result)
+       :description (:description result)
+       :database_id (:database_id result)
+       :verified (boolean (:verified result))}  ;; Fallback to False if not provided
+
+      ;; For dashboards, questions, and metrics:
+      {:id (:id result)
+       :type result-type
+       :name (:name result)                 ; Display name
+       :description (:description result)
+       :verified (boolean (or (:verified result)  ; Fallback to False if not provided
+                              (= "verified" (:moderated_status result))))})))
+
+(defn- build-search-context
+  "Build the search context for a given search term"
+  [term database-id search-models search-limit verified?]
+  (search/search-context
+   (cond-> {:search-string term
+            :models (set search-models)
+            :table-db-id database-id
+            :current-user-id api/*current-user-id*
+            :is-impersonated-user? (perms/impersonated-user?)
+            :is-sandboxed-user? (perms/sandboxed-user?)
+            :is-superuser? api/*is-superuser?*
+            :current-user-perms @api/*current-user-permissions-set*
+            :context :metabot
+            :archived false
+            :limit search-limit
+            :offset 0}
+     verified? (assoc :verified true))))
+
+(defn- filter-by-collection
+  "Filter search results by collection, always including tables"
+  [metabot-coll-id results]
+  (if metabot-coll-id
+    (filter (fn [result]
+              (or (= (:model result) "table")
+                  (= (get-in result [:collection :id]) metabot-coll-id)))
+            results)
+    results))
+
+(defn- dedupe-by-id-and-model
+  "Remove duplicate results based on id and model"
+  [results]
+  (vals (m/index-by (juxt :id :model) results)))
+
+(defn- search-for-term
+  "Execute a search for a single term"
+  [{:keys [term database-id search-models search-limit verified?]}]
+  (-> (build-search-context term database-id search-models search-limit verified?)
+      search/search
+      :data))
+
+(defn search-data-sources
+  "Search for data sources (tables, models, cards, dashboards, metrics) in Metabase.
+  Abstracted from the API endpoint logic."
+  [{:keys [keywords description database-id entity-types limit metabot-id]}]
+  (let [metabot-entity-id (get-in metabot-v3.config/metabot-config [metabot-id :entity-id] metabot-id)
+        metabot (t2/select-one :model/Metabot :entity_id metabot-entity-id)
+        metabot-coll-id (:collection_id metabot)
+
+        normalized-description (when-not (str/blank? description) description)
+        entity-types (if (seq entity-types)
+                       entity-types
+                       ["table" "model" "question" "dashboard" "metric"])
+
+        search-terms (distinct (concat (or keywords []) (when normalized-description [normalized-description])))
+        search-models (keep entity-type->search-model entity-types)
+        search-limit (or limit (if metabot-coll-id 75 50))]
+    (->> search-terms
+         (mapcat #(search-for-term {:term %
+                                    :database-id database-id
+                                    :search-models search-models
+                                    :search-limit search-limit
+                                    :verified? (:use_verified_content metabot)}))
+         (filter-by-collection metabot-coll-id)
+         dedupe-by-id-and-model
+         (map transform-search-result))))

enterprise/backend/src/metabase_enterprise/metabot_v3/tools/snippets.clj

@@ -0,0 +1,25 @@
+(ns metabase-enterprise.metabot-v3.tools.snippets
+  (:require
+   [metabase-enterprise.metabot-v3.tools.util :as metabot-v3.tools.u]
+   [metabase.native-query-snippets.core :as snippets]))
+
+(defn get-snippets
+  "Lists SQL snippets available to the current user."
+  []
+  (try
+    {:structured_output
+     (->> (snippets/list-native-query-snippets)
+          (map #(select-keys % [:id :name :description])))}
+    (catch Exception e
+      (metabot-v3.tools.u/handle-agent-error e))))
+
+(defn get-snippet-details
+  "Retrieve a specific SQL snippet by ID, including its content."
+  [snippet-id]
+  (try
+    {:structured_output
+     (when snippet-id
+       (when-let [snippet (snippets/get-native-query-snippet snippet-id)]
+         (select-keys snippet [:id :name :description :content])))}
+    (catch Exception e
+      (metabot-v3.tools.u/handle-agent-error e))))

enterprise/backend/src/metabase_enterprise/metabot_v3/tools/transforms.clj

@@ -0,0 +1,46 @@
+(ns metabase-enterprise.metabot-v3.tools.transforms
+  (:require
+   [metabase-enterprise.metabot-v3.tools.util :as metabot-v3.tools.u]
+   [metabase-enterprise.transforms.api :as api.transforms]
+   [metabase.util.malli.registry :as mr]))
+
+(mr/def ::transform-source ::api.transforms/transform-source)
+(mr/def ::transform-target ::api.transforms/transform-target)
+
+(defn- transform-source-type
+  [transform]
+  (-> transform :source :type))
+
+(defn- transform-source-query-type
+  [transform]
+  (-> transform :source :query :type))
+
+(defn- python-transform?
+  [transform]
+  (= "python" (transform-source-type transform)))
+
+(defn- native-query-transform?
+  [transform]
+  (and (= "query" (transform-source-type transform))
+       (= "native" (transform-source-query-type transform))))
+
+(defn get-transforms
+  "Get a list of all known transforms."
+  []
+  (try
+    {:structured_output
+     (->> (api.transforms/get-transforms)
+          (into [] (comp (map #(select-keys % [:id :entity_id :name :description :source]))
+                         (filter #(or (python-transform? %)
+                                      (native-query-transform? %))))))}
+    (catch Exception e
+      (metabot-v3.tools.u/handle-agent-error e))))
+
+(defn get-transform-details
+  "Get information about a transform."
+  [transform-id]
+  (try
+    {:structured_output
+     (api.transforms/get-transform transform-id)}
+    (catch Exception e
+      (metabot-v3.tools.u/handle-agent-error e))))

enterprise/backend/src/metabase_enterprise/metabot_v3/tools/util.clj

@@ -8,6 +8,8 @@
    [metabase.lib.core :as lib]
    [metabase.lib.metadata :as lib.metadata]
    [metabase.lib.types.isa :as lib.types.isa]
+   [metabase.lib.util :as lib.util]
+   [metabase.premium-features.core :as premium-features]
    [metabase.util :as u]
    [toucan2.core :as t2]))
 
@@ -30,6 +32,15 @@
       (isa? (:effective-type column) :type/Time)     :time
       (lib.types.isa/temporal? column)               :date)))
 
+(defn add-table-reference
+  "Add table-reference to columns that have FK relationships."
+  [query col]
+  (cond-> col
+    (and (:fk-field-id col)
+         (:table-id col))
+    (assoc :table-reference (-> (lib/display-name query (lib.metadata/field query (:fk-field-id col)))
+                                lib.util/strip-id))))
+
 (defn table-field-id-prefix
   "Return the field ID prefix for `table-id`."
   [table-id]
@@ -55,7 +66,8 @@
   [query column index field-id-prefix]
   (let [semantic-type (some-> (:semantic-type column) name u/->snake_case_en)]
     (-> {:field-id (str field-id-prefix index)
-         :name (lib/display-name query column)
+         :name (:name column)
+         :display_name (lib/display-name query column)
          :type (convert-field-type column)}
         (m/assoc-some :description (:description column)
                       :semantic-type semantic-type
@@ -121,74 +133,49 @@
     (let [mp (lib.metadata.jvm/application-database-metadata-provider (:db_id table))]
       (lib/query mp (lib.metadata/table mp table-id)))))
 
-(defn metabot-scope-query
+(defn metabot-metrics-and-models-query
   "Return the metric and model cards in metabot scope visible to the current user.
 
-  If provided, the filter clause `metabot-entity-condition` is used to filter the metabot_entity table. The clause can
-  refer to that table using the :mbe alias. The cards are returned with an additional field called :metabot_entity_id
-  containing the ID of the metabot entity that brought the card in scope. (Situations in which several metabot
-  entities can bring a card into scope is not supported. It is expected that at some point this query will be extended
-  to deal with tables as well.)
-
-  For example,
-    [:in [:mbe.id [1 7 42]]] can be used to consider only specific entities,
-    [:= :mbe.metabot_id 1]   can be used to consider only entities of metabot 1."
-  ([]
-   (metabot-scope-query nil))
-  ([metabot-entity-condition]
-   {:select [:*]
-    :from   [[{:union-all
-               [{:select [:card.* [:mbe.id :metabot_entity_id]]
-                 :from   [[:report_card :card]]
-                 :join   [[:metabot_entity :mbe] (cond-> [:and [:= :mbe.model [:inline "collection"]]]
-                                                   metabot-entity-condition (conj metabot-entity-condition))
-                          [:collection :ecoll]   [:= :ecoll.id :mbe.model_id]
-                          [:collection :coll]    [:or
-                                                  [:= :coll.id :ecoll.id]
-                                                  [:like :coll.location [:concat :ecoll.location :ecoll.id [:inline "/%"]]]]]
-                 :where  [:= :card.collection_id :coll.id]}
-                {:select [:card.* [:mbe.id :metabot_entity_id]]
-                 :from   [[:report_card :card]]
-                 :join   [[:metabot_entity :mbe] (cond-> [:and
-                                                          [:in :mbe.model [:inline ["dataset"  "metric"]]]]
-                                                   metabot-entity-condition (conj metabot-entity-condition))]
-                 :where  [:= :card.id :mbe.model_id]}]}
-              :card]]
-    :where  [:and
-             [:in :card.type [:inline ["metric" "model"]]]
-             [:= :card.archived false]
-             ;; check that the current user can see the card
-             (collection/visible-collection-filter-clause :card.collection_id)]}))
+  Takes a metabot-id and returns all metric and model cards in that metabot's collection
+  and its subcollections. If the metabot has use_verified_content enabled, only verified
+  content is returned."
+  [metabot-id]
+  (let [metabot (t2/select-one :model/Metabot :id metabot-id)
+        metabot-collection-id (:collection_id metabot)
+        use-verified-content? (:use_verified_content metabot)
+        collection-filter (if metabot-collection-id
+                            (let [collection (t2/select-one :model/Collection :id metabot-collection-id)
+                                  collection-ids (conj (collection/descendant-ids collection) metabot-collection-id)]
+                              [:in :collection_id collection-ids])
+                            [:and true])
+        base-query {:select [:report_card.*]
+                    :from   [[:report_card]]
+                    :where  [:and
+                             collection-filter
+                             [:in :type [:inline ["metric" "model"]]]
+                             [:= :archived false]
+                             (collection/visible-collection-filter-clause :collection_id)]}]
+    (if (and use-verified-content? (premium-features/has-feature? :content-verification))
+      (-> base-query
+          (assoc :left-join [[:moderation_review :mr] [:and
+                                                       [:= :mr.moderated_item_id :report_card.id]
+                                                       [:= :mr.moderated_item_type [:inline "card"]]
+                                                       [:= :mr.most_recent true]]])
+          (update :where conj [:= :mr.status [:inline "verified"]]))
+      base-query)))
 
 (comment
   (binding [api/*current-user-id* 2
             api/*is-superuser?* true]
-    (t2/select-fn-vec #(select-keys % [:id :name :type :metabot_entity_id])
+    (t2/select-fn-vec #(select-keys % [:id :name :type])
                       :model/Card
-                      (metabot-scope-query [:= :mbe.metabot_id 1])))
+                      (metabot-metrics-and-models-query 1)))
   -)
 
 (defn get-metrics-and-models
   "Retrieve the metric and model cards for the Metabot instance with ID `metabot-id` from the app DB.
 
-  Only card visible to the current user are returned."
+  Only cards visible to the current user are returned."
   [metabot-id]
-  (let [scope-query (metabot-scope-query [:= :mbe.metabot_id metabot-id])]
-    (t2/select :model/Card {:select   [:card.*]
-                            :from     [[scope-query :card]]
-                            :order-by [:card.id]})))
-
-(defn metabot-scope
-  "Return a map from cards (models or metrics) to the ID of metabot entity they belong to.
-
-  Warning: this function assumes that each card is bought by one entity into scope. This assumption holds
-  in the known special cases when there is but one collection entity, or when there are only dataset (model)
-  and metric type entities."
-  [entity-ids]
-  (when (seq entity-ids)
-    (let [cards (t2/select :model/Card
-                           (metabot-scope-query [:in :mbe.id entity-ids]))]
-      (reduce (fn [m c]
-                (assoc m (dissoc c :metabot_entity_id) (:metabot_entity_id c)))
-              {}
-              cards))))
+  (t2/select :model/Card (-> (metabot-metrics-and-models-query metabot-id)
+                             (assoc :order-by [:id]))))

enterprise/backend/src/metabase_enterprise/serialization/v2/models.clj

@@ -37,7 +37,6 @@
    "Dimension"
    "ParameterCard"
    "DashboardCardSeries"
-   "MetabotEntity"
    "MetabotPrompt"
    "TimelineEvent"])
 

enterprise/backend/src/metabase_enterprise/transforms/api.clj

@@ -73,15 +73,20 @@
     (api/check-400 (not (transforms.util/db-routing-enabled? database))
                    (deferred-tru "Transforms are not supported on databases with DB routing enabled."))))
 
-(api.macros/defendpoint :get "/"
+(defn get-transforms
   "Get a list of transforms."
-  [_route-params
-   _query-params]
+  []
   (api/check-superuser)
   (-> (t2/select :model/Transform)
       (t2/hydrate :last_run :transform_tag_ids)
       (->> (map #(update % :last_run transforms.util/localize-run-timestamps)))))
 
+(api.macros/defendpoint :get "/"
+  "Get a list of transforms."
+  [_route-params
+   _query-params]
+  (get-transforms))
+
 (api.macros/defendpoint :post "/"
   "Create a new transform."
   [_route-params
@@ -108,11 +113,9 @@
       ;; Return with hydrated tag_ids
       (t2/hydrate transform :transform_tag_ids))))
 
-(api.macros/defendpoint :get "/:id"
+(defn get-transform
   "Get a specific transform."
-  [{:keys [id]} :- [:map
-                    [:id ms/PositiveInt]]]
-  (log/info "get transform" id)
+  [id]
   (api/check-superuser)
   (let [{:keys [target] :as transform} (api/check-404 (t2/select-one :model/Transform id))
         database-id (source-database-id transform)
@@ -122,6 +125,13 @@
         (u/update-some :last_run transforms.util/localize-run-timestamps)
         (assoc :table target-table))))
 
+(api.macros/defendpoint :get "/:id"
+  "Get a specific transform."
+  [{:keys [id]} :- [:map
+                    [:id ms/PositiveInt]]]
+  (log/info "get transform" id)
+  (get-transform id))
+
 (api.macros/defendpoint :get "/:id/dependencies"
   "Get the dependencies of a specific transform."
   [{:keys [id]} :- [:map

enterprise/backend/test/metabase_enterprise/metabot_v3/api/metabot_test.clj

@@ -4,6 +4,7 @@
    [clojure.test :refer :all]
    [metabase-enterprise.metabot-v3.api.metabot :as metabot-v3.api.metabot]
    [metabase-enterprise.metabot-v3.client :as metabot-v3.client]
+   [metabase-enterprise.metabot-v3.config :as metabot-v3.config]
    [metabase.collections.models.collection :as collection]
    [metabase.lib.core :as lib]
    [metabase.lib.metadata :as lib.metadata]
@@ -24,323 +25,137 @@
          (when (seq original-entities#)
            (t2/insert! :model/Metabot original-entities#))))))
 
-(deftest metabot-entities-get-cards-test
-  (testing "GET /api/ee/metabot-v3/metabot/:id/entities"
-    (mt/with-premium-features #{:metabot-v3}
-      (mt/with-temp [:model/Metabot {metabot-id :id} {:name "Test Metabot"}
-                     :model/Collection {collection-id :id} {:name "Test Collection"}
-                     :model/Card {card-name-1 :name
-                                  card-id-1 :id} {:name "Test Card 1"
-                                                  :collection_id collection-id}
-                     :model/Card {card-name-2 :name
-                                  card-id-2 :id} {:name "Test Card 2"
-                                                  :collection_id nil}
-                     :model/MetabotEntity _ {:metabot_id metabot-id
-                                             :model "dataset"
-                                             :model_id card-id-1}
-                     :model/MetabotEntity _ {:metabot_id metabot-id
-                                             :model "dataset"
-                                             :model_id card-id-2}]
-
-        (testing "should return entities for a metabot"
-          (let [response (mt/user-http-request :crowberto :get 200
-                                               (format "ee/metabot-v3/metabot/%d/entities" metabot-id))]
-            (is (= 2 (:total response)))
-            (is (= 2 (count (:items response))))
-            (is (= #{card-name-1 card-name-2}
-                   (set (map :name (:items response)))))
-            (is (= #{card-id-1 card-id-2}
-                   (set (map :model_id (:items response)))))
-            (is (= #{collection-id "root"} (into #{} (map :collection_id (:items response)))))
-            (is (= #{"Test Collection" "Our analytics"} (into #{} (map :collection_name (:items response)))))))
-
-        (testing "should support pagination"
-          (let [response (mt/user-http-request :crowberto :get 200
-                                               (format "ee/metabot-v3/metabot/%d/entities?limit=1" metabot-id))]
-            (is (= 2 (:total response)))
-            (is (= 1 (count (:items response))))
-            (is (= 1 (:limit response)))))
-
-        (testing "should require superuser permissions"
-          (is (= "You don't have permissions to do that."
-                 (mt/user-http-request :rasta :get 403
-                                       (format "ee/metabot-v3/metabot/%d/entities" metabot-id)))))
-
-        (testing "should return 404 for non-existent metabot"
-          (is (= "Not found."
-                 (mt/user-http-request :crowberto :get 404
-                                       (format "ee/metabot-v3/metabot/%d/entities" Integer/MAX_VALUE)))))))))
-
-(deftest metabot-entities-get-collections-test
-  (testing "GET /api/ee/metabot-v3/metabot/:id/entities"
-    (mt/with-premium-features #{:metabot-v3}
-      (mt/with-temp [:model/Metabot {metabot-id :id} {:name "Test Metabot"}
-                     :model/Collection {collection-id :id} {:name "Test Collection"}
-                     :model/MetabotEntity _ {:metabot_id metabot-id
-                                             :model "collection"
-                                             :model_id collection-id}]
-
-        (testing "should return entities for a metabot"
-          (let [response (mt/user-http-request :crowberto :get 200
-                                               (format "ee/metabot-v3/metabot/%d/entities" metabot-id))]
-            (is (= #{collection-id}
-                   (set (map :model_id (:items response)))))
-            (is (= #{"Test Collection"} (into #{} (map :name (:items response)))))
-            (is (= #{nil} (into #{} (map :collection_id (:items response)))))
-            (is (= #{nil} (into #{} (map :collection_name (:items response)))))))
-
-        (testing "should support pagination"
-          (let [response (mt/user-http-request :crowberto :get 200
-                                               (format "ee/metabot-v3/metabot/%d/entities?limit=1" metabot-id))]
-            (is (= 1 (:total response)))
-            (is (= 1 (count (:items response))))
-            (is (= 1 (:limit response)))))
-
-        (testing "should require superuser permissions"
-          (is (= "You don't have permissions to do that."
-                 (mt/user-http-request :rasta :get 403
-                                       (format "ee/metabot-v3/metabot/%d/entities" metabot-id)))))
-
-        (testing "should return 404 for non-existent metabot"
-          (is (= "Not found."
-                 (mt/user-http-request :crowberto :get 404
-                                       (format "ee/metabot-v3/metabot/%d/entities" Integer/MAX_VALUE)))))))))
-
-(deftest metabot-entities-put-test-add-entities
-  (testing "PUT /api/ee/metabot-v3/metabot/:id/entities"
-    (mt/with-premium-features #{:metabot-v3}
-      (mt/with-temp [:model/Metabot {metabot-id :id} {:name "Test Metabot"}
-                     :model/Collection {collection-id :id} {:name "Test Collection"}
-                     :model/Card {card-id-1 :id} {:name "Test Card 1"
-                                                  :collection_id collection-id
-                                                  :type :model}
-                     :model/Card {card-id-2 :id} {:name "Test Card 2"
-                                                  :collection_id collection-id
-                                                  :type :metric}]
-
-        (testing "should add entities to metabot access list"
-          (let [entities {:items [{:id card-id-1 :model "dataset"}
-                                  {:id card-id-2 :model "metric"}]}]
-
-            ;; Make the API call to add entities
-            (with-redefs [metabot-v3.api.metabot/generate-sample-prompts identity]
-              (mt/user-http-request :crowberto :put 204
-                                    (format "ee/metabot-v3/metabot/%d/entities" metabot-id)
-                                    entities))
-
-            ;; Verify entities were added to the database
-            (let [added-entities (t2/select :model/MetabotEntity
-                                            :metabot_id metabot-id)]
-              (is (= #{card-id-1 card-id-2}
-                     (set (map :model_id added-entities))))
-              (is (= [:dataset :metric] (mapv :model added-entities))))))))))
-
-(deftest metabot-entities-put-test-add-entities-collections
-  (testing "PUT /api/ee/metabot-v3/metabot/:id/entities"
-    (mt/with-premium-features #{:metabot-v3}
-      (mt/with-temp [:model/Metabot {metabot-id :id} {:name "Test Metabot"}
-                     :model/Collection {collection-id :id} {:name "Test Collection"}]
-
-        (testing "should add entities to metabot access list"
-          (let [entities {:items [{:id collection-id :model "collection"}]}]
-
-            ;; Make the API call to add entities
-            (with-redefs [metabot-v3.api.metabot/generate-sample-prompts identity]
-              (mt/user-http-request :crowberto :put 204
-                                    (format "ee/metabot-v3/metabot/%d/entities" metabot-id)
-                                    entities))
-
-            ;; Verify entities were added to the database
-            (let [added-entities (t2/select :model/MetabotEntity
-                                            :metabot_id metabot-id)]
-              (is (= #{collection-id}
-                     (set (map :model_id added-entities))))
-              (is (= [:collection] (mapv :model added-entities))))))))))
-
-(deftest metabot-entities-put-duplicates-test
-  (testing "PUT /api/ee/metabot-v3/metabot/:id/entities"
-    (mt/with-premium-features #{:metabot-v3}
-      (mt/with-temp [:model/Metabot {metabot-id :id} {:name "Test Metabot"}
-                     :model/Collection {collection-id :id} {:name "Test Collection"}
-                     :model/Card {card-id-1 :id} {:name "Test Card 1"
-                                                  :collection_id collection-id}
-                     :model/Card {card-id-2 :id} {:name "Test Card 2"
-                                                  :collection_id collection-id}
-                     :model/MetabotEntity _ {:metabot_id metabot-id :model_id card-id-1 :model "metric"}
-                     :model/MetabotEntity _ {:metabot_id metabot-id :model_id card-id-2 :model "dataset"}]
-        (testing "should not add duplicate entities"
-          (let [entities {:items [{:id card-id-1 :model "metric"}]}] ;; This entity already exists
-
-            ;; Make the API call again with an existing entity
-            (with-redefs [metabot-v3.api.metabot/generate-sample-prompts identity]
-              (mt/user-http-request :crowberto :put 204
-                                    (format "ee/metabot-v3/metabot/%d/entities" metabot-id)
-                                    entities))
-
-            ;; Verify no new entities were added
-            (is (= 2 (t2/count :model/MetabotEntity :metabot_id metabot-id)))))))))
-
-(deftest metabot-entities-put-super-user-only-test
-  (testing "PUT /api/ee/metabot-v3/metabot/:id/entities"
-    (mt/with-premium-features #{:metabot-v3}
-      (mt/with-temp [:model/Metabot {metabot-id :id} {:name "Test Metabot"}
-                     :model/Collection {collection-id :id} {:name "Test Collection"}
-                     :model/Card {card-id-1 :id} {:name "Test Card 1"
-                                                  :collection_id collection-id}]
-        (testing "should require superuser permissions"
-          (is (= "You don't have permissions to do that."
-                 (mt/user-http-request :rasta :put 403
-                                       (format "ee/metabot-v3/metabot/%d/entities" metabot-id)
-                                       {:items [{:id (str card-id-1) :model "metric"}]}))))))))
-
-(deftest metabot-entities-put-404-non-existent-test
-  (testing "PUT /api/ee/metabot-v3/metabot/:id/entities"
-    (mt/with-premium-features #{:metabot-v3}
-      (mt/with-temp [:model/Collection {collection-id :id} {:name "Test Collection"}
-                     :model/Card {card-id-1 :id} {:name "Test Card 1"
-                                                  :collection_id collection-id}]
-        (testing "should return 404 for non-existent metabot"
-          (is (= "Not found."
-                 (mt/user-http-request :crowberto :put 404
-                                       (format "ee/metabot-v3/metabot/%d/entities" Integer/MAX_VALUE)
-                                       {:items [{:id (str card-id-1) :model "dataset"}]}))))))))
-
 (deftest sample-prompts-e2e-test
-  (let [mp (mt/metadata-provider)
-        model-source-query (lib/query mp (lib.metadata/table mp (mt/id :products)))
-        metric-source-query (-> model-source-query
-                                (lib/aggregate (lib/avg (lib.metadata/field mp (mt/id :products :rating))))
-                                (lib/breakout (lib/with-temporal-bucket
-                                                (lib.metadata/field mp (mt/id :products :created_at)) :week)))
-        metric-data  {:description "Metric description"
-                      :dataset_query (lib/->legacy-MBQL metric-source-query)
-                      :type :metric}
-        model-data   {:description "Model desc"
-                      :dataset_query (lib/->legacy-MBQL model-source-query)
-                      :type :model}]
-    (mt/with-premium-features #{:metabot-v3}
-      (mt/with-temp [:model/Collection {coll-id :id}   {}
-                     :model/Collection {child-id1 :id} {:location (collection/location-path coll-id)}
-                     :model/Collection {child-id2 :id} {:location (collection/location-path coll-id)}
-                     :model/Collection {child-id3 :id} {:location (collection/location-path coll-id child-id2)}
-                     :model/Card _ (assoc model-data  :name "Model1"  :collection_id coll-id)
-                     :model/Card _ (assoc metric-data :name "Metric1" :collection_id child-id1)
-                     :model/Card _ (assoc model-data  :name "Model2"  :collection_id child-id2)
-                     :model/Card _ (assoc metric-data :name "Metric2" :collection_id child-id3)
-                     :model/Metabot {metabot-id :id} {:name "metabot"}]
-        (let [entities {:items [{:id coll-id :model "collection"}]}
-              prompts {"Metric1" ["metric question1" "metric question2" "metric question3"]
-                       "Metric2" ["metric question4" "metric question5" "metric question6"]
-                       "Model1"  ["model question1"  "model question2"  "model question3"]
-                       "Model2"  ["model question4"  "model question5"  "model question6"]}
-              generate-prompt (fn [promptables]
-                                (map (fn [promptable]
-                                       {:questions (-> promptable :name prompts)})
-                                     promptables))
-              prompt-generator (fn [request]
-                                 (-> request
-                                     (update :metrics generate-prompt)
-                                     (update :tables  generate-prompt)
-                                     (set/rename-keys {:metrics :metric_questions
-                                                       :tables :table_questions})))]
-          ;; --------------------------- Generating sample prompts ---------------------------
-          (testing "should add entities to metabot access list and generate prompt suggestions"
-            (with-redefs [metabot-v3.client/generate-example-questions prompt-generator]
-              ;; Make the API call to add entities
-              (mt/user-http-request :crowberto :put 204
-                                    (format "ee/metabot-v3/metabot/%d/entities" metabot-id)
-                                    entities))
-            (let [entity-id     (t2/select-one-fn :id :model/MetabotEntity :metabot_id metabot-id)
-                  added-prompts (t2/select [:model/MetabotPrompt [:card.name :model_name] :prompt]
-                                           :metabot_entity_id entity-id
-                                           {:join     [[:report_card :card] [:= :card.id :card_id]]
-                                            :order-by [:metabot_prompt.id]})]
-              ;; Verify prompts were added to the database
-              (is (= prompts
-                     (-> (group-by :model_name added-prompts)
-                         (update-vals #(map :prompt %)))))))
-          ;; --------------------------- Querying sample prompts ---------------------------
-          (let [expected-prompts (into #{} (mapcat val) prompts)
-                all-prompts (mt/user-http-request :rasta :get 200
-                                                  (format "ee/metabot-v3/metabot/%d/prompt-suggestions" metabot-id))]
-            (testing "can get all prompts"
-              (is (= (count expected-prompts) (:total all-prompts)))
-              (is (= expected-prompts  (into #{} (map :prompt) (:prompts all-prompts)))))
-            (testing "can page through the prompts"
-              (let [offset 3
-                    limit 5
-                    url (format "ee/metabot-v3/metabot/%d/prompt-suggestions?offset=%d&limit=%d"
-                                metabot-id offset limit)
-                    page-prompts (mt/user-http-request :rasta :get 200 url)]
-                (is (= {:prompts (->> all-prompts :prompts (drop offset) (take limit))
-                        :limit limit, :offset offset, :total (:total all-prompts)}
-                       page-prompts))))
-            (testing "can filter by model type"
-              (doseq [model-type ["metric" "model"]]
-                (let [url (format "ee/metabot-v3/metabot/%d/prompt-suggestions?model=%s" metabot-id model-type)
-                      expected-prompts (->> all-prompts :prompts (filter (comp #{model-type} :model)))
+  (mt/dataset test-data
+    (let [mp (mt/metadata-provider)
+          model-source-query (lib/query mp (lib.metadata/table mp (mt/id :products)))
+          metric-source-query (-> model-source-query
+                                  (lib/aggregate (lib/avg (lib.metadata/field mp (mt/id :products :rating))))
+                                  (lib/breakout (lib/with-temporal-bucket
+                                                  (lib.metadata/field mp (mt/id :products :created_at)) :week)))
+          metric-data  {:description "Metric description"
+                        :dataset_query (lib/->legacy-MBQL metric-source-query)
+                        :type :metric}
+          model-data   {:description "Model desc"
+                        :dataset_query (lib/->legacy-MBQL model-source-query)
+                        :type :model}]
+      (mt/with-premium-features #{:metabot-v3}
+        (mt/with-temp [:model/Collection {coll-id :id}   {}
+                       :model/Collection {child-id1 :id} {:location (collection/location-path coll-id)}
+                       :model/Collection {child-id2 :id} {:location (collection/location-path coll-id)}
+                       :model/Collection {child-id3 :id} {:location (collection/location-path coll-id child-id2)}
+                       :model/Card _ (assoc model-data  :name "Model1"  :collection_id coll-id)
+                       :model/Card _ (assoc metric-data :name "Metric1" :collection_id child-id1)
+                       :model/Card _ (assoc model-data  :name "Model2"  :collection_id child-id2)
+                       :model/Card _ (assoc metric-data :name "Metric2" :collection_id child-id3)
+                       :model/Metabot {metabot-id :id} {:name "metabot" :collection_id coll-id}]
+          (let [prompts {"Metric1" ["metric question1" "metric question2" "metric question3"]
+                         "Metric2" ["metric question4" "metric question5" "metric question6"]
+                         "Model1"  ["model question1"  "model question2"  "model question3"]
+                         "Model2"  ["model question4"  "model question5"  "model question6"]}
+                generate-prompt (fn [promptables]
+                                  (map (fn [promptable]
+                                         {:questions (-> promptable :name prompts)})
+                                       promptables))
+                prompt-generator (fn [request]
+                                   (-> request
+                                       (update :metrics generate-prompt)
+                                       (update :tables  generate-prompt)
+                                       (set/rename-keys {:metrics :metric_questions
+                                                         :tables :table_questions})))]
+                      ;; --------------------------- Generating sample prompts ---------------------------
+            (testing "should generate prompt suggestions for metabot"
+              (with-redefs [metabot-v3.client/generate-example-questions prompt-generator]
+                          ;; Trigger prompt generation by calling the regenerate endpoint
+                (mt/user-http-request :crowberto :post 204
+                                      (format "ee/metabot-v3/metabot/%d/prompt-suggestions/regenerate" metabot-id)))
+              (let [added-prompts (t2/select [:model/MetabotPrompt [:card.name :model_name] :prompt]
+                                             :metabot_id metabot-id
+                                             {:join     [[:report_card :card] [:= :card.id :card_id]]
+                                              :order-by [:metabot_prompt.id]})]
+                          ;; Verify prompts were added to the database
+                (is (= prompts
+                       (-> (group-by :model_name added-prompts)
+                           (update-vals #(map :prompt %)))))))
+                      ;; --------------------------- Querying sample prompts ---------------------------
+            (let [expected-prompts (into #{} (mapcat val) prompts)
+                  all-prompts (mt/user-http-request :rasta :get 200
+                                                    (format "ee/metabot-v3/metabot/%d/prompt-suggestions" metabot-id))]
+              (testing "can get all prompts"
+                (is (= (count expected-prompts) (:total all-prompts)))
+                (is (= expected-prompts  (into #{} (map :prompt) (:prompts all-prompts)))))
+              (testing "can page through the prompts"
+                (let [offset 3
+                      limit 5
+                      url (format "ee/metabot-v3/metabot/%d/prompt-suggestions?offset=%d&limit=%d"
+                                  metabot-id offset limit)
+                      page-prompts (mt/user-http-request :rasta :get 200 url)]
+                  (is (= {:prompts (->> all-prompts :prompts (drop offset) (take limit))
+                          :limit limit, :offset offset, :total (:total all-prompts)}
+                         page-prompts))))
+              (testing "can filter by model type"
+                (doseq [model-type ["metric" "model"]]
+                  (let [url (format "ee/metabot-v3/metabot/%d/prompt-suggestions?model=%s" metabot-id model-type)
+                        expected-prompts (->> all-prompts :prompts (filter (comp #{model-type} :model)))
+                        filtered-prompts (mt/user-http-request :rasta :get 200 url)]
+                    (is (= {:prompts expected-prompts
+                            :limit nil, :offset nil, :total (count expected-prompts)}
+                           filtered-prompts)))))
+              (testing "can filter by model type and model ID"
+                (let [selected-prompt (rand-nth (:prompts all-prompts))
+                      url (format "ee/metabot-v3/metabot/%d/prompt-suggestions?model=%s&model_id=%d"
+                                  metabot-id (:model selected-prompt) (:model_id selected-prompt))
+                      expected-prompts (->> all-prompts :prompts (filter (comp #{(:model_id selected-prompt)} :model_id)))
                       filtered-prompts (mt/user-http-request :rasta :get 200 url)]
                   (is (= {:prompts expected-prompts
                           :limit nil, :offset nil, :total (count expected-prompts)}
-                         filtered-prompts)))))
-            (testing "can filter by model type and model ID"
-              (let [selected-prompt (rand-nth (:prompts all-prompts))
-                    url (format "ee/metabot-v3/metabot/%d/prompt-suggestions?model=%s&model_id=%d"
-                                metabot-id (:model selected-prompt) (:model_id selected-prompt))
-                    expected-prompts (->> all-prompts :prompts (filter (comp #{(:model_id selected-prompt)} :model_id)))
-                    filtered-prompts (mt/user-http-request :rasta :get 200 url)]
-                (is (= {:prompts expected-prompts
-                        :limit nil, :offset nil, :total (count expected-prompts)}
-                       filtered-prompts))))
-            (testing "can return a sample"
-              (let [limit 5
-                    url (format "ee/metabot-v3/metabot/%d/prompt-suggestions?sample=true&limit=%d"
-                                metabot-id limit)
-                    expected-prompts? (fn [prompts]
-                                        (and (= (count prompts) limit)
-                                             (set/subset? (set prompts) (set (:prompts all-prompts)))))
-                    sample-prompts (mt/user-http-request :rasta :get 200 url)]
-                (is (=? {:prompts expected-prompts?
-                         :limit limit, :offset nil, :total (:total all-prompts)}
-                        sample-prompts))))
-            ;; --------------------------- Deleting & regenerating sample prompts ---------------------------
-            (let [all-prompt-ids (into #{} (map :id) (:prompts all-prompts))
-                  current-prompt-ids #(t2/select-pks-set :model/MetabotPrompt
-                                                         {:join [[:metabot_entity :mbe]
-                                                                 [:= :mbe.id :metabot_prompt.metabot_entity_id]]
-                                                          :where [:= :mbe.metabot_id metabot-id]})
-                  selected-prompt (rand-nth (:prompts all-prompts))
-                  url (format "ee/metabot-v3/metabot/%d/prompt-suggestions/%d" metabot-id (:id selected-prompt))
-                  remaining-prompt-ids (disj all-prompt-ids (:id selected-prompt))]
+                         filtered-prompts))))
+              (testing "can return a sample"
+                (let [limit 5
+                      url (format "ee/metabot-v3/metabot/%d/prompt-suggestions?sample=true&limit=%d"
+                                  metabot-id limit)
+                      expected-prompts? (fn [prompts]
+                                          (and (= (count prompts) limit)
+                                               (set/subset? (set prompts) (set (:prompts all-prompts)))))
+                      sample-prompts (mt/user-http-request :rasta :get 200 url)]
+                  (is (=? {:prompts expected-prompts?
+                           :limit limit, :offset nil, :total (:total all-prompts)}
+                          sample-prompts))))
+                        ;; --------------------------- Deleting & regenerating sample prompts ---------------------------
+              (let [all-prompt-ids (into #{} (map :id) (:prompts all-prompts))
+                    current-prompt-ids #(t2/select-pks-set :model/MetabotPrompt :metabot_id metabot-id)
+                    selected-prompt (rand-nth (:prompts all-prompts))
+                    url (format "ee/metabot-v3/metabot/%d/prompt-suggestions/%d" metabot-id (:id selected-prompt))
+                    remaining-prompt-ids (disj all-prompt-ids (:id selected-prompt))]
 
-              (testing "deleting a specific prompt"
-                (testing "normal users cannot delete"
-                  (mt/user-http-request :rasta :delete 403 url)
-                  (is (= all-prompt-ids (current-prompt-ids))))
-                (testing "admins can delete"
-                  (mt/user-http-request :crowberto :delete 204 url)
-                  (is (= remaining-prompt-ids (current-prompt-ids)))))
+                (testing "deleting a specific prompt"
+                  (testing "normal users cannot delete"
+                    (mt/user-http-request :rasta :delete 403 url)
+                    (is (= all-prompt-ids (current-prompt-ids))))
+                  (testing "admins can delete"
+                    (mt/user-http-request :crowberto :delete 204 url)
+                    (is (= remaining-prompt-ids (current-prompt-ids)))))
 
-              (testing "generating new prompts"
-                (let [url (format "ee/metabot-v3/metabot/%d/prompt-suggestions/regenerate" metabot-id)]
-                  (testing "normal users are not allowed"
-                    (mt/user-http-request :rasta :post 403 url)
-                    (is (= remaining-prompt-ids (current-prompt-ids))))
-                  (testing "admin users are allowed"
-                    (with-redefs [metabot-v3.client/generate-example-questions prompt-generator]
-                      (mt/user-http-request :crowberto :post 204 url)))))
+                (testing "generating new prompts"
+                  (let [url (format "ee/metabot-v3/metabot/%d/prompt-suggestions/regenerate" metabot-id)]
+                    (testing "normal users are not allowed"
+                      (mt/user-http-request :rasta :post 403 url)
+                      (is (= remaining-prompt-ids (current-prompt-ids))))
+                    (testing "admin users are allowed"
+                      (with-redefs [metabot-v3.client/generate-example-questions prompt-generator]
+                        (mt/user-http-request :crowberto :post 204 url)))))
 
-              (let [new-prompt-ids (current-prompt-ids)]
-                (is (= (count all-prompt-ids) (count new-prompt-ids)))
-                (is (empty? (set/intersection all-prompt-ids new-prompt-ids)))
-                (testing "can delete all prompts"
-                  (let [url (format "ee/metabot-v3/metabot/%d/prompt-suggestions" metabot-id)]
-                    (testing "normal users cannot delete"
-                      (mt/user-http-request :rasta :delete 403 url)
-                      (is (= new-prompt-ids (current-prompt-ids))))
-                    (testing "admins can delete"
-                      (mt/user-http-request :crowberto :delete 204 url)
-                      (is (nil? (current-prompt-ids))))))))))))))
+                (let [new-prompt-ids (current-prompt-ids)]
+                  (is (= (count all-prompt-ids) (count new-prompt-ids)))
+                  (is (empty? (set/intersection all-prompt-ids new-prompt-ids)))
+                  (testing "can delete all prompts"
+                    (let [url (format "ee/metabot-v3/metabot/%d/prompt-suggestions" metabot-id)]
+                      (testing "normal users cannot delete"
+                        (mt/user-http-request :rasta :delete 403 url)
+                        (is (= new-prompt-ids (current-prompt-ids))))
+                      (testing "admins can delete"
+                        (mt/user-http-request :crowberto :delete 204 url)
+                        (is (nil? (current-prompt-ids)))))))))))))))
 
 (deftest metabot-list-test
   (testing "GET /api/ee/metabot-v3/metabot"
@@ -362,72 +177,214 @@
             (is (= "You don't have permissions to do that."
                    (mt/user-http-request :rasta :get 403 "ee/metabot-v3/metabot")))))))))
 
-(deftest metabot-entities-delete-collection-test
-  (testing "DELETE /api/ee/metabot-v3/metabot/:id/entities/:model-type/:model-id"
+(deftest metabot-get-single-test
+  (testing "GET /api/ee/metabot-v3/metabot/:id"
     (mt/with-premium-features #{:metabot-v3}
-      (mt/with-temp [:model/Metabot {metabot-id :id} {:name "Test Metabot"}
-                     :model/Collection {collection-id :id} {:name "Test Collection"}
-                     :model/MetabotEntity _ {:metabot_id metabot-id
-                                             :model_id collection-id
-                                             :model "collection"}]
+      (mt/with-temp [:model/Collection {collection-id :id} {:name "Test Collection"}
+                     :model/Metabot {metabot-id :id} {:name "Test Metabot"
+                                                      :description "Test Description"
+                                                      :use_verified_content true
+                                                      :collection_id collection-id}]
 
-        (testing "should delete the specified entity"
-          (is (= 1 (t2/count :model/MetabotEntity :metabot_id metabot-id)))
-
-          ;; Delete one entity
-          (mt/user-http-request :crowberto :delete 204
-                                (format "ee/metabot-v3/metabot/%d/entities/collection/%d"
-                                        metabot-id collection-id))
-
-          (let [remaining-entities (t2/select :model/MetabotEntity :metabot_id metabot-id)]
-            (is (= 0 (count remaining-entities)))))))))
-
-(deftest metabot-entities-delete-test
-  (testing "DELETE /api/ee/metabot-v3/metabot/:id/entities/:model-type/:model-id"
-    (mt/with-premium-features #{:metabot-v3}
-      (mt/with-temp [:model/Metabot {metabot-id :id} {:name "Test Metabot"}
-                     :model/Collection {collection-id :id} {:name "Test Collection"}
-                     :model/Card {card-id-1 :id} {:name "Test Card 1"
-                                                  :collection_id collection-id}
-                     :model/Card {card-id-2 :id} {:name "Test Card 2"
-                                                  :collection_id collection-id}
-                     :model/MetabotEntity _ {:metabot_id metabot-id
-                                             :model_id card-id-1
-                                             :model "dataset"}
-                     :model/MetabotEntity {entity-id-2 :id} {:metabot_id metabot-id
-                                                             :model_id card-id-2
-                                                             :model "metric"}]
-
-        (testing "should delete the specified entity"
-          ;; Verify both entities exist before deletion
-          (is (= 2 (t2/count :model/MetabotEntity :metabot_id metabot-id)))
-
-          ;; Delete one entity
-          (mt/user-http-request :crowberto :delete 204
-                                (format "ee/metabot-v3/metabot/%d/entities/dataset/%d"
-                                        metabot-id card-id-1))
-
-          ;; Verify only one entity remains and it's the correct one
-          (let [remaining-entities (t2/select :model/MetabotEntity :metabot_id metabot-id)]
-            (is (= 1 (count remaining-entities)))
-            (is (= entity-id-2 (:id (first remaining-entities))))
-            (is (= card-id-2 (:model_id (first remaining-entities))))
-            (is (= :metric (:model (first remaining-entities))))))
+        (testing "should return metabot with all fields"
+          (let [response (mt/user-http-request :crowberto :get 200
+                                               (format "ee/metabot-v3/metabot/%d" metabot-id))]
+            (is (= metabot-id (:id response)))
+            (is (= "Test Metabot" (:name response)))
+            (is (= "Test Description" (:description response)))
+            (is (true? (:use_verified_content response)))
+            (is (= collection-id (:collection_id response)))))
 
         (testing "should require superuser permissions"
           (is (= "You don't have permissions to do that."
-                 (mt/user-http-request :rasta :delete 403
-                                       (format "ee/metabot-v3/metabot/%d/entities/metric/%d"
-                                               metabot-id card-id-2)))))
+                 (mt/user-http-request :rasta :get 403
+                                       (format "ee/metabot-v3/metabot/%d" metabot-id)))))
 
         (testing "should return 404 for non-existent metabot"
           (is (= "Not found."
-                 (mt/user-http-request :crowberto :delete 404
-                                       (format "ee/metabot-v3/metabot/%d/entities/metric/%d"
-                                               Integer/MAX_VALUE card-id-2)))))
+                 (mt/user-http-request :crowberto :get 404
+                                       (format "ee/metabot-v3/metabot/%d" Integer/MAX_VALUE)))))))))
 
-        (testing "should return 204 even if entity doesn't exist"
-          ;; This tests idempotency - deleting a non-existent entity should still return success
-          (mt/user-http-request :crowberto :delete 204
-                                (format "ee/metabot-v3/metabot/%d/entities/dataset/%d"
-                                        metabot-id 99999)))))))
+(deftest metabot-put-test
+  (testing "PUT /api/ee/metabot-v3/metabot/:id"
+    (mt/with-premium-features #{:metabot-v3 :content-verification}
+      (mt/with-temp [:model/Collection {collection-id-1 :id} {:name "Collection 1"}
+                     :model/Collection {collection-id-2 :id} {:name "Collection 2"}
+                     :model/Metabot {metabot-id :id} {:name "Test Metabot"
+                                                      :use_verified_content false
+                                                      :collection_id collection-id-1}]
+
+        (testing "should update use_verified_content field"
+          (with-redefs [metabot-v3.api.metabot/generate-sample-prompts (constantly nil)]
+            (let [response (mt/user-http-request :crowberto :put 200
+                                                 (format "ee/metabot-v3/metabot/%d" metabot-id)
+                                                 {:use_verified_content true})]
+              (is (true? (:use_verified_content response)))
+              (is (= collection-id-1 (:collection_id response))) ; Should remain unchanged
+              ;; Verify in database
+              (let [updated-metabot (t2/select-one :model/Metabot :id metabot-id)]
+                (is (true? (:use_verified_content updated-metabot)))))))
+
+        (testing "should update collection_id field"
+          (with-redefs [metabot-v3.api.metabot/generate-sample-prompts (constantly nil)]
+            (let [response (mt/user-http-request :crowberto :put 200
+                                                 (format "ee/metabot-v3/metabot/%d" metabot-id)
+                                                 {:collection_id collection-id-2})]
+              (is (= collection-id-2 (:collection_id response)))
+              (is (true? (:use_verified_content response))) ; Should remain from previous test
+              ;; Verify in database
+              (let [updated-metabot (t2/select-one :model/Metabot :id metabot-id)]
+                (is (= collection-id-2 (:collection_id updated-metabot)))))))
+
+        (testing "should update collection_id to null"
+          (with-redefs [metabot-v3.api.metabot/generate-sample-prompts (constantly nil)]
+            (let [response (mt/user-http-request :crowberto :put 200
+                                                 (format "ee/metabot-v3/metabot/%d" metabot-id)
+                                                 {:collection_id nil})]
+              (is (= nil (:collection_id response)))
+              ;; Verify in database
+              (let [updated-metabot (t2/select-one :model/Metabot :id metabot-id)]
+                (is (= nil (:collection_id updated-metabot)))))))
+
+        (testing "should update all fields simultaneously"
+          (with-redefs [metabot-v3.api.metabot/generate-sample-prompts (constantly nil)]
+            (let [response (mt/user-http-request :crowberto :put 200
+                                                 (format "ee/metabot-v3/metabot/%d" metabot-id)
+                                                 {:use_verified_content false
+                                                  :collection_id collection-id-1})]
+              (is (= false (:use_verified_content response)))
+              (is (= collection-id-1 (:collection_id response))))))
+
+        (testing "should require superuser permissions"
+          (is (= "You don't have permissions to do that."
+                 (mt/user-http-request :rasta :put 403
+                                       (format "ee/metabot-v3/metabot/%d" metabot-id)
+                                       {:use_verified_content true}))))
+
+        (testing "should return 404 for non-existent metabot"
+          (is (= "Not found."
+                 (mt/user-http-request :crowberto :put 404
+                                       (format "ee/metabot-v3/metabot/%d" Integer/MAX_VALUE)
+                                       {:use_verified_content true}))))
+
+        (testing "should prevent enabling verified content without premium feature"
+          (mt/with-premium-features #{:metabot-v3}  ; Only metabot-v3, no content-verification
+            (is (= "Content verification is a paid feature not currently available to your instance. Please upgrade to use it. Learn more at metabase.com/upgrade/"
+                   (:message (mt/user-http-request :crowberto :put 402
+                                                   (format "ee/metabot-v3/metabot/%d" metabot-id)
+                                                   {:use_verified_content true}))))))
+
+        (testing "should prevent updating collection_id on primary metabot instance"
+          (let [metabot-id (metabot-v3.config/normalize-metabot-id metabot-v3.config/internal-metabot-id)]
+            (is (= "Cannot update collection_id for the primary metabot instance."
+                   (mt/user-http-request :crowberto :put 400
+
+                                         (format "ee/metabot-v3/metabot/%d" metabot-id)
+                                         {:collection_id collection-id-1})))
+            ;; Verify the collection_id was not updated
+            (let [unchanged-metabot (t2/select-one :model/Metabot :id metabot-id)]
+              (is (= nil (:collection_id unchanged-metabot))))
+
+            ;; Verify that updating other fields still works
+            (with-redefs [metabot-v3.api.metabot/generate-sample-prompts (constantly nil)]
+              (let [response (mt/user-http-request :crowberto :put 200
+                                                   (format "ee/metabot-v3/metabot/%d" metabot-id)
+                                                   {:use_verified_content true})]
+                (is (true? (:use_verified_content response)))
+                (is (= nil (:collection_id response)))))))))))
+
+(deftest metabot-prompt-regeneration-on-config-change-test
+  (mt/dataset test-data
+    (testing "PUT /api/ee/metabot-v3/metabot/:id should regenerate prompts when config changes"
+      (mt/with-premium-features #{:metabot-v3 :content-verification}
+        (let [metric-query {:type :query, :database (mt/id), :query {:source-table (mt/id :products)}}
+              model-query {:type :query, :database (mt/id), :query {:source-table (mt/id :products)}}]
+          (mt/with-temp [:model/Collection {collection-id-1 :id} {:name "Collection 1"}
+                         :model/Collection {collection-id-2 :id} {:name "Collection 2"}
+                         :model/Card {card-id-1 :id} {:name "Test Metric Card"
+                                                      :type :metric
+                                                      :dataset_query metric-query}
+                         :model/Card {card-id-2 :id} {:name "Test Model Card"
+                                                      :type :model
+                                                      :dataset_query model-query}
+                         :model/Card {card-id-3 :id} {:name "Test Card for New Prompt"
+                                                      :type :metric
+                                                      :dataset_query metric-query}
+                         :model/Card {card-id-4 :id} {:name "Test Card for Collection Change"
+                                                      :type :model
+                                                      :dataset_query model-query}
+                         :model/Card {card-id-5 :id} {:name "Test Card for Baseline"
+                                                      :type :metric
+                                                      :dataset_query metric-query}
+                         :model/Metabot {metabot-id :id} {:name "Test Metabot"
+                                                          :use_verified_content false
+                                                          :collection_id collection-id-1}
+                         :model/MetabotPrompt {prompt-id-1 :id} {:metabot_id metabot-id
+                                                                 :prompt "old prompt 1"
+                                                                 :model :metric
+                                                                 :card_id card-id-1}
+                         :model/MetabotPrompt {prompt-id-2 :id} {:metabot_id metabot-id
+                                                                 :prompt "old prompt 2"
+                                                                 :model :model
+                                                                 :card_id card-id-2}]
+
+            (let [original-prompt-ids #{prompt-id-1 prompt-id-2}]
+
+              (testing "should regenerate prompts when use_verified_content changes"
+                (with-redefs [metabot-v3.api.metabot/generate-sample-prompts
+                              (fn [metabot-id]
+                                (t2/insert! :model/MetabotPrompt {:metabot_id metabot-id
+                                                                  :prompt "new prompt after verified change"
+                                                                  :model :metric
+                                                                  :card_id card-id-3}))]
+                  (mt/user-http-request :crowberto :put 200
+                                        (format "ee/metabot-v3/metabot/%d" metabot-id)
+                                        {:use_verified_content true})
+
+                            ;; Verify old prompts were deleted and new ones created
+                  (let [current-prompts (t2/select :model/MetabotPrompt :metabot_id metabot-id)
+                        current-prompt-ids (set (map :id current-prompts))]
+                    (is (= 1 (count current-prompts)))
+                    (is (empty? (set/intersection original-prompt-ids current-prompt-ids)))
+                    (is (= "new prompt after verified change" (:prompt (first current-prompts)))))))
+
+              (testing "should regenerate prompts when collection_id changes"
+                (with-redefs [metabot-v3.api.metabot/generate-sample-prompts
+                              (fn [metabot-id]
+                                (t2/insert! :model/MetabotPrompt {:metabot_id metabot-id
+                                                                  :prompt "new prompt after collection change"
+                                                                  :model :model
+                                                                  :card_id card-id-4}))]
+                  (mt/user-http-request :crowberto :put 200
+                                        (format "ee/metabot-v3/metabot/%d" metabot-id)
+                                        {:collection_id collection-id-2})
+
+                            ;; Verify prompts were regenerated again
+                  (let [current-prompts (t2/select :model/MetabotPrompt :metabot_id metabot-id)]
+                    (is (= 1 (count current-prompts)))
+                    (is (= "new prompt after collection change" (:prompt (first current-prompts)))))))
+
+              (testing "should NOT regenerate prompts when no relevant fields change"
+                          ;; First, establish a baseline
+                (t2/delete! :model/MetabotPrompt :metabot_id metabot-id)
+                (t2/insert! :model/MetabotPrompt {:metabot_id metabot-id
+                                                  :prompt "baseline prompt"
+                                                  :model :metric
+                                                  :card_id card-id-5})
+                (let [baseline-prompts (t2/select :model/MetabotPrompt :metabot_id metabot-id)
+                      baseline-ids (set (map :id baseline-prompts))]
+
+                            ;; Make a PUT request that doesn't change verified content or collection_id
+                            ;; (This would be if we add other fields to update in the future)
+                  (with-redefs [metabot-v3.api.metabot/generate-sample-prompts
+                                (fn [_] (throw (Exception. "Should not be called")))]
+                    (mt/user-http-request :crowberto :put 200
+                                          (format "ee/metabot-v3/metabot/%d" metabot-id)
+                                          {:use_verified_content true  ; Same as current value
+                                           :collection_id collection-id-2})) ; Same as current value
+
+                            ;; Verify prompts were NOT changed
+                  (let [current-prompts (t2/select :model/MetabotPrompt :metabot_id metabot-id)
+                        current-ids (set (map :id current-prompts))]
+                    (is (= baseline-ids current-ids))
+                    (is (= "baseline prompt" (:prompt (first current-prompts))))))))))))))

enterprise/backend/test/metabase_enterprise/metabot_v3/config_test.clj

@@ -36,9 +36,9 @@
                    (metabot-v3.config/resolve-dynamic-profile-id nil "any-metabot-id")))))
 
         (testing "metabot-id mapping takes third precedence"
-          (is (= "experimental"
+          (is (= "internal"
                  (metabot-v3.config/resolve-dynamic-profile-id nil metabot-v3.config/internal-metabot-id)))
-          (is (= "default"
+          (is (= "embedding"
                  (metabot-v3.config/resolve-dynamic-profile-id nil metabot-v3.config/embedded-metabot-id))))
 
         (testing "falls back to default when no matches"
@@ -47,7 +47,7 @@
 
         (testing "single arity version uses dynamic metabot resolution"
           (mt/with-temporary-setting-values [metabot-v3.settings/metabot-id metabot-v3.config/embedded-metabot-id]
-            (is (= "default"
+            (is (= "embedding"
                    (metabot-v3.config/resolve-dynamic-profile-id nil)))))))))
 
 (deftest integrated-resolution-test
@@ -75,4 +75,4 @@
           (let [metabot-id (metabot-v3.config/resolve-dynamic-metabot-id nil)
                 profile-id (metabot-v3.config/resolve-dynamic-profile-id nil metabot-id)]
             (is (= metabot-v3.config/internal-metabot-id metabot-id))
-            (is (= "experimental" profile-id))))))))
+            (is (= "internal" profile-id))))))))

enterprise/backend/test/metabase_enterprise/metabot_v3/context_test.clj

@@ -10,13 +10,16 @@
         extra [{:id 3 :name "extra1"} {:id 4 :name "extra2"}]
         all-tables (concat used extra)]
     (with-redefs [table-utils/used-tables (fn [_] used)
-                  table-utils/database-tables (fn [_ opts]
-                                                (take (:all-tables-limit opts 100) all-tables))]
-      (let [result (#'context/database-tables-for-context 123 {:query {:native {:query "SELECT *"}} :all-tables-limit 3})]
-        (is (= (take 2 result) used) "Used tables should be first")
-        (is (= 3 (count result)) "Should respect limit")
-        (is (some #(= 3 (:id %)) result) "Should include extra table if room")
-        (is (not (some #(= 4 (:id %)) result)) "Should not exceed limit")
+                  table-utils/enhanced-database-tables (fn [_ {:keys [priority-tables all-tables-limit] :or {all-tables-limit 100}}]
+                                                         (let [priority-ids (set (map :id priority-tables))
+                                                               non-priority (remove #(priority-ids (:id %)) all-tables)
+                                                               result (concat priority-tables (take (- all-tables-limit (count priority-tables)) non-priority))]
+                                                           (take all-tables-limit result)))]
+      (let [result (#'context/database-tables-for-context {:query {:database 123 :native {:query "SELECT *"}}, :all-tables-limit 3})]
+        (is (= used result) "Should only return used tables, in order")
+        (is (= (count used) (count result)) "Should return all used tables")
+        (is (not-any? #(= 3 (:id %)) result) "Should not include extra tables")
+        (is (not-any? #(= 4 (:id %)) result) "Should not include extra tables")
         (is (= (count (distinct (map :id result))) (count result))) ; no duplicates
         ))))
 
@@ -25,74 +28,101 @@
         extra (mapv #(hash-map :id % :name (str "extra" %)) (range 6 10))
         all-tables (concat used extra)]
     (with-redefs [table-utils/used-tables (fn [_] used)
-                  table-utils/database-tables (fn [_ opts]
-                                                (take (:all-tables-limit opts 100) all-tables))]
-      (let [result (#'context/database-tables-for-context 123 {:query {:native {:query "SELECT *"}} :all-tables-limit 3})]
-        (is (= 3 (count result)) "Should respect limit")
-        (is (every? (set used) result) "Only used tables should be present if they exceed limit")
+                  table-utils/enhanced-database-tables (fn [_ {:keys [priority-tables all-tables-limit] :or {all-tables-limit 100}}]
+                                                         (let [priority-ids (set (map :id priority-tables))
+                                                               non-priority (remove #(priority-ids (:id %)) all-tables)
+                                                               result (concat priority-tables (take (- all-tables-limit (count priority-tables)) non-priority))]
+                                                           (take all-tables-limit result)))]
+      (let [result (#'context/database-tables-for-context {:query {:database 123 :native {:query "SELECT *"}}, :all-tables-limit 3})]
+        (is (= used result) "Should return all used tables")
+        (is (= (count used) (count result)) "Should return all used tables")
         (is (= (count (distinct (map :id result))) (count result))) ; no duplicates
         ))))
 
 (deftest database-tables-for-context-no-used-tables
   (let [extra (mapv #(hash-map :id % :name (str "extra" %)) (range 1 5))]
     (with-redefs [table-utils/used-tables (fn [_] [])
-                  table-utils/database-tables (fn [_ opts]
-                                                (take (:all-tables-limit opts 100) extra))]
-      (let [result (#'context/database-tables-for-context 123 {:query {:native {:query "SELECT *"}} :all-tables-limit 3})]
-        (is (= 3 (count result)))
-        (is (every? (set extra) result))
-        (is (= (count (distinct (map :id result))) (count result))) ; no duplicates
-        ))))
+                  table-utils/enhanced-database-tables (fn [_ opts]
+                                                         (take (:all-tables-limit opts 100) extra))]
+      (let [result (#'context/database-tables-for-context {:query {:database 123 :native {:query "SELECT *"}}, :all-tables-limit 3})]
+        (is (empty? result) "Should return empty when no used tables")))))
 
 (deftest database-tables-for-context-used-tables-exact-limit
   (let [used (mapv #(hash-map :id % :name (str "used" %)) (range 1 4)) ; 3 used tables
         extra (mapv #(hash-map :id % :name (str "extra" %)) (range 4 7))
         all-tables (concat used extra)]
     (with-redefs [table-utils/used-tables (fn [_] used)
-                  table-utils/database-tables (fn [_ opts]
-                                                (take (:all-tables-limit opts 100) all-tables))]
-      (let [result (#'context/database-tables-for-context 123 {:query {:native {:query "SELECT *"}} :all-tables-limit 3})]
-        (is (= 3 (count result)))
+                  table-utils/enhanced-database-tables (fn [_ {:keys [priority-tables all-tables-limit] :or {all-tables-limit 100}}]
+                                                         (let [priority-ids (set (map :id priority-tables))
+                                                               non-priority (remove #(priority-ids (:id %)) all-tables)
+                                                               result (concat priority-tables (take (- all-tables-limit (count priority-tables)) non-priority))]
+                                                           (take all-tables-limit result)))]
+      (let [result (#'context/database-tables-for-context {:query {:database 123 :native {:query "SELECT *"}}, :all-tables-limit 3})]
         (is (= used result) "Should be exactly the used tables, in order")
+        (is (= (count used) (count result)))
         (is (= (count (distinct (map :id result))) (count result))) ; no duplicates
         ))))
 
+(deftest database-tables-for-context-exception-handling
+  (testing "Returns empty when table-utils/enhanced-database-tables throws"
+    (with-redefs [table-utils/used-tables (fn [_] [{:id 1}])
+                  table-utils/enhanced-database-tables (fn [_ _] (throw (Exception. "boom")))]
+      (let [result (#'context/database-tables-for-context {:query {:database 123 :native {:query "SELECT *"}}})]
+        (is (empty? result))))))
+
+(deftest database-tables-for-context-nil-used-tables
+  (testing "Handles nil used-tables gracefully"
+    (with-redefs [table-utils/used-tables (fn [_] nil)
+                  table-utils/enhanced-database-tables (fn [_ opts]
+                                                         (take (:all-tables-limit opts 100) [{:id 1} {:id 2}]))]
+      (let [result (#'context/database-tables-for-context {:query {:database 123 :native {:query "SELECT *"}}})]
+        (is (empty? result) "Should return empty when used-tables is nil")))))
+
+(deftest database-tables-for-context-duplicate-ids
+  (testing "Deduplicates tables returned by database-tables"
+    (let [used [{:id 1}]
+          dup [{:id 1} {:id 1} {:id 2}]]
+      (with-redefs [table-utils/used-tables (fn [_] used)
+                    table-utils/enhanced-database-tables (fn [_ _] dup)]
+        (let [result (#'context/database-tables-for-context {:query {:database 123 :native {:query "SELECT *"}}})]
+          (is (= [1 2] (map :id result)) "Should preserve order and remove duplicates"))))))
+
 ;; --- enhance-context-with-schema tests ---
 
 (deftest enhance-context-with-schema-native
   (let [mock-tables [{:id 1 :name "table1"} {:id 2 :name "table2"}]]
-    (with-redefs [context/database-tables-for-context (fn [_ _] mock-tables)]
+    (with-redefs [context/database-tables-for-context (fn [_] mock-tables)]
       (testing "Enhances context with schema for native queries"
         (let [input {:user_is_viewing [{:query {:type :native :database 123}}]}
               result (#'context/enhance-context-with-schema input)]
-          (is (= (get-in result [:user_is_viewing 0 :database_schema]) mock-tables)))))))
+          (is (= (get-in result [:user_is_viewing 0 :used_tables]) mock-tables)))))))
 
 (deftest enhance-context-with-schema-non-native
   (testing "Does not enhance context for non-native queries"
     (let [input {:user_is_viewing [{:query {:type :query :database 123}}]}
           result (#'context/enhance-context-with-schema input)]
-      (is (nil? (get-in result [:user_is_viewing 0 :database_schema]))))))
+      (is (nil? (get-in result [:user_is_viewing 0 :used_tables]))))))
 
 (deftest enhance-context-with-schema-no-database
   (testing "Does not enhance context if no database present"
     (let [input {:user_is_viewing [{:query {:type :native}}]}
           result (#'context/enhance-context-with-schema input)]
-      (is (nil? (get-in result [:user_is_viewing 0 :database_schema]))))))
+      (is (nil? (get-in result [:user_is_viewing 0 :used_tables]))))))
 
 (deftest enhance-context-with-schema-complete-native-query
   (testing "Enhances context with schema for complete native query structure"
     (let [mock-tables [{:id 1 :name "users"} {:id 2 :name "orders"}]]
-      (with-redefs [context/database-tables-for-context (fn [_ _] mock-tables)]
+      (with-redefs [context/database-tables-for-context (fn [_] mock-tables)]
         (let [input {:user_is_viewing [{:query {:type :native
                                                 :database 123
                                                 :native {:query "SELECT * FROM users WHERE active = true"}}}]}
               result (#'context/enhance-context-with-schema input)]
-          (is (= (get-in result [:user_is_viewing 0 :database_schema]) mock-tables)))))))
+          (is (= (get-in result [:user_is_viewing 0 :used_tables]) mock-tables)))))))
 
 (deftest enhance-context-with-schema-complete-mbql-query
   (testing "Does not enhance context for complete MBQL query structure and doesn't call database-tables-for-context"
     (let [call-count (atom 0)]
-      (with-redefs [context/database-tables-for-context (fn [_ _]
+      (with-redefs [context/database-tables-for-context (fn [_]
                                                           (swap! call-count inc)
                                                           [{:id 1 :name "should-not-be-used"}])]
         (let [input {:user_is_viewing [{:query {:type :query
@@ -101,7 +131,7 @@
                                                         :aggregation [[:count]]
                                                         :breakout [[:field 789 nil]]}}}]}
               result (#'context/enhance-context-with-schema input)]
-          (is (nil? (get-in result [:user_is_viewing 0 :database_schema]))
+          (is (nil? (get-in result [:user_is_viewing 0 :used_tables]))
               "Should not add database_schema for MBQL queries")
           (is (= 0 @call-count)
               "Should not call database-tables-for-context for MBQL queries"))))))

enterprise/backend/test/metabase_enterprise/metabot_v3/models/metabot_test.clj

@@ -2,30 +2,41 @@
   (:require
    [clojure.test :refer :all]
    [metabase-enterprise.metabot-v3.models.metabot]
+   [metabase-enterprise.metabot-v3.models.metabot-prompt]
    [metabase.test :as mt]
    [metabase.test.fixtures :as fixtures]
    [toucan2.core :as t2]))
 
 (use-fixtures :once (fixtures/initialize :db))
 
-(deftest metabot-entities-hydration-test
-  (testing "batched hydration of metabot entities"
-    (mt/with-temp
-      [:model/Database {db-id :id} {}
-       :model/Card card1 {:name "Test Card 1" :type :model :database_id db-id}
-       :model/Card card2 {:name "Test Card 2" :type :model :database_id db-id}
-       :model/Metabot {metabot1-id :id} {:name "Test Metabot 1"}
-       :model/Metabot {metabot2-id :id} {:name "Test Metabot 2"}
-       :model/MetabotEntity _ {:metabot_id metabot1-id :model "dataset" :model_id (:id card1)}
-       :model/MetabotEntity _ {:metabot_id metabot1-id :model "dataset" :model_id (:id card2)}]
+(deftest metabot-prompts-hydration-test
+  (mt/dataset test-data
+    (testing "batched hydration of metabot prompts"
+      (mt/with-temp
+        [:model/Card {card1-id :id} {:name "Test Card 1"
+                                     :type :metric
+                                     :dataset_query {:database (mt/id)
+                                                     :type :query
+                                                     :query {:source-table (mt/id :products)}}}
+         :model/Card {card2-id :id} {:name "Test Card 2"
+                                     :type :model
+                                     :dataset_query {:database (mt/id)
+                                                     :type :query
+                                                     :query {:source-table (mt/id :products)}}}
+         :model/Metabot {metabot1-id :id} {:name "Test Metabot 1"}
+         :model/Metabot {metabot2-id :id} {:name "Test Metabot 2"}
+         :model/MetabotPrompt _ {:metabot_id metabot1-id :prompt "Prompt 1" :model :metric :card_id card1-id}
+         :model/MetabotPrompt _ {:metabot_id metabot1-id :prompt "Prompt 2" :model :model :card_id card2-id}]
 
-      (let [hydrated-metabots (t2/hydrate (t2/select :model/Metabot :id [:in [metabot1-id metabot2-id]]) :entities)]
-        (testing "should hydrate entities for metabots with entities"
-          (let [metabot1 (first (filter #(= (:id %) metabot1-id) hydrated-metabots))]
-            (is (= 2 (count (:entities metabot1))))
-            (is (= #{(:id card1) (:id card2)}
-                   (set (map :model_id (:entities metabot1)))))))
+        (let [hydrated-metabots (t2/hydrate (t2/select :model/Metabot :id [:in [metabot1-id metabot2-id]]) :prompts)]
+          (testing "should hydrate prompts for metabots with prompts"
+            (let [metabot1 (first (filter #(= (:id %) metabot1-id) hydrated-metabots))]
+              (is (= 2 (count (:prompts metabot1))))
+              (is (= #{"Prompt 1" "Prompt 2"}
+                     (set (map :prompt (:prompts metabot1)))))
+              (is (= #{:metric :model}
+                     (set (map :model (:prompts metabot1)))))))
 
-        (testing "should return empty list for metabots without entities"
-          (let [metabot2 (first (filter #(= (:id %) metabot2-id) hydrated-metabots))]
-            (is (= [] (:entities metabot2)))))))))
+          (testing "should return empty list for metabots without prompts"
+            (let [metabot2 (first (filter #(= (:id %) metabot2-id) hydrated-metabots))]
+              (is (= [] (:prompts metabot2))))))))))

enterprise/backend/test/metabase_enterprise/metabot_v3/table_utils_test.clj

@@ -263,6 +263,67 @@
           (is (every? #(not= "hidden_table" (:name %)) tables))
           (is (some #(= "visible_table" (:name %)) tables)))))))
 
+(deftest enhanced-database-tables-test
+  (testing "enhanced-database-tables function with new format"
+    (mt/with-temp [:model/Database {db-id :id} {}
+                   :model/Table    {table1-id :id} {:db_id db-id, :name "users", :schema "public", :active true, :visibility_type nil}
+                   :model/Table    {table2-id :id} {:db_id db-id, :name "orders", :schema "public", :active true, :visibility_type nil}
+                   :model/Field    {user-id-field :id} {:table_id table1-id, :name "id", :database_type "INTEGER", :base_type :type/Integer, :semantic_type :type/PK}
+                   :model/Field    {} {:table_id table1-id, :name "name", :database_type "VARCHAR", :base_type :type/Text}
+                   :model/Field    {} {:table_id table2-id, :name "id", :database_type "INTEGER", :base_type :type/Integer, :semantic_type :type/PK}
+                   :model/Field    {} {:table_id table2-id, :name "user_id", :database_type "INTEGER", :base_type :type/Integer, :semantic_type :type/FK, :fk_target_field_id user-id-field}
+                   :model/Field    {} {:table_id table2-id, :name "total", :database_type "DECIMAL", :base_type :type/Decimal}]
+
+      (testing "returns tables with new enhanced formatting"
+        (mt/with-current-user (mt/user->id :crowberto)
+          (let [tables (table-utils/enhanced-database-tables db-id)]
+            (is (vector? tables))
+            (is (every? #(contains? % :name) tables))
+            (is (every? #(contains? % :database_schema) tables))
+            (is (every? #(contains? % :fields) tables))
+            (is (every? #(contains? % :type) tables))
+            (is (every? #(contains? % :display_name) tables))
+            (is (every? #(= :table (:type %)) tables))
+            (is (every? #(every? (fn [field] (contains? field :field-id)) (:fields %)) tables))
+            (is (every? #(every? (fn [field] (contains? field :name)) (:fields %)) tables))
+            (is (every? #(every? (fn [field] (contains? field :type)) (:fields %)) tables))
+            (is (every? #(contains? % :metrics) tables)))))
+
+      (testing "includes table-reference for implicitly joined fields"
+        (mt/dataset test-data
+          (mt/with-current-user (mt/user->id :crowberto)
+            (let [test-db-id (mt/id)
+                  tables (table-utils/enhanced-database-tables test-db-id)
+                  orders-table (first (filter #(= "ORDERS" (:name %)) tables))
+                  all-fields (:fields orders-table)
+                  user-fields (filter #(= "User" (:table-reference %)) all-fields)
+                  product-fields (filter #(= "Product" (:table-reference %)) all-fields)]
+              (is (some? orders-table) "Expected to find ORDERS table")
+              (is (seq user-fields) "Expected to find fields with table-reference 'User' from implicit join")
+              (is (some #(= "NAME" (:name %)) user-fields) "Expected to find User NAME field from implicit join")
+              (is (seq product-fields) "Expected to find fields with table-reference 'Product' from implicit join")
+              (is (some #(= "TITLE" (:name %)) product-fields) "Expected to find Product TITLE field from implicit join")))))
+
+      (testing "enhanced format respects all-tables-limit option"
+        (mt/with-current-user (mt/user->id :crowberto)
+          (let [tables (table-utils/enhanced-database-tables db-id {:all-tables-limit 1})]
+            (is (<= (count tables) 1)))))
+
+      (testing "enhanced format excludes specified table IDs"
+        (mt/with-current-user (mt/user->id :crowberto)
+          (let [all-tables (table-utils/enhanced-database-tables db-id)
+                filtered-tables (table-utils/enhanced-database-tables db-id {:exclude-table-ids #{table1-id}})]
+            (is (< (count filtered-tables) (count all-tables)))
+            (is (not-any? #(= table1-id (:id %)) filtered-tables)))))
+
+      (testing "enhanced format prioritizes specified tables"
+        (mt/with-current-user (mt/user->id :crowberto)
+          (let [priority-table {:id table2-id :name "orders" :schema "public"}
+                tables (table-utils/enhanced-database-tables db-id {:priority-tables [priority-table]})]
+            (is (seq tables))
+            ;; Priority table should appear first (if it appears at all)
+            (is (= "orders" (-> tables first :name)))))))))
+
 (deftest ^:parallel format-escaped-test
   (are [in out] (= out
                    (with-out-str (#'table-utils/format-escaped in *out*)))

enterprise/backend/test/metabase_enterprise/metabot_v3/tools/filters_test.clj

@@ -13,9 +13,9 @@
 (defn- by-name
   [dimensions dimension-name]
   (let [pred (if (string? dimension-name)
-               (comp #{dimension-name} :name)
+               (comp #{dimension-name} :display_name)
                (every-pred (comp #{(first dimension-name)} :table-reference)
-                           (comp #{(second dimension-name)} :name)))]
+                           (comp #{(second dimension-name)} :display_name)))]
     (m/find-first pred dimensions)))
 
 (deftest ^:parallel query-metric-test
@@ -489,3 +489,227 @@
           (is (=? expected
                   (metabot-v3.tools.filters/filter-records
                    (assoc input :data-source (select-keys query-details [:query]))))))))))
+
+(deftest ^:parallel query-datasource-table-test
+  (mt/with-current-user (mt/user->id :crowberto)
+    (let [mp (mt/metadata-provider)
+          table-id (mt/id :orders)
+          table-details (#'metabot-v3.dummy-tools/table-details table-id {:metadata-provider mp})
+          ->field-id #(u/prog1 (-> table-details :fields (by-name %) :field-id)
+                        (when-not <>
+                          (throw (ex-info (str "Column " % " not found") {:column %}))))]
+      (testing "Basic table query works"
+        (is (=? {:structured-output {:type :query
+                                     :query-id string?
+                                     :query {:database (mt/id)
+                                             :type :query
+                                             :query {:source-table table-id}}}}
+                (metabot-v3.tools.filters/query-datasource
+                 {:table-id table-id}))))
+
+      (testing "Table query with fields selection"
+        (is (=? {:structured-output {:type :query
+                                     :query-id string?
+                                     :query {:database (mt/id)
+                                             :type :query
+                                             :query {:source-table table-id
+                                                     :fields [[:field (mt/id :orders :created_at) {:base-type :type/DateTimeWithLocalTZ}]
+                                                              [:field (mt/id :orders :total) {:base-type :type/Float}]]}}}}
+                (metabot-v3.tools.filters/query-datasource
+                 {:table-id table-id
+                  :fields [{:field-id (->field-id "Created At")}
+                           {:field-id (->field-id "Total")}]}))))
+
+      (testing "Table query with filters"
+        (is (=? {:structured-output {:type :query
+                                     :query-id string?
+                                     :query {:database (mt/id)
+                                             :type :query
+                                             :query {:source-table table-id
+                                                     :filter [:and
+                                                              [:> [:field (mt/id :orders :discount) {:base-type :type/Float}] 3]
+                                                              [:= [:field (mt/id :orders :user_id) {:base-type :type/Integer}] 10]]}}}}
+                (metabot-v3.tools.filters/query-datasource
+                 {:table-id table-id
+                  :filters [{:field-id (->field-id "Discount")
+                             :operation :number-greater-than
+                             :value 3}
+                            {:field-id (->field-id "User ID")
+                             :operation :equals
+                             :value 10}]}))))
+
+      (testing "Table query with aggregations and grouping"
+        (is (=? {:structured-output {:type :query
+                                     :query-id string?
+                                     :query {:database (mt/id)
+                                             :type :query
+                                             :query {:source-table table-id
+                                                     :aggregation [[:sum [:field (mt/id :orders :subtotal) {:base-type :type/Float}]]
+                                                                   [:avg [:field (mt/id :orders :discount) {:base-type :type/Float}]]]
+                                                     :breakout [[:field (mt/id :orders :product_id) {:base-type :type/Integer}]]}}}}
+                (metabot-v3.tools.filters/query-datasource
+                 {:table-id table-id
+                  :aggregations [{:field-id (->field-id "Subtotal")
+                                  :function :sum}
+                                 {:field-id (->field-id "Discount")
+                                  :function :avg}]
+                  :group-by [{:field-id (->field-id "Product ID")}]}))))
+
+      (testing "Table query with order by and limit"
+        (is (=? {:structured-output {:type :query
+                                     :query-id string?
+                                     :query {:database (mt/id)
+                                             :type :query
+                                             :query {:source-table table-id
+                                                     :order-by [[:asc [:field (mt/id :orders :created_at) {:base-type :type/DateTimeWithLocalTZ}]]
+                                                                [:desc [:field (mt/id :orders :total) {:base-type :type/Float}]]]
+                                                     :limit 100}}}}
+                (metabot-v3.tools.filters/query-datasource
+                 {:table-id table-id
+                  :order-by [{:field {:field-id (->field-id "Created At")}
+                              :direction :asc}
+                             {:field {:field-id (->field-id "Total")}
+                              :direction :desc}]
+                  :limit 100}))))
+
+      (testing "Table query with temporal bucketing"
+        (is (=? {:structured-output {:type :query
+                                     :query-id string?
+                                     :query {:database (mt/id)
+                                             :type :query
+                                             :query {:source-table table-id
+                                                     :aggregation [[:count]]
+                                                     :breakout [[:field (mt/id :orders :created_at)
+                                                                 {:base-type :type/DateTimeWithLocalTZ
+                                                                  :temporal-unit :month}]]}}}}
+                (metabot-v3.tools.filters/query-datasource
+                 {:table-id table-id
+                  :aggregations [{:field-id (->field-id "Product ID")
+                                  :function :count}]
+                  :group-by [{:field-id (->field-id "Created At")
+                              :field-granularity :month}]})))))))
+
+(deftest ^:parallel query-datasource-model-test
+  (mt/with-temp [:model/Card {model-id :id} {:dataset_query (mt/mbql-query orders {})
+                                             :database_id (mt/id)
+                                             :name "Orders Model"
+                                             :description "Test model for orders"
+                                             :type :model}]
+    (mt/with-current-user (mt/user->id :crowberto)
+      (let [model-details (-> (metabot-v3.dummy-tools/get-table-details {:model-id model-id})
+                              :structured-output)
+            model-card-id (str "card__" model-id)
+            ->field-id #(u/prog1 (-> model-details :fields (by-name %) :field-id)
+                          (when-not <>
+                            (throw (ex-info (str "Column " % " not found") {:column %}))))]
+        (testing "Basic model query works"
+          (is (=? {:structured-output {:type :query
+                                       :query-id string?
+                                       :query (mt/mbql-query orders {:source-table model-card-id})}}
+                  (metabot-v3.tools.filters/query-datasource
+                   {:model-id model-id}))))
+
+        (testing "Model query with fields selection"
+          (is (=? {:structured-output {:type :query
+                                       :query-id string?
+                                       :query (mt/mbql-query orders
+                                                {:source-table model-card-id
+                                                 :fields [[:field "CREATED_AT" {:base-type :type/DateTimeWithLocalTZ}]
+                                                          [:field "TOTAL" {:base-type :type/Float}]]})}}
+                  (metabot-v3.tools.filters/query-datasource
+                   {:model-id model-id
+                    :fields [{:field-id (->field-id "Created At")}
+                             {:field-id (->field-id "Total")}]}))))
+
+        (testing "Model query with filters"
+          (is (=? {:structured-output {:type :query
+                                       :query-id string?
+                                       :query (mt/mbql-query orders
+                                                {:source-table model-card-id
+                                                 :filter [:and
+                                                          [:> [:field "DISCOUNT" {:base-type :type/Float}] 5]
+                                                          [:< [:field "SUBTOTAL" {:base-type :type/Float}] 100]]})}}
+                  (metabot-v3.tools.filters/query-datasource
+                   {:model-id model-id
+                    :filters [{:field-id (->field-id "Discount")
+                               :operation :number-greater-than
+                               :value 5}
+                              {:field-id (->field-id "Subtotal")
+                               :operation :number-less-than
+                               :value 100}]}))))
+
+        (testing "Model query with aggregations and grouping"
+          (is (=? {:structured-output {:type :query
+                                       :query-id string?
+                                       :query (mt/mbql-query orders
+                                                {:source-table model-card-id
+                                                 :aggregation [[:count]
+                                                               [:max [:field "TOTAL" {:base-type :type/Float}]]]
+                                                 :breakout [[:field "USER_ID" {:base-type :type/Integer}]]})}}
+                  (metabot-v3.tools.filters/query-datasource
+                   {:model-id model-id
+                    :aggregations [{:field-id (->field-id "Product ID")
+                                    :function :count}
+                                   {:field-id (->field-id "Total")
+                                    :function :max}]
+                    :group-by [{:field-id (->field-id "User ID")}]}))))
+
+        (testing "Model query with order by and limit"
+          (is (=? {:structured-output {:type :query
+                                       :query-id string?
+                                       :query (mt/mbql-query orders
+                                                {:source-table model-card-id
+                                                 :order-by [[:desc [:field "CREATED_AT" {:base-type :type/DateTimeWithLocalTZ}]]]
+                                                 :limit 50})}}
+                  (metabot-v3.tools.filters/query-datasource
+                   {:model-id model-id
+                    :order-by [{:field {:field-id (->field-id "Created At")}
+                                :direction :desc}]
+                    :limit 50}))))))))
+
+(deftest ^:parallel query-datasource-validation-test
+  (mt/with-current-user (mt/user->id :crowberto)
+    (testing "Error when neither table-id nor model-id provided"
+      (is (= {:output "Either table_id or model_id must be provided"}
+             (metabot-v3.tools.filters/query-datasource {}))))
+
+    (testing "Error when both table-id and model-id provided"
+      (is (= {:output "Cannot provide both table_id and model_id"}
+             (metabot-v3.tools.filters/query-datasource
+              {:table-id (mt/id :orders)
+               :model-id 123}))))
+
+    (testing "Error with invalid table-id"
+      (is (= {:output "Invalid table_id not-a-number"}
+             (metabot-v3.tools.filters/query-datasource
+              {:table-id "not-a-number"}))))
+
+    (testing "Error with invalid model-id"
+      (is (= {:output "Invalid model_id not-a-number"}
+             (metabot-v3.tools.filters/query-datasource
+              {:model-id "not-a-number"}))))
+
+    (testing "Error with non-existent table"
+      (is (= {:output (str "No table found with table_id " Integer/MAX_VALUE)}
+             (metabot-v3.tools.filters/query-datasource
+              {:table-id Integer/MAX_VALUE}))))
+
+    (testing "Error with non-existent model"
+      (is (= {:output (str "No model found with model_id " Integer/MAX_VALUE)}
+             (metabot-v3.tools.filters/query-datasource
+              {:model-id Integer/MAX_VALUE}))))))
+
+(deftest ^:parallel query-datasource-permissions-test
+  (mt/with-temp [:model/Card {model-id :id} {:dataset_query (mt/mbql-query orders {})
+                                             :database_id (mt/id)
+                                             :name "Orders Model"
+                                             :type :model}]
+    (testing "User without permissions gets error for table"
+      (is (thrown-with-msg? clojure.lang.ExceptionInfo #"You don't have permissions to do that."
+                            (metabot-v3.tools.filters/query-datasource
+                             {:table-id (mt/id :orders)}))))
+
+    (testing "User without permissions gets error for model"
+      (is (thrown-with-msg? clojure.lang.ExceptionInfo #"You don't have permissions to do that."
+                            (metabot-v3.tools.filters/query-datasource
+                             {:model-id model-id}))))))

enterprise/backend/test/metabase_enterprise/metabot_v3/tools/find_outliers_test.clj

@@ -13,7 +13,7 @@
 
 (defn- by-name
   [dimensions dimension-name]
-  (m/find-first (comp #{dimension-name} :name) dimensions))
+  (m/find-first (comp #{dimension-name} :display_name) dimensions))
 
 (defn- test-card
   []

enterprise/backend/test/metabase_enterprise/metabot_v3/tools/search_data_sources_test.clj

@@ -0,0 +1,152 @@
+(ns metabase-enterprise.metabot-v3.tools.search-data-sources-test
+  (:require
+   [clojure.test :refer :all]
+   [metabase-enterprise.metabot-v3.tools.search-data-sources :as sds]
+   [metabase.api.common :as api]
+   [metabase.permissions.core :as perms]
+   [metabase.search.core :as search]
+   [metabase.test :as mt]
+   [toucan2.core :as t2]))
+
+(deftest transform-search-result-test
+  (testing "table result transformation"
+    (let [result {:model "table"
+                  :id 1
+                  :table_name "orders"
+                  :name "Orders"
+                  :description "Order table"
+                  :database_id 42
+                  :table_schema "public"}
+          expected {:id 1
+                    :type :table
+                    :name "orders"
+                    :display_name "Orders"
+                    :description "Order table"
+                    :database_id 42
+                    :database_schema "public"}]
+      (is (= expected (sds/transform-search-result result)))))
+
+  (testing "model result transformation"
+    (let [result {:model "dataset"
+                  :id 2
+                  :name "Sales Model"
+                  :description "Model for sales"
+                  :database_id 43
+                  :verified true}
+          expected {:id 2
+                    :type :model
+                    :name "Sales Model"
+                    :display_name "Sales Model"
+                    :description "Model for sales"
+                    :database_id 43
+                    :verified true}]
+      (is (= expected (sds/transform-search-result result)))))
+
+  (testing "dashboard result transformation"
+    (let [result {:model "dashboard"
+                  :id 3
+                  :name "Main Dashboard"
+                  :description "Dashboard desc"
+                  :verified false}
+          expected {:id 3
+                    :type :dashboard
+                    :name "Main Dashboard"
+                    :description "Dashboard desc"
+                    :verified false}]
+      (is (= expected (sds/transform-search-result result)))))
+
+  (testing "question result transformation with moderated_status"
+    (let [result {:model "card"
+                  :id 4
+                  :name "Q1"
+                  :description "Question desc"
+                  :moderated_status "verified"}
+          expected {:id 4
+                    :type :question
+                    :name "Q1"
+                    :description "Question desc"
+                    :verified true}]
+      (is (= expected (sds/transform-search-result result)))))
+
+  (testing "metric result transformation"
+    (let [result {:model "metric"
+                  :id 5
+                  :name "Revenue"
+                  :description "Metric desc"
+                  :verified nil}
+          expected {:id 5
+                    :type :metric
+                    :name "Revenue"
+                    :description "Metric desc"
+                    :verified false}]
+      (is (= expected (sds/transform-search-result result))))))
+
+(deftest search-data-sources-test
+  (mt/with-premium-features #{:content-verification}
+    (mt/with-test-user :rasta
+      (let [order-table {:id 1
+                         :model "table"
+                         :table_name "orders"
+                         :name "Orders"
+                         :description "Order table"
+                         :database_id 42
+                         :table_schema "public"}
+            dashboard {:id 2
+                       :model "dashboard"
+                       :name "Sales Dashboard"
+                       :description "Dashboard for sales"
+                       :verified true}]
+
+        (with-redefs [perms/impersonated-user? (fn [] false)
+                      perms/sandboxed-user? (fn [] false)
+                      api/*current-user-id* 1]
+          (testing "search-data-sources returns transformed and deduplicated results for multiple keywords"
+            (with-redefs [search/search (fn [_] {:data [order-table]})]
+              (let [args {:keywords ["orders" "sales"]
+                          :entity-types ["table"]}
+                    results (sds/search-data-sources args)
+                    expected [(sds/transform-search-result order-table)]]
+                (is (= expected results)))))
+
+          (testing "search-data-sources handles multiple results and deduplication"
+            (with-redefs [search/search (fn [_] {:data [order-table
+                                                        order-table
+                                                        dashboard]})]
+              (let [args {:keywords ["orders"]
+                          :entity-types ["table" "dashboard"]}
+                    results (sds/search-data-sources args)
+                    expected [(sds/transform-search-result order-table)
+                              (sds/transform-search-result dashboard)]]
+                (is (= expected results)))))
+
+          (testing "search-data-sources handles empty results"
+            (with-redefs [search/search (fn [_] {:data []})]
+              (let [args {:keywords ["none"]
+                          :entity-types ["table"]}
+                    results (sds/search-data-sources args)]
+                (is (empty? results)))))
+
+          (testing "search-data-sources with metabot collection and verified content"
+            (let [metabot {:entity_id "test-bot"
+                           :collection_id 100
+                           :use_verified_content true}
+                  table-result {:id 1 :model "table" :name "orders"}
+                  card-in-collection {:id 2 :model "card" :name "Q1" :collection {:id 100} :verified true}
+                  card-other-collection {:id 3 :model "card" :name "Q2" :collection {:id 200} :verified true}]
+
+              (with-redefs [t2/select-one (fn [model & _]
+                                            (is (= :model/Metabot model) "Should query for Metabot model")
+                                            metabot)
+                            search/search (fn [context]
+                                            ;; Verify that verified flag is set when metabot has use_verified_content
+                                            (is (true? (:verified context)))
+                                            {:data [table-result card-in-collection card-other-collection]})]
+
+                (let [results (sds/search-data-sources {:keywords ["test"]
+                                                        :metabot-id "test-bot"
+                                                        :entity-types ["table" "question"]})]
+                  ;; Should only return table and card in collection 100
+                  (is (= 2 (count results)))
+                  (is (some #(= (:id %) 1) results) "Table should be included")
+                  (is (some #(= (:id %) 2) results) "Card in metabot collection should be included")
+                  (is (not (some #(= (:id %) 3) results)) "Card in different collection should be excluded"))))))))))

enterprise/backend/test/metabase_enterprise/metabot_v3/tools/search_test.clj

@@ -0,0 +1,305 @@
+(ns metabase-enterprise.metabot-v3.tools.search-test
+  (:require
+   [clojure.test :refer :all]
+   [metabase-enterprise.metabot-v3.tools.search :as search]
+   [metabase.api.common :as api]
+   [metabase.permissions.core :as perms]
+   [metabase.search.core :as search-core]
+   [metabase.test :as mt]
+   [toucan2.core :as t2]))
+
+(deftest reciprocal-rank-fusion-test
+  (testing "Basic RRF with single list"
+    (let [single-list [[{:id 1 :model "card" :name "Card 1"}
+                        {:id 2 :model "dashboard" :name "Dashboard 1"}
+                        {:id 3 :model "table" :name "Table 1"}]]
+          result (#'search/reciprocal-rank-fusion single-list)]
+      (is (= 3 (count result)))
+      (is (= 1 (-> result first :id)))
+      (is (= 2 (-> result second :id)))
+      (is (= 3 (-> result last :id)))))
+
+  (testing "RRF with multiple lists - no overlap"
+    (let [list1 [{:id 1 :model "card" :name "Card 1"}
+                 {:id 2 :model "dashboard" :name "Dashboard 1"}]
+          list2 [{:id 3 :model "table" :name "Table 1"}
+                 {:id 4 :model "metric" :name "Metric 1"}]
+          result (#'search/reciprocal-rank-fusion [list1 list2])]
+      (is (= 4 (count result)))
+      (is (every? #(contains? #{1 2 3 4} (:id %)) result))))
+
+  (testing "RRF with overlapping results - should boost common items"
+    (let [list1 [{:id 1 :model "card" :name "Revenue Report"}
+                 {:id 2 :model "dashboard" :name "Sales Dashboard"}
+                 {:id 3 :model "table" :name "Orders"}]
+          list2 [{:id 2 :model "dashboard" :name "Sales Dashboard"}
+                 {:id 1 :model "card" :name "Revenue Report"}
+                 {:id 4 :model "metric" :name "Total Revenue"}]
+          result (#'search/reciprocal-rank-fusion [list1 list2])]
+      (is (= 4 (count result)))
+      ;; Items appearing in both lists should rank higher
+      (let [top-two-ids (set (map :id (take 2 result)))]
+        (is (contains? top-two-ids 1))
+        (is (contains? top-two-ids 2)))))
+
+  (testing "RRF with identical items at different positions"
+    (let [list1 [{:id 1 :model "card" :name "First"}
+                 {:id 2 :model "dashboard" :name "Second"}
+                 {:id 3 :model "table" :name "Third"}]
+          list2 [{:id 3 :model "table" :name "Third"}
+                 {:id 2 :model "dashboard" :name "Second"}
+                 {:id 1 :model "card" :name "First"}]
+          list3 [{:id 2 :model "dashboard" :name "Second"}
+                 {:id 3 :model "table" :name "Third"}
+                 {:id 1 :model "card" :name "First"}]
+          result (#'search/reciprocal-rank-fusion [list1 list2 list3])]
+      (is (= 3 (count result)))
+      ;; Item 2 appears first in list3, second in list1 and list2, so should rank highest
+      (is (= 2 (-> result first :id)))))
+
+  (testing "RRF with empty lists"
+    (let [list1 []
+          list2 [{:id 1 :model "card" :name "Card 1"}]
+          result (#'search/reciprocal-rank-fusion [list1 list2])]
+      (is (= 1 (count result)))
+      (is (= 1 (-> result first :id)))))
+
+  (testing "RRF with all empty lists"
+    (let [result (#'search/reciprocal-rank-fusion [[] [] []])]
+      (is (empty? result))))
+
+  (testing "RRF score calculation correctness"
+    ;; Test that the RRF formula 1/(k+r) where k=60 is correctly applied
+    (let [list1 [{:id 1 :model "card" :name "Rank 1"}]  ; rank=1, score=1/61
+          list2 [{:id 2 :model "dashboard" :name "Other"}
+                 {:id 1 :model "card" :name "Rank 1"}]  ; rank=2, score=1/62
+          result (#'search/reciprocal-rank-fusion [list1 list2])
+          first-item (first result)
+          second-item (second result)]
+      ;; Item 1 appears at rank 1 in list1 (score=1/61) and rank 2 in list2 (score=1/62)
+      ;; Total score = 1/61 + 1/62 ≈ 0.0164 + 0.0161 = 0.0325
+      ;; Item 2 appears only at rank 1 in list2 (score=1/61 ≈ 0.0164)
+      ;; So item 1 should rank higher than item 2
+      (is (= 1 (:id first-item)))
+      (is (= 2 (:id second-item)))))
+
+  (testing "RRF preserves item data"
+    (let [complex-item {:id 42
+                        :model "dataset"
+                        :name "Complex Dataset"
+                        :description "A detailed description"
+                        :database_id 1
+                        :created_at "2024-01-01"
+                        :extra_field "preserved"}
+          result (#'search/reciprocal-rank-fusion [[complex-item]])]
+      (is (= 1 (count result)))
+      (is (= complex-item (first result)))))
+
+  (testing "RRF with many lists"
+    (let [lists (for [i (range 5)]
+                  [{:id (inc i) :model "card" :name (str "Card " (inc i))}
+                   {:id 99 :model "dashboard" :name "Common Dashboard"}
+                   {:id (+ i 10) :model "table" :name (str "Table " (+ i 10))}])
+          result (#'search/reciprocal-rank-fusion lists)]
+      ;; Item 99 appears in all 5 lists at position 2, so should rank very high
+      (is (= 99 (:id (first result)))))))
+
+(deftest transform-search-result-test
+  (testing "table result transformation"
+    (let [result {:model "table"
+                  :id 1
+                  :table_name "orders"
+                  :name "Orders"
+                  :description "Order table"
+                  :database_id 42
+                  :table_schema "public"
+                  :updated_at "2024-01-01"
+                  :created_at "2024-01-01"}
+          expected {:id 1
+                    :type "table"
+                    :name "orders"
+                    :display_name "Orders"
+                    :description "Order table"
+                    :database_id 42
+                    :database_schema "public"
+                    :updated_at "2024-01-01"
+                    :created_at "2024-01-01"}]
+      (is (= expected (#'search/transform-search-result result)))))
+
+  (testing "model (dataset) result transformation"
+    (let [result {:model "dataset"
+                  :id 2
+                  :name "Sales Model"
+                  :description "Model for sales"
+                  :database_id 43
+                  :verified true
+                  :collection nil
+                  :updated_at "2024-01-02"
+                  :created_at "2024-01-02"}
+          expected {:id 2
+                    :type "model"
+                    :name "Sales Model"
+                    :description "Model for sales"
+                    :database_id 43
+                    :verified true
+                    :collection {}
+                    :updated_at "2024-01-02"
+                    :created_at "2024-01-02"}]
+      (is (= expected (#'search/transform-search-result result)))))
+
+  (testing "dashboard result transformation"
+    (let [result {:model "dashboard"
+                  :id 3
+                  :name "Main Dashboard"
+                  :description "Dashboard desc"
+                  :verified false
+                  :collection {:name "Finance" :authority_level "official"}
+                  :updated_at "2024-01-03"
+                  :created_at "2024-01-03"}
+          expected {:id 3
+                    :type "dashboard"
+                    :name "Main Dashboard"
+                    :description "Dashboard desc"
+                    :verified false
+                    :collection {:name "Finance" :authority_level "official"}
+                    :updated_at "2024-01-03"
+                    :created_at "2024-01-03"}]
+      (is (= expected (#'search/transform-search-result result)))))
+
+  (testing "question (card) result transformation with moderated_status"
+    (let [result {:model "card"
+                  :id 4
+                  :name "Q1"
+                  :description "Question desc"
+                  :moderated_status "verified"
+                  :collection {:name "Analytics" :authority_level nil}
+                  :updated_at "2024-01-04"
+                  :created_at "2024-01-04"}
+          expected {:id 4
+                    :type "question"
+                    :name "Q1"
+                    :description "Question desc"
+                    :database_id nil
+                    :verified true
+                    :collection {:name "Analytics" :authority_level nil}
+                    :updated_at "2024-01-04"
+                    :created_at "2024-01-04"}]
+      (is (= expected (#'search/transform-search-result result)))))
+
+  (testing "metric result transformation"
+    (let [result {:model "metric"
+                  :id 5
+                  :name "Revenue"
+                  :description "Metric desc"
+                  :verified nil
+                  :updated_at "2024-01-05"
+                  :created_at "2024-01-05"}
+          expected {:id 5
+                    :type "metric"
+                    :name "Revenue"
+                    :description "Metric desc"
+                    :database_id nil
+                    :verified false
+                    :collection {}
+                    :updated_at "2024-01-05"
+                    :created_at "2024-01-05"}]
+      (is (= expected (#'search/transform-search-result result)))))
+
+  (testing "database result transformation"
+    (let [result {:model "database"
+                  :id 6
+                  :name "Production DB"
+                  :description "Main database"
+                  :updated_at "2024-01-06"
+                  :created_at "2024-01-06"}
+          expected {:id 6
+                    :type "database"
+                    :name "Production DB"
+                    :description "Main database"
+                    :updated_at "2024-01-06"
+                    :created_at "2024-01-06"}]
+      (is (= expected (#'search/transform-search-result result))))))
+
+(deftest search-test
+  (mt/with-premium-features #{:content-verification}
+    (mt/with-test-user :rasta
+      (let [order-table {:id 1
+                         :model "table"
+                         :table_name "orders"
+                         :name "Orders"
+                         :description "Order table"
+                         :database_id 42
+                         :table_schema "public"}
+            dashboard {:id 2
+                       :model "dashboard"
+                       :name "Sales Dashboard"
+                       :description "Dashboard for sales"
+                       :verified true}]
+
+        (with-redefs [perms/impersonated-user? (fn [] false)
+                      perms/sandboxed-user? (fn [] false)
+                      api/*current-user-id* 1]
+
+          (testing "search returns transformed results for term queries"
+            (with-redefs [search-core/search (fn [_] {:data [order-table]})]
+              (let [args {:term-queries ["orders"]
+                          :entity-types ["table"]}
+                    results (search/search args)
+                    expected [(#'search/transform-search-result order-table)]]
+                (is (= expected results)))))
+
+          (testing "search returns transformed results for semantic queries"
+            (with-redefs [search-core/search (fn [_] {:data [dashboard]})]
+              (let [args {:semantic-queries ["sales metrics"]
+                          :entity-types ["dashboard"]}
+                    results (search/search args)
+                    expected [(#'search/transform-search-result dashboard)]]
+                (is (= expected results)))))
+
+          (testing "search combines term and semantic queries using RRF"
+            (with-redefs [search-core/search (fn [context]
+                                               (if (= (:search-string context) "orders")
+                                                 {:data [order-table]}
+                                                 {:data [dashboard]}))]
+              (let [args {:term-queries ["orders"]
+                          :semantic-queries ["sales"]
+                          :entity-types ["table" "dashboard"]}
+                    results (search/search args)]
+                ;; Should return both results combined via RRF
+                (is (= 2 (count results)))
+                (is (some #(= (:id %) 1) results))
+                (is (some #(= (:id %) 2) results)))))
+
+          (testing "search applies RRF to overlapping results"
+            (with-redefs [search-core/search (fn [_]
+                                               {:data [order-table dashboard]})]
+              (let [args {:term-queries ["orders" "sales"]
+                          :entity-types ["table" "dashboard"]}
+                    results (search/search args)]
+                ;; Both queries return same results, RRF should boost them
+                (is (= 2 (count results)))
+                (is (some #(= (:id %) 1) results))
+                (is (some #(= (:id %) 2) results)))))
+
+          (testing "search handles empty results"
+            (with-redefs [search-core/search (fn [_] {:data []})]
+              (let [args {:term-queries ["nonexistent"]
+                          :entity-types ["table"]}
+                    results (search/search args)]
+                (is (empty? results)))))
+
+          (testing "search with metabot verified content flag"
+            (let [metabot {:entity_id "test-bot"
+                           :use_verified_content true}]
+              (with-redefs [t2/select-one (fn [model & _]
+                                            (is (= :model/Metabot model) "Should query for Metabot model")
+                                            metabot)
+                            search-core/search (fn [context]
+                                                ;; Verify that verified flag is set when metabot has use_verified_content
+                                                 (is (true? (:verified context)))
+                                                 {:data [dashboard]})]
+                (let [results (search/search {:term-queries ["test"]
+                                              :metabot-id "test-bot"
+                                              :entity-types ["dashboard"]})]
+                  (is (= 1 (count results)))
+                  (is (= 2 (:id (first results)))))))))))))

enterprise/backend/test/metabase_enterprise/metabot_v3/tools/util_test.clj

@@ -3,64 +3,165 @@
    [clojure.test :refer :all]
    [metabase-enterprise.metabot-v3.tools.util :as metabot-v3.tools.util]
    [metabase.collections.models.collection :as collection]
+   [metabase.content-verification.core :as moderation]
+   [metabase.lib-be.metadata.jvm :as lib.metadata.jvm]
+   [metabase.lib.core :as lib]
+   [metabase.lib.metadata :as lib.metadata]
    [metabase.permissions.models.permissions :as perms]
    [metabase.test :as mt]))
 
-(deftest metabot-scope-test
-  (mt/dataset test-data
-    (mt/with-non-admin-groups-no-root-collection-perms
-      (mt/with-temp [:model/PermissionsGroup {group-id :id} {}
-                     :model/PermissionsGroupMembership _ {:user_id (mt/user->id :rasta)
-                                                          :group_id group-id}
-                     :model/Collection container-coll {:name "container coll"}
-                     :model/Collection metabot-coll   {:name "mb coll"
-                                                       :location (collection/location-path container-coll)}
-                     :model/Collection mb-child-coll1 {:name "mbc1"
-                                                       :location (collection/location-path container-coll metabot-coll)}
-                     :model/Collection mb-child-coll2 {:name "mbc2"
-                                                       :location (collection/location-path container-coll metabot-coll)}
-                     :model/Collection mb-child-coll3 {:name "mbc3"
-                                                       :location (collection/location-path
-                                                                  container-coll metabot-coll mb-child-coll2)}
-                     :model/Collection non-mb-coll    {:name "non-mbc"
-                                                       :location (collection/location-path container-coll)}
-                     :model/Card mb-model1  {:type :model,  :collection_id (:id metabot-coll)}
-                     :model/Card mb-model2  {:type :model,  :collection_id (:id mb-child-coll1)}
-                     :model/Card mb-model3  {:type :model,  :collection_id (:id mb-child-coll1)}
-                     :model/Card mb-model4  {:type :model,  :collection_id (:id non-mb-coll)}
-                     :model/Card _          {:type :model,  :collection_id (:id non-mb-coll)}
-                     :model/Card mb-metric1 {:type :metric, :collection_id (:id metabot-coll)}
-                     :model/Card mb-metric2 {:type :metric, :collection_id (:id mb-child-coll2)}
-                     :model/Card mb-metric3 {:type :metric, :collection_id (:id mb-child-coll3)}
-                     :model/Card mb-metric4 {:type :metric, :collection_id (:id non-mb-coll)}
-                     :model/Card _          {:type :metric, :collection_id (:id non-mb-coll)}
-                     :model/Metabot metabot {:name "metabot"}
-                     :model/MetabotEntity {mb-coll-entity-id :id}   {:metabot_id (:id metabot)
-                                                                     :model :collection
-                                                                     :model_id (:id metabot-coll)}
-                     :model/MetabotEntity {mb-model-entity-id :id}  {:metabot_id (:id metabot)
-                                                                     :model :dataset
-                                                                     :model_id (:id mb-model4)}
-                     :model/MetabotEntity {mb-metric-entity-id :id} {:metabot_id (:id metabot)
-                                                                     :model :metric
-                                                                     :model_id (:id mb-metric4)}]
-        (perms/grant-collection-read-permissions! group-id mb-child-coll2)
-        (testing "admins can see all cards"
-          (let [admin-result (mt/with-test-user :crowberto
-                               (metabot-v3.tools.util/metabot-scope
-                                [mb-coll-entity-id mb-model-entity-id mb-metric-entity-id]))]
-            (is (= {mb-model1 mb-coll-entity-id
-                    mb-model2 mb-coll-entity-id
-                    mb-model3 mb-coll-entity-id
-                    mb-model4 mb-model-entity-id
-                    mb-metric1 mb-coll-entity-id
-                    mb-metric2 mb-coll-entity-id
-                    mb-metric3 mb-coll-entity-id
-                    mb-metric4 mb-metric-entity-id}
-                   admin-result))))
-        (testing "normal users can see permitted cards"
-          (let [user-result (mt/with-test-user :rasta
-                              (metabot-v3.tools.util/metabot-scope
-                               [mb-coll-entity-id mb-model-entity-id mb-metric-entity-id]))]
-            (is (= {mb-metric2 mb-coll-entity-id}
-                   user-result))))))))
+(deftest metabot-scope-query-test
+  (testing "metabot-scope-query with collection hierarchy"
+    (mt/dataset test-data
+      (mt/with-non-admin-groups-no-root-collection-perms
+        (mt/with-temp [:model/PermissionsGroup {group-id :id} {}
+                       :model/PermissionsGroupMembership _ {:user_id (mt/user->id :rasta)
+                                                            :group_id group-id}
+                       :model/Collection container-coll {:name "container coll"}
+                       :model/Collection metabot-coll   {:name "mb coll"
+                                                         :location (collection/location-path container-coll)}
+                       :model/Collection mb-child-coll1 {:name "mbc1"
+                                                         :location (collection/location-path container-coll metabot-coll)}
+                       :model/Collection mb-child-coll2 {:name "mbc2"
+                                                         :location (collection/location-path container-coll metabot-coll)}
+                       :model/Collection non-mb-coll    {:name "non-mbc"
+                                                         :location (collection/location-path container-coll)}
+                       :model/Card mb-model1  {:type :model,  :collection_id (:id metabot-coll)}
+                       :model/Card mb-model2  {:type :model,  :collection_id (:id mb-child-coll1)}
+                       :model/Card mb-metric1 {:type :metric, :collection_id (:id metabot-coll)}
+                       :model/Card mb-metric2 {:type :metric, :collection_id (:id mb-child-coll2)}
+                       :model/Card outside-model {:type :model, :collection_id (:id non-mb-coll)}
+                       :model/Card outside-metric {:type :metric, :collection_id (:id non-mb-coll)}
+                       :model/Metabot metabot {:name "metabot"
+                                               :collection_id (:id metabot-coll)
+                                               :use_verified_content false}]
+          (perms/grant-collection-read-permissions! group-id metabot-coll)
+          (perms/grant-collection-read-permissions! group-id mb-child-coll1)
+          (perms/grant-collection-read-permissions! group-id mb-child-coll2)
+
+          (testing "admin can see all cards in metabot collection and subcollections"
+            (let [admin-result (mt/with-test-user :crowberto
+                                 (metabot-v3.tools.util/get-metrics-and-models (:id metabot)))
+                  card-ids (set (map :id admin-result))]
+              (is (contains? card-ids (:id mb-model1)))
+              (is (contains? card-ids (:id mb-model2)))
+              (is (contains? card-ids (:id mb-metric1)))
+              (is (contains? card-ids (:id mb-metric2)))
+              (is (not (contains? card-ids (:id outside-model))))
+              (is (not (contains? card-ids (:id outside-metric))))))
+
+          (testing "normal user sees only permitted cards"
+            (let [user-result (mt/with-test-user :rasta
+                                (metabot-v3.tools.util/get-metrics-and-models (:id metabot)))
+                  card-ids (set (map :id user-result))]
+              (is (contains? card-ids (:id mb-model1)))
+              (is (contains? card-ids (:id mb-model2)))
+              (is (contains? card-ids (:id mb-metric1)))
+              (is (contains? card-ids (:id mb-metric2))))))))))
+
+(deftest metabot-scope-query-root-collection-test
+  (testing "metabot-scope-query with root collection (null collection_id)"
+    (mt/dataset test-data
+      (mt/with-temp [:model/Card root-model  {:type :model, :collection_id nil}
+                     :model/Card root-metric {:type :metric, :collection_id nil}
+                     :model/Collection some-coll {:name "some collection"}
+                     :model/Card coll-model {:type :model, :collection_id (:id some-coll)}
+                     :model/Metabot metabot {:name "root metabot"
+                                             :collection_id nil
+                                             :use_verified_content false}]
+        (testing "metabot with root collection sees all content"
+          (let [result (mt/with-test-user :crowberto
+                         (metabot-v3.tools.util/get-metrics-and-models (:id metabot)))
+                card-ids (set (map :id result))]
+            (is (contains? card-ids (:id root-model)))
+            (is (contains? card-ids (:id root-metric)))
+            (is (contains? card-ids (:id coll-model)))))))))
+
+(deftest add-table-reference-test
+  (testing "add-table-reference function adds table-reference for FK fields"
+    (mt/dataset test-data
+      (mt/with-current-user (mt/user->id :crowberto)
+        (let [test-db-id (mt/id)
+              mp (lib.metadata.jvm/application-database-metadata-provider test-db-id)
+              orders-query (lib/query mp (lib.metadata/table mp (mt/id :orders)))
+              columns (lib/visible-columns orders-query)]
+
+          (testing "adds table-reference for implicitly joined columns"
+            (let [processed-columns (map #(metabot-v3.tools.util/add-table-reference orders-query %) columns)
+                  user-name-column (first (filter #(and (= "NAME" (:name %))
+                                                        (:fk-field-id %)) processed-columns))]
+              (is (some? user-name-column) "Expected to find implicitly joined User NAME column")
+              (is (contains? user-name-column :table-reference))
+              (is (string? (:table-reference user-name-column)))
+              (is (seq (:table-reference user-name-column)))
+              (is (= "User" (:table-reference user-name-column)))))
+
+          (testing "does not add table-reference for direct table columns"
+            (let [processed-columns (map #(metabot-v3.tools.util/add-table-reference orders-query %) columns)
+                  id-column (first (filter #(and (= "ID" (:name %))
+                                                 (not (:fk-field-id %))) processed-columns))]
+              (is (some? id-column) "Expected to find direct ORDERS ID column")
+              (is (not (contains? id-column :table-reference)))))
+
+          (testing "handles columns without fk-field-id or table-id gracefully"
+            (let [mock-column {:name "test-column" :type :string}
+                  result (metabot-v3.tools.util/add-table-reference orders-query mock-column)]
+              (is (= mock-column result))
+              (is (not (contains? result :table-reference)))))
+
+          (testing "handles columns with fk-field-id but no table-id"
+            (let [mock-column {:name "test-fk" :fk-field-id 123}
+                  result (metabot-v3.tools.util/add-table-reference orders-query mock-column)]
+              (is (= mock-column result))
+              (is (not (contains? result :table-reference)))))
+
+          (testing "handles columns with table-id but no fk-field-id"
+            (let [mock-column {:name "test-field" :table-id (mt/id :orders)}
+                  result (metabot-v3.tools.util/add-table-reference orders-query mock-column)]
+              (is (= mock-column result))
+              (is (not (contains? result :table-reference))))))))))
+
+(deftest metabot-verified-content-test
+  (testing "metabot-scope-query with verified content filtering"
+    (mt/dataset test-data
+      (mt/with-premium-features #{:content-verification}
+        (mt/with-temp [:model/Collection metabot-coll {:name "mb coll"}
+                       :model/Card verified-model {:type :model, :collection_id (:id metabot-coll)}
+                       :model/Card unverified-model {:type :model, :collection_id (:id metabot-coll)}
+                       :model/Card verified-metric {:type :metric, :collection_id (:id metabot-coll)}
+                       :model/Card unverified-metric {:type :metric, :collection_id (:id metabot-coll)}
+                       :model/Metabot verified-metabot {:name "verified metabot"
+                                                        :collection_id (:id metabot-coll)
+                                                        :use_verified_content true}
+                       :model/Metabot unverified-metabot {:name "unverified metabot"
+                                                          :collection_id (:id metabot-coll)
+                                                          :use_verified_content false}]
+        ;; Mark some content as verified
+          (moderation/create-review! {:moderated_item_id (:id verified-model)
+                                      :moderated_item_type "card"
+                                      :moderator_id (mt/user->id :crowberto)
+                                      :status "verified"
+                                      :text "This is verified"})
+          (moderation/create-review! {:moderated_item_id (:id verified-metric)
+                                      :moderated_item_type "card"
+                                      :moderator_id (mt/user->id :crowberto)
+                                      :status "verified"
+                                      :text "This is verified"})
+
+          (testing "metabot with use_verified_content=true sees only verified content"
+            (let [result (mt/with-test-user :crowberto
+                           (metabot-v3.tools.util/get-metrics-and-models (:id verified-metabot)))
+                  card-ids (set (map :id result))]
+              (is (contains? card-ids (:id verified-model)))
+              (is (contains? card-ids (:id verified-metric)))
+              (is (not (contains? card-ids (:id unverified-model))))
+              (is (not (contains? card-ids (:id unverified-metric))))))
+
+          (testing "metabot with use_verified_content=false sees all content"
+            (let [result (mt/with-test-user :crowberto
+                           (metabot-v3.tools.util/get-metrics-and-models (:id unverified-metabot)))
+                  card-ids (set (map :id result))]
+              (is (contains? card-ids (:id verified-model)))
+              (is (contains? card-ids (:id verified-metric)))
+              (is (contains? card-ids (:id unverified-model)))
+              (is (contains? card-ids (:id unverified-metric))))))))))

enterprise/backend/test/metabase_enterprise/models/entity_id_test.clj

@@ -110,7 +110,7 @@
     (testing "All exported models should get entity id except those with other unique property (like name)"
       (is (= (set (concat serdes.models/exported-models
                           ;; those are inline models which still have entity_id
-                          ["DashboardCard" "DashboardTab" "Dimension" "MetabotEntity" "MetabotPrompt"]))
+                          ["DashboardCard" "DashboardTab" "Dimension" "MetabotPrompt"]))
              (set (->> (concat entity-id-models
                                entities-external-name)
                        (map name))))))

enterprise/backend/test/metabase_enterprise/serialization/v2/extract_test.clj

@@ -1737,16 +1737,16 @@
                     model-eid :entity_id} {:name "AI Model"
                                            :type :model}
 
+       :model/Collection {coll-id :id
+                          coll-eid :entity_id} {:name "Metabot Collection"}
+
        :model/Metabot {metabot-id :id
                        metabot-eid :entity_id} {:name "Test Metabot"
-                                                :description "A test metabot"}
+                                                :description "A test metabot"
+                                                :use_verified_content false
+                                                :collection_id coll-id}
 
-       :model/MetabotEntity {metabot-entity-id :id
-                             metabot-entity-eid :entity_id} {:metabot_id metabot-id
-                                                             :model :dataset
-                                                             :model_id model-id}
-
-       :model/MetabotPrompt {metabot-prompt-eid :entity_id} {:metabot_entity_id metabot-entity-id
+       :model/MetabotPrompt {metabot-prompt-eid :entity_id} {:metabot_id metabot-id
                                                              :prompt "A sample prompt"
                                                              :model :model
                                                              :card_id model-id}]
@@ -1757,20 +1757,15 @@
                    :name "Test Metabot"
                    :description "A test metabot"
                    :entity_id metabot-eid
-                   :entities [{:model "dataset"
-                               :model_id model-eid
-                               :entity_id metabot-entity-eid
-                               :serdes/meta [{:model "Metabot" :id metabot-eid}
-                                             {:model "MetabotEntity" :id metabot-entity-eid}]
-                               :prompts [{:prompt "A sample prompt"
-                                          :model "model"
-                                          :entity_id metabot-prompt-eid
-                                          :card_id model-eid
-                                          :serdes/meta [{:model "Metabot" :id metabot-eid}
-                                                        {:model "MetabotEntity" :id metabot-entity-eid}
-                                                        {:model "MetabotPrompt" :id metabot-prompt-eid}]
-                                          :created_at string?}]
-                               :created_at string?}]
+                   :collection_id coll-eid
+                   :use_verified_content false
+                   :prompts [{:prompt "A sample prompt"
+                              :model "model"
+                              :entity_id metabot-prompt-eid
+                              :card_id model-eid
+                              :serdes/meta [{:model "Metabot" :id metabot-eid}
+                                            {:model "MetabotPrompt" :id metabot-prompt-eid}]
+                              :created_at string?}]
                    :created_at string?}
                   ser))
           (is (not (contains? ser :id)))
@@ -1782,24 +1777,23 @@
 (deftest metabot-collection-test
   (mt/with-empty-h2-app-db!
     (ts/with-temp-dpc
-      [:model/Collection {model-id :id
-                          model-eid :entity_id} {:name "AI Model"}
+      [:model/Collection {model-id :id} {:name "AI Model"}
 
        :model/Card {card-id :id
                     card-eid :entity_id} {:name "AI Model"
                                           :type :model
                                           :collection_id model-id}
 
+       :model/Collection {coll-id :id
+                          coll-eid :entity_id} {:name "Metabot Collection"}
+
        :model/Metabot {metabot-id :id
                        metabot-eid :entity_id} {:name "Test Metabot"
-                                                :description "A test metabot"}
+                                                :description "A test metabot"
+                                                :use_verified_content false
+                                                :collection_id coll-id}
 
-       :model/MetabotEntity {metabot-entity-id :id
-                             metabot-entity-eid :entity_id} {:metabot_id metabot-id
-                                                             :model :collection
-                                                             :model_id model-id}
-
-       :model/MetabotPrompt {metabot-prompt-eid :entity_id} {:metabot_entity_id metabot-entity-id
+       :model/MetabotPrompt {metabot-prompt-eid :entity_id} {:metabot_id metabot-id
                                                              :prompt "A sample prompt"
                                                              :model :model
                                                              :card_id card-id}]
@@ -1810,26 +1804,21 @@
                    :name "Test Metabot"
                    :description "A test metabot"
                    :entity_id metabot-eid
-                   :entities [{:model "collection"
-                               :model_id model-eid
-                               :entity_id metabot-entity-eid
-                               :serdes/meta [{:model "Metabot" :id metabot-eid}
-                                             {:model "MetabotEntity" :id metabot-entity-eid}]
-                               :prompts [{:prompt "A sample prompt"
-                                          :model "model"
-                                          :entity_id metabot-prompt-eid
-                                          :card_id card-eid
-                                          :serdes/meta [{:model "Metabot" :id metabot-eid}
-                                                        {:model "MetabotEntity" :id metabot-entity-eid}
-                                                        {:model "MetabotPrompt" :id metabot-prompt-eid}]
-                                          :created_at string?}]
-                               :created_at string?}]
+                   :collection_id coll-eid
+                   :use_verified_content false
+                   :prompts [{:prompt "A sample prompt"
+                              :model "model"
+                              :entity_id metabot-prompt-eid
+                              :card_id card-eid
+                              :serdes/meta [{:model "Metabot" :id metabot-eid}
+                                            {:model "MetabotPrompt" :id metabot-prompt-eid}]
+                              :created_at string?}]
                    :created_at string?}
                   ser))
           (is (not (contains? ser :id)))
 
-          (testing "metabot depends on its model entities and their prompts"
-            (is (= #{[{:model "Collection" :id model-eid}] [{:model "Card" :id card-eid}]}
+          (testing "metabot depends on its prompts' cards"
+            (is (= #{[{:model "Card" :id card-eid}]}
                    (set (serdes/dependencies ser))))))))))
 
 (deftest document-test

enterprise/backend/test/metabase_enterprise/serialization/v2/round_trip_test.clj

@@ -171,7 +171,7 @@
                 (is (nil? delta)
                     (str "Content mismatch for file: " (strip-base-path source-dir file)))
 
-              ;; Leave behind files for developers to inspect
+                ;; Leave behind files for developers to inspect
                 (when (and (.exists dev-inspect-dir) delta)
                   (vreset! wrote-files? true)
                   (create-files-to-diff! ref-file out-file))))

enterprise/frontend/src/metabase-enterprise/api/ai-streaming/schemas.ts

@@ -1,16 +1,25 @@
 import * as Yup from "yup";
 
+import type { MetabotTodoItem, SuggestedTransform } from "metabase-types/api";
+
 export const dataPartSchema = Yup.object({
   type: Yup.string().required(),
   version: Yup.number().required(),
   value: Yup.mixed(),
 });
 
-export const knownDataPartTypes = ["state", "navigate_to"];
+export const knownDataPartTypes = [
+  "navigate_to",
+  "state",
+  "todo_list",
+  "transform_suggestion",
+];
 
 export type KnownDataPart =
+  | { type: "navigate_to"; version: 1; value: string }
   | { type: "state"; version: 1; value: Record<string, any> }
-  | { type: "navigate_to"; version: 1; value: string };
+  | { type: "todo_list"; version: 1; value: MetabotTodoItem[] }
+  | { type: "transform_suggestion"; version: 1; value: SuggestedTransform };
 
 export const toolCallPartSchema = Yup.object({
   toolCallId: Yup.string().required(),

enterprise/frontend/src/metabase-enterprise/api/metabot.ts

@@ -1,19 +1,15 @@
 import type {
   DeleteSuggestedMetabotPromptRequest,
-  MetabotApiEntity,
-  MetabotEntity,
   MetabotGenerateContentRequest,
   MetabotGenerateContentResponse,
   MetabotId,
   MetabotInfo,
-  PaginationRequest,
-  PaginationResponse,
   SuggestedMetabotPromptsRequest,
   SuggestedMetabotPromptsResponse,
 } from "metabase-types/api";
 
 import { EnterpriseApi } from "./api";
-import { idTag } from "./tags";
+import { idTag, invalidateTags, listTag } from "./tags";
 
 export const metabotApi = EnterpriseApi.injectEndpoints({
   endpoints: (builder) => ({
@@ -22,64 +18,28 @@ export const metabotApi = EnterpriseApi.injectEndpoints({
         method: "GET",
         url: "/api/ee/metabot-v3/metabot",
       }),
+      providesTags: (result) => [
+        listTag("metabot"),
+        ...(result?.items || []).map((metabot) => idTag("metabot", metabot.id)),
+      ],
     }),
-    listMetabotsEntities: builder.query<
-      { items: MetabotEntity[] } & PaginationResponse,
-      { id: MetabotId } & PaginationRequest
+    updateMetabot: builder.mutation<
+      MetabotInfo,
+      { id: MetabotId } & Partial<
+        Pick<MetabotInfo, "use_verified_content" | "collection_id">
+      >
     >({
-      query: ({ id, ...paginationProps }) => ({
-        method: "GET",
-        url: `/api/ee/metabot-v3/metabot/${id}/entities`,
-        body: paginationProps,
-      }),
-      providesTags: ["metabot-entities-list"],
-      transformResponse: (
-        response: { items: MetabotApiEntity[] } & PaginationResponse,
-      ) => ({
-        ...response,
-        // transform model_id to id in items
-        items: response.items.map(({ model_id: id, ...item }) => ({
-          ...item,
-          id,
-        })),
-      }),
-    }),
-    updateMetabotEntities: builder.mutation<
-      void,
-      {
-        id: MetabotId;
-        entities: Pick<MetabotEntity, "model" | "id">[];
-      }
-    >({
-      query: ({ id, entities }) => ({
+      query: ({ id, ...updates }) => ({
         method: "PUT",
-        url: `/api/ee/metabot-v3/metabot/${id}/entities`,
-        body: { items: entities },
+        url: `/api/ee/metabot-v3/metabot/${id}`,
+        body: updates,
       }),
       invalidatesTags: (_, error, { id }) =>
-        !error
-          ? ["metabot-entities-list", idTag("metabot-prompt-suggestions", id)]
-          : [],
-    }),
-    deleteMetabotEntities: builder.mutation<
-      void,
-      {
-        metabotId: MetabotId;
-        entityModel: MetabotEntity["model"];
-        entityId: MetabotEntity["id"];
-      }
-    >({
-      query: ({ metabotId, entityModel, entityId }) => ({
-        method: "DELETE",
-        url: `/api/ee/metabot-v3/metabot/${metabotId}/entities/${entityModel}/${entityId}`,
-      }),
-      invalidatesTags: (_, error, { metabotId }) =>
-        !error
-          ? [
-              "metabot-entities-list",
-              idTag("metabot-prompt-suggestions", metabotId),
-            ]
-          : [],
+        invalidateTags(error, [
+          listTag("metabot"),
+          idTag("metabot", id),
+          idTag("metabot-prompt-suggestions", id),
+        ]),
     }),
     getSuggestedMetabotPrompts: builder.query<
       SuggestedMetabotPromptsResponse,
@@ -103,7 +63,9 @@ export const metabotApi = EnterpriseApi.injectEndpoints({
         url: `/api/ee/metabot-v3/metabot/${metabot_id}/prompt-suggestions/${prompt_id}`,
       }),
       invalidatesTags: (_, error, { metabot_id }) =>
-        !error ? [idTag("metabot-prompt-suggestions", metabot_id)] : [],
+        invalidateTags(error, [
+          idTag("metabot-prompt-suggestions", metabot_id),
+        ]),
     }),
     regenerateSuggestedMetabotPrompts: builder.mutation<void, MetabotId>({
       query: (metabot_id) => ({
@@ -111,7 +73,9 @@ export const metabotApi = EnterpriseApi.injectEndpoints({
         url: `/api/ee/metabot-v3/metabot/${metabot_id}/prompt-suggestions/regenerate`,
       }),
       invalidatesTags: (_, error, metabot_id) =>
-        !error ? [idTag("metabot-prompt-suggestions", metabot_id)] : [],
+        invalidateTags(error, [
+          idTag("metabot-prompt-suggestions", metabot_id),
+        ]),
     }),
     metabotGenerateContent: builder.query<
       MetabotGenerateContentResponse,
@@ -128,9 +92,7 @@ export const metabotApi = EnterpriseApi.injectEndpoints({
 
 export const {
   useListMetabotsQuery,
-  useListMetabotsEntitiesQuery,
-  useUpdateMetabotEntitiesMutation,
-  useDeleteMetabotEntitiesMutation,
+  useUpdateMetabotMutation,
   useGetSuggestedMetabotPromptsQuery,
   useDeleteSuggestedMetabotPromptMutation,
   useRegenerateSuggestedMetabotPromptsMutation,

enterprise/frontend/src/metabase-enterprise/metabot/components/Metabot.tsx

@@ -12,9 +12,13 @@ import { MetabotChat } from "./MetabotChat";
 
 export interface MetabotProps {
   hide?: boolean;
+  emptyConvoText?: string;
 }
 
-export const MetabotAuthenticated = ({ hide }: MetabotProps) => {
+export const MetabotAuthenticated = ({
+  hide,
+  emptyConvoText,
+}: MetabotProps) => {
   const { visible, setVisible } = useMetabotAgent();
 
   useEffect(() => {
@@ -44,7 +48,7 @@ export const MetabotAuthenticated = ({ hide }: MetabotProps) => {
 
   return (
     <ErrorBoundary errorComponent={() => null}>
-      <MetabotChat />
+      <MetabotChat emptyConvoText={emptyConvoText} />
     </ErrorBoundary>
   );
 };

enterprise/frontend/src/metabase-enterprise/metabot/components/MetabotAdmin/MetabotAdminPage.tsx

@@ -1,8 +1,8 @@
 import { useDisclosure } from "@mantine/hooks";
 import { useEffect, useMemo } from "react";
 import { push } from "react-router-redux";
-import { match } from "ts-pattern";
-import { c, t } from "ttag";
+import { P, match } from "ts-pattern";
+import { c, jt, t } from "ttag";
 import _ from "underscore";
 
 import ErrorBoundary from "metabase/ErrorBoundary";
@@ -13,6 +13,7 @@ import {
 import { SettingsSection } from "metabase/admin/components/SettingsSection";
 import { SettingHeader } from "metabase/admin/settings/components/SettingHeader";
 import { skipToken, useGetCollectionQuery } from "metabase/api";
+import { canonicalCollectionId } from "metabase/collections/utils";
 import { AdminSettingsLayout } from "metabase/common/components/AdminLayout/AdminSettingsLayout";
 import { LoadingAndErrorWrapper } from "metabase/common/components/LoadingAndErrorWrapper";
 import { CollectionPickerModal } from "metabase/common/components/Pickers/CollectionPicker";
@@ -20,75 +21,84 @@ import { useToast } from "metabase/common/hooks";
 import { color } from "metabase/lib/colors";
 import { getIcon } from "metabase/lib/icon";
 import { useDispatch } from "metabase/lib/redux";
-import { Box, Button, Flex, Icon, Loader, Text } from "metabase/ui";
 import {
-  useDeleteMetabotEntitiesMutation,
-  useListMetabotsEntitiesQuery,
+  Box,
+  Button,
+  Flex,
+  Icon,
+  Loader,
+  Stack,
+  Switch,
+  Text,
+} from "metabase/ui";
+import {
   useListMetabotsQuery,
-  useUpdateMetabotEntitiesMutation,
+  useUpdateMetabotMutation,
 } from "metabase-enterprise/api";
+import {
+  FIXED_METABOT_ENTITY_IDS,
+  FIXED_METABOT_IDS,
+} from "metabase-enterprise/metabot/constants";
+import { hasPremiumFeature } from "metabase-enterprise/settings";
 import type {
+  Collection,
   CollectionEssentials,
-  MetabotEntity,
-  MetabotId,
+  MetabotInfo,
 } from "metabase-types/api";
 
 import { MetabotPromptSuggestionPane } from "./MetabotAdminSuggestedPrompts";
 import { useMetabotIdPath } from "./utils";
 
 export function MetabotAdminPage() {
-  const metabotId = useMetabotIdPath();
+  const metabotId = useMetabotIdPath() ?? FIXED_METABOT_IDS.DEFAULT;
   const { data, isLoading, error } = useListMetabotsQuery();
-  const metabotName =
-    data?.items?.find((bot) => bot.id === metabotId)?.name ?? t`Metabot`;
-  const isEmbeddedMetabot = metabotName.toLowerCase().includes("embed");
 
-  const { data: entityList } = useListMetabotsEntitiesQuery(
-    metabotId ? { id: metabotId } : skipToken,
-  );
-  const hasEntities = (entityList?.items?.length ?? 0) > 0;
+  const metabot = data?.items?.find((bot) => bot.id === metabotId);
 
-  if (isLoading || !data) {
+  if (isLoading || !metabot) {
     return (
       <LoadingAndErrorWrapper
-        loading={isLoading}
-        error={error ? t`Error fetching Metabots` : null}
+        loading={isLoading || !data}
+        error={match({ isLoading, error, metabot })
+          .with(
+            { isLoading: false, error: P.not(null) },
+            () => t`Error fetching Metabots`,
+          )
+          .with({ isLoading: false, metabot: undefined }, () => t`Not found.`)
+          .otherwise(() => null)}
       />
     );
   }
 
+  const isEmbedMetabot =
+    metabot.entity_id === FIXED_METABOT_ENTITY_IDS.EMBEDDED;
+
   return (
     <AdminSettingsLayout sidebar={<MetabotNavPane />}>
       <ErrorBoundary>
-        <SettingsSection>
+        <SettingsSection key={metabotId}>
           <Box>
             <SettingHeader
               id="configure-metabot"
               title={c("{0} is the name of an AI assistant")
-                .t`Configure ${metabotName}`}
+                .t`Configure ${metabot.name}`}
               description={c("{0} is the name of an AI assistant") // eslint-disable-next-line no-literal-metabase-strings -- admin ui
-                .t`${metabotName} is Metabase's AI agent. To help ${metabotName} more easily find and focus on the data you care about most, select the collection containing the models and metrics it should be able to use to create queries.`}
+                .t`${metabot.name} is Metabase's AI agent. To help ${metabot.name} more easily find and focus on the data you care about most, configure what content it should be able to access or use to create queries.`}
             />
-            {isEmbeddedMetabot && (
+            {isEmbedMetabot && (
               <Text c="text-medium" maw="40rem">
                 {t`If you're embedding the Metabot component in an app, you can specify a different collection that embedded Metabot is allowed to use for creating queries.`}
               </Text>
             )}
           </Box>
-          {metabotId && (
-            <>
-              <MetabotConfigurationPane
-                metabotId={metabotId}
-                metabotName={metabotName}
-              />
-              {hasEntities && (
-                <MetabotPromptSuggestionPane
-                  key={metabotId}
-                  metabotId={metabotId}
-                />
-              )}
-            </>
+
+          <MetabotVerifiedContentConfigurationPane metabot={metabot} />
+
+          {isEmbedMetabot && (
+            <MetabotCollectionConfigurationPane metabot={metabot} />
           )}
+
+          <MetabotPromptSuggestionPane metabotId={metabotId} />
         </SettingsSection>
       </ErrorBoundary>
     </AdminSettingsLayout>
@@ -130,27 +140,72 @@ function MetabotNavPane() {
   );
 }
 
-function MetabotConfigurationPane({
-  metabotId,
-  metabotName,
+function MetabotVerifiedContentConfigurationPane({
+  metabot,
 }: {
-  metabotId: MetabotId;
-  metabotName: string;
+  metabot: MetabotInfo;
 }) {
+  const [updateMetabot, { isLoading: isUpdating }] = useUpdateMetabotMutation();
+  const [sendToast] = useToast();
+
+  const handleVerifiedContentToggle = async (checked: boolean) => {
+    const result = await updateMetabot({
+      id: metabot.id,
+      use_verified_content: checked,
+    });
+
+    if (result.error) {
+      sendToast({
+        message: t`Error updating Metabot`,
+        icon: "warning",
+      });
+    }
+  };
+
+  if (!hasPremiumFeature("content_verification")) {
+    return null;
+  }
+
+  return (
+    <Stack gap="sm">
+      <SettingHeader
+        id="verified-content"
+        title={t`Verified content`}
+        description={jt`When enabled, Metabot will only use models and metrics marked as Verified.`}
+      />
+      <Switch
+        label={t`Only use Verified content`}
+        checked={!!metabot.use_verified_content}
+        onChange={(e) => handleVerifiedContentToggle(e.target.checked)}
+        disabled={isUpdating}
+        w="auto"
+        size="sm"
+      />
+    </Stack>
+  );
+}
+
+function MetabotCollectionConfigurationPane({
+  metabot,
+}: {
+  metabot: MetabotInfo;
+}) {
+  const metabotId = metabot.id;
+  const metabotName = metabot.name;
+
   const {
-    data: entityList,
+    data: collection,
     isLoading,
     error,
-  } = useListMetabotsEntitiesQuery({ id: metabotId });
-  const [updateEntities, { isLoading: isUpdating }] =
-    useUpdateMetabotEntitiesMutation();
-  const [deleteEntity, { isLoading: isDeleting }] =
-    useDeleteMetabotEntitiesMutation();
-  const isMutating = isUpdating || isDeleting;
+  } = useGetCollectionQuery({
+    id: metabot.collection_id ?? "root",
+  });
+
+  const [updateMetabot, { isLoading: isUpdating }] = useUpdateMetabotMutation();
   const [isOpen, { open, close }] = useDisclosure(false);
   const [sendToast] = useToast();
 
-  if (isLoading || !entityList || error) {
+  if (isLoading || !collection || error) {
     return (
       <LoadingAndErrorWrapper
         loading={isLoading}
@@ -159,37 +214,18 @@ function MetabotConfigurationPane({
     );
   }
 
-  const collection: MetabotEntity | undefined = entityList?.items?.[0];
-  const handleDelete = async () => {
-    if (collection) {
-      const result = await deleteEntity({
-        metabotId,
-        entityModel: "collection",
-        entityId: collection.id,
-      });
-
-      if (result.error) {
-        sendToast({
-          message: t`Error removing folder`,
-          icon: "warning",
-        });
-      }
-    }
-  };
-
-  const handleAddEntity = async (
-    newEntity: Pick<MetabotEntity, "model" | "id" | "name">,
+  const handleUpdateCollectionId = async (
+    newEntity: Pick<MetabotInfo, "id" | "name">,
   ) => {
     close();
-    await handleDelete();
-    const result = await updateEntities({
+    const result = await updateMetabot({
       id: metabotId,
-      entities: [_.pick(newEntity, "model", "id")],
+      collection_id: canonicalCollectionId(newEntity.id),
     });
 
     if (result.error) {
       sendToast({
-        message: t`Error adding ${newEntity.name}`,
+        message: t`Error setting ${newEntity.name}`,
         icon: "warning",
       });
     }
@@ -204,30 +240,23 @@ function MetabotConfigurationPane({
       />
       <CollectionInfo collection={collection} />
       <Flex gap="md" mt="md">
-        <Button onClick={open} leftSection={isMutating && <Loader size="xs" />}>
-          {match({ isMutating, collection })
+        <Button onClick={open} leftSection={isUpdating && <Loader size="xs" />}>
+          {match({ isMutating: isUpdating, collection })
             .with({ isMutating: true }, () => t`Updating collection...`)
             .with({ collection: undefined }, () => t`Pick a collection`)
             .otherwise(() => t`Pick a different collection`)}
         </Button>
-        {collection && (
-          <Button onClick={handleDelete}>
-            <Icon name="trash" />
-          </Button>
-        )}
       </Flex>
       {isOpen && (
         <CollectionPickerModal
           title={t`Select a collection`}
-          shouldDisableItem={(item) => item.id === "root"}
-          canSelectItem={(item) => item && item.id !== "root"}
           value={{
             id: collection?.id ?? null,
             model: "collection",
           }}
           onChange={(item) =>
-            handleAddEntity(
-              item as unknown as Pick<MetabotEntity, "model" | "id" | "name">,
+            handleUpdateCollectionId(
+              item as unknown as Pick<MetabotInfo, "id" | "name">,
             )
           }
           onClose={close}
@@ -241,7 +270,7 @@ function MetabotConfigurationPane({
   );
 }
 
-function CollectionInfo({ collection }: { collection: MetabotEntity | null }) {
+function CollectionInfo({ collection }: { collection: Collection }) {
   const { data: collectionInfo } = useGetCollectionQuery(
     collection?.id ? { id: collection.id } : skipToken,
   );

enterprise/frontend/src/metabase-enterprise/metabot/components/MetabotAdmin/MetabotAdminPage.unit.spec.tsx

@@ -5,112 +5,93 @@ import { Route } from "react-router";
 import {
   findRequests,
   setupCollectionByIdEndpoint,
-  setupCollectionItemsEndpoint,
-  setupCollectionsEndpoints,
   setupRecentViewsAndSelectionsEndpoints,
 } from "__support__/server-mocks";
 import {
-  setupMetabotAddEntitiesEndpoint,
-  setupMetabotDeleteEntitiesEndpoint,
-  setupMetabotEntitiesEndpoint,
   setupMetabotPromptSuggestionsEndpoint,
-  setupMetabotsEndpoint,
+  setupMetabotsEndpoints,
 } from "__support__/server-mocks/metabot";
-import { renderWithProviders, screen, waitFor, within } from "__support__/ui";
-import type {
-  MetabotApiEntity,
-  MetabotEntity,
-  MetabotId,
-  RecentItem,
-} from "metabase-types/api";
-import { createMockCollection } from "metabase-types/api/mocks";
+import { renderWithProviders, screen, waitFor } from "__support__/ui";
+import {
+  FIXED_METABOT_ENTITY_IDS,
+  FIXED_METABOT_IDS,
+} from "metabase-enterprise/metabot/constants";
+import { hasPremiumFeature } from "metabase-enterprise/settings";
+import type { MetabotId, RecentItem } from "metabase-types/api";
+import {
+  createMockCollection,
+  createMockMetabotInfo,
+} from "metabase-types/api/mocks";
 
 import { MetabotAdminPage } from "./MetabotAdminPage";
 import * as hooks from "./utils";
 
+jest.mock("metabase-enterprise/settings");
+const mockHasPremiumFeature = hasPremiumFeature as jest.MockedFunction<
+  typeof hasPremiumFeature
+>;
+
 const mockPathParam = (id: MetabotId) => {
   jest.spyOn(hooks, "useMetabotIdPath").mockReturnValue(id);
 };
 
-const metabots = [
-  {
-    id: 1,
-    name: "Metabot One",
-  },
-  {
-    id: 2,
-    name: "Metabot Two",
-  },
-  {
-    id: 3,
-    name: "Embedded Metabot Three",
-  },
+const defaultMetabots = [
+  createMockMetabotInfo({
+    id: FIXED_METABOT_IDS.DEFAULT,
+    entity_id: FIXED_METABOT_ENTITY_IDS.DEFAULT,
+  }),
+  createMockMetabotInfo({
+    id: FIXED_METABOT_IDS.EMBEDDED,
+    name: "Embedded Metabot",
+    entity_id: FIXED_METABOT_ENTITY_IDS.EMBEDDED,
+    collection_id: 21,
+  }),
 ];
 
-const entities = {
-  1: [
-    {
-      model_id: 11,
-      name: "Collection One",
-      model: "collection",
-      collection_id: 12,
-      collection_name: "Collection One Prime",
+const defaultSeedCollections = [
+  createMockCollection({ id: "root", name: "Our Analytics" }),
+  {
+    id: 21,
+    name: "Collection Two",
+    model: "collection",
+    collection_name: "Collection Two Prime",
+    parent_collection: {
+      id: 3,
+      name: "Collection Two Prime",
     },
-  ],
-  2: [
-    {
-      model_id: 21,
-      name: "Collection Two",
-      model: "collection",
-      collection_id: 22,
-      collection_name: "Collection Two Prime",
+  },
+  {
+    id: 31,
+    name: "Collection Three",
+    model: "collection",
+    parent_collection: {
+      id: 3,
+      name: "Collection Three Prime",
     },
-  ] as MetabotApiEntity[],
-  3: [],
-  recents: [
-    {
-      id: 31,
-      name: "Collection Three",
-      model: "collection",
-      parent_collection: {
-        id: 3,
-        name: "Collection Three Prime",
-      },
+  },
+  {
+    id: 32,
+    name: "Collection Four",
+    model: "collection",
+    parent_collection: {
+      id: 3,
+      name: "Collection Four Prime",
     },
-    {
-      id: 32,
-      name: "Collection Four",
-      model: "collection",
-      parent_collection: {
-        id: 3,
-        name: "Collection Four Prime",
-      },
-    },
-  ],
-};
-
+  },
+];
 const setup = async (
   initialPathParam: MetabotId = 1,
-  seedData = entities,
+  metabots = defaultMetabots,
+  seedCollections = defaultSeedCollections,
   error = false,
 ) => {
   mockPathParam(initialPathParam);
-  setupMetabotsEndpoint(metabots, error ? 500 : undefined);
-  const collections = [...seedData[1], ...seedData[2], ...seedData.recents];
+  setupMetabotsEndpoints(metabots, error ? 500 : undefined);
   setupCollectionByIdEndpoint({
-    collections: collections.map((c: any) => ({ id: c.model_id, ...c })),
+    collections: seedCollections.map((c: any) => ({ id: c.model_id, ...c })),
   });
 
-  metabots.forEach((metabot) => {
-    setupMetabotEntitiesEndpoint(
-      metabot.id,
-      seedData[metabot.id as 1 | 2] as MetabotApiEntity[],
-    );
-    setupMetabotAddEntitiesEndpoint(metabot.id);
-  });
-  setupMetabotDeleteEntitiesEndpoint();
-
-  setupRecentViewsAndSelectionsEndpoints(seedData.recents as RecentItem[]);
+  setupRecentViewsAndSelectionsEndpoints(seedCollections as RecentItem[]);
 
   metabots.forEach((mb) =>
     setupMetabotPromptSuggestionsEndpoint({
@@ -145,135 +126,80 @@ describe("MetabotAdminPage", () => {
 
   it("should render the metabots list", async () => {
     await setup();
-    expect(await screen.findByText("Metabot One")).toBeInTheDocument();
-    expect(screen.getByText("Metabot Two")).toBeInTheDocument();
+    expect(await screen.findByText("Metabot")).toBeInTheDocument();
+    expect(screen.getByText("Embedded Metabot")).toBeInTheDocument();
   });
 
-  it("should render a selected collection", async () => {
+  it("should not be able to select a collection for default metabot", async () => {
     await setup();
-    expect(await screen.findByText("Collection One")).toBeInTheDocument();
+    expect(await screen.findByText("Configure Metabot")).toBeInTheDocument();
+    expect(
+      screen.queryByText("Collection Metabot can use"),
+    ).not.toBeInTheDocument();
   });
 
-  it("should be able to switch between metabots", async () => {
-    await setup(1);
-    await screen.findByText("Collection One");
-
-    mockPathParam(2);
-    await userEvent.click(await screen.findByText("Metabot Two"));
+  it("should render a selected collection for embedded metabot", async () => {
+    await setup(FIXED_METABOT_IDS.EMBEDDED);
     expect(await screen.findByText("Collection Two")).toBeInTheDocument();
   });
 
-  it("should change selected collection", async () => {
-    await setup(1);
+  it("should render a root collection if collection_id is null for metabot", async () => {
+    await setup(FIXED_METABOT_IDS.EMBEDDED, [
+      createMockMetabotInfo({
+        id: FIXED_METABOT_IDS.EMBEDDED,
+        name: "Embedded Metabot",
+        entity_id: FIXED_METABOT_ENTITY_IDS.EMBEDDED,
+        collection_id: null,
+      }),
+    ]);
+    expect(await screen.findByText("Our Analytics")).toBeInTheDocument();
+  });
+
+  it("should be able to switch between metabots", async () => {
+    await setup(FIXED_METABOT_IDS.DEFAULT);
+    expect(await screen.findByText("Configure Metabot")).toBeInTheDocument();
+
+    mockPathParam(FIXED_METABOT_IDS.EMBEDDED);
+    await userEvent.click(await screen.findByText("Embedded Metabot"));
+    expect(await screen.findByText("Collection Two")).toBeInTheDocument();
+  });
+
+  it("should change selected collection for embedded metabot", async () => {
+    await setup(FIXED_METABOT_IDS.EMBEDDED);
 
     expect(
       fetchMock.callHistory.calls(
-        `path:/api/ee/metabot-v3/metabot/1/prompt-suggestions`,
+        `path:/api/ee/metabot-v3/metabot/${FIXED_METABOT_IDS.EMBEDDED}/prompt-suggestions?limit=10&offset=0`,
       ).length,
     ).toEqual(1); // should have loaded prompt suggestions
 
-    expect(await screen.findByText("Collection One")).toBeInTheDocument();
+    expect(await screen.findByText("Collection Two")).toBeInTheDocument();
     await userEvent.click(screen.getByText("Pick a different collection"));
 
     await screen.findByText("Select a collection");
     await userEvent.click(screen.getByText("Collection Three"));
     await userEvent.click(screen.getByText("Select"));
 
-    await waitFor(async () => {
-      const deletes = await findRequests("DELETE");
-      expect(deletes.length).toBe(1);
-    });
-
     await waitFor(async () => {
       const puts = await findRequests("PUT");
       expect(puts.length).toBe(1);
     });
 
-    const [{ url, body }] = await findRequests("PUT");
-    expect(url).toMatch(/\/api\/ee\/metabot-v3\/metabot\/1\/entities/);
-    expect(body.items).toHaveLength(1); // 1 new
-
-    expect(
-      body.items.find((item: MetabotEntity) => item.id === 31),
-    ).toBeTruthy();
+    const puts = await findRequests("PUT");
+    expect(puts[0].url).toMatch(
+      new RegExp(`/api/ee/metabot-v3/metabot/${FIXED_METABOT_IDS.EMBEDDED}`),
+    );
+    expect(puts[0].body).toEqual({ collection_id: 31 });
 
     expect(
       fetchMock.callHistory.calls(
-        `path:/api/ee/metabot-v3/metabot/1/prompt-suggestions`,
+        `path:/api/ee/metabot-v3/metabot/${FIXED_METABOT_IDS.EMBEDDED}/prompt-suggestions?limit=10&offset=0`,
       ).length,
-    ).toEqual(3); // +1 refetch for DELETE, +1 for PUT
-  });
-
-  it("should not allow selecting the root collection", async () => {
-    // setup entity picker endpoints
-    const rootCollection = createMockCollection({ id: "root" });
-    setupCollectionsEndpoints({ collections: [rootCollection] });
-    setupCollectionItemsEndpoint({
-      collection: rootCollection,
-      collectionItems: [],
-      models: [],
-    });
-    setupCollectionItemsEndpoint({
-      collection: createMockCollection({ id: 1 }),
-      collectionItems: [],
-      models: [],
-    });
-
-    // default to no entities for default metabot - this will default the
-    // entity picker's initial value to be the root collection
-    await setup(1, { ...entities, 1: [] });
-
-    await userEvent.click(screen.getByText("Pick a collection"));
-
-    const entityPicker = await screen.findByTestId("entity-picker-modal");
-    const entityPickerTabs =
-      await within(entityPicker).findByTestId("tabs-view");
-    await userEvent.click(
-      await within(entityPickerTabs).findByText(/Collections/),
-    );
-
-    // should not be able to select the default Our analytics option
-    expect(
-      await screen.findByRole("button", { name: /Select/ }),
-    ).toBeDisabled();
-  });
-
-  it("should delete the selected collection", async () => {
-    const metabotId = 1;
-    await setup(metabotId);
-
-    expect(await screen.findByText("Collection One")).toBeInTheDocument();
-    expect(await screen.findByText("Prompt suggestions")).toBeInTheDocument();
-    const [deleteButton] = await screen.findAllByLabelText("trash icon");
-
-    // simulte the response after a delete
-    fetchMock.modifyRoute(`metabot-${metabotId}-entities-get`, {
-      response: { items: [] },
-    });
-    await userEvent.click(deleteButton);
-
-    const [{ url: deleteUrl }, ...rest] = await findRequests("DELETE");
-    expect(deleteUrl).toContain("metabot/1/entities/collection/11");
-    expect(rest).toHaveLength(0); // only 1 delete
-    await waitFor(() => {
-      expect(screen.queryByText("Prompt suggestions")).not.toBeInTheDocument();
-    });
-  });
-
-  it("should show an empty state when no entities", async () => {
-    await setup(1, {
-      1: [],
-      2: [],
-      3: [],
-      recents: [],
-    });
-    expect(await screen.findByText("Pick a collection")).toBeInTheDocument();
-
-    expect(screen.queryByLabelText("trash icon")).not.toBeInTheDocument();
+    ).toEqual(2); // +1 refetch for DELETE, +1 for PUT
   });
 
   it("should show special copy for embedded metabot", async () => {
-    await setup(3);
+    await setup(FIXED_METABOT_IDS.EMBEDDED);
 
     expect(
       await screen.findByText(/embedding the metabot component/i),
@@ -281,10 +207,76 @@ describe("MetabotAdminPage", () => {
   });
 
   it("should show an error message when a request fails", async () => {
-    await setup(3, entities, true);
+    await setup(404, defaultMetabots, defaultSeedCollections, true);
 
     expect(
       await screen.findByText("Error fetching Metabots"),
     ).toBeInTheDocument();
   });
+
+  describe("MetabotVerifiedContentConfigurationPane", () => {
+    const mockContentVerificationEnabled = (enabled: boolean) => {
+      mockHasPremiumFeature.mockImplementation((feature) => {
+        if (feature === "content_verification") {
+          return enabled;
+        }
+        return true; // Mock other features as enabled by default
+      });
+    };
+
+    it("should not show verification switch without content_verification feature", async () => {
+      mockContentVerificationEnabled(false);
+
+      await setup();
+
+      // First ensure the page has loaded
+      await screen.findByText(/Configure Metabot/);
+
+      expect(screen.queryByText("Verified content")).not.toBeInTheDocument();
+      expect(
+        screen.queryByText("Only use Verified content"),
+      ).not.toBeInTheDocument();
+    });
+
+    it("should show verification switch with content_verification feature", async () => {
+      mockContentVerificationEnabled(true);
+
+      await setup();
+
+      expect(await screen.findByText("Verified content")).toBeInTheDocument();
+      expect(
+        await screen.findByText("Only use Verified content"),
+      ).toBeInTheDocument();
+      expect(
+        await screen.findByRole("switch", {
+          name: "Only use Verified content",
+        }),
+      ).toBeInTheDocument();
+    });
+
+    it("should allow enabling/disabling verified switch affecting use_verified_content", async () => {
+      mockContentVerificationEnabled(true);
+
+      await setup();
+
+      const verifiedSwitch = await screen.findByRole("switch", {
+        name: "Only use Verified content",
+      });
+
+      // Verify switch is initially unchecked (default metabot has use_verified_content: false)
+      expect(verifiedSwitch).not.toBeChecked();
+
+      // Click to enable
+      await userEvent.click(verifiedSwitch);
+
+      // Verify API call was made with correct payload
+      await waitFor(async () => {
+        const putRequests = await findRequests("PUT");
+        expect(putRequests.length).toBe(1);
+      });
+
+      const putRequests = await findRequests("PUT");
+      expect(putRequests[0].body).toEqual({ use_verified_content: true });
+    });
+  });
 });

enterprise/frontend/src/metabase-enterprise/metabot/components/MetabotChat/MetabotChat.module.css

@@ -104,6 +104,14 @@
   background: var(--mb-color-bg-medium);
 }
 
+.messageUserAction {
+  padding: 0.5rem 1rem;
+  border-radius: 0.5rem;
+  background: var(--mb-color-bg-light);
+  border: 1px solid var(--mb-color-bg-medium);
+  font-weight: bold;
+}
+
 .messageActions {
   opacity: 0;
   transition: all 150ms linear;

enterprise/frontend/src/metabase-enterprise/metabot/components/MetabotChat/MetabotChat.tsx

@@ -27,7 +27,11 @@ import { Messages } from "./MetabotChatMessage";
 import { MetabotThinking } from "./MetabotThinking";
 import { useScrollManager } from "./hooks";
 
-export const MetabotChat = () => {
+export const MetabotChat = ({
+  emptyConvoText,
+}: {
+  emptyConvoText?: string;
+}) => {
   const metabot = useMetabotAgent();
   const { handleSubmitInput, handleRetryMessage, handleResetInput } =
     useMetabotChatHandlers();
@@ -64,7 +68,9 @@ export const MetabotChat = () => {
         <Box ref={headerRef} className={Styles.header}>
           <Flex align-items="center">
             <Text lh={1} fz="sm" c="text-secondary">
-              {t`Metabot isn't perfect. Double-check results.`}
+              {metabot.profileOverride
+                ? t`Using profile: ${metabot.profileOverride}`
+                : t`Metabot isn't perfect. Double-check results.`}
             </Text>
           </Flex>
 
@@ -101,33 +107,34 @@ export const MetabotChat = () => {
                 data-testid="metabot-empty-chat-info"
               >
                 <Box component={EmptyDashboardBot} w="6rem" />
-                <Text
-                  c="text-light"
-                  maw="12rem"
-                  ta="center"
-                >{t`I can help you explore your metrics and models.`}</Text>
+                <Text c="text-light" maw="12rem" ta="center">
+                  {emptyConvoText ??
+                    t`I can help you explore your metrics and models.`}
+                </Text>
               </Flex>
-              {/* empty state with suggested prompts */}
-              <Stack
-                gap="sm"
-                className={Styles.promptSuggestionsContainer}
-                data-testid="metabot-prompt-suggestions"
-              >
-                <>
-                  {suggestedPrompts.map(({ prompt }, index) => (
-                    <Box key={index}>
-                      <Button
-                        fz="sm"
-                        size="xs"
-                        onClick={() => handleSubmitInput(prompt)}
-                        className={Styles.promptSuggestionButton}
-                      >
-                        {prompt}
-                      </Button>
-                    </Box>
-                  ))}
-                </>
-              </Stack>
+              {/* empty state with suggested prompts - TODO: temporarily disabled */}
+              {!!false && (
+                <Stack
+                  gap="sm"
+                  className={Styles.promptSuggestionsContainer}
+                  data-testid="metabot-prompt-suggestions"
+                >
+                  <>
+                    {suggestedPrompts.map(({ prompt }, index) => (
+                      <Box key={index}>
+                        <Button
+                          fz="sm"
+                          size="xs"
+                          onClick={() => handleSubmitInput(prompt)}
+                          className={Styles.promptSuggestionButton}
+                        >
+                          {prompt}
+                        </Button>
+                      </Box>
+                    ))}
+                  </>
+                </Stack>
+              )}
             </>
           )}
 
@@ -144,7 +151,6 @@ export const MetabotChat = () => {
                 isDoingScience={metabot.isDoingScience}
                 showFeedbackButtons
               />
-
               {/* loading */}
               {metabot.isDoingScience && (
                 <MetabotThinking
@@ -154,10 +160,8 @@ export const MetabotChat = () => {
                   }
                 />
               )}
-
               {/* filler - height gets set via ref mutation */}
               <div ref={fillerRef} data-testid="metabot-message-filler" />
-
               {/* long convo warning */}
               {metabot.isLongConversation && <MetabotResetLongChatButton />}
             </Box>

enterprise/frontend/src/metabase-enterprise/metabot/components/MetabotChat/MetabotChatMessage.tsx

@@ -1,20 +1,33 @@
 import { useClipboard } from "@mantine/hooks";
+import type { UnknownAction } from "@reduxjs/toolkit";
 import cx from "classnames";
 import { useCallback, useState } from "react";
+import { push } from "react-router-redux";
+import { useLocation } from "react-use";
+import { t } from "ttag";
 
 import { useToast } from "metabase/common/hooks";
 import { downloadObjectAsJson } from "metabase/lib/download";
+import { useDispatch } from "metabase/lib/redux";
 import {
   ActionIcon,
+  Button,
   Flex,
   type FlexProps,
   Icon,
   type IconName,
+  Stack,
   Text,
 } from "metabase/ui";
-import type {
-  MetabotChatMessage,
-  MetabotErrorMessage,
+import {
+  type MetabotAgentChatMessage,
+  type MetabotAgentEditSuggestionChatMessage,
+  type MetabotAgentTextChatMessage,
+  type MetabotAgentTodoListChatMessage,
+  type MetabotChatMessage,
+  type MetabotErrorMessage,
+  type MetabotUserChatMessage,
+  setSuggestedTransform,
 } from "metabase-enterprise/metabot/state";
 import type { MetabotFeedback } from "metabase-types/api";
 
@@ -49,7 +62,8 @@ export const MessageContainer = ({
   />
 );
 
-interface UserMessageProps extends BaseMessageProps {
+interface UserMessageProps extends Omit<BaseMessageProps, "message"> {
+  message: MetabotUserChatMessage;
   onCopy: () => void;
 }
 
@@ -61,14 +75,22 @@ export const UserMessage = ({
   ...props
 }: UserMessageProps) => (
   <MessageContainer chatRole={message.role} {...props}>
-    <Text
-      className={cx(
-        Styles.message,
-        message.role === "user" && Styles.messageUser,
-      )}
-    >
-      {message.message}
-    </Text>
+    {message.type === "text" && (
+      <Text className={cx(Styles.message, Styles.messageUser)}>
+        {message.message}
+      </Text>
+    )}
+
+    {message.type === "action" && (
+      <Flex direction="column" gap="xs">
+        <Flex align="center" gap="xs">
+          <Text className={cx(Styles.message, Styles.messageUserAction)}>
+            {message.userMessage}
+          </Text>
+        </Flex>
+      </Flex>
+    )}
+
     <Flex className={Styles.messageActions}>
       {!hideActions && (
         <ActionIcon
@@ -104,7 +126,8 @@ const FeedbackButton = ({
   </ActionIcon>
 );
 
-interface AgentMessageProps extends BaseMessageProps {
+interface AgentMessageProps extends Omit<BaseMessageProps, "message"> {
+  message: MetabotAgentChatMessage;
   onRetry: (messageId: string) => void;
   onCopy: (messageId: string) => void;
   showFeedbackButtons: boolean;
@@ -113,6 +136,81 @@ interface AgentMessageProps extends BaseMessageProps {
   onInternalLinkClick?: (link: string) => void;
 }
 
+const AgentSuggestionMessage = ({
+  message,
+}: {
+  message: MetabotAgentEditSuggestionChatMessage;
+}) => {
+  const dispatch = useDispatch();
+
+  const transform = message.payload;
+
+  const url = useLocation();
+  const disabled = url.pathname?.startsWith(
+    `/admin/transforms/${transform.id}/query`,
+  );
+
+  return (
+    <Flex direction="column" gap="sm">
+      <Button
+        disabled={disabled}
+        onClick={() => {
+          const url = transform.id
+            ? `/admin/transforms/${transform.id}/query`
+            : "/admin/transforms/new/native";
+          dispatch(push(url) as UnknownAction);
+          dispatch(setSuggestedTransform(transform));
+        }}
+      >
+        {message.payload.id
+          ? t`Proposed changes to: ${transform.name}`
+          : t`Create Transform: ${transform.name}`}
+      </Button>
+    </Flex>
+  );
+};
+
+const AgentTodoListMessage = ({
+  list,
+}: {
+  list: MetabotAgentTodoListChatMessage["payload"];
+}) => {
+  const statusConfig = {
+    completed: { icon: "check", color: "success" },
+    in_progress: { icon: "play", color: "brand" },
+    cancelled: { icon: "close", color: "text-light" },
+    pending: { icon: "circle" as IconName, color: "text-medium" },
+  } as const;
+
+  return (
+    <Stack gap="sm" w="100%">
+      <Text fw="bold">{t`Todo List`}</Text>
+      {list.map((item) => (
+        <Flex
+          key={item.id}
+          style={{ borderRadius: "2px" }}
+          gap="sm"
+          align="center"
+        >
+          <Icon
+            name={statusConfig[item.status].icon}
+            size="1rem"
+            c={statusConfig[item.status].color}
+            mt="2px"
+          />
+          <Text
+            size="md"
+            td={item.status === "completed" ? "line-through" : undefined}
+            c={item.status === "cancelled" ? "text-light" : undefined}
+          >
+            {item.content}
+          </Text>
+        </Flex>
+      ))}
+    </Stack>
+  );
+};
+
 export const AgentMessage = ({
   message,
   className,
@@ -124,66 +222,77 @@ export const AgentMessage = ({
   onInternalLinkClick,
   hideActions,
   ...props
-}: AgentMessageProps) => (
-  <MessageContainer chatRole={message.role} {...props}>
-    <AIMarkdown
-      className={Styles.message}
-      onInternalLinkClick={onInternalLinkClick}
-    >
-      {message.message}
-    </AIMarkdown>
-
-    <Flex className={Styles.messageActions}>
-      {!hideActions && (
-        <>
-          <ActionIcon
-            h="sm"
-            data-testid="metabot-chat-message-copy"
-            onClick={() => onCopy(message.id)}
-          >
-            <Icon name="copy" size="1rem" />
-          </ActionIcon>
-          {showFeedbackButtons && setFeedbackMessage && (
-            <>
-              <FeedbackButton
-                data-testid="metabot-chat-message-thumbs-up"
-                icon="thumbs_up"
-                hasBeenClicked={submittedFeedback === "positive"}
-                disabled={!!submittedFeedback}
-                onClick={() =>
-                  setFeedbackMessage({
-                    messageId: message.id,
-                    positive: true,
-                  })
-                }
-              />
-              <FeedbackButton
-                data-testid="metabot-chat-message-thumbs-down"
-                icon="thumbs_down"
-                hasBeenClicked={submittedFeedback === "negative"}
-                disabled={!!submittedFeedback}
-                onClick={() =>
-                  setFeedbackMessage({
-                    messageId: message.id,
-                    positive: false,
-                  })
-                }
-              />
-            </>
-          )}
-
-          <ActionIcon
-            onClick={() => onRetry(message.id)}
-            h="sm"
-            data-testid="metabot-chat-message-retry"
-          >
-            <Icon name="revert" size="1rem" />
-          </ActionIcon>
-        </>
+}: AgentMessageProps) => {
+  return (
+    <MessageContainer chatRole={message.role} {...props}>
+      {message.type === "text" && (
+        <AIMarkdown
+          className={Styles.message}
+          onInternalLinkClick={onInternalLinkClick}
+        >
+          {message.message}
+        </AIMarkdown>
       )}
-    </Flex>
-  </MessageContainer>
-);
+
+      {message.type === "edit_suggestion" && (
+        <AgentSuggestionMessage message={message} />
+      )}
+
+      {message.type === "todo_list" && (
+        <AgentTodoListMessage list={message.payload} />
+      )}
+      <Flex className={Styles.messageActions}>
+        {!hideActions && (
+          <>
+            <ActionIcon
+              h="sm"
+              data-testid="metabot-chat-message-copy"
+              onClick={() => onCopy(message.id)}
+            >
+              <Icon name="copy" size="1rem" />
+            </ActionIcon>
+            {showFeedbackButtons && setFeedbackMessage && (
+              <>
+                <FeedbackButton
+                  data-testid="metabot-chat-message-thumbs-up"
+                  icon="thumbs_up"
+                  hasBeenClicked={submittedFeedback === "positive"}
+                  disabled={!!submittedFeedback}
+                  onClick={() =>
+                    setFeedbackMessage({
+                      messageId: message.id,
+                      positive: true,
+                    })
+                  }
+                />
+                <FeedbackButton
+                  data-testid="metabot-chat-message-thumbs-down"
+                  icon="thumbs_down"
+                  hasBeenClicked={submittedFeedback === "negative"}
+                  disabled={!!submittedFeedback}
+                  onClick={() =>
+                    setFeedbackMessage({
+                      messageId: message.id,
+                      positive: false,
+                    })
+                  }
+                />
+              </>
+            )}
+
+            <ActionIcon
+              onClick={() => onRetry(message.id)}
+              h="sm"
+              data-testid="metabot-chat-message-retry"
+            >
+              <Icon name="revert" size="1rem" />
+            </ActionIcon>
+          </>
+        )}
+      </Flex>
+    </MessageContainer>
+  );
+};
 
 export const AgentErrorMessage = ({
   message,
@@ -278,7 +387,11 @@ export const Messages = ({
   const onAgentMessageCopy = useCallback(
     (messageId: string) => {
       const allMessages = getFullAgentReply(messages, messageId);
-      clipboard.copy(allMessages.map((msg) => msg.message).join("\n\n"));
+      const textMessages = allMessages.filter(
+        (msg): msg is MetabotAgentTextChatMessage =>
+          msg.role === "agent" && msg.type === "text",
+      );
+      clipboard.copy(textMessages.map((msg) => msg.message).join("\n\n"));
     },
     [messages, clipboard],
   );
@@ -316,7 +429,14 @@ export const Messages = ({
             data-testid="metabot-chat-message"
             message={message}
             hideActions={isDoingScience && messages.length === index + 1}
-            onCopy={() => clipboard.copy(message.message)}
+            onCopy={() => {
+              // TODO: probably want to make this not show up long-term
+              const copyText =
+                message.type === "action"
+                  ? `${message.userMessage}: ${message.message}`
+                  : message.message;
+              clipboard.copy(copyText);
+            }}
           />
         ),
       )}

enterprise/frontend/src/metabase-enterprise/metabot/components/MetabotChat/MetabotThinking.tsx

@@ -24,7 +24,7 @@ const ThoughtProcess = ({ toolCalls }: { toolCalls: MetabotToolCall[] }) => {
   return (
     <Stack gap={0} className={S.toolCalls}>
       {slots.map((_, i) => {
-        const tc = toolCalls[i];
+        const tc: MetabotToolCall | undefined = toolCalls[i];
         return (
           <Transition
             key={i}
@@ -35,9 +35,9 @@ const ThoughtProcess = ({ toolCalls }: { toolCalls: MetabotToolCall[] }) => {
             {(styles) => (
               <Text
                 style={styles}
-                className={cx(tc.status === "started" && S.toolCallStarted)}
+                className={cx(tc?.status === "started" && S.toolCallStarted)}
               >
-                {tc.message}
+                {tc?.message}
               </Text>
             )}
           </Transition>

enterprise/frontend/src/metabase-enterprise/metabot/constants.ts

@@ -12,6 +12,11 @@ export const METABOT_REQUEST_IDS = {
   EMBEDDED: "c61bf5f5-1025-47b6-9298-bf1827105bb6",
 };
 
+export const FIXED_METABOT_ENTITY_IDS = {
+  DEFAULT: "metabotmetabotmetabot" as const,
+  EMBEDDED: "embeddedmetabotmetabo" as const,
+};
+
 export const METABOT_ERR_MSG = {
   get default() {
     return t`Sorry, I ran into an error. Could you please try that again?`;
@@ -37,7 +42,31 @@ export const TOOL_CALL_MESSAGES: Record<string, string | undefined> = {
   get search_data_sources() {
     return t`Checking available data sources`;
   },
+  get search() {
+    return t`Searching`;
+  },
   get search_metabase_documentation() {
     return t`Consulting the docs`;
   },
+  get write_transform_sql() {
+    return t`Writing SQL`;
+  },
+  get todo_write() {
+    return t`Planning`;
+  },
+  get todo_read() {
+    return t`Planning`;
+  },
+  get search_transforms() {
+    return t`Searching transforms`;
+  },
+  get get_transform_details() {
+    return t`Getting transform details`;
+  },
+  get get_field_values() {
+    return t`Retrieving table metadata`;
+  },
+  get search_tables() {
+    return t`Searching database tables`;
+  },
 };

enterprise/frontend/src/metabase-enterprise/metabot/hooks/use-metabot-agent.ts

@@ -7,14 +7,16 @@ import { useMetabotContext } from "metabase/metabot";
 import { trackMetabotRequestSent } from "../analytics";
 import {
   type MetabotPromptSubmissionResult,
+  type MetabotUserChatMessage,
   getAgentErrorMessages,
   getIsLongMetabotConversation,
   getIsProcessing,
-  getLastAgentMessagesByType,
   getMessages,
   getMetabotId,
+  getMetabotReactionsState,
   getMetabotRequestId,
   getMetabotVisible,
+  getProfileOverride,
   getToolCalls,
   resetConversation as resetConversationAction,
   retryPrompt,
@@ -35,9 +37,6 @@ export const useMetabotAgent = () => {
   const messages = useSelector(getMessages as any) as ReturnType<
     typeof getMessages
   >;
-  const lastAgentMessages = useSelector(
-    getLastAgentMessagesByType as any,
-  ) as ReturnType<typeof getLastAgentMessagesByType>;
 
   const errorMessages = useSelector(getAgentErrorMessages as any) as ReturnType<
     typeof getAgentErrorMessages
@@ -67,6 +66,14 @@ export const useMetabotAgent = () => {
     [dispatch],
   );
 
+  const profileOverride = useSelector(getProfileOverride as any) as ReturnType<
+    typeof getProfileOverride
+  >;
+
+  const reactions = useSelector(getMetabotReactionsState as any) as ReturnType<
+    typeof getMetabotReactionsState
+  >;
+
   const resetConversation = useCallback(
     () => dispatch(resetConversationAction()),
     [dispatch],
@@ -83,18 +90,27 @@ export const useMetabotAgent = () => {
   );
 
   const submitInput = useCallback(
-    async (prompt: string) => {
+    async (prompt: string | Omit<MetabotUserChatMessage, "id" | "role">) => {
       if (!visible) {
         setVisible(true);
       }
 
       const context = await getChatContext();
       const action = await dispatch(
-        submitInputAction({
-          message: prompt,
-          context,
-          metabot_id: metabotRequestId,
-        }),
+        submitInputAction(
+          typeof prompt === "string"
+            ? {
+                type: "text",
+                message: prompt,
+                context,
+                metabot_id: metabotRequestId,
+              }
+            : {
+                ...prompt,
+                context,
+                metabot_id: metabotRequestId,
+              },
+        ),
       );
 
       trackMetabotRequestSent();
@@ -152,7 +168,6 @@ export const useMetabotAgent = () => {
     metabotId,
     visible,
     messages,
-    lastAgentMessages,
     errorMessages,
     isLongConversation,
     resetConversation,
@@ -161,5 +176,7 @@ export const useMetabotAgent = () => {
     submitInput,
     retryMessage,
     toolCalls,
+    profileOverride,
+    reactions,
   };
 };

enterprise/frontend/src/metabase-enterprise/metabot/metabot.unit.spec.tsx

@@ -986,7 +986,9 @@ describe("metabot-streaming", () => {
 
       // adding messages this long via the ui's input makes the test hang
       act(() => {
-        store.dispatch(addUserMessage({ id: "1", message: longMsg }));
+        store.dispatch(
+          addUserMessage({ id: "1", type: "text", message: longMsg }),
+        );
       });
       expect(await screen.findByText(/xxxxxxx/)).toBeInTheDocument();
       expect(
@@ -994,7 +996,9 @@ describe("metabot-streaming", () => {
       ).not.toBeInTheDocument();
 
       act(() => {
-        store.dispatch(addUserMessage({ id: "2", message: longMsg }));
+        store.dispatch(
+          addUserMessage({ id: "2", type: "text", message: longMsg }),
+        );
       });
       expect(
         await screen.findByText(/This chat is getting long/),

enterprise/frontend/src/metabase-enterprise/metabot/reactions/errors.ts

@@ -1,22 +0,0 @@
-import { t } from "ttag";
-
-import { addAgentErrorMessage, stopProcessing } from "../state";
-
-import type { ReactionHandler } from "./types";
-
-type UnknownReaction = { type: string } & Record<string, unknown>;
-
-export const notifyUnknownReaction: ReactionHandler<UnknownReaction> = (
-  reaction,
-) => {
-  return ({ dispatch }) => {
-    console.error("Unknown reaction recieved", reaction);
-    dispatch(
-      addAgentErrorMessage({
-        type: "message",
-        message: t`Whoops, I actually can’t do this. Sorry.`,
-      }),
-    );
-    dispatch(stopProcessing());
-  };
-};

enterprise/frontend/src/metabase-enterprise/metabot/reactions/index.ts

@@ -1,18 +0,0 @@
-import type { MetabotReaction } from "metabase-types/api";
-
-import { showMessage } from "./messages";
-import { redirect } from "./metabot";
-import type { ReactionHandler } from "./types";
-
-export * from "./errors";
-
-type ReactionHandlers = {
-  [key in MetabotReaction["type"]]: ReactionHandler<
-    Extract<MetabotReaction, { type: key }>
-  >;
-};
-
-export const reactionHandlers: ReactionHandlers = {
-  "metabot.reaction/message": showMessage,
-  "metabot.reaction/redirect": redirect,
-};

enterprise/frontend/src/metabase-enterprise/metabot/reactions/messages.ts

@@ -1,13 +0,0 @@
-import type { MetabotMessageReaction } from "metabase-types/api";
-
-import { addAgentMessage } from "../state";
-
-import type { ReactionHandler } from "./types";
-
-export const showMessage: ReactionHandler<MetabotMessageReaction> = (
-  reaction,
-) => {
-  return ({ dispatch }) => {
-    dispatch(addAgentMessage({ message: reaction.message }));
-  };
-};

enterprise/frontend/src/metabase-enterprise/metabot/reactions/metabot.ts

@@ -1,21 +0,0 @@
-import { push } from "react-router-redux";
-
-import type { MetabotRedirectReaction } from "metabase-types/api";
-
-import type { ReactionHandler } from "./types";
-
-export const redirect: ReactionHandler<MetabotRedirectReaction> = (
-  reaction,
-) => {
-  const redirectUrl = new URL(`${window.location.origin}${reaction.url}`);
-
-  return async ({ dispatch }) => {
-    await dispatch(
-      push({
-        pathname: redirectUrl.pathname,
-        hash: redirectUrl.hash,
-        search: redirectUrl.search,
-      }),
-    );
-  };
-};

enterprise/frontend/src/metabase-enterprise/metabot/reactions/types.ts

@@ -1,8 +0,0 @@
-import type { Dispatch, GetState } from "metabase-types/store";
-
-export type ReactionHandler<Reaction> = (
-  reaction: Reaction,
-) => (reduxApis: {
-  dispatch: Dispatch;
-  getState: GetState;
-}) => void | Promise<void>;

enterprise/frontend/src/metabase-enterprise/metabot/state/actions.ts

@@ -3,6 +3,7 @@ import { push } from "react-router-redux";
 import { P, match } from "ts-pattern";
 
 import { createAsyncThunk } from "metabase/lib/redux";
+import { addUndo } from "metabase/redux/undo";
 import { getIsEmbedding } from "metabase/selectors/embed";
 import { getUser } from "metabase/selectors/user";
 import { EnterpriseApi } from "metabase-enterprise/api";
@@ -15,14 +16,18 @@ import type {
   MetabotAgentRequest,
   MetabotAgentResponse,
   MetabotChatContext,
-  MetabotReaction,
 } from "metabase-types/api";
 import type { Dispatch } from "metabase-types/store";
 
 import { METABOT_ERR_MSG } from "../constants";
-import { notifyUnknownReaction, reactionHandlers } from "../reactions";
 
-import { type MetabotErrorMessage, metabot } from "./reducer";
+import {
+  type MetabotAgentEditSuggestionChatMessage,
+  type MetabotAgentTodoListChatMessage,
+  type MetabotErrorMessage,
+  type MetabotUserChatMessage,
+  metabot,
+} from "./reducer";
 import {
   getAgentErrorMessages,
   getAgentRequestMetadata,
@@ -32,7 +37,8 @@ import {
   getMetabotConversationId,
   getUserPromptForMessageId,
 } from "./selectors";
-import { createMessageId } from "./utils";
+import type { SlashCommand } from "./types";
+import { createMessageId, parseSlashCommand } from "./utils";
 
 export const {
   addAgentTextDelta,
@@ -44,6 +50,9 @@ export const {
   setNavigateToPath,
   toolCallStart,
   toolCallEnd,
+  setProfileOverride,
+  setMetabotReqIdOverride,
+  setSuggestedTransform,
 } = metabot.actions;
 
 type PromptErrorOutcome = {
@@ -90,6 +99,30 @@ export const setVisible =
     dispatch(metabot.actions.setVisible(isVisible));
   };
 
+export const executeSlashCommand = createAsyncThunk<void, SlashCommand>(
+  "metabase-enterprise/metabot/executeSlashCommand",
+  async (slashCommand, { dispatch }) => {
+    match(slashCommand)
+      .with({ cmd: "profile" }, ({ args }) => {
+        if (args.length <= 1) {
+          dispatch(setProfileOverride(args[0]));
+        } else {
+          dispatch(addUndo({ message: "/profile <name>" }));
+        }
+      })
+      .with({ cmd: "metabot" }, ({ args }) => {
+        if (args.length <= 1) {
+          dispatch(setMetabotReqIdOverride(args[0]));
+        } else {
+          dispatch(addUndo({ message: "/metabot <name>" }));
+        }
+      })
+      .otherwise(() => {
+        dispatch(addUndo({ message: "Unknown command" }));
+      });
+  },
+);
+
 export type MetabotPromptSubmissionResult =
   | { prompt: string; success: true; shouldRetry?: void }
   | { prompt: string; success: false; shouldRetry: false }
@@ -97,8 +130,7 @@ export type MetabotPromptSubmissionResult =
 
 export const submitInput = createAsyncThunk<
   MetabotPromptSubmissionResult,
-  {
-    message: string;
+  Omit<MetabotUserChatMessage, "id" | "role"> & {
     context: MetabotChatContext;
     metabot_id?: string;
   }
@@ -120,11 +152,18 @@ export const submitInput = createAsyncThunk<
         dispatch(rewindConversation(lastMessageId));
       }
 
+      const slashCommand = parseSlashCommand(data.message);
+      if (slashCommand) {
+        await dispatch(executeSlashCommand(slashCommand));
+        return { prompt: data.message, success: true };
+      }
+
       // it's important that we get the current metadata containing the history before
       // altering it by adding the current message the user is wanting to send
       const agentMetadata = getAgentRequestMetadata(getState() as any);
       const messageId = createMessageId();
-      dispatch(addUserMessage({ id: messageId, message: data.message }));
+      dispatch(addUserMessage({ id: messageId, ...data }));
+
       const sendMessageRequestPromise = dispatch(
         sendAgentRequest({
           ...data,
@@ -198,6 +237,17 @@ export const sendAgentRequest = createAsyncThunk<
             match(part)
               // only update the convo state if the request is successful
               .with({ type: "state" }, (part) => (state = part.value))
+              .with({ type: "todo_list" }, (part) => {
+                const message: Omit<
+                  MetabotAgentTodoListChatMessage,
+                  "id" | "role"
+                > = {
+                  type: "todo_list",
+                  payload: part.value,
+                };
+
+                dispatch(addAgentMessage(message));
+              })
               .with({ type: "navigate_to" }, (part) => {
                 dispatch(setNavigateToPath(part.value));
 
@@ -205,11 +255,25 @@ export const sendAgentRequest = createAsyncThunk<
                   dispatch(push(part.value) as UnknownAction);
                 }
               })
+              .with(
+                { type: "transform_suggestion" },
+                ({ value: transform }) => {
+                  dispatch(setSuggestedTransform(transform));
+                  const message: Omit<
+                    MetabotAgentEditSuggestionChatMessage,
+                    "id" | "role"
+                  > = {
+                    type: "edit_suggestion",
+                    model: "transform",
+                    payload: transform,
+                  };
+
+                  dispatch(addAgentMessage(message));
+                },
+              )
               .exhaustive();
           },
-          onTextPart: (part) => {
-            dispatch(addAgentTextDelta(String(part)));
-          },
+          onTextPart: (part) => dispatch(addAgentTextDelta(String(part))),
           onToolCallPart: (part) => dispatch(toolCallStart(part)),
           onToolResultPart: (part) => dispatch(toolCallEnd(part)),
           onError: (part) => (error = part),
@@ -274,7 +338,13 @@ export const retryPrompt = createAsyncThunk<
     dispatch(metabot.actions.rewindStateToMessageId(messageId));
 
     return await dispatch(
-      submitInput({ message: prompt.message, context, metabot_id }),
+      // TODO: fix with new action type
+      submitInput({
+        type: "text",
+        message: prompt.message,
+        context,
+        metabot_id,
+      }),
     ).unwrap();
   },
 );
@@ -291,37 +361,6 @@ export const resetConversation = createAsyncThunk(
   },
 );
 
-export const processMetabotReactions = createAsyncThunk(
-  "metabase-enterprise/metabot/processMetabotReactions",
-  async (reactions: MetabotReaction[], { dispatch, getState }) => {
-    dispatch(setIsProcessing(true));
-
-    for (const reaction of reactions) {
-      try {
-        const reactionHandler =
-          reactionHandlers[reaction.type] ?? notifyUnknownReaction;
-        // TS isn't smart enough to know the reaction matches the handler
-        await reactionHandler(reaction as any)({ dispatch, getState });
-      } catch (error: any) {
-        console.error("Halting processing of reactions.", error);
-        dispatch(stopProcessingAndNotify());
-        break;
-      }
-
-      // TODO: make an EnterpriseStore
-      const isProcessing = getIsProcessing(getState() as any);
-      if (!isProcessing) {
-        console.warn(
-          "A handler has stopped further procesing of metabot reactions",
-        );
-        break;
-      }
-    }
-
-    dispatch(setIsProcessing(false));
-  },
-);
-
 export const stopProcessing = () => (dispatch: Dispatch) => {
   dispatch(setIsProcessing(false));
 };

enterprise/frontend/src/metabase-enterprise/metabot/state/reducer.ts

@@ -3,19 +3,67 @@ import _ from "underscore";
 
 import { logout } from "metabase/auth/actions";
 import { uuid } from "metabase/lib/uuid";
-import type { MetabotHistory } from "metabase-types/api";
+import type {
+  MetabotHistory,
+  MetabotTodoItem,
+  SuggestedTransform,
+} from "metabase-types/api";
 
 import { TOOL_CALL_MESSAGES } from "../constants";
 
 import { sendAgentRequest } from "./actions";
 import { createMessageId } from "./utils";
 
-export type MetabotChatMessage = {
+export type MetabotUserTextChatMessage = {
   id: string;
-  role: "user" | "agent";
+  role: "user";
+  type: "text";
   message: string;
 };
 
+export type MetabotUserActionChatMessage = {
+  id: string;
+  role: "user";
+  type: "action";
+  message: string;
+  userMessage: string;
+};
+
+export type MetabotAgentTextChatMessage = {
+  id: string;
+  role: "agent";
+  type: "text";
+  message: string;
+};
+
+export type MetabotAgentTodoListChatMessage = {
+  id: string;
+  role: "agent";
+  type: "todo_list";
+  payload: MetabotTodoItem[];
+};
+
+export type MetabotAgentEditSuggestionChatMessage = {
+  id: string;
+  role: "agent";
+  type: "edit_suggestion";
+  model: "transform";
+  payload: SuggestedTransform;
+};
+
+export type MetabotAgentChatMessage =
+  | MetabotAgentTextChatMessage
+  | MetabotAgentTodoListChatMessage
+  | MetabotAgentEditSuggestionChatMessage;
+
+export type MetabotUserChatMessage =
+  | MetabotUserTextChatMessage
+  | MetabotUserActionChatMessage;
+
+export type MetabotChatMessage =
+  | MetabotUserChatMessage
+  | MetabotAgentChatMessage;
+
 export type MetabotErrorMessage = {
   type: "message" | "alert";
   message: string;
@@ -30,6 +78,7 @@ export type MetabotToolCall = {
 
 export type MetabotReactionsState = {
   navigateToPath: string | null;
+  suggestedTransform: SuggestedTransform | undefined;
 };
 
 export interface MetabotState {
@@ -42,6 +91,10 @@ export interface MetabotState {
   state: any;
   reactions: MetabotReactionsState;
   toolCalls: MetabotToolCall[];
+  experimental: {
+    metabotReqIdOverride: string | undefined;
+    profileOverride: string | undefined;
+  };
 }
 
 export const getMetabotInitialState = (): MetabotState => ({
@@ -54,8 +107,13 @@ export const getMetabotInitialState = (): MetabotState => ({
   state: {},
   reactions: {
     navigateToPath: null,
+    suggestedTransform: undefined,
   },
   toolCalls: [],
+  experimental: {
+    metabotReqIdOverride: undefined,
+    profileOverride: undefined,
+  },
 });
 
 export const metabot = createSlice({
@@ -64,24 +122,25 @@ export const metabot = createSlice({
   reducers: {
     addUserMessage: (
       state,
-      action: PayloadAction<Omit<MetabotChatMessage, "role">>,
+      action: PayloadAction<Omit<MetabotUserChatMessage, "role">>,
     ) => {
-      const { id, message } = action.payload;
+      const { id, message, ...rest } = action.payload;
 
       state.errorMessages = [];
-      state.messages.push({ id, role: "user", message });
+      state.messages.push({ id, role: "user", message, ...rest } as any);
       state.history.push({ id, role: "user", content: message });
     },
     addAgentMessage: (
       state,
-      action: PayloadAction<Omit<MetabotChatMessage, "id" | "role">>,
+      action: PayloadAction<Omit<MetabotAgentChatMessage, "id" | "role">>,
     ) => {
       state.toolCalls = [];
+      // @ts-expect-error - TODO: figure out why this type causes issues
       state.messages.push({
         id: createMessageId(),
         role: "agent",
-        message: action.payload.message,
-      });
+        ...action.payload,
+      } as MetabotAgentChatMessage);
     },
     addAgentErrorMessage: (
       state,
@@ -92,14 +151,18 @@ export const metabot = createSlice({
     addAgentTextDelta: (state, action: PayloadAction<string>) => {
       const hasToolCalls = state.toolCalls.length > 0;
       const lastMessage = _.last(state.messages);
-      const canAppend = !hasToolCalls && lastMessage?.role === "agent";
+      const canAppend =
+        !hasToolCalls &&
+        lastMessage?.role === "agent" &&
+        lastMessage.type === "text";
 
       if (canAppend) {
-        lastMessage!.message = lastMessage!.message + action.payload;
+        lastMessage.message = lastMessage.message + action.payload;
       } else {
         state.messages.push({
           id: createMessageId(),
           role: "agent",
+          type: "text",
           message: action.payload,
         });
       }
@@ -148,6 +211,8 @@ export const metabot = createSlice({
       state.isProcessing = false;
       state.toolCalls = [];
       state.conversationId = uuid();
+      state.reactions.suggestedTransform = undefined;
+      state.experimental.metabotReqIdOverride = undefined;
     },
     resetConversationId: (state) => {
       state.conversationId = uuid();
@@ -158,9 +223,24 @@ export const metabot = createSlice({
     setNavigateToPath: (state, action: PayloadAction<string>) => {
       state.reactions.navigateToPath = action.payload;
     },
+    setSuggestedTransform: (
+      state,
+      action: PayloadAction<SuggestedTransform | undefined>,
+    ) => {
+      state.reactions.suggestedTransform = action.payload;
+    },
     setVisible: (state, action: PayloadAction<boolean>) => {
       state.visible = action.payload;
     },
+    setMetabotReqIdOverride: (
+      state,
+      action: PayloadAction<string | undefined>,
+    ) => {
+      state.experimental.metabotReqIdOverride = action.payload;
+    },
+    setProfileOverride: (state, action: PayloadAction<string | undefined>) => {
+      state.experimental.profileOverride = action.payload;
+    },
   },
   extraReducers: (builder) => {
     builder

enterprise/frontend/src/metabase-enterprise/metabot/state/selectors.ts

@@ -2,6 +2,8 @@ import { createSelector } from "@reduxjs/toolkit";
 import _ from "underscore";
 
 import { getIsEmbedding } from "metabase/selectors/embed";
+import { getLocation } from "metabase/selectors/routing";
+import type { TransformId } from "metabase-types/api";
 
 import {
   FIXED_METABOT_IDS,
@@ -9,6 +11,7 @@ import {
   METABOT_REQUEST_IDS,
 } from "../constants";
 
+import type { MetabotUserChatMessage } from "./reducer";
 import type { MetabotStoreState } from "./types";
 
 export const getMetabot = (state: MetabotStoreState) => {
@@ -16,8 +19,9 @@ export const getMetabot = (state: MetabotStoreState) => {
 };
 
 export const getMetabotVisible = createSelector(
-  getMetabot,
-  (metabot) => metabot.visible,
+  [getMetabot, getLocation],
+  (metabot, location) =>
+    location.pathname.startsWith("/admin/transforms") ? true : metabot.visible,
 );
 
 export const getMessages = createSelector(
@@ -25,6 +29,11 @@ export const getMessages = createSelector(
   (metabot) => metabot.messages,
 );
 
+export const getToolCalls = createSelector(
+  getMetabot,
+  (metabot) => metabot.toolCalls,
+);
+
 export const getLastMessage = createSelector(getMessages, (messages) =>
   _.last(messages),
 );
@@ -39,7 +48,7 @@ export const getAgentErrorMessages = createSelector(
 // that exact message will be returned.
 export const getUserPromptForMessageId = createSelector(
   [getMessages, (_, messageId: string) => messageId],
-  (messages, messageId) => {
+  (messages, messageId): MetabotUserChatMessage | undefined => {
     const messageIndex = messages.findLastIndex((m) => m.id === messageId);
     const message = messages[messageIndex];
     if (!message) {
@@ -49,28 +58,16 @@ export const getUserPromptForMessageId = createSelector(
     if (message.role === "user") {
       return message;
     } else {
-      return messages.slice(0, messageIndex).findLast((m) => m.role === "user");
+      // TODO: avoid type cast
+      return messages
+        .slice(0, messageIndex)
+        .findLast((m) => m.role === "user") as
+        | MetabotUserChatMessage
+        | undefined;
     }
   },
 );
 
-export const getLastAgentMessagesByType = createSelector(
-  [getMessages, getAgentErrorMessages],
-  (messages, errorMessages) => {
-    if (errorMessages.length > 0) {
-      return errorMessages.map(({ message }) => message);
-    }
-
-    const start = messages.findLastIndex((msg) => msg.role !== "agent") + 1;
-    return messages.slice(start).map(({ message }) => message);
-  },
-);
-
-export const getToolCalls = createSelector(
-  getMetabot,
-  (metabot) => metabot.toolCalls,
-);
-
 export const getIsProcessing = createSelector(
   getMetabot,
   (metabot) => metabot.isProcessing,
@@ -100,7 +97,8 @@ export const getIsLongMetabotConversation = createSelector(
   getMessages,
   (messages) => {
     const totalMessageLength = messages.reduce((sum, msg) => {
-      return sum + msg.message.length;
+      // TODO: fix this to be any message w/ a message property
+      return sum + (msg.type === "text" ? msg.message.length : 0);
     }, 0);
     return totalMessageLength >= LONG_CONVO_MSG_LENGTH_THRESHOLD;
   },
@@ -110,20 +108,49 @@ export const getMetabotId = createSelector(getIsEmbedding, (isEmbedding) =>
   isEmbedding ? FIXED_METABOT_IDS.EMBEDDED : FIXED_METABOT_IDS.DEFAULT,
 );
 
+export const getMetabotReqIdOverride = createSelector(
+  getMetabot,
+  (metabot) => metabot.experimental.metabotReqIdOverride,
+);
+
 export const getMetabotRequestId = createSelector(
+  getMetabotReqIdOverride,
   getIsEmbedding,
-  (isEmbedding) => (isEmbedding ? METABOT_REQUEST_IDS.EMBEDDED : undefined),
+  (metabotReqIdOverride, isEmbedding) =>
+    metabotReqIdOverride ??
+    (isEmbedding ? METABOT_REQUEST_IDS.EMBEDDED : undefined),
+);
+
+export const getProfileOverride = createSelector(
+  getMetabot,
+  (metabot) => metabot.experimental.profileOverride,
+);
+
+export const getProfile = createSelector(
+  getProfileOverride,
+  getLocation,
+  (profileOverride, location) => {
+    if (profileOverride) {
+      return profileOverride;
+    }
+
+    return location.pathname.startsWith("/admin/transforms")
+      ? "transforms_codegen"
+      : undefined;
+  },
 );
 
 export const getAgentRequestMetadata = createSelector(
   getHistory,
   getMetabotState,
-  (history, state) => ({
+  getProfile,
+  (history, state, profile) => ({
     state,
     // NOTE: need end to end support for ids on messages as BE will error if ids are present
     history: history.map((h) =>
       h.id && h.id.startsWith(`msg_`) ? _.omit(h, "id") : h,
     ),
+    ...(profile ? { profile_id: profile } : {}),
   }),
 );
 
@@ -131,3 +158,16 @@ export const getNavigateToPath = createSelector(
   getMetabotReactionsState,
   (reactionsState) => reactionsState.navigateToPath,
 );
+
+export const getMetabotSuggestedTransform = createSelector(
+  [getMetabotReactionsState, (_, transformId?: TransformId) => transformId],
+  (reactionsState, transformId) => {
+    if (!transformId) {
+      return reactionsState.suggestedTransform;
+    }
+
+    return reactionsState.suggestedTransform?.id === transformId
+      ? reactionsState.suggestedTransform
+      : undefined;
+  },
+);

enterprise/frontend/src/metabase-enterprise/metabot/state/selectors.unit.spec.ts

@@ -3,7 +3,6 @@ import { createMockState } from "metabase-types/store/mocks";
 
 import {
   type MetabotState,
-  getLastAgentMessagesByType,
   getMetabotInitialState,
   getUserPromptForMessageId,
 } from "./index";
@@ -23,71 +22,41 @@ function setup(metabotState: Partial<MetabotState>) {
 }
 
 describe("metabot selectors", () => {
-  describe("getLastAgentMessagesByType", () => {
-    it("should not return any messages if there are none", () => {
-      const state = setup({ messages: [] });
-      const messages = getLastAgentMessagesByType(state as any);
-      expect(messages).toEqual([]);
-    });
-
-    it("should not return any messages if latest message is from a user", () => {
-      const state = setup({
-        messages: [{ id: "1", role: "user", message: "bleh" }],
-      });
-      const messages = getLastAgentMessagesByType(state as any);
-      expect(messages).toEqual([]);
-    });
-
-    it("should return latest agent reply messages only", () => {
-      const state = setup({
-        messages: [
-          { id: "1", role: "user", message: "bleh" },
-          { id: "2", role: "agent", message: "blah" },
-          { id: "3", role: "agent", message: "blah" },
-        ],
-      });
-      const messages = getLastAgentMessagesByType(state as any);
-      expect(messages).toEqual(["blah", "blah"]);
-    });
-
-    it("should return latest agent error messages only", () => {
-      const state = setup({
-        messages: [{ id: "1", role: "agent", message: "blah" }],
-        errorMessages: [
-          { type: "message", message: "BLAH" },
-          { type: "message", message: "BLAH" },
-        ],
-      });
-      const messages = getLastAgentMessagesByType(state as any);
-      expect(messages).toEqual(["BLAH", "BLAH"]);
-    });
-  });
-
   describe("getUserPromptForMessageId", () => {
     it("should return the message with the matching id if id is for a user message", () => {
       const state = setup({
         messages: [
-          { id: "1", role: "user", message: "bleh" },
-          { id: "2", role: "agent", message: "blah" },
+          { id: "1", role: "user", type: "text", message: "bleh" },
+          { id: "2", role: "agent", type: "text", message: "blah" },
         ],
       });
       const message = getUserPromptForMessageId(state as any, "1");
-      expect(message).toEqual({ id: "1", role: "user", message: "bleh" });
+      expect(message).toEqual({
+        id: "1",
+        role: "user",
+        type: "text",
+        message: "bleh",
+      });
     });
 
     it("should return the message with the matching id if id is for an agent message", () => {
       const state = setup({
         messages: [
-          { id: "1", role: "user", message: "bleh" },
-          { id: "2", role: "agent", message: "blah" },
-          { id: "3", role: "user", message: "bleh bleh" },
-          { id: "4", role: "agent", message: "blah blah" },
+          { id: "1", type: "text", role: "user", message: "bleh" },
+          { id: "2", type: "text", role: "agent", message: "blah" },
+          { id: "3", type: "text", role: "user", message: "bleh bleh" },
+          { id: "4", type: "text", role: "agent", message: "blah blah" },
         ],
       });
       const message1 = getUserPromptForMessageId(state as any, "2");
       expect(message1).toEqual({ id: "1", role: "user", message: "bleh" });
       const message2 = getUserPromptForMessageId(state as any, "4");
-      expect(message2).toEqual({ id: "3", role: "user", message: "bleh bleh" });
+      expect(message2).toEqual({
+        id: "3",
+        role: "user",
+        type: "text",
+        message: "bleh bleh",
+      });
     });
   });
 });

enterprise/frontend/src/metabase-enterprise/metabot/state/types.ts

@@ -9,3 +9,8 @@ export interface MetabotStoreState extends EnterpriseState {
     metabotPlugin: MetabotState;
   };
 }
+
+export interface SlashCommand {
+  cmd: string;
+  args: string[];
+}

enterprise/frontend/src/metabase-enterprise/metabot/state/utils.ts

@@ -1,5 +1,23 @@
 import { nanoid } from "@reduxjs/toolkit";
 
+import type { SlashCommand } from "./types";
+
 export const createMessageId = () => {
   return `msg_${nanoid()}`;
 };
+
+export const parseSlashCommand = (
+  message: string,
+): SlashCommand | undefined => {
+  const { cmd, args } =
+    message.match(/^\/(?<cmd>\w+)(?:\s(?<args>.+))?/)?.groups || {};
+
+  if (!cmd) {
+    return undefined;
+  }
+
+  return {
+    cmd,
+    args: args ? args.split(" ") : [],
+  };
+};

enterprise/frontend/src/metabase-enterprise/transforms/components/QueryEditor/EditorBody/EditorBody.tsx

@@ -15,6 +15,7 @@ import type Database from "metabase-lib/v1/metadata/Database";
 import type NativeQuery from "metabase-lib/v1/queries/NativeQuery";
 import type {
   Database as ApiDatabase,
+  DatasetQuery,
   NativeQuerySnippet,
   RecentItem,
 } from "metabase-types/api";
@@ -53,6 +54,9 @@ type EditorBodyProps = {
   nativeEditorSelectedText?: string | null;
 
   databases: ApiDatabase[];
+  proposedQuestion?: Question;
+  onRejectProposed?: () => void;
+  onAcceptProposed?: (query: DatasetQuery) => void;
 };
 
 export function EditorBody({
@@ -73,6 +77,9 @@ export function EditorBody({
   onChangeModalSnippet,
   onChangeNativeEditorSelection,
   nativeEditorSelectedText,
+  proposedQuestion,
+  onRejectProposed,
+  onAcceptProposed,
 }: EditorBodyProps) {
   const [isResizing, setIsResizing] = useState(false);
   const reportTimezone = useSetting("report-timezone-long");
@@ -135,6 +142,9 @@ export function EditorBody({
       }}
       setNativeEditorSelectedRange={onChangeNativeEditorSelection}
       nativeEditorSelectedText={nativeEditorSelectedText}
+      proposedQuestion={proposedQuestion}
+      onRejectProposed={onRejectProposed}
+      onAcceptProposed={onAcceptProposed}
     />
   ) : (
     <ResizableBox

enterprise/frontend/src/metabase-enterprise/transforms/components/QueryEditor/EditorHeader/EditorHeader.tsx

@@ -4,6 +4,7 @@ import Button from "metabase/common/components/Button";
 import EditBar from "metabase/common/components/EditBar";
 
 type EditorHeaderProps = {
+  name?: string;
   canSave: boolean;
   isNew: boolean;
   isSaving: boolean;
@@ -12,6 +13,7 @@ type EditorHeaderProps = {
 };
 
 export function EditorHeader({
+  name,
   canSave,
   isNew,
   isSaving,
@@ -20,7 +22,7 @@ export function EditorHeader({
 }: EditorHeaderProps) {
   return (
     <EditBar
-      title={getTitle(isNew)}
+      title={getTitle(isNew, name)}
       admin
       buttons={[
         <Button key="cancel" small onClick={onCancel}>{t`Cancel`}</Button>,
@@ -38,11 +40,11 @@ export function EditorHeader({
   );
 }
 
-function getTitle(isNew: boolean) {
+function getTitle(isNew: boolean, name?: string) {
   if (isNew) {
     return t`You’re creating a new transform`;
   } else {
-    return t`You’re editing a transform`;
+    return t`You’re editing the "${name}" transform`;
   }
 }
 

enterprise/frontend/src/metabase-enterprise/transforms/components/QueryEditor/QueryEditor.tsx

@@ -6,7 +6,7 @@ import type { SelectionRange } from "metabase/query_builder/components/NativeQue
 import { Center, Flex, Loader, Stack } from "metabase/ui";
 import * as Lib from "metabase-lib";
 import type Question from "metabase-lib/v1/Question";
-import type { DatasetQuery, NativeQuerySnippet } from "metabase-types/api";
+import type { DatasetQuery, NativeQuerySnippet, Transform } from "metabase-types/api";
 
 import { useQueryMetadata } from "../../hooks/use-query-metadata";
 import { useQueryResults } from "../../hooks/use-query-results";
@@ -20,21 +20,32 @@ import S from "./QueryEditor.module.css";
 import { useInsertSnippetHandler, useSelectedText } from "./util";
 
 type QueryEditorProps = {
+  transform?: Transform;
   initialQuery: DatasetQuery;
+  proposedQuery?: DatasetQuery;
   isNew?: boolean;
   isSaving?: boolean;
   onSave: (newQuery: DatasetQuery) => void;
+  onChange?: (newQuery: DatasetQuery) => void;
   onCancel: () => void;
+  onRejectProposed?: () => void;
+  onAcceptProposed?: (query: DatasetQuery) => void;
 };
 
 export function QueryEditor({
   initialQuery,
+  proposedQuery,
+  transform,
   isNew = true,
   isSaving = false,
   onSave,
+  onChange,
   onCancel,
+  onRejectProposed,
+  onAcceptProposed,
 }: QueryEditorProps) {
-  const { question, isQueryDirty, setQuestion } = useQueryState(initialQuery);
+  const { question, proposedQuestion, isQueryDirty, setQuestion } =
+    useQueryState(initialQuery, proposedQuery);
   const { isInitiallyLoaded } = useQueryMetadata(question);
   const {
     result,
@@ -44,12 +55,13 @@ export function QueryEditor({
     isResultDirty,
     runQuery,
     cancelQuery,
-  } = useQueryResults(question);
+  } = useQueryResults(question, proposedQuestion);
   const canSave = Lib.canSave(question.query(), question.type());
   const { isNative } = Lib.queryDisplayInfo(question.query());
 
   const handleChange = async (newQuestion: Question) => {
     setQuestion(newQuestion);
+    onChange?.(newQuestion.datasetQuery());
   };
 
   const handleSave = () => {
@@ -119,6 +131,7 @@ export function QueryEditor({
       gap={0}
     >
       <EditorHeader
+        name={transform?.name}
         isNew={isNew}
         isSaving={isSaving}
         canSave={canSave && (isNew || isQueryDirty)}
@@ -129,6 +142,7 @@ export function QueryEditor({
         <Stack flex="2 1 100%">
           <EditorBody
             question={question}
+            proposedQuestion={proposedQuestion}
             isNative={isNative}
             isRunnable={isRunnable}
             isRunning={isRunning}
@@ -138,6 +152,8 @@ export function QueryEditor({
             onChange={handleChange}
             onRunQuery={runQuery}
             onCancelQuery={cancelQuery}
+            onRejectProposed={onRejectProposed}
+            onAcceptProposed={onAcceptProposed}
             databases={databases?.data ?? []}
             onToggleDataReference={handleToggleDataReference}
             onToggleSnippetSidebar={handleToggleSnippetSidebar}

enterprise/frontend/src/metabase-enterprise/transforms/hooks/use-query-results/use-query-results.ts

@@ -8,13 +8,18 @@ import * as Lib from "metabase-lib";
 import Question from "metabase-lib/v1/Question";
 import type { DatasetQuery } from "metabase-types/api";
 
-export function useQueryResults(question: Question) {
+export function useQueryResults(
+  question: Question,
+  proposedQuestion?: Question,
+) {
   const metadata = useSelector(getMetadata);
   const [lastRunQuery, setLastRunQuery] = useState<DatasetQuery | null>(null);
 
+  const currentQuestion = proposedQuestion ?? question;
+
   const [{ value: results = null, loading: isRunning }, runQuery] = useAsyncFn(
-    () => runQuestionQuery(question),
-    [question],
+    () => runQuestionQuery(currentQuestion),
+    [currentQuestion],
   );
 
   const { result, rawSeries, isRunnable, isResultDirty } = useMemo(() => {
@@ -26,9 +31,13 @@ export function useQueryResults(question: Question) {
       lastRunQuestion && result
         ? [{ card: lastRunQuestion.card(), data: result.data }]
         : null;
-    const isRunnable = Lib.canRun(question.query(), question.type());
+    const isRunnable = Lib.canRun(
+      currentQuestion.query(),
+      currentQuestion.type(),
+    );
     const isResultDirty =
-      lastRunQuestion == null || question.isDirtyComparedTo(lastRunQuestion);
+      lastRunQuestion == null ||
+      currentQuestion.isDirtyComparedTo(lastRunQuestion);
 
     return {
       result,
@@ -36,12 +45,12 @@ export function useQueryResults(question: Question) {
       isRunnable,
       isResultDirty,
     };
-  }, [question, results, metadata, lastRunQuery]);
+  }, [currentQuestion, results, metadata, lastRunQuery]);
 
   const handleRunQuery = useCallback(async () => {
     await runQuery();
-    setLastRunQuery(question.datasetQuery());
-  }, [question, runQuery]);
+    setLastRunQuery(currentQuestion.datasetQuery());
+  }, [currentQuestion, runQuery]);
 
   const handleCancelQuery = useCallback(() => {
     return null;

enterprise/frontend/src/metabase-enterprise/transforms/hooks/use-query-state/use-query-state.ts

@@ -6,7 +6,10 @@ import * as Lib from "metabase-lib";
 import Question from "metabase-lib/v1/Question";
 import type { DatasetQuery } from "metabase-types/api";
 
-export function useQueryState(initialQuery: DatasetQuery) {
+export function useQueryState(
+  initialQuery: DatasetQuery,
+  proposedQuery?: DatasetQuery,
+) {
   const [query, setQuery] = useState(initialQuery);
   const metadata = useSelector(getMetadata);
 
@@ -15,6 +18,30 @@ export function useQueryState(initialQuery: DatasetQuery) {
     [query, metadata],
   );
 
+  const proposedQuestion = useMemo(() => {
+    if (!proposedQuery) {
+      return undefined;
+    }
+
+    const question = Question.create({
+      dataset_query: proposedQuery,
+      metadata,
+    });
+    const { isNative } = Lib.queryDisplayInfo(question.query());
+
+    if (isNative) {
+      const nativeQuery = question.legacyNativeQuery();
+      if (nativeQuery) {
+        const queryText = nativeQuery.queryText();
+        // For native queries, ensure template tags are processed
+        const updatedQuery = nativeQuery.setQueryText(queryText);
+        return updatedQuery.question();
+      }
+    }
+
+    return question;
+  }, [proposedQuery, metadata]);
+
   const isQueryDirty = useMemo(
     () => !Lib.areLegacyQueriesEqual(query, initialQuery),
     [query, initialQuery],
@@ -24,5 +51,5 @@ export function useQueryState(initialQuery: DatasetQuery) {
     setQuery(newQuestion.datasetQuery());
   };
 
-  return { question, isQueryDirty, setQuestion };
+  return { question, proposedQuestion, isQueryDirty, setQuestion };
 }

enterprise/frontend/src/metabase-enterprise/transforms/pages/NewTransformQueryPage/CreateTransformModal/CreateTransformModal.tsx

@@ -28,14 +28,25 @@ import type {
   Transform,
 } from "metabase-types/api";
 
+const NEW_TRANSFORM_SCHEMA = Yup.object({
+  name: Yup.string().required(Errors.required),
+  description: Yup.string().nullable(),
+  targetName: Yup.string().required(Errors.required),
+  targetSchema: Yup.string().nullable(),
+});
+
+type NewTransformValues = Yup.InferType<typeof NEW_TRANSFORM_SCHEMA>;
+
 type CreateTransformModalProps = {
   query: DatasetQuery;
+  initValues?: Partial<NewTransformValues>;
   onCreate: (transform: Transform) => void;
   onClose: () => void;
 };
 
 export function CreateTransformModal({
   query,
+  initValues,
   onCreate,
   onClose,
 }: CreateTransformModalProps) {
@@ -44,6 +55,7 @@ export function CreateTransformModal({
       <FocusTrap.InitialFocus />
       <CreateTransformForm
         query={query}
+        initValues={initValues}
         onCreate={onCreate}
         onClose={onClose}
       />
@@ -53,26 +65,14 @@ export function CreateTransformModal({
 
 type CreateTransformFormProps = {
   query: DatasetQuery;
+  initValues?: Partial<NewTransformValues>;
   onCreate: (transform: Transform) => void;
   onClose: () => void;
 };
 
-type NewTransformValues = {
-  name: string;
-  description: string | null;
-  targetName: string;
-  targetSchema: string | null;
-};
-
-const NEW_TRANSFORM_SCHEMA = Yup.object({
-  name: Yup.string().required(Errors.required),
-  description: Yup.string().nullable(),
-  targetName: Yup.string().required(Errors.required),
-  targetSchema: Yup.string().nullable(),
-});
-
 function CreateTransformForm({
   query,
+  initValues,
   onCreate,
   onClose,
 }: CreateTransformFormProps) {
@@ -99,8 +99,8 @@ function CreateTransformForm({
   const supportsSchemas = database && hasFeature(database, "schemas");
 
   const initialValues: NewTransformValues = useMemo(
-    () => getInitialValues(schemas),
-    [schemas],
+    () => getInitialValues(schemas, initValues),
+    [schemas, initValues],
   );
 
   if (isLoading || error != null) {
@@ -161,12 +161,16 @@ function CreateTransformForm({
   );
 }
 
-function getInitialValues(schemas: string[]): NewTransformValues {
+function getInitialValues(
+  schemas: string[],
+  initValues?: Partial<NewTransformValues>,
+): NewTransformValues {
   return {
     name: "",
     description: null,
     targetName: "",
     targetSchema: schemas?.[0] || null,
+    ...initValues,
   };
 }
 

enterprise/frontend/src/metabase-enterprise/transforms/pages/NewTransformQueryPage/NewTransformQueryPage.tsx

@@ -1,11 +1,16 @@
 import { useDisclosure } from "@mantine/hooks";
-import { useState } from "react";
+import { useMemo, useState } from "react";
 import { push } from "react-router-redux";
 
 import { skipToken, useGetCardQuery } from "metabase/api";
+import { AdminSettingsLayout } from "metabase/common/components/AdminLayout/AdminSettingsLayout";
 import { LoadingAndErrorWrapper } from "metabase/common/components/LoadingAndErrorWrapper";
-import { useDispatch } from "metabase/lib/redux";
+import { useDispatch, useSelector } from "metabase/lib/redux";
 import * as Urls from "metabase/lib/urls";
+import {
+  getMetabotSuggestedTransform,
+  setSuggestedTransform,
+} from "metabase-enterprise/metabot/state";
 import Question from "metabase-lib/v1/Question";
 import type { Card, CardId, DatasetQuery, Transform } from "metabase-types/api";
 
@@ -36,11 +41,20 @@ export function NewTransformQueryPage({ params }: NewTransformQueryPageProps) {
     error,
   } = useGetCardQuery(cardId ? { id: cardId } : skipToken);
 
+  const suggestedTransform = useSelector(
+    getMetabotSuggestedTransform as any,
+  ) as ReturnType<typeof getMetabotSuggestedTransform>;
+  const suggestedQuery = suggestedTransform?.source.query;
+
+  const [initialSuggestedQuery] = useState(suggestedQuery);
+
+  const initialQuery = initialSuggestedQuery || getInitialQuery(card, type);
+
   if (isLoading || error) {
     return <LoadingAndErrorWrapper loading={isLoading} error={error} />;
   }
 
-  return <NewTransformPageBody initialQuery={getInitialQuery(card, type)} />;
+  return <NewTransformPageBody initialQuery={initialQuery} />;
 }
 
 type NewTransformPageBodyProps = {
@@ -66,22 +80,63 @@ function NewTransformPageBody({ initialQuery }: NewTransformPageBodyProps) {
     dispatch(push(getTransformListUrl()));
   };
 
+  const suggestedTransform = useSelector(
+    getMetabotSuggestedTransform as any,
+  ) as ReturnType<typeof getMetabotSuggestedTransform>;
+
+  const onRejectProposed = () => dispatch(setSuggestedTransform(undefined));
+  const onAcceptProposed = (query: DatasetQuery) => {
+    handleSaveClick(query);
+  };
+
+  const suggestedQuery = suggestedTransform?.source.query;
+
+  const initQuery =
+    initialQuery.type === "native" && initialQuery.native.query.length > 0
+      ? initialQuery
+      : (suggestedQuery ?? initialQuery);
+
+  const proposedQuery =
+    suggestedQuery?.type === "native" &&
+    initQuery.type === "native" &&
+    suggestedQuery.native.query === initQuery.native.query
+      ? undefined
+      : suggestedQuery;
+
+  const createTransformInitValues = useMemo(
+    () =>
+      suggestedTransform
+        ? {
+            name: suggestedTransform.name,
+            description: suggestedTransform.description,
+            targetName: suggestedTransform.target.name,
+            // TODO: enabling this breaks the UI for some reason...
+            // targetSchema: suggestedTransform.target.schema,
+          }
+        : undefined,
+    [suggestedTransform],
+  );
+
   return (
-    <>
+    <AdminSettingsLayout fullWidthContent>
       <QueryEditor
-        initialQuery={initialQuery}
+        initialQuery={initQuery}
         isNew
         onSave={handleSaveClick}
         onCancel={handleCancelClick}
+        proposedQuery={proposedQuery}
+        onRejectProposed={onRejectProposed}
+        onAcceptProposed={onAcceptProposed}
       />
       {isModalOpened && (
         <CreateTransformModal
           query={query}
+          initValues={createTransformInitValues}
           onCreate={handleCreate}
           onClose={closeModal}
         />
       )}
-    </>
+    </AdminSettingsLayout>
   );
 }
 

enterprise/frontend/src/metabase-enterprise/transforms/pages/TransformPage/TransformPage.tsx

@@ -4,6 +4,7 @@ import { t } from "ttag";
 import { skipToken } from "metabase/api";
 import { LoadingAndErrorWrapper } from "metabase/common/components/LoadingAndErrorWrapper";
 import * as Urls from "metabase/lib/urls";
+import { useRegisterMetabotContextProvider } from "metabase/metabot";
 import { Stack } from "metabase/ui";
 import { useGetTransformQuery } from "metabase-enterprise/api";
 import type { Transform, TransformId } from "metabase-types/api";
@@ -49,6 +50,12 @@ export function TransformPage({ params }: TransformPageProps) {
     setIsPolling(isPollingNeeded(transform));
   }
 
+  useRegisterMetabotContextProvider(async () => {
+    return transform
+      ? { user_is_viewing: [{ type: "transform", ...transform }] }
+      : {};
+  }, [transform]);
+
   if (isLoading || error != null) {
     return <LoadingAndErrorWrapper loading={isLoading} error={error} />;
   }

enterprise/frontend/src/metabase-enterprise/transforms/pages/TransformQueryPage/TransformQueryPage.tsx

@@ -1,15 +1,23 @@
+import { useState } from "react";
 import { push } from "react-router-redux";
 import { t } from "ttag";
 
 import { skipToken } from "metabase/api";
+import { AdminSettingsLayout } from "metabase/common/components/AdminLayout/AdminSettingsLayout";
 import { LoadingAndErrorWrapper } from "metabase/common/components/LoadingAndErrorWrapper";
-import { useDispatch } from "metabase/lib/redux";
+import { useDispatch, useSelector } from "metabase/lib/redux";
 import * as Urls from "metabase/lib/urls";
+import { useRegisterMetabotContextProvider } from "metabase/metabot";
 import { useMetadataToasts } from "metabase/metadata/hooks";
 import {
   useGetTransformQuery,
   useUpdateTransformMutation,
 } from "metabase-enterprise/api";
+import { useMetabotAgent } from "metabase-enterprise/metabot/hooks";
+import {
+  getMetabotSuggestedTransform,
+  setSuggestedTransform,
+} from "metabase-enterprise/metabot/state";
 import type { DatasetQuery, Transform } from "metabase-types/api";
 
 import { QueryEditor } from "../../components/QueryEditor";
@@ -53,7 +61,50 @@ export function TransformQueryPageBody({
   const dispatch = useDispatch();
   const { sendErrorToast, sendSuccessToast } = useMetadataToasts();
 
-  const handleSave = async (query: DatasetQuery) => {
+  const metabot = useMetabotAgent();
+
+  const initialQuery = transform.source.query;
+  const [latestQuery, setLatestQuery] = useState<DatasetQuery>(initialQuery);
+
+  const suggestedTransform = useSelector(
+    (state) => getMetabotSuggestedTransform(state, transform.id) as any,
+  ) as ReturnType<typeof getMetabotSuggestedTransform>;
+  const proposedQuery = suggestedTransform?.source.query;
+
+  useRegisterMetabotContextProvider(async () => {
+    const viewedTransform = suggestedTransform ?? {
+      ...transform,
+      source: { ...transform.source, query: latestQuery },
+    };
+    return { user_is_viewing: [{ type: "transform", ...viewedTransform }] };
+  }, [transform, latestQuery, suggestedTransform]);
+
+  const onRejectProposed = () => {
+    dispatch(setSuggestedTransform(undefined));
+    metabot.submitInput({
+      type: "action",
+      message:
+        "HIDDEN MESSAGE: the user has rejected your changes, ask for clarification on what they'd like to do instead.",
+      // @ts-expect-error -- TODO
+      userMessage: "❌ You rejected the change",
+    });
+  };
+  const onAcceptProposed = async (query: DatasetQuery) => {
+    await handleSave(query, { leaveEditor: false });
+    dispatch(setSuggestedTransform(undefined));
+    metabot.submitInput({
+      type: "action",
+      message:
+        "HIDDEN MESSAGE: user has accepted your changes, move to the next step!",
+      // @ts-expect-error -- TODO
+      userMessage: "✅ You accepted the change",
+    });
+  };
+
+  const handleSave = async (
+    query: DatasetQuery,
+    { leaveEditor } = { leaveEditor: true },
+  ) => {
     const { error } = await updateTransform({
       id: transform.id,
       source: {
@@ -66,7 +117,9 @@ export function TransformQueryPageBody({
       sendErrorToast(t`Failed to update transform query`);
     } else {
       sendSuccessToast(t`Transform query updated`);
-      dispatch(push(getTransformUrl(transform.id)));
+      if (leaveEditor) {
+        dispatch(push(getTransformUrl(transform.id)));
+      }
     }
   };
 
@@ -75,12 +128,19 @@ export function TransformQueryPageBody({
   };
 
   return (
-    <QueryEditor
-      initialQuery={transform.source.query}
-      isNew={false}
-      isSaving={isLoading}
-      onSave={handleSave}
-      onCancel={handleCancel}
-    />
+    <AdminSettingsLayout fullWidthContent key={transform.id}>
+      <QueryEditor
+        initialQuery={initialQuery}
+        transform={transform}
+        isNew={false}
+        isSaving={isLoading}
+        onSave={handleSave}
+        onChange={setLatestQuery}
+        onCancel={handleCancel}
+        proposedQuery={proposedQuery}
+        onRejectProposed={onRejectProposed}
+        onAcceptProposed={onAcceptProposed}
+      />
+    </AdminSettingsLayout>
   );
 }

frontend/src/metabase-types/api/metabot.ts

@@ -2,12 +2,12 @@ import type {
   CardDisplayType,
   CardId,
   CardType,
-  CollectionId,
   DashboardId,
   DatasetQuery,
   PaginationResponse,
   RowValue,
-  SearchModel,
+  SuggestedTransform,
+  Transform,
   UnsavedCard,
   Version,
 } from ".";
@@ -57,18 +57,6 @@ export type MetabotHistory = any[];
 
 export type MetabotStateContext = Record<string, any>;
 
-export type MetabotMessageReaction = {
-  type: "metabot.reaction/message";
-  message: string;
-};
-
-export type MetabotRedirectReaction = {
-  type: "metabot.reaction/redirect";
-  url: string;
-};
-
-export type MetabotReaction = MetabotMessageReaction | MetabotRedirectReaction;
-
 export type MetabotColumnType =
   | "number"
   | "string"
@@ -135,11 +123,16 @@ export type MetabotDocumentInfo = {
   id: number;
 };
 
+export type MetabotTransformInfo = (Transform | SuggestedTransform) & {
+  type: "transform";
+};
+
 export type MetabotEntityInfo =
   | MetabotCardInfo
   | MetabotDashboardInfo
   | MetabotAdhocQueryInfo
-  | MetabotDocumentInfo;
+  | MetabotDocumentInfo
+  | MetabotTransformInfo;
 
 /* Metabot v3 - API Request Types */
 
@@ -150,10 +143,10 @@ export type MetabotAgentRequest = {
   state: MetabotStateContext;
   conversation_id: string; // uuid
   metabot_id?: string;
+  profile_id?: string;
 };
 
 export type MetabotAgentResponse = {
-  reactions: MetabotReaction[];
   history: MetabotHistory[];
   conversation_id: string;
   state: any;
@@ -218,19 +211,13 @@ export type MetabotName = string;
 
 export type MetabotInfo = {
   id: MetabotId;
+  entity_id: string;
   name: MetabotName;
-};
-
-export type MetabotEntity = {
-  name: string;
-  id: CollectionId;
-  model: Extract<SearchModel, "collection">;
-  collection_id: CollectionId;
-  collection_name: string;
-};
-
-export type MetabotApiEntity = Omit<MetabotEntity, "id"> & {
-  model_id: MetabotEntity["id"];
+  description: string;
+  use_verified_content: boolean;
+  collection_id: number | null;
+  created_at: string;
+  updated_at: string;
 };
 
 /* Metabot v3 - Document Types */
@@ -251,3 +238,12 @@ export interface MetabotGenerateContentResponse {
 export interface PurchaseMetabotAddOnRequest {
   terms_of_service: boolean;
 }
+
+/* Metabot v3 - Data Part Types */
+
+export type MetabotTodoItem = {
+  id: string;
+  content: string;
+  status: "pending" | "in_progress" | "completed" | "cancelled";
+  priority: "high" | "medium" | "low";
+};

frontend/src/metabase-types/api/mocks/index.ts

@@ -14,6 +14,7 @@ export * from "./field";
 export * from "./geojson";
 export * from "./group";
 export * from "./logger";
+export * from "./metabot";
 export * from "./modelIndexes";
 export * from "./models";
 export * from "./notification";

frontend/src/metabase-types/api/mocks/metabot.ts

@@ -0,0 +1,15 @@
+import type { MetabotInfo } from "../metabot";
+
+export const createMockMetabotInfo = (
+  opts?: Partial<MetabotInfo>,
+): MetabotInfo => ({
+  id: 1,
+  name: "Metabot",
+  entity_id: "metabot",
+  description: "",
+  use_verified_content: false,
+  collection_id: null,
+  created_at: new Date().toISOString(),
+  updated_at: new Date().toISOString(),
+  ...opts,
+});

frontend/src/metabase-types/api/transform.ts

@@ -22,6 +22,9 @@ export type Transform = {
   last_run?: TransformRun | null;
 };
 
+export type SuggestedTransform = Partial<Pick<Transform, "id">> &
+  Pick<Transform, "name" | "description" | "source" | "target">;
+
 export type TransformSource = {
   type: "query";
   query: DatasetQuery;

frontend/src/metabase/common/components/AdminLayout/AdminSettingsLayout.module.css

@@ -22,7 +22,6 @@
   display: flex;
   justify-content: center;
   width: 100%;
-  padding: 2rem;
   position: relative;
   background-color: var(--mb-color-bg-light);
 }

frontend/src/metabase/common/components/AdminLayout/AdminSettingsLayout.tsx

@@ -1,6 +1,10 @@
 import type React from "react";
+import { t } from "ttag";
 
 import ErrorBoundary from "metabase/ErrorBoundary";
+import { useSelector } from "metabase/lib/redux";
+import { PLUGIN_METABOT } from "metabase/plugins";
+import { getLocation } from "metabase/selectors/routing";
 import { Box } from "metabase/ui";
 
 import { NotFound } from "../ErrorPages";
@@ -10,12 +14,20 @@ import S from "./AdminSettingsLayout.module.css";
 export const AdminSettingsLayout = ({
   sidebar,
   children,
-  maw = "50rem",
+  fullWidthContent = false,
+  maw = fullWidthContent ? "100%" : "50rem",
+  padding = fullWidthContent ? "0" : "2rem",
 }: {
   sidebar?: React.ReactNode;
   children?: React.ReactNode;
+  fullWidthContent?: boolean;
   maw?: string;
+  padding?: string;
 }) => {
+  const location = useSelector(getLocation);
+  const isMetabotEnabledForRoute =
+    location.pathname.startsWith("/admin/transforms");
+
   return (
     <Box className={S.Wrapper}>
       <Box className={S.Main}>
@@ -24,13 +36,22 @@ export const AdminSettingsLayout = ({
             {sidebar}
           </Box>
         )}
-        <Box className={S.Content} data-testid="admin-layout-content">
+        <Box
+          className={S.Content}
+          style={{ padding }}
+          data-testid="admin-layout-content"
+        >
           <Box maw={maw} w="100%">
-            <Box pb="2rem">
+            <Box {...(fullWidthContent ? { h: "100%" } : { pb: "2rem" })}>
               <ErrorBoundary>{children ?? <NotFound />}</ErrorBoundary>
             </Box>
           </Box>
         </Box>
+
+        <PLUGIN_METABOT.Metabot
+          hide={!isMetabotEnabledForRoute}
+          emptyConvoText={t`Let's transform your data together!`}
+        />
       </Box>
     </Box>
   );

frontend/src/metabase/plugins/index.ts

@@ -700,7 +700,8 @@ export const PLUGIN_AI_ENTITY_ANALYSIS: PluginAIEntityAnalysis = {
 
 export const PLUGIN_METABOT = {
   isEnabled: () => false,
-  Metabot: (_props: { hide?: boolean }) => null as React.ReactElement | null,
+  Metabot: (_props: { hide?: boolean; emptyConvoText?: string }) =>
+    null as React.ReactElement | null,
   defaultMetabotContextValue,
   MetabotContext: React.createContext(defaultMetabotContextValue),
   getMetabotProvider: () => {

frontend/src/metabase/query_builder/components/NativeQueryEditor/CodeMirrorEditor/CodeMirrorEditor.tsx

@@ -21,6 +21,7 @@ import type { SelectionRange } from "../types";
 
 export type CodeMirrorEditorProps = {
   query: Lib.Query;
+  proposedQuery?: Lib.Query;
   highlightedLineNumbers?: number[];
   placeholder?: string;
   readOnly?: boolean;
@@ -51,6 +52,7 @@ export const CodeMirrorEditor = forwardRef<
 >(function CodeMirrorEditor(
   {
     query,
+    proposedQuery,
     highlightedLineNumbers,
     placeholder = getPlaceholderText(Lib.engine(query)),
     readOnly,
@@ -64,7 +66,11 @@ export const CodeMirrorEditor = forwardRef<
   ref,
 ) {
   const editorRef = useRef<CodeMirrorRef>(null);
-  const extensions = useExtensions({ query, onRunQuery });
+  const extensions = useExtensions({
+    query,
+    diff: !!proposedQuery,
+    onRunQuery,
+  });
 
   useImperativeHandle(ref, () => {
     return {
@@ -128,13 +134,17 @@ export const CodeMirrorEditor = forwardRef<
     [highlightedLineNumbers],
   );
 
+  const value = useMemo(() => {
+    return Lib.rawNativeQuery(proposedQuery ?? query);
+  }, [proposedQuery, query]);
+
   return (
     <CodeMirror
       ref={editorRef}
       data-testid="native-query-editor"
       className={S.editor}
       extensions={extensions}
-      value={Lib.rawNativeQuery(query)}
+      value={value}
       readOnly={readOnly}
       onChange={onChange}
       height="100%"

frontend/src/metabase/query_builder/components/NativeQueryEditor/CodeMirrorEditor/extensions.tsx

@@ -1,4 +1,5 @@
 import { autocompletion } from "@codemirror/autocomplete";
+import { unifiedMergeView } from "@codemirror/merge";
 import { type Extension, Prec } from "@codemirror/state";
 import {
   Decoration,
@@ -26,10 +27,15 @@ import { getReferencedCardIds } from "./util";
 
 type Options = {
   query: Lib.Query;
+  diff: boolean;
   onRunQuery?: () => void;
 };
 
-export function useExtensions({ query, onRunQuery }: Options): Extension[] {
+export function useExtensions({
+  query,
+  diff,
+  onRunQuery,
+}: Options): Extension[] {
   const { databaseId, engine, referencedCardIds } = useMemo(
     () => ({
       databaseId: Lib.databaseID(query),
@@ -78,10 +84,18 @@ export function useExtensions({ query, onRunQuery }: Options): Extension[] {
           },
         ]),
       ),
+      diff
+        ? unifiedMergeView({
+            original: Lib.rawNativeQuery(query),
+            mergeControls: false,
+          })
+        : null,
     ]
       .flat()
       .filter(isNotNull);
   }, [
+    query,
+    diff,
     engine,
     schemaCompletion,
     snippetCompletion,

frontend/src/metabase/query_builder/components/NativeQueryEditor/NativeQueryEditor.tsx

@@ -7,6 +7,7 @@ import {
   useCallback,
 } from "react";
 import { ResizableBox, type ResizableBoxProps } from "react-resizable";
+import { t } from "ttag";
 import _ from "underscore";
 
 import ExplicitSize from "metabase/common/components/ExplicitSize";
@@ -23,7 +24,7 @@ import {
 import SnippetFormModal from "metabase/query_builder/components/template_tags/SnippetFormModal";
 import type { QueryModalType } from "metabase/query_builder/constants";
 import { useNotebookScreenSize } from "metabase/query_builder/hooks/use-notebook-screen-size";
-import { Flex } from "metabase/ui";
+import { Button, Flex, Icon, Stack, Tooltip } from "metabase/ui";
 import * as Lib from "metabase-lib";
 import type Question from "metabase-lib/v1/Question";
 import type Database from "metabase-lib/v1/metadata/Database";
@@ -32,6 +33,7 @@ import type {
   CardId,
   Collection,
   DatabaseId,
+  DatasetQuery,
   NativeQuerySnippet,
   ParameterId,
 } from "metabase-types/api";
@@ -60,6 +62,10 @@ type OwnProps = {
   question: Question;
   query: NativeQuery;
 
+  proposedQuestion?: Question | undefined;
+  onRejectProposed?: () => void;
+  onAcceptProposed?: (query: DatasetQuery) => void;
+
   nativeEditorSelectedText?: string;
   modalSnippet?: NativeQuerySnippet;
   viewHeight: number;
@@ -240,6 +246,9 @@ class NativeQueryEditor extends Component<Props, NativeQueryEditorState> {
       resizableBoxProps = {},
       snippetCollections = [],
       question,
+      proposedQuestion,
+      onRejectProposed,
+      onAcceptProposed,
       query,
       readOnly,
       isNativeEditorOpen,
@@ -337,6 +346,7 @@ class NativeQueryEditor extends Component<Props, NativeQueryEditorState> {
               <CodeMirrorEditor
                 ref={this.editor}
                 query={question.query()}
+                proposedQuery={proposedQuestion?.query()}
                 readOnly={readOnly}
                 placeholder={placeholder}
                 highlightedLineNumbers={highlightedLineNumbers}
@@ -350,16 +360,54 @@ class NativeQueryEditor extends Component<Props, NativeQueryEditorState> {
                 }
               />
 
-              {hasRunButton && !readOnly && (
-                <NativeQueryEditorRunButton
-                  cancelQuery={this.props.cancelQuery}
-                  isResultDirty={this.props.isResultDirty}
-                  isRunnable={this.props.isRunnable}
-                  isRunning={this.props.isRunning}
-                  nativeEditorSelectedText={this.props.nativeEditorSelectedText}
-                  runQuery={this.props.runQuery}
-                />
-              )}
+              <Stack m="1rem" gap="md" mt="auto">
+                {proposedQuestion && onRejectProposed && onAcceptProposed && (
+                  <>
+                    <Tooltip label={t`Accept proposed changes`} position="top">
+                      <Button
+                        variant="filled"
+                        bg="success"
+                        px="0"
+                        w="2.5rem"
+                        onClick={() => {
+                          // TODO: clean this up so the condition is not needed
+                          const proposedQuery =
+                            proposedQuestion.legacyNativeQuery();
+                          if (proposedQuery) {
+                            this.onChange(proposedQuery.queryText());
+                            onAcceptProposed(proposedQuery.datasetQuery());
+                          }
+                        }}
+                      >
+                        <Icon name="check" />
+                      </Button>
+                    </Tooltip>
+                    <Tooltip label={t`Reject proposed changes`} position="top">
+                      <Button
+                        w="2.5rem"
+                        px="0"
+                        variant="filled"
+                        bg="danger"
+                        onClick={onRejectProposed}
+                      >
+                        <Icon name="close" />
+                      </Button>
+                    </Tooltip>
+                  </>
+                )}
+                {hasRunButton && !readOnly && (
+                  <NativeQueryEditorRunButton
+                    cancelQuery={this.props.cancelQuery}
+                    isResultDirty={this.props.isResultDirty}
+                    isRunnable={this.props.isRunnable}
+                    isRunning={this.props.isRunning}
+                    nativeEditorSelectedText={
+                      this.props.nativeEditorSelectedText
+                    }
+                    runQuery={this.props.runQuery}
+                  />
+                )}
+              </Stack>
             </Flex>
           </ResizableBox>
         </div>

frontend/src/metabase/query_builder/components/NativeQueryEditor/NativeQueryEditorRunButton/NativeQueryEditorRunButton.module.css

@@ -1,6 +1,4 @@
 .RunButtonWithTooltipStyled {
-  margin: 1rem;
-  margin-top: auto;
   height: 40px;
   width: 40px;
 }

frontend/test/__support__/server-mocks/collection.ts

@@ -192,11 +192,12 @@ export function setupCollectionByIdEndpoint({
     return;
   }
 
-  fetchMock.get(/api\/collection\/\d+$/, (call) => {
+  fetchMock.get(/api\/collection\/(\d+|root)$/, (call) => {
     const urlString = call.url;
     const parts = urlString.split("/");
     const collectionIdParam = parts[parts.length - 1];
-    const collectionId = Number(collectionIdParam);
+    const collectionId =
+      collectionIdParam === "root" ? "root" : Number(collectionIdParam);
 
     const collection = collections.find(
       (collection) => collection.id === collectionId,

frontend/test/__support__/server-mocks/metabot.ts

@@ -1,14 +1,13 @@
 import fetchMock, { type UserRouteConfig } from "fetch-mock";
 
 import type {
-  MetabotApiEntity,
   MetabotId,
   MetabotInfo,
   SuggestedMetabotPrompt,
   SuggestedMetabotPromptsResponse,
 } from "metabase-types/api";
 
-export function setupMetabotsEndpoint(
+export function setupMetabotsEndpoints(
   metabots: MetabotInfo[],
   statusCode?: number,
 ) {
@@ -16,30 +15,10 @@ export function setupMetabotsEndpoint(
     "path:/api/ee/metabot-v3/metabot",
     statusCode ? { status: statusCode } : { items: metabots },
   );
-}
-
-export function setupMetabotEntitiesEndpoint(
-  metabotId: MetabotId,
-  entities: MetabotApiEntity[],
-) {
-  fetchMock.get(
-    `path:/api/ee/metabot-v3/metabot/${metabotId}/entities`,
-    {
-      items: entities,
-    },
-    { name: `metabot-${metabotId}-entities-get` },
-  );
-}
-
-export function setupMetabotAddEntitiesEndpoint(metabotId: MetabotId) {
-  fetchMock.put(`path:/api/ee/metabot-v3/metabot/${metabotId}/entities`, {
-    status: 204,
-  });
-}
-
-export function setupMetabotDeleteEntitiesEndpoint() {
-  fetchMock.delete(/api\/ee\/metabot-v3\/metabot\/\d+\/entities/, {
-    status: 204,
+  metabots.forEach((metabot) => {
+    fetchMock.put(`path:/api/ee/metabot-v3/metabot/${metabot.id}`, (call) => {
+      return { ...metabot, ...JSON.parse(call.options?.body as string) };
+    });
   });
 }
 
@@ -52,6 +31,12 @@ export function setupMetabotPromptSuggestionsEndpointError(
   );
 }
 
+export function setupMetabotAddEntitiesEndpoint(metabotId: MetabotId) {
+  fetchMock.put(`path:/api/ee/metabot-v3/metabot/${metabotId}/entities`, {
+    status: 204,
+  });
+}
+
 type SuggestionsEndpointOptions = {
   metabotId: MetabotId;
   prompts: SuggestedMetabotPromptsResponse["prompts"];

package.json

@@ -17,6 +17,7 @@
     "@codemirror/lang-python": "^6.1.7",
     "@codemirror/lang-sql": "^6.8.0",
     "@codemirror/legacy-modes": "^6.4.2",
+    "@codemirror/merge": "^6.10.2",
     "@dnd-kit/core": "^6.0.8",
     "@dnd-kit/modifiers": "^6.0.1",
     "@dnd-kit/sortable": "^8.0.0",

resources/migrations/056_update_migrations.yaml

@@ -373,6 +373,288 @@ databaseChangeLog:
                     AND mt2.name = SUBSTRING(metabase_table.name, 1, POSITION('__mbarchiv__' IN metabase_table.name) - 1)
                     AND mt2.id != metabase_table.id);
 
+  - changeSet:
+      id: v56.2025-08-22T12:00:00
+      author: sloansparger
+      comment: Add use_verified_content column to metabot
+      changes:
+        - addColumn:
+            tableName: metabot
+            columns:
+              - column:
+                  name: use_verified_content
+                  type: ${boolean.type}
+                  remarks: Whether this metabot should only use verified content
+                  constraints:
+                    nullable: true
+                  defaultValueBoolean: false
+
+  - changeSet:
+      id: v56.2025-08-22T12:00:01
+      author: sloansparger
+      comment: Add collection_id column to metabot
+      changes:
+        - addColumn:
+            tableName: metabot
+            columns:
+              - column:
+                  name: collection_id
+                  type: int
+                  remarks: ID of the collection this metabot can access
+                  constraints:
+                    nullable: true
+
+  - changeSet:
+      id: v56.2025-08-22T12:00:02
+      author: sloansparger
+      comment: Add metabot_id column to metabot_prompt
+      changes:
+        - addColumn:
+            tableName: metabot_prompt
+            columns:
+              - column:
+                  name: metabot_id
+                  type: int
+                  remarks: The metabot this prompt is associated with
+                  constraints:
+                    nullable: true
+
+  - changeSet:
+      id: v56.2025-08-22T12:00:03
+      author: sloansparger
+      comment: Migrate collection references from metabot_entity to metabot.collection_id with exception to the default metabot instance
+      changes:
+        - sql:
+            comment: Migrate collection references from metabot_entity to metabot.collection_id with exception to the default metabot instance
+            sql: >-
+              UPDATE metabot
+              SET collection_id = (
+                SELECT me.model_id
+                FROM metabot_entity me
+                WHERE me.metabot_id = metabot.id
+                  AND me.model = 'collection'
+                LIMIT 1
+              )
+              WHERE EXISTS (
+                SELECT 1
+                FROM metabot_entity me
+                WHERE me.metabot_id = metabot.id
+                  AND me.model = 'collection'
+              )
+              AND metabot.entity_id != 'metabotmetabotmetabot';
+      rollback:
+        - sql:
+            comment: Clear collection_id values set by the migration
+            sql: >-
+              UPDATE metabot
+              SET collection_id = NULL
+              WHERE collection_id IS NOT NULL
+              AND metabot.entity_id != 'metabotmetabotmetabot';
+
+  - changeSet:
+      id: v56.2025-08-22T12:00:04
+      author: sloansparger
+      comment: Migrate metabot_prompt.metabot_entity_id to metabot_prompt.metabot_id
+      changes:
+        - sql:
+            comment: Migrate metabot_prompt.metabot_entity_id to metabot_prompt.metabot_id
+            sql: >-
+              UPDATE metabot_prompt
+              SET metabot_id = (
+                SELECT me.metabot_id
+                FROM metabot_entity me
+                WHERE me.id = metabot_prompt.metabot_entity_id
+              );
+      rollback:
+        - sql:
+            comment: Clear metabot_id values set by the migration
+            sql: >-
+              UPDATE metabot_prompt
+              SET metabot_id = NULL
+              WHERE metabot_id IS NOT NULL;
+
+  - changeSet:
+      id: v56.2025-08-22T12:00:05
+      author: sloansparger
+      comment: Add not null constraint to metabot_prompt.metabot_id
+      changes:
+        - addNotNullConstraint:
+            tableName: metabot_prompt
+            columnName: metabot_id
+            columnDataType: int
+
+  - changeSet:
+      id: v56.2025-08-22T12:00:06
+      author: sloansparger
+      comment: Add index on metabot_prompt.metabot_id
+      changes:
+        - createIndex:
+            tableName: metabot_prompt
+            indexName: idx_metabot_prompt_metabot_id
+            columns:
+              - column:
+                  name: metabot_id
+
+  - changeSet:
+      id: v56.2025-08-22T12:00:07
+      author: sloansparger
+      comment: Add foreign key constraint metabot_prompt.metabot_id -> metabot.id
+      changes:
+        - addForeignKeyConstraint:
+            baseTableName: metabot_prompt
+            baseColumnNames: metabot_id
+            referencedTableName: metabot
+            referencedColumnNames: id
+            constraintName: fk_metabot_prompt_metabot_id
+            onDelete: CASCADE
+
+  - changeSet:
+      id: v56.2025-08-22T12:00:08
+      author: sloansparger
+      comment: Drop foreign key constraint metabot_prompt.metabot_entity_id
+      changes:
+        - dropForeignKeyConstraint:
+            baseTableName: metabot_prompt
+            constraintName: fk_metabot_prompt_metabot_entity_id
+      rollback:
+        - addForeignKeyConstraint:
+            baseTableName: metabot_prompt
+            baseColumnNames: metabot_entity_id
+            referencedTableName: metabot_entity
+            referencedColumnNames: id
+            constraintName: fk_metabot_prompt_metabot_entity_id
+            onDelete: CASCADE
+
+  - changeSet:
+      id: v56.2025-08-22T12:00:09
+      author: sloansparger
+      comment: Drop index on metabot_prompt.metabot_entity_id
+      changes:
+        - dropIndex:
+            tableName: metabot_prompt
+            indexName: idx_metabot_prompt_metabot_entity_id
+      rollback:
+        - createIndex:
+            tableName: metabot_prompt
+            indexName: idx_metabot_prompt_metabot_entity_id
+            columns:
+              - column:
+                  name: metabot_entity_id
+
+  - changeSet:
+      id: v56.2025-08-22T12:00:10
+      author: sloansparger
+      comment: Drop metabot_entity_id column from metabot_prompt
+      changes:
+        - dropColumn:
+            tableName: metabot_prompt
+            columnName: metabot_entity_id
+      rollback:
+        - addColumn:
+            tableName: metabot_prompt
+            columns:
+              - column:
+                  name: metabot_entity_id
+                  type: int
+                  remarks: The metabot entity this prompt is associated with
+                  constraints:
+                    nullable: true
+            preConditions:
+              - onFail: MARK_RAN
+              - not:
+                  - columnExists:
+                      tableName: metabot_prompt
+                      columnName: metabot_entity_id
+        - sql:
+            comment: Restore metabot_prompt.metabot_entity_id values from metabot_prompt.metabot_id
+            sql: >-
+              UPDATE metabot_prompt
+              SET metabot_entity_id = (
+                SELECT me.id
+                FROM metabot_entity me
+                WHERE me.metabot_id = metabot_prompt.metabot_id
+                  AND me.model = 'collection'
+                LIMIT 1
+              );
+        - addNotNullConstraint:
+            tableName: metabot_prompt
+            columnName: metabot_entity_id
+            columnDataType: int
+
+  - changeSet:
+      id: v56.2025-08-22T12:00:11
+      author: sloansparger
+      comment: Remove metabot_entity table as collection association is now direct
+      changes:
+        - dropTable:
+            tableName: metabot_entity
+            cascadeConstraints: true
+            comment: Remove metabot_entity table as collection association is now direct
+      rollback:
+        - createTable:
+            tableName: metabot_entity
+            remarks: Entities associated with a metabot
+            columns:
+              - column:
+                  name: id
+                  type: int
+                  autoIncrement: true
+                  constraints:
+                    primaryKey: true
+                    nullable: false
+              - column:
+                  name: metabot_id
+                  type: int
+                  remarks: The metabot this entity is associated with
+                  constraints:
+                    nullable: false
+              - column:
+                  name: model
+                  type: varchar(32)
+                  remarks: The type of model this entity references
+                  constraints:
+                    nullable: false
+              - column:
+                  name: model_id
+                  type: int
+                  remarks: The ID of the model this entity references
+                  constraints:
+                    nullable: false
+              - column:
+                  name: entity_id
+                  type: char(21)
+                  remarks: Random NanoID tag for unique identity
+                  constraints:
+                    nullable: true
+                    unique: true
+              - column:
+                  remarks: The timestamp of when the entity was created
+                  name: created_at
+                  type: ${timestamp_type}
+                  defaultValueComputed: current_timestamp
+                  constraints:
+                    nullable: false
+        - createIndex:
+            tableName: metabot_entity
+            indexName: idx_metabot_entity_metabot_id
+            columns:
+              - column:
+                  name: metabot_id
+        - addForeignKeyConstraint:
+            baseTableName: metabot_entity
+            baseColumnNames: metabot_id
+            referencedTableName: metabot
+            referencedColumnNames: id
+            constraintName: fk_metabot_entity_metabot_id
+            onDelete: CASCADE
+        - sql:
+            comment: Restore collection references to metabot_entity from metabot.collection_id
+            sql: >-
+              INSERT INTO metabot_entity (metabot_id, model, model_id)
+              SELECT id, 'collection', collection_id
+              FROM metabot
+              WHERE collection_id IS NOT NULL;
+
   - changeSet:
       id: v56.2025-09-03T17:15:00
       author: piranha

src/metabase/cmd/copy.clj

@@ -118,7 +118,6 @@
       :model/ConnectionImpersonation
       :model/Metabot
       :model/MetabotConversation
-      :model/MetabotEntity
       :model/MetabotMessage
       :model/MetabotPrompt
       :model/Document

src/metabase/models/resolution.clj

@@ -50,7 +50,6 @@
     :model/LoginHistory                      metabase.login-history.models.login-history
     :model/Metabot                           metabase-enterprise.metabot-v3.models.metabot
     :model/MetabotConversation               metabase-enterprise.metabot-v3.models.metabot-conversation
-    :model/MetabotEntity                     metabase-enterprise.metabot-v3.models.metabot-entity
     :model/MetabotMessage                    metabase-enterprise.metabot-v3.models.metabot-message
     :model/MetabotPrompt                     metabase-enterprise.metabot-v3.models.metabot-prompt
     :model/ModelIndex                        metabase.indexed-entities.models.model-index

src/metabase/native_query_snippets/api.clj

@@ -14,26 +14,33 @@
 
 (set! *warn-on-reflection* true)
 
-(mu/defn- hydrated-native-query-snippet :- [:maybe (ms/InstanceOf :model/NativeQuerySnippet)]
-  [id :- ms/PositiveInt]
-  (-> (api/read-check (t2/select-one :model/NativeQuerySnippet :id id))
-      (t2/hydrate :creator)))
+(mu/defn list-native-query-snippets :- [:sequential (ms/InstanceOf :model/NativeQuerySnippet)]
+  "List all native query snippets the current user has read access to."
+  ([]
+   (list-native-query-snippets false))
+  ([archived :- ms/BooleanValue]
+   (let [snippets (t2/select :model/NativeQuerySnippet
+                             :archived archived
+                             {:order-by [[:%lower.name :asc]]})]
+     (t2/hydrate (filter mi/can-read? snippets) :creator))))
 
 (api.macros/defendpoint :get "/"
   "Fetch all snippets"
   [_route-params
    {:keys [archived]} :- [:map
                           [:archived {:default false} [:maybe ms/BooleanValue]]]]
-  (let [snippets (t2/select :model/NativeQuerySnippet
-                            :archived archived
-                            {:order-by [[:%lower.name :asc]]})]
-    (t2/hydrate (filter mi/can-read? snippets) :creator)))
+  (list-native-query-snippets (boolean archived)))
+
+(mu/defn get-native-query-snippet :- [:maybe (ms/InstanceOf :model/NativeQuerySnippet)]
+  [id :- ms/PositiveInt]
+  (-> (api/read-check (t2/select-one :model/NativeQuerySnippet :id id))
+      (t2/hydrate :creator)))
 
 (api.macros/defendpoint :get "/:id"
   "Fetch native query snippet with ID."
   [{:keys [id]} :- [:map
                     [:id ms/PositiveInt]]]
-  (hydrated-native-query-snippet id))
+  (get-native-query-snippet id))
 
 (defn- check-snippet-name-is-unique [snippet-name]
   (when (t2/exists? :model/NativeQuerySnippet :name snippet-name)
@@ -72,7 +79,7 @@
       (when-let [new-name (:name changes)]
         (check-snippet-name-is-unique new-name))
       (t2/update! :model/NativeQuerySnippet id changes))
-    (hydrated-native-query-snippet id)))
+    (get-native-query-snippet id)))
 
 (api.macros/defendpoint :put "/:id"
   "Update an existing `NativeQuerySnippet`."

src/metabase/native_query_snippets/core.clj

@@ -1,10 +1,14 @@
 (ns metabase.native-query-snippets.core
   (:require
+   [metabase.native-query-snippets.api]
    [metabase.native-query-snippets.models.native-query-snippet.permissions]
    [potemkin :as p]))
 
 (comment metabase.native-query-snippets.models.native-query-snippet.permissions/keep-me)
 
 (p/import-vars
+ [metabase.native-query-snippets.api
+  list-native-query-snippets
+  get-native-query-snippet]
  [metabase.native-query-snippets.models.native-query-snippet.permissions
   has-any-native-permissions?])

test_resources/serialization_baseline/collections/gQuHAIM4C-7fToHzHyNNQ_ai_model/metabots/gr1EkP92whbdA_TKQljLU.yaml

@@ -0,0 +1,22 @@
+name: Test Metabot
+description: A test metabot
+created_at: '2025-06-09T20:12:13.816701Z'
+entity_id: gr1EkP92whbdA_TKQljLU
+updated_at: '2025-06-09T20:12:13.816701Z'
+use_verified_content: false
+collection_id: gQuHAIM4C-7fToHzHyNNQ
+prompts:
+- card_id: Mr5p0TE2FSxGeUFaWd9yu
+  created_at: '2025-06-09T20:12:13.816701Z'
+  entity_id: U4V5PzzpaMasl1Kt9BvNy
+  model: model
+  prompt: A sample prompt
+  updated_at: '2025-06-09T20:12:13.816701Z'
+  serdes/meta:
+  - id: gr1EkP92whbdA_TKQljLU
+    model: Metabot
+  - id: U4V5PzzpaMasl1Kt9BvNy
+    model: MetabotPrompt
+serdes/meta:
+- id: gr1EkP92whbdA_TKQljLU
+  model: Metabot

test_resources/serialization_baseline/collections/metabots/CyLr7tYRKiOoTWlcbl36q.yaml

@@ -1,19 +1,11 @@
-created_at: '2025-05-21T13:50:15.894507Z'
-description: A test metabot
-entities:
-- created_at: '2025-05-21T13:50:15.894507Z'
-  entity_id: uv52JUhq5eLSSjBAJCK1k
-  model: dataset
-  model_id: RC0EJ8I00lOUnvbnGJJNf
-  serdes/meta:
-  - id: CyLr7tYRKiOoTWlcbl36q
-    model: Metabot
-  - id: uv52JUhq5eLSSjBAJCK1k
-    model: MetabotEntity
-  prompts: []
-entity_id: CyLr7tYRKiOoTWlcbl36q
 name: Test Metabot
+description: A test metabot
+created_at: '2025-05-21T13:50:15.894507Z'
+entity_id: CyLr7tYRKiOoTWlcbl36q
 updated_at: '2025-05-21T13:50:15.894507Z'
+use_verified_content: false
+collection_id: null
+prompts: []
 serdes/meta:
 - id: CyLr7tYRKiOoTWlcbl36q
   model: Metabot

test_resources/serialization_baseline/collections/metabots/OYeHvLi8mMGwwMwJpetYw.yaml

@@ -3,30 +3,20 @@ description: A test metabot
 created_at: '2025-06-09T20:12:07.978442Z'
 entity_id: OYeHvLi8mMGwwMwJpetYw
 updated_at: '2025-06-09T20:12:07.978442Z'
-serdes/meta:
-- id: OYeHvLi8mMGwwMwJpetYw
-  model: Metabot
-entities:
-- created_at: '2025-06-09T20:12:07.978442Z'
-  entity_id: aaBQ4EZe97-T5cVdiz1uD
-  model: dataset
-  model_id: 4ahdpFn43jybRFVQTRjX_
-  prompts:
-  - card_id: 4ahdpFn43jybRFVQTRjX_
-    created_at: '2025-06-09T20:12:07.978442Z'
-    entity_id: SNoFh7dZXzjvgeNp5aJJY
-    model: model
-    prompt: A sample prompt
-    updated_at: '2025-06-09T20:12:07.978442Z'
-    serdes/meta:
-    - id: OYeHvLi8mMGwwMwJpetYw
-      model: Metabot
-    - id: aaBQ4EZe97-T5cVdiz1uD
-      model: MetabotEntity
-    - id: SNoFh7dZXzjvgeNp5aJJY
-      model: MetabotPrompt
+use_verified_content: false
+collection_id: null
+prompts:
+- card_id: 4ahdpFn43jybRFVQTRjX_
+  created_at: '2025-06-09T20:12:07.978442Z'
+  entity_id: SNoFh7dZXzjvgeNp5aJJY
+  model: model
+  prompt: A sample prompt
+  updated_at: '2025-06-09T20:12:07.978442Z'
   serdes/meta:
   - id: OYeHvLi8mMGwwMwJpetYw
     model: Metabot
-  - id: aaBQ4EZe97-T5cVdiz1uD
-    model: MetabotEntity
+  - id: SNoFh7dZXzjvgeNp5aJJY
+    model: MetabotPrompt
+serdes/meta:
+- id: OYeHvLi8mMGwwMwJpetYw
+  model: Metabot

test_resources/serialization_baseline/collections/metabots/embeddedmetabotmetabo.yaml

@@ -1,9 +1,11 @@
-created_at: '2025-05-28T17:49:20.649159Z'
+name: Embedded Metabot
 description: Metabot instance for embedded metabase users.
-entities: []
+created_at: '2025-05-28T17:49:20.649159Z'
 entity_id: embeddedmetabotmetabo
-name: Embeddeded Metabot
 updated_at: '2025-05-28T17:49:20.649159Z'
 serdes/meta:
 - id: embeddedmetabotmetabo
   model: Metabot
+collection_id: null
+prompts: []
+use_verified_content: false

test_resources/serialization_baseline/collections/metabots/gr1EkP92whbdA_TKQljLU.yaml

@@ -1,32 +0,0 @@
-name: Test Metabot
-description: A test metabot
-created_at: '2025-06-09T20:12:13.816701Z'
-entity_id: gr1EkP92whbdA_TKQljLU
-updated_at: '2025-06-09T20:12:13.816701Z'
-serdes/meta:
-- id: gr1EkP92whbdA_TKQljLU
-  model: Metabot
-entities:
-- created_at: '2025-06-09T20:12:13.816701Z'
-  entity_id: lOPH7ciO8ndcpNCzkOCbH
-  model: collection
-  model_id: gQuHAIM4C-7fToHzHyNNQ
-  prompts:
-  - card_id: Mr5p0TE2FSxGeUFaWd9yu
-    created_at: '2025-06-09T20:12:13.816701Z'
-    entity_id: U4V5PzzpaMasl1Kt9BvNy
-    model: model
-    prompt: A sample prompt
-    updated_at: '2025-06-09T20:12:13.816701Z'
-    serdes/meta:
-    - id: gr1EkP92whbdA_TKQljLU
-      model: Metabot
-    - id: lOPH7ciO8ndcpNCzkOCbH
-      model: MetabotEntity
-    - id: U4V5PzzpaMasl1Kt9BvNy
-      model: MetabotPrompt
-  serdes/meta:
-  - id: gr1EkP92whbdA_TKQljLU
-    model: Metabot
-  - id: lOPH7ciO8ndcpNCzkOCbH
-    model: MetabotEntity

test_resources/serialization_baseline/collections/metabots/metabotmetabotmetabot.yaml

@@ -1,9 +1,11 @@
-created_at: '2025-05-28T17:49:20.649159Z'
-description: Metabot instance for internal users.
-entities: []
-entity_id: metabotmetabotmetabot
 name: Metabot
+description: Metabot instance for internal users.
+created_at: '2025-05-28T17:49:20.649159Z'
+entity_id: metabotmetabotmetabot
 updated_at: '2025-05-28T17:49:20.649159Z'
 serdes/meta:
 - id: metabotmetabotmetabot
   model: Metabot
+collection_id: null
+prompts: []
+use_verified_content: false

test_resources/serialization_baseline/collections/metabots/zXxMBaxni6nRGRhD7vmPc.yaml

@@ -1,19 +1,11 @@
-created_at: '2025-05-28T17:49:20.813848Z'
-description: A test metabot
-entities:
-- created_at: '2025-05-28T17:49:20.813848Z'
-  entity_id: _Jc9eftJesJVUU9gFKAXV
-  model: dataset
-  model_id: kDd3G__y2cqdZu7A1SUsU
-  serdes/meta:
-  - id: zXxMBaxni6nRGRhD7vmPc
-    model: Metabot
-  - id: _Jc9eftJesJVUU9gFKAXV
-    model: MetabotEntity
-  prompts: []
-entity_id: zXxMBaxni6nRGRhD7vmPc
 name: Test Metabot
+description: A test metabot
+created_at: '2025-05-28T17:49:20.813848Z'
+entity_id: zXxMBaxni6nRGRhD7vmPc
 updated_at: '2025-05-28T17:49:20.813848Z'
+use_verified_content: false
+collection_id: null
+prompts: []
 serdes/meta:
 - id: zXxMBaxni6nRGRhD7vmPc
   model: Metabot

yarn.lock

@@ -1300,6 +1300,17 @@
     "@codemirror/view" "^6.35.0"
     crelt "^1.0.5"
 
+"@codemirror/merge@^6.10.2":
+  version "6.10.2"
+  resolved "https://registry.yarnpkg.com/@codemirror/merge/-/merge-6.10.2.tgz#b87b27fff45f38c3e3c7a463d8a0c2f2e81c4263"
+  integrity sha512-rmHzVkt5FnCtsi0IgvDIDjh/J4LmbfOboB7FMvVl21IHO0p1QM6jSwjkBjBD3D+c+T79OabEqoduCqvJCBV8Yg==
+  dependencies:
+    "@codemirror/language" "^6.0.0"
+    "@codemirror/state" "^6.0.0"
+    "@codemirror/view" "^6.17.0"
+    "@lezer/highlight" "^1.0.0"
+    style-mod "^4.1.0"
+
 "@codemirror/search@^6.0.0":
   version "6.5.9"
   resolved "https://registry.yarnpkg.com/@codemirror/search/-/search-6.5.9.tgz#08829cf1db9d093dd4822bb22ece93da3ebffefc"