Compare commits

...

35 Commits

Author SHA1 Message Date
Noah Moss
e882da2eed Add debug mode to metabot chat 2025-09-11 15:31:14 -04:00
Sloan Sparger
c3abfc359f half hacky way to use current editor state for context in transform editing 2025-09-11 13:34:31 -05:00
Sloan Sparger
d94134da52 fix transform context on subsequent edits for the same transform 2025-09-11 13:18:15 -05:00
Sloan Sparger
47edcb46e8 misc. fe improvements 2025-09-11 11:55:29 -05:00
Sloan Sparger
20a8644a09 quick fix preview error not displaying correctly 2025-09-10 21:58:24 -05:00
Sloan Sparger
2cde7d4db7 hack in a message to metabot on save 2025-09-10 18:51:43 -05:00
Sloan Sparger
d10a687c4f wip on creation flow 2025-09-10 15:38:54 -05:00
Noah Moss
e2ba497bce Parse out tables from transforms in metabot user context 2025-09-10 16:01:03 -04:00
Sloan Sparger
856815885e use transforms instead of data queries 2025-09-10 12:43:54 -05:00
Sloan Sparger
3eb87b838f add transform context on edit query page 2025-09-09 18:11:19 -05:00
Mike Appleby
5ce546eccd Cleanup malli schemas for get-transforms and get-transform-details 2025-09-09 15:20:58 -07:00
Sloan Sparger
f8fb30a0ef accept tranform suggestions from metabot 2025-09-09 17:20:19 -05:00
Sloan Sparger
a0241f1f67 show metabot on edit + create transforms page 2025-09-09 16:58:02 -05:00
Sloan Sparger
78e998d238 send transform in context 2025-09-09 15:53:35 -05:00
Sloan Sparger
ebfa374011 use transforms_codegen profile for metabot on transforms page 2025-09-09 15:42:31 -05:00
Sloan Sparger
d8e219e7c4 allow opening metabot on transforms admin pages 2025-09-09 15:27:44 -05:00
Mike Appleby
0ef9eeea7a Add /ee/metabot-tools/get-transform-details endpoint 2025-09-09 13:17:58 -07:00
Mike Appleby
c717a99716 Add test that get-transforms does not include non-native MBQL transforms 2025-09-09 12:10:50 -07:00
Mike Appleby
ae1a51758f Add /ee/metabot-tools/get-transforms endpoint
Returns a sequence of all python or native query transforms.

BOT-391 [Backend] Create transform index tool API
2025-09-08 17:15:20 -07:00
Sloan Sparger
e77b0cc5b7 FE - add codemirror diffing, add diff support to native query editor, fake a proposed query 2025-09-08 17:55:17 -05:00
Sloan Sparger
eaa0b647b5 init branch 2025-09-08 11:48:14 -05:00
Sloan Sparger
0fb925ba33 [Answer blending] Respect collection id in search (#63150)
* remove find metric tool

* adjustments to search endpoint

* add tests

* fix failing test in ci

* add way to override the metabot id

* add way to override metabot request id

* fix whitespace lint

* missing paren...

* fix failing test
2025-09-08 11:29:21 -05:00
Tyler Plude
2362eadf28 Merge branch 'master' into tsp-bot-292-metabot-pro-search 2025-09-06 12:26:36 -07:00
Tyler Plude
b16b5b79ba [BOT-284] BE for notebook query tool table support (#62710)
* init

* moved query-table endpoint to query-datasource; add support for both model and table ids

* unit test fixes

* use existing card-query and table-query utils

* adding unit tests
2025-09-05 16:20:36 -07:00
Sloan Sparger
29493ce8ce Revert "fix failing test"
This reverts commit b1c7e9efb9.
2025-09-05 15:54:05 -05:00
Sloan Sparger
b1c7e9efb9 fix failing test 2025-09-05 15:43:32 -05:00
Sloan Sparger
675cdf1b2a [Metabot] Quick slash command + dynamic profile ui (#62909)
* quick slash command + dyanmic profile ui

* add back code i did not mean to remove
2025-08-29 17:27:44 +02:00
Sloan Sparger
6709cd1052 [Answer blending] Updates to the metabot admin configs (#62551)
* remove ability to configure the selected collection for default metabot

* add verified toggle

* add migrations for metabot collection configuration changes

* add metabot PUT endpoint, rip out metabot entities endpoints and usages, convert FE over to new FE, fix tests

* remove metabot_entity table, model, usage and refactor code to use updated metabot and metabot_prompts schemas

* enforce verification and regenerate prompts on verification + collection id change

* fix bad references

* fix bad regexp for fetch-mock collections by id endpoints breaking unit tests

* fix kondo

* fix + increase tests for enterprise/backend/test/metabase_enterprise/metabot_v3/tools/api_test.clj

* fix migration boolean default value syntax

* maybe fix be tests

* fix some tests

* enforce verified filter on the search endpoint

* factor in verified check into find-metric

* hide verification setting behind token feature check, fix migration lint, do not migrate collection id for default metabot, fix kondo errors

* fix core tests

* fix kondo

* fix most of the failing serialization tests

* fix migration lints

* fix one test file

* maybe fix migration failures

* move test resource file

* prevent default metabot from getting a collection id configured

* yolo

* add cascade delete to entity drop table migration

* increase unit tests for verification config UI

* remove extra add column

* fix migration... i think for real this time

* for real for real fix migrations

* no really

* tests fixes and add preventing verification toggle if the instance does not have the content_verification feature

* fix failing feature tests in ci

* fix one more

* fix cljfmt issue

---------

Co-authored-by: Tyler Plude <tsplude@gmail.com>
2025-08-29 17:26:30 +02:00
Thomas Schmidt
dc01024f24 Update Metabot profile IDs (#62864) 2025-08-29 17:26:12 +02:00
Thomas Schmidt
3dd3554f34 Add table reference for implicitly joinable fields in context injection (#62753) 2025-08-29 17:26:12 +02:00
Thomas Schmidt
6a32cb4cf8 Answer blending sql querying via search tables (#62415)
* Only put tables used in SQL into context
* Change format of injected database content by reusing existing models
* Create separate function for new schema enchancement
Co-authored-by: Sloan Sparger <sloansparger@gmail.com>
2025-08-29 17:26:12 +02:00
Sloan Sparger
e66a3b8ac8 fix bad merge conflict 2025-08-29 17:26:12 +02:00
Sloan Sparger
f32fe2dde3 dedupe search terms and description if they happen to be the same 2025-08-29 17:26:12 +02:00
Thomas Schmidt
a63858ecaf More tests and better split out of logic into separate files 2025-08-29 17:26:12 +02:00
Thomas Schmidt
a505563b06 Add search endpoint for Metabot and adjusted metadata 2025-08-29 17:26:12 +02:00
78 changed files with 3776 additions and 1806 deletions

View File

@@ -1784,6 +1784,7 @@
pulse
query-processor
request
search
server
settings
sync
@@ -1792,7 +1793,8 @@
util
warehouse-schema
warehouses
enterprise/documents}}
enterprise/documents
enterprise/transforms}}
enterprise/permission-debug
{:team "UX West"

View File

@@ -81,7 +81,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]

View File

@@ -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

View File

@@ -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

View File

@@ -53,42 +53,49 @@
: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))]
(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))

View File

@@ -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)
@@ -291,7 +306,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)

View File

@@ -18,15 +18,16 @@
;;; --------------------------------------------------- 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 []}))
;;; ------------------------------------------------ Serdes Hashing -------------------------------------------------
(defmethod serdes/hash-fields :model/Metabot
@@ -36,14 +37,15 @@
;;; ------------------------------------------------- Serialization -------------------------------------------------
(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)}})

View File

@@ -1,56 +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})
;;; ------------------------------------------------ Serdes Hashing -------------------------------------------------
(defmethod serdes/hash-fields :model/MetabotEntity
[_table]
[:metabot_id :model :model_id])
;;; ------------------------------------------------- Serialization -------------------------------------------------
(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)}})

View File

@@ -22,12 +22,12 @@
(defmethod serdes/hash-fields :model/MetabotPrompt
[_table]
[:metabot_entity_id :model :card_id :prompt])
[:metabot_id :model :card_id :prompt])
;;; ------------------------------------------------- Serialization -------------------------------------------------
(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]
@@ -35,8 +35,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)}})

View File

@@ -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]

View File

@@ -18,9 +18,10 @@
: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-data-sources :as metabot-v3.tools.search-data-sources]
[metabase-enterprise.metabot-v3.tools.transforms :as metabot-v3.tools.transforms]
[metabase-enterprise.metabot-v3.util :as metabot-v3.u]
[metabase.api.common :as api]
[metabase.api.macros :as api.macros]
@@ -241,6 +242,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
@@ -402,7 +420,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
@@ -466,7 +485,6 @@
[:name :string]
[:email_address :string]]]]
[:map [:output :string]]])
(mr/def ::get-dashboard-details-result
[:or
[:map
@@ -475,7 +493,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
@@ -530,7 +549,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
@@ -604,6 +624,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]]])
@@ -620,6 +643,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
@@ -629,6 +689,79 @@
[: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-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-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-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-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-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-table-result
::search-dashboard-result
::search-question-result
::search-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-result-item]]
[:total_count :int]]]]
[:map [:output :string]]])
(api.macros/defendpoint :post "/answer-sources" :- [:merge ::answer-sources-result ::tool-request]
"Return top level meta information about available information sources."
[_route-params
@@ -697,20 +830,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
@@ -864,6 +983,34 @@
(assoc :conversation_id conversation_id))
(metabot-v3.context/log :llm.log/be->llm))))
(api.macros/defendpoint :get "/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
@@ -896,6 +1043,51 @@
(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)))))
(defn- enforce-authentication
"Middleware that returns a 401 response if no `ai-session` can be found for `request`."
[handler]

View File

@@ -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

View File

@@ -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."}))))

View File

@@ -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))))

View File

@@ -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))))

View File

@@ -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]))))

View File

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

View File

@@ -71,15 +71,20 @@
(api/check-400 (driver.u/supports? (:engine database) feature database)
(deferred-tru "The database does not support the requested transform target type."))))
(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
@@ -106,11 +111,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)
@@ -120,6 +123,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

View File

@@ -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))))))))))))))

View File

@@ -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))))))))

View File

@@ -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"))))))

View File

@@ -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))))))))))

View File

@@ -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*)))

View File

@@ -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}))))))

View File

@@ -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
[]

View File

@@ -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"))))))))))

View File

@@ -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))))))))))

View File

@@ -105,7 +105,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))))))

View File

@@ -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

View File

@@ -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))))

View File

@@ -1,16 +1,23 @@
import * as Yup from "yup";
import type { 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 = [
"state",
"navigate_to",
"transform_suggestion",
];
export type KnownDataPart =
| { type: "state"; version: 1; value: Record<string, any> }
| { type: "navigate_to"; version: 1; value: string };
| { type: "navigate_to"; version: 1; value: string }
| { type: "transform_suggestion"; version: 1; value: SuggestedTransform };
export const toolCallPartSchema = Yup.object({
toolCallId: Yup.string().required(),

View File

@@ -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,

View File

@@ -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,
);

View File

@@ -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 });
});
});
});

View File

@@ -63,7 +63,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.profile
? t`Using profile: ${metabot.profile}`
: t`Metabot isn't perfect. Double-check results.`}
</Text>
</Flex>
@@ -140,6 +142,8 @@ export const MetabotChat = () => {
onRetryMessage={handleRetryMessage}
isDoingScience={metabot.isDoingScience}
showFeedbackButtons
debugMode={metabot.debugMode}
toolCallsByMessage={metabot.toolCallsByMessage}
/>
{/* loading */}

View File

@@ -15,12 +15,14 @@ import {
import type {
MetabotChatMessage,
MetabotErrorMessage,
MetabotToolCall,
} from "metabase-enterprise/metabot/state";
import type { MetabotFeedback } from "metabase-types/api";
import { AIMarkdown } from "../AIMarkdown/AIMarkdown";
import Styles from "./MetabotChat.module.css";
import { MetabotDebugToolCall } from "./MetabotDebugToolCall";
import { MetabotFeedbackModal } from "./MetabotFeedbackModal";
interface BaseMessageProps extends Omit<FlexProps, "onCopy"> {
@@ -111,6 +113,8 @@ interface AgentMessageProps extends BaseMessageProps {
setFeedbackMessage?: (data: { messageId: string; positive: boolean }) => void;
submittedFeedback: "positive" | "negative" | undefined;
onInternalLinkClick?: (link: string) => void;
debugMode?: boolean;
toolCalls?: MetabotToolCall[];
}
export const AgentMessage = ({
@@ -123,6 +127,8 @@ export const AgentMessage = ({
submittedFeedback,
onInternalLinkClick,
hideActions,
debugMode,
toolCalls,
...props
}: AgentMessageProps) => (
<MessageContainer chatRole={message.role} {...props}>
@@ -133,6 +139,15 @@ export const AgentMessage = ({
{message.message}
</AIMarkdown>
{/* Debug tool calls inline */}
{debugMode && toolCalls && toolCalls.length > 0 && (
<div style={{ marginTop: "12px" }}>
{toolCalls.map((toolCall) => (
<MetabotDebugToolCall key={toolCall.id} toolCall={toolCall} />
))}
</div>
)}
<Flex className={Styles.messageActions}>
{!hideActions && (
<>
@@ -241,6 +256,8 @@ export const Messages = ({
isDoingScience,
showFeedbackButtons,
onInternalLinkClick,
debugMode,
toolCallsByMessage,
}: {
messages: MetabotChatMessage[];
errorMessages: MetabotErrorMessage[];
@@ -248,6 +265,8 @@ export const Messages = ({
isDoingScience: boolean;
showFeedbackButtons: boolean;
onInternalLinkClick?: (navigateToPath: string) => void;
debugMode?: boolean;
toolCallsByMessage?: Record<string, MetabotToolCall[]>;
}) => {
const clipboard = useClipboard();
const [sendToast] = useToast();
@@ -309,6 +328,8 @@ export const Messages = ({
submittedFeedback={feedbackState.submitted[message.id]}
hideActions={messages[index + 1]?.role === "agent"}
onInternalLinkClick={onInternalLinkClick}
debugMode={debugMode}
toolCalls={toolCallsByMessage?.[message.id] || []}
/>
) : (
<UserMessage

View File

@@ -0,0 +1,118 @@
import { useState } from "react";
import { t } from "ttag";
import { color } from "metabase/lib/colors";
import { Accordion, Box, Code, Stack, Text } from "metabase/ui";
import type { MetabotToolCall } from "metabase-enterprise/metabot/state";
interface MetabotDebugToolCallProps {
toolCall: MetabotToolCall;
}
export const MetabotDebugToolCall = ({
toolCall,
}: MetabotDebugToolCallProps) => {
const [isExpanded, setIsExpanded] = useState(false);
return (
<Box
p="sm"
mb="xs"
style={{
backgroundColor: color("brand-light"),
borderLeft: `4px solid ${color("brand")}`,
borderRadius: "8px",
}}
data-testid="metabot-debug-tool-call"
>
<Stack gap="xs">
<Text size="sm" fw="bold" c="brand">
{t`🔧 Tool Call: ${toolCall.name}`}
</Text>
<Text size="sm" c="text-medium">
{t`Status: ${toolCall.status === "started" ? t`Running...` : t`Completed`}`}
</Text>
{toolCall.message && (
<Text size="sm" c="text-medium">
{toolCall.message}
</Text>
)}
<Accordion
value={isExpanded ? "details" : null}
onChange={(value) => setIsExpanded(value === "details")}
>
<Accordion.Item value="details">
<Accordion.Control>
<Text size="xs" c="text-medium">
{t`Show Request/Response Details`}
</Text>
</Accordion.Control>
<Accordion.Panel>
<Stack gap="md">
{toolCall.args && (
<Box>
<Text size="xs" fw="bold" c="text-medium" mb="xs">
{t`Request Arguments:`}
</Text>
<Code
block
style={{
backgroundColor: color("bg-light"),
maxHeight: "150px",
overflow: "auto",
fontSize: "11px",
}}
>
{typeof toolCall.args === "string"
? toolCall.args
: JSON.stringify(toolCall.args, null, 2)}
</Code>
</Box>
)}
{toolCall.result !== undefined && (
<Box>
<Text size="xs" fw="bold" c="text-medium" mb="xs">
{t`Response Result:`}
</Text>
<Code
block
style={{
backgroundColor: color("bg-light"),
maxHeight: "150px",
overflow: "auto",
fontSize: "11px",
}}
>
{typeof toolCall.result === "string"
? toolCall.result
: JSON.stringify(toolCall.result, null, 2)}
</Code>
</Box>
)}
<Box>
<Text size="xs" fw="bold" c="text-medium" mb="xs">
{t`Full Tool Call Object:`}
</Text>
<Code
block
style={{
backgroundColor: color("bg-light"),
maxHeight: "200px",
overflow: "auto",
fontSize: "11px",
}}
>
{JSON.stringify(toolCall, null, 2)}
</Code>
</Box>
</Stack>
</Accordion.Panel>
</Accordion.Item>
</Accordion>
</Stack>
</Box>
);
};

View File

@@ -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?`;

View File

@@ -7,6 +7,8 @@ import { useMetabotContext } from "metabase/metabot";
import {
type MetabotPromptSubmissionResult,
getAgentErrorMessages,
getAllToolCalls,
getDebugMode,
getIsLongMetabotConversation,
getIsProcessing,
getLastAgentMessagesByType,
@@ -14,9 +16,12 @@ import {
getMetabotId,
getMetabotRequestId,
getMetabotVisible,
getProfile,
getToolCalls,
getToolCallsByMessage,
resetConversation as resetConversationAction,
retryPrompt,
setDebugMode,
setVisible as setVisibleAction,
submitInput as submitInputAction,
} from "../state";
@@ -66,6 +71,22 @@ export const useMetabotAgent = () => {
[dispatch],
);
const profile = useSelector(getProfile as any) as ReturnType<
typeof getProfile
>;
const debugMode = useSelector(getDebugMode as any) as ReturnType<
typeof getDebugMode
>;
const allToolCalls = useSelector(getAllToolCalls as any) as ReturnType<
typeof getAllToolCalls
>;
const toolCallsByMessage = useSelector(
getToolCallsByMessage as any,
) as ReturnType<typeof getToolCallsByMessage>;
const resetConversation = useCallback(
() => dispatch(resetConversationAction()),
[dispatch],
@@ -141,6 +162,11 @@ export const useMetabotAgent = () => {
[resetConversation, setVisible, promptInputRef, submitInput],
);
const toggleDebugMode = useCallback(
() => dispatch(setDebugMode(!debugMode)),
[dispatch, debugMode],
);
return {
isDoingScience,
prompt,
@@ -158,5 +184,10 @@ export const useMetabotAgent = () => {
submitInput,
retryMessage,
toolCalls,
profile,
debugMode,
allToolCalls,
toolCallsByMessage,
toggleDebugMode,
};
};

View File

@@ -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";
@@ -32,7 +33,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 +46,10 @@ export const {
setNavigateToPath,
toolCallStart,
toolCallEnd,
setProfileOverride,
setMetabotReqIdOverride,
setSuggestedTransform,
setDebugMode,
} = metabot.actions;
type PromptErrorOutcome = {
@@ -90,6 +96,43 @@ export const setVisible =
dispatch(metabot.actions.setVisible(isVisible));
};
export const executeSlashCommand = createAsyncThunk<void, SlashCommand>(
"metabase-enterprise/metabot/executeSlashCommand",
async (slashCommand, { dispatch, getState }) => {
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>" }));
}
})
.with({ cmd: "debug" }, () => {
const state = getState() as any;
const currentDebugMode =
state.plugins?.metabotPlugin?.debugMode ?? false;
dispatch(setDebugMode(!currentDebugMode));
dispatch(
addUndo({
message: currentDebugMode
? "Debug mode disabled"
: "Debug mode enabled",
}),
);
})
.otherwise(() => {
dispatch(addUndo({ message: "Unknown command" }));
});
},
);
export type MetabotPromptSubmissionResult =
| { prompt: string; success: true; shouldRetry?: void }
| { prompt: string; success: false; shouldRetry: false }
@@ -120,11 +163,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 }));
const sendMessageRequestPromise = dispatch(
sendAgentRequest({
...data,
@@ -205,6 +255,16 @@ export const sendAgentRequest = createAsyncThunk<
dispatch(push(part.value) as UnknownAction);
}
})
.with(
{ type: "transform_suggestion" },
({ value: transform }) => {
dispatch(setSuggestedTransform(transform));
const url = transform.id
? `/admin/transforms/${transform.id}/query`
: `/admin/transforms/new/native`;
dispatch(push(url) as UnknownAction);
},
)
.exhaustive();
},
onTextPart: (part) => {

View File

@@ -3,7 +3,7 @@ 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, SuggestedTransform } from "metabase-types/api";
import { TOOL_CALL_MESSAGES } from "../constants";
@@ -26,10 +26,14 @@ export type MetabotToolCall = {
name: string;
message: string | undefined;
status: "started" | "ended";
args?: string;
result?: any;
messageId?: string; // Associates tool call with a specific message
};
export type MetabotReactionsState = {
navigateToPath: string | null;
suggestedTransform: SuggestedTransform | undefined;
};
export interface MetabotState {
@@ -42,6 +46,12 @@ export interface MetabotState {
state: any;
reactions: MetabotReactionsState;
toolCalls: MetabotToolCall[];
debugMode: boolean;
allToolCalls: MetabotToolCall[];
experimental: {
metabotReqIdOverride: string | undefined;
profileOverride: string | undefined;
};
}
export const getMetabotInitialState = (): MetabotState => ({
@@ -54,8 +64,15 @@ export const getMetabotInitialState = (): MetabotState => ({
state: {},
reactions: {
navigateToPath: null,
suggestedTransform: undefined,
},
toolCalls: [],
debugMode: false,
allToolCalls: [],
experimental: {
metabotReqIdOverride: undefined,
profileOverride: undefined,
},
});
export const metabot = createSlice({
@@ -108,19 +125,40 @@ export const metabot = createSlice({
},
toolCallStart: (
state,
action: PayloadAction<{ toolCallId: string; toolName: string }>,
action: PayloadAction<{
toolCallId: string;
toolName: string;
args?: string;
}>,
) => {
const { toolCallId, toolName } = action.payload;
state.toolCalls.push({
const { toolCallId, toolName, args } = action.payload;
// Find the last agent message to associate this tool call with
const lastAgentMessage = [...state.messages]
.reverse()
.find((m) => m.role === "agent");
const toolCall = {
id: toolCallId,
name: toolName,
message: TOOL_CALL_MESSAGES[toolName],
status: "started",
});
status: "started" as const,
args,
messageId: lastAgentMessage?.id,
};
state.toolCalls.push(toolCall);
state.allToolCalls.push(toolCall);
},
toolCallEnd: (state, action: PayloadAction<{ toolCallId: string }>) => {
toolCallEnd: (
state,
action: PayloadAction<{ toolCallId: string; result?: any }>,
) => {
const { toolCallId, result } = action.payload;
state.toolCalls = state.toolCalls.map((tc) =>
tc.id === action.payload.toolCallId ? { ...tc, status: "ended" } : tc,
tc.id === toolCallId ? { ...tc, status: "ended", result } : tc,
);
state.allToolCalls = state.allToolCalls.map((tc) =>
tc.id === toolCallId ? { ...tc, status: "ended", result } : tc,
);
},
// NOTE: this reducer fn should be made smarter if/when we want to have
@@ -148,6 +186,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 +198,27 @@ 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;
},
setDebugMode: (state, action: PayloadAction<boolean>) => {
state.debugMode = action.payload;
},
},
extraReducers: (builder) => {
builder

View File

@@ -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,
@@ -25,6 +27,11 @@ export const getMessages = createSelector(
(metabot) => metabot.messages,
);
export const getToolCalls = createSelector(
getMetabot,
(metabot) => metabot.toolCalls,
);
export const getLastMessage = createSelector(getMessages, (messages) =>
_.last(messages),
);
@@ -66,11 +73,6 @@ export const getLastAgentMessagesByType = createSelector(
},
);
export const getToolCalls = createSelector(
getMetabot,
(metabot) => metabot.toolCalls,
);
export const getIsProcessing = createSelector(
getMetabot,
(metabot) => metabot.isProcessing,
@@ -110,20 +112,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 +162,42 @@ 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;
},
);
export const getDebugMode = createSelector(
getMetabot,
(metabot) => metabot.debugMode,
);
export const getAllToolCalls = createSelector(
getMetabot,
(metabot) => metabot.allToolCalls,
);
export const getToolCallsByMessage = createSelector(
getAllToolCalls,
(toolCalls) => {
const toolCallsByMessage: Record<string, typeof toolCalls> = {};
toolCalls.forEach((toolCall) => {
if (toolCall.messageId) {
if (!toolCallsByMessage[toolCall.messageId]) {
toolCallsByMessage[toolCall.messageId] = [];
}
toolCallsByMessage[toolCall.messageId].push(toolCall);
}
});
return toolCallsByMessage;
},
);

View File

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

View File

@@ -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(" ") : [],
};
};

View File

@@ -28,6 +28,9 @@ type EditorBodyProps = {
onRunQuery: () => Promise<void>;
onCancelQuery: () => void;
databases: ApiDatabase[];
proposedQuestion?: Question;
onRejectProposed?: () => void;
onAcceptProposed?: (query: DatasetQuery) => void;
};
export function EditorBody({
@@ -40,6 +43,9 @@ export function EditorBody({
onRunQuery,
onCancelQuery,
databases,
proposedQuestion,
onRejectProposed,
onAcceptProposed,
}: EditorBodyProps) {
const [isResizing, setIsResizing] = useState(false);
const reportTimezone = useSetting("report-timezone-long");
@@ -94,6 +100,9 @@ export function EditorBody({
const database = databases.find((database) => database.id === db.id);
return !doesDatabaseSupportTransforms(database);
}}
proposedQuestion={proposedQuestion}
onRejectProposed={onRejectProposed}
onAcceptProposed={onAcceptProposed}
/>
) : (
<ResizableBox

View File

@@ -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
@@ -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`;
}
}

View File

@@ -4,7 +4,7 @@ import { useListDatabasesQuery } from "metabase/api";
import { Center, Flex, Loader } from "metabase/ui";
import * as Lib from "metabase-lib";
import type Question from "metabase-lib/v1/Question";
import type { DatasetQuery } from "metabase-types/api";
import type { DatasetQuery, Transform } from "metabase-types/api";
import { useQueryMetadata } from "../../hooks/use-query-metadata";
import { useQueryResults } from "../../hooks/use-query-results";
@@ -16,21 +16,32 @@ import { EditorVisualization } from "./EditorVisualization";
import S from "./QueryEditor.module.css";
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,
@@ -40,12 +51,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 = () => {
@@ -84,6 +96,7 @@ export function QueryEditor({
data-testid="transform-query-editor"
>
<EditorHeader
name={transform?.name}
isNew={isNew}
isSaving={isSaving}
canSave={canSave && (isNew || isQueryDirty)}
@@ -92,6 +105,7 @@ export function QueryEditor({
/>
<EditorBody
question={question}
proposedQuestion={proposedQuestion}
isNative={isNative}
isRunnable={isRunnable}
isRunning={isRunning}
@@ -100,6 +114,8 @@ export function QueryEditor({
onRunQuery={runQuery}
onCancelQuery={cancelQuery}
databases={databases?.data ?? []}
onRejectProposed={onRejectProposed}
onAcceptProposed={onAcceptProposed}
/>
<EditorVisualization
question={question}

View File

@@ -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;

View File

@@ -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,14 @@ export function useQueryState(initialQuery: DatasetQuery) {
[query, metadata],
);
const proposedQuestion = useMemo(
() =>
proposedQuery
? Question.create({ dataset_query: proposedQuery, metadata })
: undefined,
[proposedQuery, metadata],
);
const isQueryDirty = useMemo(
() => !Lib.areLegacyQueriesEqual(query, initialQuery),
[query, initialQuery],
@@ -24,5 +35,5 @@ export function useQueryState(initialQuery: DatasetQuery) {
setQuery(newQuestion.datasetQuery());
};
return { question, isQueryDirty, setQuestion };
return { question, proposedQuestion, isQueryDirty, setQuestion };
}

View File

@@ -27,14 +27,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) {
@@ -43,6 +54,7 @@ export function CreateTransformModal({
<FocusTrap.InitialFocus />
<CreateTransformForm
query={query}
initValues={initValues}
onCreate={onCreate}
onClose={onClose}
/>
@@ -52,26 +64,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) {
@@ -98,8 +98,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) {
@@ -158,12 +158,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,
};
}

View File

@@ -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";
@@ -66,22 +71,62 @@ 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,
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>
);
}

View File

@@ -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";
@@ -44,6 +45,12 @@ export function TransformPage({ params }: TransformPageProps) {
setIsPolling(!isPolling);
}
useRegisterMetabotContextProvider(async () => {
return transform
? { user_is_viewing: [{ type: "transform", ...transform }] }
: {};
}, [transform]);
if (isLoading || error != null) {
return <LoadingAndErrorWrapper loading={isLoading} error={error} />;
}

View File

@@ -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,6 +61,38 @@ export function TransformQueryPageBody({
const dispatch = useDispatch();
const { sendErrorToast, sendSuccessToast } = useMetadataToasts();
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(
"HIDDEN MESSAGE: user has rejected your changes, ask for clarification on what they'd like to do instead.",
);
};
const onAcceptProposed = async (query: DatasetQuery) => {
await handleSave(query);
dispatch(setSuggestedTransform(undefined));
metabot.submitInput(
"HIDDEN MESSAGE: user has accepted your changes, move to the next step!",
);
};
const handleSave = async (query: DatasetQuery) => {
const { error } = await updateTransform({
id: transform.id,
@@ -75,12 +115,19 @@ export function TransformQueryPageBody({
};
return (
<QueryEditor
initialQuery={transform.source.query}
isNew={false}
isSaving={isLoading}
onSave={handleSave}
onCancel={handleCancel}
/>
<AdminSettingsLayout fullWidthContent>
<QueryEditor
initialQuery={initialQuery}
transform={transform}
isNew={false}
isSaving={isLoading}
onSave={handleSave}
onChange={setLatestQuery}
onCancel={handleCancel}
proposedQuery={proposedQuery}
onRejectProposed={onRejectProposed}
onAcceptProposed={onAcceptProposed}
/>
</AdminSettingsLayout>
);
}

View File

@@ -2,12 +2,12 @@ import type {
CardDisplayType,
CardId,
CardType,
CollectionId,
DashboardId,
DatasetQuery,
PaginationResponse,
RowValue,
SearchModel,
SuggestedTransform,
Transform,
UnsavedCard,
Version,
} from ".";
@@ -135,11 +135,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,6 +155,7 @@ export type MetabotAgentRequest = {
state: MetabotStateContext;
conversation_id: string; // uuid
metabot_id?: string;
profile_id?: string;
};
export type MetabotAgentResponse = {
@@ -218,19 +224,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 */

View File

@@ -15,6 +15,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";

View File

@@ -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,
});

View File

@@ -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;

View File

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

View File

@@ -1,6 +1,9 @@
import type React from "react";
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 +13,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 +35,19 @@ 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} />
</Box>
</Box>
);

View File

@@ -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%"

View File

@@ -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,

View File

@@ -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,
@@ -338,6 +347,7 @@ class NativeQueryEditor extends Component<Props, NativeQueryEditorState> {
<CodeMirrorEditor
ref={this.editor}
query={question.query()}
proposedQuery={proposedQuestion?.query()}
readOnly={readOnly}
placeholder={placeholder}
highlightedLineNumbers={highlightedLineNumbers}
@@ -351,16 +361,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) {
setDatasetQuery(proposedQuery);
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>

View File

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

View File

@@ -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,

View File

@@ -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"];

View File

@@ -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",

View File

@@ -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;
# >>>>>>>>>> DO NOT ADD NEW MIGRATIONS BELOW THIS LINE! ADD THEM ABOVE <<<<<<<<<<
########################################################################################################################

View File

@@ -116,7 +116,6 @@
[:model/GroupTableAccessPolicy
:model/ConnectionImpersonation
:model/Metabot
:model/MetabotEntity
:model/MetabotPrompt
:model/Document
:model/DocumentBookmark])))

View File

@@ -47,7 +47,6 @@
:model/ImplicitAction metabase.actions.models
:model/LoginHistory metabase.login-history.models.login-history
:model/Metabot metabase-enterprise.metabot-v3.models.metabot
:model/MetabotEntity metabase-enterprise.metabot-v3.models.metabot-entity
:model/MetabotPrompt metabase-enterprise.metabot-v3.models.metabot-prompt
:model/ModelIndex metabase.indexed-entities.models.model-index
:model/ModelIndexValue metabase.indexed-entities.models.model-index

View File

@@ -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

View File

@@ -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

View File

@@ -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

View File

@@ -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

View File

@@ -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

View File

@@ -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

View File

@@ -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

View File

@@ -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"