Kimirock/metabase
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Retrieve Library Entities tool using OSI metadata (#76394)

Browse Source
This commit is contained in:
Chris Truter
2026-06-30 21:15:53 +02:00
committed by GitHub
parent b113cf9527
commit 78382c9a96
72 changed files with 3655 additions and 1532 deletions
Download Patch File Download Diff File Expand all files Collapse all files

5
.clj-kondo/config.edn
View File

@@ -448,7 +448,7 @@
instaparse.core/defparser clojure.core/def
metabase-enterprise.action-v2.test-util/with-test-tables! clojure.core/let
metabase-enterprise.content-translation.routes-test/with-new-secret-key! clojure.core/fn
metabase-enterprise.curated-search.reconcile-test/with-isolated-mirror clojure.core/fn
metabase-enterprise.entity-retrieval.reconcile-test/with-isolated-index clojure.core/fn
metabase-enterprise.serialization.metadata-file-import.wire-format-test/with-tiny-fixture! clojure.core/fn
metabase-enterprise.serialization.test-util/with-dbs clojure.core/fn
metabase-enterprise.serialization.test-util/with-random-dump-dir clojure.core/let
@@ -465,7 +465,7 @@
metabase.channel.impl.http-test/with-captured-http-requests clojure.core/fn
metabase.channel.impl.http-test/with-server clojure.core/let
metabase.collections.models.collection-test/with-collection-in-location clojure.core/let
metabase.curated-search.api-test/with-test-entry clojure.core/let
metabase.collections.test-utils/with-library clojure.core/fn
metabase.driver.mongo.connection/with-mongo-client clojure.core/let
metabase.driver.mongo.connection/with-mongo-database clojure.core/let
metabase.driver.mongo.query-processor/mongo-let clojure.core/let
@@ -479,6 +479,7 @@
metabase.notification.test-util/with-card-notification clojure.core/let
metabase.notification.test-util/with-system-event-notification! clojure.core/let
metabase.notification.test-util/with-temp-notification clojure.core/let
metabase.osi.ai-context.api-test/with-test-entry clojure.core/let
metabase.premium-features.core/defenterprise-schema clj-kondo.lint-as/def-catch-all
metabase.premium-features.defenterprise/defenterprise-schema clj-kondo.lint-as/def-catch-all
metabase.premium-features.settings/define-premium-feature clojure.core/def

60
.clj-kondo/config/modules/config.edn
View File

@@ -687,18 +687,6 @@
:api #{}
:uses :any}
curated-search
{:team "Metabot"
:api #{metabase.curated-search.api
metabase.curated-search.core}
:uses #{api
models
premium-features
request
util}
:model-exports #{:model/CuratedSearchEntry}
:model-imports #{:model/Card}}
dashboards
{:team "UX West"
:api #{metabase.dashboards.autoplace
@@ -975,6 +963,13 @@
util}
:model-imports #{:model/Card :model/Dashboard :model/DashboardCard :model/EmbeddingTheme}}
entity-retrieval
{:team "Metabot"
:api #{metabase.entity-retrieval.core
metabase.entity-retrieval.mirror}
:uses #{premium-features}
:model-imports #{:model/OsiAiContext}}
events
{:team "DevEx"
:api #{metabase.events.core
@@ -1283,9 +1278,9 @@
channel
collections
config
curated-search
documents
driver
entity-retrieval
interestingness
legacy-mbql
lib
@@ -1489,6 +1484,20 @@
tracing
util}}
osi
{:team "Metabot"
:api #{metabase.osi.ai-context.api}
:uses #{api
app-db
entity-retrieval
models
request
util}
:model-exports #{:model/OsiAiContext}
:model-imports #{:model/Card
:model/Measure
:model/Segment}}
;; TODO -- too many API namespaces
parameters
{:team "Graphy"
@@ -2982,15 +2991,6 @@
:api #{metabase-enterprise.core.init}
:uses :any}
enterprise/curated-search
{:team "Metabot"
:api #{metabase-enterprise.curated-search.init}
:uses #{premium-features
task
util
enterprise/semantic-search}
:model-imports #{:model/CuratedSearchEntry}}
enterprise/custom-viz-plugin
{:team "Gadget"
:api #{metabase-enterprise.custom-viz-plugin.api
@@ -3167,6 +3167,22 @@
;; make token tracking a hook). Drop this dep once the relocation lands.
:uses #{enterprise/semantic-search}}
enterprise/entity-retrieval
{:team "Metabot"
:api #{metabase-enterprise.entity-retrieval.init}
:uses #{analytics-interface
collections
entity-retrieval
premium-features
task
util
enterprise/semantic-search}
:model-imports #{:model/Card
:model/Measure
:model/OsiAiContext
:model/Segment
:model/Table}}
enterprise/erd
{:team "Gadget"
:api #{metabase-enterprise.erd.api}

6
.gitattributes vendored
View File

@@ -7,3 +7,9 @@
#
# TODO (Bryan 12/12/25) -- reenable this when we can completely verify correctness
# resources/migrations/*_update_migrations.yaml merge=yaml-migrations
# Generated API documentation/specs — regenerated by the "Auto-update documentation for master"
# CI automation, not edited by hand. Mark generated so tools/humans don't review or diff them as source.
docs/api.json linguist-generated
docs/api.html linguist-generated
resources/openapi/openapi.json linguist-generated

8
.github/workflows/semantic-search.yml vendored
View File

@@ -34,7 +34,7 @@ jobs:
- name: Semantic Search Tests
build-static-viz: false
test-args: >-
:only '["enterprise/backend/test/metabase_enterprise/semantic_search" "enterprise/backend/test/metabase_enterprise/curated_search" "enterprise/backend/test/metabase_enterprise/metabot_v3/tools/semantic_search_test.clj"]'
:only '["enterprise/backend/test/metabase_enterprise/semantic_search" "enterprise/backend/test/metabase_enterprise/entity_retrieval" "enterprise/backend/test/metabase_enterprise/metabot_v3/tools/semantic_search_test.clj"]'
exclude-tag: ':mb/driver-tests'
services:
@@ -93,7 +93,7 @@ jobs:
- name: Semantic Search Tests
build-static-viz: false
test-args: >-
:only '["enterprise/backend/test/metabase_enterprise/semantic_search" "enterprise/backend/test/metabase_enterprise/curated_search" "enterprise/backend/test/metabase_enterprise/metabot_v3/tools/semantic_search_test.clj"]'
:only '["enterprise/backend/test/metabase_enterprise/semantic_search" "enterprise/backend/test/metabase_enterprise/entity_retrieval" "enterprise/backend/test/metabase_enterprise/metabot_v3/tools/semantic_search_test.clj"]'
exclude-tag: ':mb/driver-tests'
services:
@@ -166,7 +166,7 @@ jobs:
- name: Semantic Search Tests
build-static-viz: false
test-args: >-
:only '["enterprise/backend/test/metabase_enterprise/semantic_search" "enterprise/backend/test/metabase_enterprise/curated_search" "enterprise/backend/test/metabase_enterprise/metabot_v3/tools/semantic_search_test.clj"]'
:only '["enterprise/backend/test/metabase_enterprise/semantic_search" "enterprise/backend/test/metabase_enterprise/entity_retrieval" "enterprise/backend/test/metabase_enterprise/metabot_v3/tools/semantic_search_test.clj"]'
exclude-tag: ':mb/driver-tests'
services:
@@ -239,7 +239,7 @@ jobs:
- name: Semantic Search Tests
build-static-viz: false
test-args: >-
:only '["enterprise/backend/test/metabase_enterprise/semantic_search" "enterprise/backend/test/metabase_enterprise/curated_search" "enterprise/backend/test/metabase_enterprise/metabot_v3/tools/semantic_search_test.clj"]'
:only '["enterprise/backend/test/metabase_enterprise/semantic_search" "enterprise/backend/test/metabase_enterprise/entity_retrieval" "enterprise/backend/test/metabase_enterprise/metabot_v3/tools/semantic_search_test.clj"]'
exclude-tag: ':mb/driver-tests'
services:

2
enterprise/backend/src/metabase_enterprise/core/init.clj
View File

@@ -10,11 +10,11 @@
[metabase-enterprise.advanced-config.init]
[metabase-enterprise.audit-app.init]
[metabase-enterprise.cache.init]
[metabase-enterprise.curated-search.init]
[metabase-enterprise.custom-viz-plugin.init]
[metabase-enterprise.data-complexity-score.init]
[metabase-enterprise.database-replication.init]
[metabase-enterprise.dependencies.init]
[metabase-enterprise.entity-retrieval.init]
[metabase-enterprise.gsheets.init]
[metabase-enterprise.memoize-monitor.init]
[metabase-enterprise.metabot.init]

120
enterprise/backend/src/metabase_enterprise/curated_search/core.clj
View File

@@ -1,120 +0,0 @@
(ns metabase-enterprise.curated-search.core
"Enterprise implementation of the curated search mirror: pgvector-backed similarity search over the
curated `curated_search_entries` table, and the write-path nudge that keeps the mirror fresh.
OSS shims live in [[metabase.curated-search.mirror]]; the background sync they nudge is the
[[metabase-enterprise.curated-search.task.sync]] Quartz job."
(:require
[clojurewerkz.quartzite.jobs :as jobs]
[honey.sql :as sql]
[honey.sql.helpers :as sql.helpers]
[metabase-enterprise.curated-search.index-table :as index-table]
[metabase-enterprise.semantic-search.db.datasource :as semantic.db.datasource]
[metabase-enterprise.semantic-search.embedding :as embedding]
[metabase.premium-features.core :as premium-features :refer [defenterprise]]
[metabase.task.core :as task]
[metabase.util.json :as json]
[next.jdbc :as jdbc]
[next.jdbc.result-set :as jdbc.rs])
(:import
(java.sql SQLException)
(org.postgresql.util PGobject)))
(set! *warn-on-reflection* true)
(def sync-job-key
"Quartz job key of the background sync; [[request-sync!]] triggers it and
[[metabase-enterprise.curated-search.task.sync]] schedules it."
(jobs/key "metabase-enterprise.curated-search.sync.job"))
(defn pgvector-configured?
"True when a pgvector store is configured (read once at boot from MB_PGVECTOR_DB_URL).
The sync task gates scheduling on this rather than [[available?]], so the periodic safety net exists
even when the semantic-search feature is turned on after startup (a common onboarding flow)."
[]
(string? (not-empty semantic.db.datasource/db-url)))
(defn available?
"Whether the curated search mirror can run right now: a pgvector store is configured and the license
includes semantic search.
The feature check can flip at runtime (token entered post-boot), so callers re-evaluate per use."
[]
(and (pgvector-configured?)
(premium-features/has-feature? :semantic-search)))
(defenterprise request-sync!
"Trigger the background sync job to reconcile the mirror soon.
Fire-and-forget: [[metabase.task.core/trigger-now!]] swallows scheduler errors, and the job's periodic
schedule covers anything a lost trigger misses."
:feature :semantic-search
[]
(when (available?)
(task/trigger-now! sync-job-key)))
;; Weighted scoring shaped like regular search (metabase.search.scoring): each factor contributes
;; weight * score; similarity (1 - cosine distance) is primary, verified a flat 0/1 boost.
(def ^:private similarity-weight 1.0)
(def ^:private verified-weight 0.1)
(def ^:private default-limit 10)
(defn- scorer [nm score weight]
(let [score (double score) weight (double weight)]
{:name nm :score score :weight weight :contribution (* score weight)}))
(defn- score
"Build a weighted-scorer breakdown for one row, shaped like regular search's scoring.
Returns a `:scores` vector of `{:name :score :weight :contribution}` plus the weighted-sum
`:total_score`."
[{:keys [distance verified]}]
(let [scores [(scorer :similarity (- 1.0 (double distance)) similarity-weight)
(scorer :verified (if verified 1.0 0.0) verified-weight)]]
{:scores scores
:total_score (reduce + (map :contribution scores))}))
(defn- decode-entity [v]
(cond
(instance? PGobject v) (json/decode (.getValue ^PGobject v) true)
(string? v) (json/decode v true)
:else v))
(defenterprise search
"Find the saved prompts nearest to `user-search-prompt` by cosine distance, up to `limit`.
Results are ranked by blended score (similarity + verified boost), each shaped
`{:saved_search_prompt :usage_instructions :entity :score}`.
Returns [] when the pgvector store is unconfigured."
:feature :semantic-search
[user-search-prompt limit]
(if-not (available?)
[]
(let [pgvector (semantic.db.datasource/ensure-initialized-data-source!)
limit (or limit default-limit)
model (embedding/get-configured-model)
embedding (embedding/get-embedding model user-search-prompt
{:type :query :record-tokens? true})
lit (index-table/format-embedding embedding)
;; Order by the blended score, not raw distance, so a verified row whose boost lifts it into
;; the top N survives the LIMIT. No HNSW index could accelerate this expression — see
;; index-table for why the table has none.
ranking (format "(embedding <=> %s) - (CASE WHEN verified THEN %s ELSE 0.0 END)"
lit verified-weight)
rows (try
(jdbc/execute!
pgvector
(-> (sql.helpers/select :search_prompt :usage_instructions :entity :verified
[[:raw (str "embedding <=> " lit)] :distance])
(sql.helpers/from (keyword index-table/*vectors-table*))
(sql.helpers/order-by [[:raw ranking] :asc])
(sql.helpers/limit limit)
(sql/format {:quoted true}))
{:builder-fn jdbc.rs/as-unqualified-lower-maps})
;; 42P01 = undefined table: the background sync hasn't created the mirror yet.
(catch SQLException e
(if (= "42P01" (.getSQLState e)) [] (throw e))))]
(->> rows
(map (fn [row]
{:saved_search_prompt (:search_prompt row)
:usage_instructions (:usage_instructions row)
:entity (decode-entity (:entity row))
:score (score row)}))
(sort-by (comp :total_score :score) >)
vec))))

5
enterprise/backend/src/metabase_enterprise/curated_search/init.clj
View File

@@ -1,5 +0,0 @@
(ns metabase-enterprise.curated-search.init
"Loads the curated-search module's side-effecting namespaces (the background sync task) at
system startup. See [[metabase-enterprise.core.init]]."
(:require
[metabase-enterprise.curated-search.task.sync]))

112
enterprise/backend/src/metabase_enterprise/curated_search/reconcile.clj
View File

@@ -1,112 +0,0 @@
(ns metabase-enterprise.curated-search.reconcile
"Reconciles the pgvector mirror with the authoritative appdb `curated_search_entries` table.
Rather than streaming individual writes (the table is curated and small — hundreds to low thousands of
rows), [[reconcile!]] diffs the whole table against the mirror by content hash on every run: rows whose
hash is missing or stale are re-embedded and upserted, mirror rows with no appdb counterpart are
deleted, unchanged rows cost nothing.
Because each run re-derives the mirror from the appdb, any missed write — pgvector downtime, a crashed
node, an import that bypassed the model hooks — is repaired on the next run with no operator action.
Runs from the [[metabase-enterprise.curated-search.task.sync]] Quartz job; callers supply the
datasource and embedding model so this namespace reads no settings itself."
(:require
[buddy.core.hash :as buddy-hash]
[honey.sql :as sql]
[honey.sql.helpers :as sql.helpers]
[metabase-enterprise.curated-search.index-table :as index-table]
[metabase-enterprise.semantic-search.embedding :as embedding]
[metabase.util :as u]
[metabase.util.json :as json]
[metabase.util.log :as log]
[next.jdbc :as jdbc]
[next.jdbc.result-set :as jdbc.rs]
[toucan2.core :as t2]))
(set! *warn-on-reflection* true)
(def ^:private embed-batch-size
"Rows embedded per embedding-service request while backfilling stale mirror rows."
50)
(defn content-hash
"Hash of an appdb row's mirror-relevant fields.
Stored on the mirror row at upsert time; a mismatch (or absence) marks the row stale."
[{:keys [search_prompt usage_instructions entity verified]}]
(u/encode-base64-bytes
(buddy-hash/sha1
(json/encode (sorted-map :search_prompt search_prompt
:usage_instructions (or usage_instructions "")
:entity entity
:verified (boolean verified))))))
(defn- mirror-hashes
"Map of index_id -> content_hash for every row currently in the mirror."
[pgvector]
(into {}
(map (juxt :index_id :content_hash))
(jdbc/execute! pgvector
[(format "SELECT index_id, content_hash FROM \"%s\"" index-table/*vectors-table*)]
{:builder-fn jdbc.rs/as-unqualified-lower-maps})))
(defn- delete-rows! [pgvector ids]
(when (seq ids)
(jdbc/execute! pgvector
(-> (sql.helpers/delete-from (keyword index-table/*vectors-table*))
(sql.helpers/where [:in :index_id ids])
(sql/format {:quoted true})))))
(defn- row->record [row hash embedding]
{:index_id (:id row)
:search_prompt (:search_prompt row)
:usage_instructions (or (:usage_instructions row) "")
:entity [:cast (json/encode (:entity row)) :jsonb]
:verified (boolean (:verified row))
:content_hash hash
:embedding [:raw (index-table/format-embedding embedding)]})
(defn- upsert-batch!
"Embed one batch of stale appdb rows and upsert them into the mirror.
Returns the number of rows upserted."
[pgvector embedding-model rows]
(let [embeddings (embedding/get-embeddings-batch embedding-model (map :search_prompt rows)
{:type :index :record-tokens? true})
records (map (fn [row embedding] (row->record row (content-hash row) embedding))
rows embeddings)]
(jdbc/execute! pgvector
(-> (sql.helpers/insert-into (keyword index-table/*vectors-table*))
(sql.helpers/values (vec records))
(sql.helpers/on-conflict :index_id)
(sql.helpers/do-update-set :search_prompt :usage_instructions :entity
:verified :content_hash :embedding)
(sql/format {:quoted true})))
(count rows)))
(defn reconcile!
"Bring the pgvector mirror in line with the appdb table; see the namespace docstring for the approach.
Embedding-service calls are scoped to changed rows only.
A failed batch is logged and skipped — its rows stay stale and are retried on the next run.
Returns {:upserted n :deleted n :unchanged n}."
[pgvector embedding-model]
(index-table/ensure-tables! pgvector embedding-model)
(let [appdb-rows (t2/select [:model/CuratedSearchEntry :id :search_prompt :usage_instructions
:entity :verified])
mirrored (mirror-hashes pgvector)
stale (remove #(= (content-hash %) (get mirrored (:id %))) appdb-rows)
orphans (remove (set (map :id appdb-rows)) (keys mirrored))
upserted (transduce
(partition-all embed-batch-size)
(completing
(fn [n batch]
(+ n (try
(upsert-batch! pgvector embedding-model batch)
(catch Exception e
(log/error e "curated search mirror: failed to upsert batch of"
(count batch) "rows; will retry next run")
0)))))
0
stale)]
(delete-rows! pgvector orphans)
{:upserted upserted
:deleted (count orphans)
:unchanged (- (count appdb-rows) (count stale))}))

59
enterprise/backend/src/metabase_enterprise/curated_search/task/sync.clj
View File

@@ -1,59 +0,0 @@
(ns metabase-enterprise.curated-search.task.sync
"Quartz job running the curated search mirror reconciliation.
Scheduled periodically as a safety net; appdb writes also trigger it immediately via
[[metabase-enterprise.curated-search.core/request-sync!]], so curator edits become searchable
within seconds.
`DisallowConcurrentExecution` plus Quartz's trigger coalescing debounce edit bursts to at most one
queued extra run, and Quartz clustering ensures a single node runs it."
(:require
[clojurewerkz.quartzite.jobs :as jobs]
[clojurewerkz.quartzite.schedule.simple :as simple]
[clojurewerkz.quartzite.triggers :as triggers]
[metabase-enterprise.curated-search.core :as curated-search.core]
[metabase-enterprise.curated-search.reconcile :as reconcile]
[metabase-enterprise.semantic-search.db.datasource :as semantic.db.datasource]
[metabase-enterprise.semantic-search.embedding :as embedding]
[metabase.task.core :as task]
[metabase.util.log :as log])
(:import
(java.time Duration Instant)
(java.util Date)))
(set! *warn-on-reflection* true)
(task/defjob ^{org.quartz.DisallowConcurrentExecution true
:doc "Reconciles the curated search pgvector mirror with the appdb table."}
CuratedSearchEntrySync [_ctx]
(when (curated-search.core/available?)
(try
(let [result (reconcile/reconcile! (semantic.db.datasource/ensure-initialized-data-source!)
(embedding/get-configured-model))]
(when (pos? (+ (:upserted result) (:deleted result)))
(log/info "curated search mirror reconciled" result)))
(catch Throwable e
;; Log and move on: the next periodic run retries from the authoritative appdb table.
(log/error e "curated search mirror reconciliation failed")))))
(def ^:private ^Duration startup-delay (Duration/parse "PT10S"))
(def ^:private ^Duration run-frequency (Duration/parse "PT30S"))
(defmethod task/init! ::CuratedSearchEntrySync [_]
;; Gate scheduling on pgvector being configured (boot-fixed), NOT on available? — the job body
;; self-gates on the feature flag, so scheduling here lets the periodic safety net (and the
;; write-path trigger's target job) exist even when the license is enabled after startup.
(when (curated-search.core/pgvector-configured?)
(let [job (jobs/build
(jobs/of-type CuratedSearchEntrySync)
(jobs/store-durably)
(jobs/with-identity curated-search.core/sync-job-key))
trigger (triggers/build
(triggers/with-identity
(triggers/key "metabase-enterprise.curated-search.sync.trigger"))
(triggers/for-job curated-search.core/sync-job-key)
(triggers/start-at (Date/from (.plus (Instant/now) startup-delay)))
(triggers/with-schedule
(simple/schedule
(simple/with-interval-in-milliseconds (.toMillis run-frequency))
(simple/repeat-forever))))]
(task/schedule-task! job trigger))))

292
enterprise/backend/src/metabase_enterprise/entity_retrieval/core.clj Normal file
View File

@@ -0,0 +1,292 @@
(ns metabase-enterprise.entity-retrieval.core
"Enterprise implementation of the library entity index: pgvector-backed similarity search over the
per-value `library_entity_index` documents, and the write-path nudge that keeps the index fresh.
OSS shims live in [[metabase.entity-retrieval.mirror]]; the background sync they nudge is the
[[metabase-enterprise.entity-retrieval.task.sync]] Quartz job."
(:require
[clojure.string :as str]
[clojurewerkz.quartzite.jobs :as jobs]
[honey.sql :as sql]
[honey.sql.helpers :as sql.helpers]
[metabase-enterprise.entity-retrieval.index-table :as index-table]
[metabase-enterprise.entity-retrieval.reconcile :as reconcile]
[metabase-enterprise.semantic-search.db.datasource :as semantic.db.datasource]
[metabase-enterprise.semantic-search.embedding :as embedding]
[metabase.analytics-interface.core :as analytics]
[metabase.premium-features.core :as premium-features :refer [defenterprise]]
[metabase.util :as u]
[metabase.util.log :as log]
[next.jdbc :as jdbc]
[next.jdbc.result-set :as jdbc.rs])
(:import
(java.sql SQLException)))
(set! *warn-on-reflection* true)
(def sync-job-key
"Quartz job key of the periodic full-reconcile backstop, scheduled by
[[metabase-enterprise.entity-retrieval.task.sync]]."
(jobs/key "metabase-enterprise.entity-retrieval.sync.job"))
(defn pgvector-configured?
"True when a pgvector store is configured (read once at boot from MB_PGVECTOR_DB_URL).
The sync task gates scheduling on this rather than [[available?]], so the periodic safety net exists
even when the semantic-search feature is turned on after startup (a common onboarding flow)."
[]
(string? (not-empty semantic.db.datasource/db-url)))
(defn available?
"Whether the entity-retrieval mirror can run right now: a pgvector store is configured and the license
includes semantic search.
The feature check can flip at runtime (token entered post-boot), so callers re-evaluate per use."
[]
(and (pgvector-configured?)
(premium-features/has-feature? :semantic-search)))
(defn- index-ready?
"Whether the library entity index can serve a query right now: its meta row matches the configured
embedding model and schema version, and it holds at least one document.
A mismatched meta row (a model/dimension/format change whose rebuild hasn't run yet), a missing or empty
vectors table, or any query error all read as not-ready.
The compatibility check is what keeps a stale index from being offered: querying a vectors table built for
a different model returns nothing (or errors) rather than answering, so the agent must get the
general-search fallback instead."
[]
(boolean
(try
(let [ds (semantic.db.datasource/ensure-initialized-data-source!)]
(and (index-table/index-compatible? ds (embedding/get-configured-model))
(seq (jdbc/execute! ds [(format "SELECT 1 FROM \"%s\" LIMIT 1" index-table/*vectors-table*)]))))
(catch Throwable _ false))))
;; OSS-callable surface used to decide whether to OFFER the curated retrieve_library_entities tool: it must
;; be able to actually answer, so beyond config + license the index has to be built for the current model and
;; populated. (The write/reconcile path gates on the looser [[available?]] instead, so an empty or stale
;; index can still be rebuilt.) Runs unconditionally rather than gating on :feature — the OSS fallback
;; `false` already covers the unlicensed case.
(defenterprise entity-retrieval-available?
"EE impl: pgvector configured + semantic-search licensed AND the index is ready (built for the current
model and populated), so the curated tool is offered only when it can serve a query (otherwise the agent
gets the general-search fallback)."
:feature :none
[]
(and (available?) (index-ready?)))
;;; ----------------------------------------- Reconcile scheduling -----------------------------------------
;;;
;;; Two independent two-slot coalescing schedules, both mutated only under run-lock and both serialized on
;;; the appdb side by reconcile's pg advisory lock so they never race on the index:
;;; - full-* drives full reconciles (force-reconcile API, periodic backstop, startup); blocking.
;;; - targeted-* drives the per-entity write path; fire-and-forget, draining `dirty-entities`.
;;; Each is the "current run + one queued follow-up" pattern: a caller never joins the in-flight run (it may
;;; predate the caller's write); it starts a run when idle, or queues (or joins) the single follow-up, which
;;; begins only after the in-flight run finishes. So every write is covered by a run that starts after it,
;;; and a burst collapses to one extra run. Keeping the schedules separate (rather than one with a
;;; full-vs-targeted flag) means a blocking full caller always derefs a run that did a full reconcile.
(defonce ^:private run-lock (Object.))
(defonce ^:private full-current (atom nil))
(defonce ^:private full-next (atom nil))
(defonce ^:private targeted-current (atom nil))
(defonce ^:private targeted-next (atom nil))
(defonce ^:private dirty-entities (atom #{}))
(defn- record-run!
"Emit metrics and a log line for one completed reconcile. `scope` is \"full\" or \"targeted\"; the index
size gauges come only from a full run, which counts the whole index."
[scope {:keys [inserted deleted unchanged rebuilt? documents entities]} ran-ms]
(analytics/observe! :metabase-entity-retrieval/reconcile-duration-ms {:scope scope} ran-ms)
(when (pos? (long (or inserted 0))) (analytics/inc! :metabase-entity-retrieval/docs-inserted (long inserted)))
(when (pos? (long (or deleted 0))) (analytics/inc! :metabase-entity-retrieval/docs-deleted (long deleted)))
(when documents (analytics/set-gauge! :metabase-entity-retrieval/index-documents documents))
(when entities (analytics/set-gauge! :metabase-entity-retrieval/index-entities entities))
(when (or (pos? (long (or inserted 0))) (pos? (long (or deleted 0))) rebuilt?)
(log/info "library entity index reconciled"
{:scope scope :inserted inserted :deleted deleted :unchanged unchanged
:rebuilt? rebuilt? :documents documents :entities entities :ran_ms ran-ms})))
(defn- run-loop
"Future for the two-slot schedule held in `cur`/`nxt`: await `predecessor` (nil to start now), call
`(work scheduled-timer)`, then promote the queued follow-up. Returns `work`'s value (the full path's
result map; nil for the targeted path)."
[cur nxt predecessor work]
(future
(let [scheduled (u/start-timer)]
(when predecessor
;; A failed predecessor must not block the follow-up — its writes still need reconciling.
(try @predecessor (catch Throwable _ nil)))
(try
(work scheduled)
(finally
(locking run-lock
(reset! cur @nxt)
(reset! nxt nil)))))))
(defn- schedule!
"Schedule (or join) a run of `work` on the `cur`/`nxt` two-slot schedule, returning its future. Must be
called holding run-lock; deref the future outside the lock."
[cur nxt work]
(cond
(nil? @cur) (let [f (run-loop cur nxt nil work)] (reset! cur f) f)
(nil? @nxt) (let [f (run-loop cur nxt @cur work)] (reset! nxt f) f)
:else @nxt))
(defn- elapsed-ms ^long [timer]
(Math/round ^double (u/since-ms timer)))
(defn- do-full-run
"Run a full reconcile, record its metrics, and return {:index {...} :execution {...}}. The model is
resolved under the reconcile lock (see [[reconcile/reconcile!]]), so a run that waited out a concurrent
node uses the current model rather than one captured before the wait."
[scheduled]
(let [ds (semantic.db.datasource/ensure-initialized-data-source!)
waited-ms (elapsed-ms scheduled)
started (u/start-timer)
diff (reconcile/reconcile! ds embedding/get-configured-model)
ran-ms (elapsed-ms started)]
(record-run! "full" diff ran-ms)
{:index (select-keys diff [:inserted :deleted :unchanged])
:execution {:waited_ms waited-ms :ran_ms ran-ms}}))
(defn- do-targeted-run
"Targeted-reconcile each currently-dirty entity, recording per-entity metrics.
The datasource is resolved *before* the dirty set is snapshot-and-cleared, so a datasource failure leaves
the entries in place for the next run rather than dropping them. Clearing the set up front (vs after the
loop) means a write arriving mid-run re-dirties and is picked up by the next run. A per-entity reconcile
failure re-dirties that entity, so a later run (or the periodic backstop) retries it instead of losing
the write to the slow backstop."
[_scheduled]
(let [ds (semantic.db.datasource/ensure-initialized-data-source!)
dirty (locking run-lock (let [d @dirty-entities] (reset! dirty-entities #{}) d))]
(doseq [[entity-type entity-local-id :as entity-key] dirty]
(try
(let [started (u/start-timer)
diff (reconcile/reconcile-entity! ds embedding/get-configured-model entity-type entity-local-id)]
(record-run! "targeted" diff (elapsed-ms started)))
(catch Throwable e
(log/error e "library entity index: targeted reconcile failed; re-queuing"
entity-type entity-local-id)
(locking run-lock (swap! dirty-entities conj entity-key)))))))
(defn reconcile-full-coalesced!
"Run a full reconcile through the full-reconcile schedule, blocking until a run covering this call
finishes; returns {:index {:inserted n :deleted n :unchanged n} :execution {:waited_ms _ :ran_ms _}}.
Used by the force-reconcile API and the periodic backstop job. Callers must have checked [[available?]]."
[]
@(locking run-lock (schedule! full-current full-next do-full-run)))
(defenterprise force-reconcile!
"Force a full reconcile of the `library_entity_index` against the appdb, blocking until a run covering
this call finishes. Returns {:index {:inserted n :deleted n :unchanged n} :execution {:waited_ms _
:ran_ms _}}, or nil when entity retrieval isn't available (no pgvector store configured).
See [[reconcile-full-coalesced!]] for the coalescing semantics it shares with the periodic backstop."
:feature :semantic-search
[]
(when (available?)
(reconcile-full-coalesced!)))
(defenterprise request-entity-sync!
"Mark one library entity dirty and ensure the targeted write path reconciles its index slice soon.
Fire-and-forget: never blocks, throws, or does embedding/pgvector work on the calling (appdb-write)
thread — the reconcile runs later on a future. Covers `osi_ai_context` edits; membership / name /
description changes to the underlying entity are caught by the periodic full reconcile."
:feature :semantic-search
[entity-type entity-local-id]
(when (available?)
(try
(locking run-lock
(swap! dirty-entities conj [entity-type entity-local-id])
(schedule! targeted-current targeted-next do-targeted-run))
(catch Throwable _ nil))
nil))
(defenterprise library-entity-keys
"Live set of `[entity_type entity_local_id]` for entities currently in the library (see the OSS shim)."
:feature :semantic-search
[]
(reconcile/library-entity-keys))
(def ^:private default-limit 10)
(def ^:private weights
"Blended-score weights, keyed like regular search's scorer weights (a namespaced keyword per variant).
Similarity is the base; the slight per-doc_type bump lets a name match edge out a synonym on a tie."
{:similarity 1.0
:doc-type/name 0.02
:doc-type/synonym 0.01})
(defn- doc-type-boost ^double [doc-type]
(double (get weights (keyword "doc-type" doc-type) 0.0)))
(defn- scorer [nm score weight]
(let [score (double score) weight (double weight)]
{:name nm :score score :weight weight :contribution (* score weight)}))
(defn- score
"Weighted-scorer breakdown for one row: a `:scores` vector of `{:name :score :weight :contribution}`
plus the weighted-sum `:total_score`."
[{:keys [distance doc_type]}]
(let [scores [(scorer :similarity (- 1.0 (double distance)) (:similarity weights))
(scorer :doc-type 1.0 (doc-type-boost doc_type))]]
{:scores scores
:total_score (reduce + (map :contribution scores))}))
(defn- ranking-sql
"`distance - doc_type_boost`: a boosted doc_type sorts earlier even when slightly farther."
[distance-expr]
(let [cases (str/join " " (for [[k w] weights :when (= "doc-type" (namespace k))]
(format "WHEN doc_type = '%s' THEN %s" (name k) w)))]
(format "(%s) - (CASE %s ELSE 0.0 END)" distance-expr cases)))
(defenterprise search
"Find the library-entity documents best matching `user-search-prompt`, up to `limit`, ranked by a
blended score (cosine similarity plus a slight doc_type bump).
Each result is shaped `{:entity {:model :id} :doc_type :doc_text :score}`, best score first; the caller
dedupes the (many-per-entity) docs down to distinct entities.
Returns [] when the pgvector store is unconfigured."
:feature :semantic-search
[user-search-prompt limit]
(if-not (available?)
[]
(let [pgvector (semantic.db.datasource/ensure-initialized-data-source!)
limit (or limit default-limit)
model (embedding/get-configured-model)
embedding (embedding/get-embedding model user-search-prompt
{:type :query :record-tokens? true})
lit (index-table/format-embedding embedding)
distance (str "doc_embedding <=> " lit)
rows (try
(jdbc/execute!
pgvector
(-> (sql.helpers/select :entity_type :entity_local_id :doc_type :doc_text
[[:raw distance] :distance])
(sql.helpers/from (keyword index-table/*vectors-table*))
;; Exact scan, no HNSW: the blended order-by can't use an ANN index, and the
;; curated set is tiny.
(sql.helpers/order-by [[:raw (ranking-sql distance)] :asc])
(sql.helpers/limit limit)
(sql/format {:quoted true}))
{:builder-fn jdbc.rs/as-unqualified-lower-maps})
(catch SQLException e
;; Only the index-not-ready states degrade to no results — the index, not the agent,
;; is at fault and the next reconcile heals it:
;; 42P01 vectors table doesn't exist yet (pre-first-build; expected at boot, stay quiet)
;; 22000 stored vectors incompatible with the query vector (a dimension/format change
;; awaiting rebuild — e.g. a vector(OLD) column vs a vector(NEW) literal)
;; Any other SQL error — a connection loss or pgvector outage — propagates so the tool
;; reports the search as unavailable rather than as an empty library.
(case (.getSQLState e)
"42P01" []
"22000" (do (log/warn e "library entity index incompatible with the query vector; returning no results")
[])
(do (analytics/inc! :metabase-entity-retrieval/search-failed)
(throw e)))))]
(->> rows
(map (fn [row]
{:entity {:model (:entity_type row) :id (:entity_local_id row)}
:doc_type (:doc_type row)
:doc_text (:doc_text row)
:score (score row)}))
(sort-by (comp :total_score :score) >)
vec))))

88
enterprise/backend/src/metabase_enterprise/curated_search/index_table.clj → enterprise/backend/src/metabase_enterprise/entity_retrieval/index_table.clj
View File

@@ -1,22 +1,24 @@
(ns metabase-enterprise.curated-search.index-table
"Lifecycle of the pgvector tables backing the curated search mirror.
(ns metabase-enterprise.entity-retrieval.index-table
"Lifecycle of the pgvector tables backing the library entity index.
Two tables live in the pgvector store:
- the vectors table: one row per appdb `curated_search_entries` row, carrying the embedding, a copy of
the searchable fields, and a `content_hash` used by the reconciler to detect stale rows;
- the vectors table (`library_entity_index`): many rows per library entity — one *embedded* document
per value (the entity's name, its description, each ai_context synonym, each ai_context example).
Each row carries `doc_embedding`, the `doc_text` it was embedded from, and flat `entity_type` /
`entity_local_id` columns.
The primary key `doc_id` is content-addressed over the embedded text (see [[reconcile]]), so a text
edit mints a new row and the reconciler GCs the old one.
Curator `instructions` are NOT stored here — the tool reads them live from `osi_ai_context` at query
time.
- a single-row meta table recording the embedding model identity and schema version the vectors table
was built for.
There is deliberately no HNSW index: the search path orders by the blended score (distance minus the
verified boost), which an HNSW index can't accelerate anyway, and an exact scan over a curated table
of this size is microseconds.
[[ensure-tables!]] is the only entry point: it idempotently creates both tables and, when the
configured embedding model or the schema version no longer matches the meta row, drops and recreates
the vectors table so the next reconcile re-embeds everything.
That drop-and-rebuild is the entire model-change story — the appdb table is authoritative and small,
so rebuilding from it is cheap."
That drop-and-rebuild is the entire model-change story — the index is rebuilt from the authoritative
appdb (library membership + `osi_ai_context`), so rebuilding is cheap."
(:require
[clojure.string :as str]
[honey.sql :as sql]
@@ -30,17 +32,24 @@
(def ^:dynamic *vectors-table*
"Vectors table name. Dynamic so tests can rebind it to an isolated table."
"curated_search_index")
"library_entity_index")
(def ^:dynamic *meta-table*
"Meta table name. Dynamic so tests can rebind it to an isolated table."
"curated_search_index_meta")
"library_entity_index_meta")
(def schema-version
"Version of the vectors table schema.
Bump it to force a drop-and-rebuild on instances whose meta row records an older version."
;; 2: dropped the HNSW index — the rebuild sheds it from tables built under version 1.
2)
"Canonical version of the index's *document format* — both the vectors table schema and the
doc-derivation contract (doc_id scheme, doc_type set, doc_text source, dedup/key rules).
It's part of the meta row's [[model-identity]], so bumping it makes [[ensure-tables!]] drop and rebuild
the vectors table; the post-upgrade startup reconcile then repopulates from the appdb under the new
format. Bump on ANY format-affecting change in [[metabase-enterprise.entity-retrieval.reconcile]] or
the table schema: a vectors-table column/type change, a new or renamed doc_type, a changed doc_text
source, or a changed doc_id / dedup / key scheme — anything that makes old rows incomparable to newly
derived desired docs. A bump forces a full re-embed of the library on every instance at upgrade, so do
it only when the format truly moved, never as a refresh convenience."
;; v1 — initial schema.
1)
;; Advisory lock serializing concurrent ensure-tables! calls (e.g. several cluster nodes starting at
;; once). Arbitrary app-wide-unique constant; semantic-search's migration lock uses 19991.
@@ -73,13 +82,12 @@
(defn- create-vectors-table-sql [dims]
(-> (sql.helpers/create-table (keyword *vectors-table*) :if-not-exists)
(sql.helpers/with-columns
[[:index_id :bigint [:primary-key]]
[:search_prompt :text :not-null]
[:usage_instructions :text :not-null]
[:entity :jsonb :not-null]
[:verified :boolean [:default false] :not-null]
[:content_hash :text :not-null]
[:embedding [:raw (format "vector(%d)" dims)] :not-null]])
[[:doc_id :text [:primary-key]]
[:entity_type :text :not-null]
[:entity_local_id :bigint :not-null]
[:doc_type :text :not-null]
[:doc_text :text :not-null]
[:doc_embedding [:raw (format "vector(%d)" dims)] :not-null]])
sql-format-quoted))
(defn- model-identity
@@ -91,7 +99,9 @@
:schema_version schema-version})
(defn- read-meta [tx]
(jdbc/execute-one! tx [(format "SELECT provider, model_name, vector_dimensions, schema_version FROM \"%s\" WHERE id = 1" *meta-table*)]
(jdbc/execute-one! tx
[(format "SELECT provider, model_name, vector_dimensions, schema_version FROM \"%s\" WHERE id = 1"
*meta-table*)]
{:builder-fn jdbc.rs/as-unqualified-lower-maps}))
(defn- write-meta! [tx embedding-model]
@@ -112,12 +122,31 @@
(defn- create-tables! [tx dims]
(jdbc/execute! tx (create-vectors-table-sql dims)))
(defn- vectors-table-exists? [tx]
(some? (:exists (jdbc/execute-one! tx [(format "SELECT to_regclass('%s') AS exists" *vectors-table*)]
{:builder-fn jdbc.rs/as-unqualified-lower-maps}))))
(defn index-compatible?
"Whether the meta row matches `embedding-model` and the current [[schema-version]] — i.e. the vectors
table holds embeddings the configured model can be queried against.
False when the meta row is missing or describes a different provider/model/dimensions/format, meaning the
index is stale and a rebuild is still pending; the same comparison [[ensure-tables!]] uses to decide a
rebuild.
Reads the meta table directly, so a caller must guard against the table not existing yet (it throws a SQL
error before the first [[ensure-tables!]] of the process)."
[pgvector embedding-model]
(= (read-meta pgvector) (model-identity embedding-model)))
(defn ensure-tables!
"Idempotently create the vectors + meta tables for `embedding-model`, rebuilding on model change.
When the meta row's model identity (provider, model name, dimensions) or schema version no longer
matches, the vectors table is dropped and recreated empty — the next reconcile re-embeds every row.
Serialized across nodes with an advisory lock. Returns :created, :rebuilt or :ok."
Serialized across nodes with an advisory lock.
Returns :created (the vectors table was just created empty — first build, or healing a manual drop),
:rebuilt (dropped and recreated empty for a model/format change), or :ok (already present, untouched).
A caller doing a targeted reconcile must treat :created and :rebuilt alike: both leave an empty table
that a one-entity write can't fill."
[pgvector embedding-model]
(jdbc/with-transaction [tx pgvector]
(jdbc/execute! tx [(format "SELECT pg_advisory_xact_lock(%d)" ensure-lock-id)])
@@ -139,6 +168,9 @@
:rebuilt)
:else
;; Re-issue the IF NOT EXISTS DDL so a manually dropped vectors table heals itself.
(do (create-tables! tx dims)
:ok)))))
;; Meta matches. Re-issue the IF NOT EXISTS DDL so a manually dropped vectors table heals itself,
;; and report :created when it had actually gone missing — the recreated table is empty, so a
;; targeted reconcile must repopulate the whole library rather than fill it with one entity.
(let [existed? (vectors-table-exists? tx)]
(create-tables! tx dims)
(if existed? :ok :created))))))

5
enterprise/backend/src/metabase_enterprise/entity_retrieval/init.clj Normal file
View File

@@ -0,0 +1,5 @@
(ns metabase-enterprise.entity-retrieval.init
"Loads the entity-retrieval module's side-effecting namespaces (the background sync task) at
system startup. See [[metabase-enterprise.core.init]]."
(:require
[metabase-enterprise.entity-retrieval.task.sync]))

455
enterprise/backend/src/metabase_enterprise/entity_retrieval/reconcile.clj Normal file
View File

@@ -0,0 +1,455 @@
(ns metabase-enterprise.entity-retrieval.reconcile
"Reconciles the pgvector `library_entity_index` with the authoritative appdb.
The index holds one embedded document per library-entity value — the entity's name, its description, and
each `ai_context` synonym/example.
The primary key `doc_id` is content-addressed over `entity_type|entity_local_id|doc_type|doc_text`, so
editing text mints a new row.
Curator `instructions` are not stored here — the tool reads them live from `osi_ai_context` at query time.
Two reconcile entry points, both run via the coalescing schedule in
[[metabase-enterprise.entity-retrieval.core]] and serialized across nodes by one advisory lock:
- [[reconcile!]] — a full diff: derive the full desired doc set from library membership (published
Tables, library metric/model Cards, and the Measures/Segments on those Tables) plus `osi_ai_context`,
insert the docs whose `doc_id` is absent (embedding them), and delete the docs with no desired
counterpart. That delete is the GC, and re-deriving from the appdb every run self-heals any missed
write. Runs on a schedule as the backstop and from the force-reconcile API.
- [[reconcile-entity!]] — the same diff scoped to one entity's doc slice, driven by the
`osi_ai_context` write hooks so a curator edit becomes searchable without a full scan.
Callers pass the datasource and a model resolver, so this namespace reads no settings."
(:require
[buddy.core.hash :as buddy-hash]
[clojure.string :as str]
[honey.sql :as sql]
[honey.sql.helpers :as sql.helpers]
[medley.core :as m]
[metabase-enterprise.entity-retrieval.index-table :as index-table]
[metabase-enterprise.semantic-search.embedding :as embedding]
[metabase.collections.core :as collections]
[metabase.entity-retrieval.core :as entity-retrieval]
[metabase.util :as u]
[metabase.util.log :as log]
[next.jdbc :as jdbc]
[next.jdbc.result-set :as jdbc.rs]
[toucan2.core :as t2])
(:import
(java.sql Connection)))
(set! *warn-on-reflection* true)
(def ^:private embed-batch-size
"Documents embedded per embedding-service request while backfilling new index rows.
Larger than semantic-search's indexer default (150): our docs are short (names, single
synonyms/examples), so item count, not the per-request token budget, is the binding constraint."
512)
;; Advisory lock serializing whole reconcile runs across cluster nodes (insert-vs-orphan-delete races).
;; Arbitrary app-wide-unique constant, distinct from index-table's ensure lock (20011) and
;; semantic-search's migration lock (19991).
(def ^:private reconcile-lock-id 20012)
(defn doc-id
"Content-addressed primary key for an index document.
`instructions` is intentionally not an input: editing it must not re-embed an entity's name/synonyms."
[entity-type entity-local-id doc-type doc-text]
(u/encode-base64-bytes
(buddy-hash/sha1 (str entity-type "|" entity-local-id "|" doc-type "|" doc-text))))
;;; ------------------------------------------------- Desired docs -------------------------------------------------
(def ^:private max-doc-chars
"Char cap on a doc's text, applied before [[doc-id]] and embedding.
The embedding layer skips a single text over the per-item token budget, which drops the doc and
under-indexes the entity; truncating keeps it indexed.
Coarse (chars-per-token varies) — a residual over-budget outlier is dropped by [[insert-batch!]], never
inserted as an empty vector."
8000)
(def ^:private max-values-per-kind
"Cap on synonym docs (and, separately, example docs) indexed per entity, mirroring the API's per-list cap.
Bounds index bloat from rows that bypass the API schema — SerDes, direct appdb writes, or rows predating
the cap — since reconcile reads osi_ai_context directly."
50)
(defn- make-doc [entity-type entity-local-id doc-type doc-text]
(let [doc-text (cond-> doc-text
(and (string? doc-text) (> (count doc-text) max-doc-chars)) (subs 0 max-doc-chars))]
{:doc_id (doc-id entity-type entity-local-id doc-type doc-text)
:entity_type entity-type
:entity_local_id entity-local-id
:doc_type doc-type
:doc_text doc-text}))
(defn- entity->docs
"All desired docs for one library entity: a `name` doc (always), a `description` doc (non-blank), and a
`synonym`/`example` doc per `ai_context` value.
Instructions are not indexed — the tool reads them live from `osi_ai_context`."
[{:keys [entity_type entity_local_id name description]} ai_context]
(let [doc #(make-doc entity_type entity_local_id %1 %2)]
(concat
[(doc "name" name)]
(when-not (str/blank? description) [(doc "description" description)])
;; cap the list length (each value is also char-capped in make-doc) so a row that skipped the API's
;; bounds can't bloat the index with an unbounded number of synonym/example docs.
(map #(doc "synonym" %) (take max-values-per-kind (remove str/blank? (:synonyms ai_context))))
(map #(doc "example" %) (take max-values-per-kind (remove str/blank? (:examples ai_context)))))))
(defn- ->library-entity [entity-type id nm description]
{:entity_type entity-type
:entity_local_id id
:name nm
:description description})
;;; ------------------------------------------- Library membership ------------------------------------------------
;;;
;;; The four per-type selects are the single source of membership truth, shared by the full-scan
;;; [[library-entities]] and the point [[library-entity]]. Each takes an optional `id` to restrict to one
;;; entity. Headless: the reconcile job runs with no current user, so these selects must NOT permission-filter.
(defn- library-ids
"Collection ids that count as library content: the Library root and its descendants. Content usually
lives in the Data/Metrics sub-collections, but an entity placed directly in the root is library content
too, so the root id is included."
[lib]
(vec (distinct (cons (:id lib) (collections/descendant-ids lib)))))
(defn- library-cards [lib-ids id]
;; :card_schema is mandatory in any column-scoped Card SELECT (toucan guard). Card :type is keywordized
;; (:metric / :model); the entity_type is its name string.
(->> (apply t2/select [:model/Card :id :name :description :type :card_schema]
(cond-> [:collection_id [:in lib-ids], :archived false, :type [:in ["metric" "model"]]]
id (conj :id id)))
(map (fn [c] (->library-entity (name (:type c)) (:id c) (:name c) (:description c))))))
(defn- library-tables [lib-ids id]
;; A published table's user-facing label is its display_name; fall back to the raw name.
(->> (apply t2/select [:model/Table :id :name :display_name :description]
(cond-> [:collection_id [:in lib-ids], :is_published true, :active true]
id (conj :id id)))
(map (fn [t] (->library-entity "table" (:id t) (or (:display_name t) (:name t)) (:description t))))))
(defn- library-measures [table-ids id]
(->> (apply t2/select [:model/Measure :id :name :description]
(cond-> [:table_id [:in table-ids], :archived false]
id (conj :id id)))
(map (fn [mv] (->library-entity "measure" (:id mv) (:name mv) (:description mv))))))
(defn- library-segments [table-ids id]
(->> (apply t2/select [:model/Segment :id :name :description]
(cond-> [:table_id [:in table-ids], :archived false]
id (conj :id id)))
(map (fn [s] (->library-entity "segment" (:id s) (:name s) (:description s))))))
(defn- library-entities
"Uniform `{:entity_type :entity_local_id :name :description}` maps for every entity in the library."
[]
(when-let [lib (collections/library-collection)]
(let [lib-ids (library-ids lib)
cards (library-cards lib-ids nil)
tables (library-tables lib-ids nil)
table-ids (not-empty (mapv :entity_local_id tables))
measures (when table-ids (library-measures table-ids nil))
segments (when table-ids (library-segments table-ids nil))]
(concat cards tables measures segments))))
(defn library-entity
"The `{:entity_type :entity_local_id :name :description}` map for one entity if it is currently a library
member, else nil. Shares [[library-entities]]' membership rules via the per-type selects.
The hook may pass either Card label; the returned map carries the entity's *stored* type, so a
metric↔model relabel keeps `doc_id`s stable."
[entity-type entity-local-id]
(when-let [lib (collections/library-collection)]
(let [lib-ids (library-ids lib)]
(cond
(entity-retrieval/card-entity-type? entity-type)
(first (library-cards lib-ids entity-local-id))
(= "table" entity-type)
(first (library-tables lib-ids entity-local-id))
(#{"measure" "segment"} entity-type)
;; A measure/segment is a member only when its parent table is a current library table.
(when-let [table-id (t2/select-one-fn :table_id
(if (= entity-type "measure") :model/Measure :model/Segment)
:id entity-local-id)]
(when (seq (library-tables lib-ids table-id))
(first (if (= entity-type "measure")
(library-measures [table-id] entity-local-id)
(library-segments [table-id] entity-local-id)))))
:else nil))))
(defn library-entity-keys
"Set of `[entity_type entity_local_id]` for every entity currently in the library.
The tool post-filters its index hits against this so a stale index never surfaces an entity that has
since left the library, the same way it post-filters on read permissions."
[]
(into #{} (map (juxt :entity_type :entity_local_id)) (library-entities)))
(defn- entity-class
"Map-taking adapter over [[entity-retrieval/entity-class]] for the `{:entity_type :entity_local_id}` maps
this namespace passes around (rows, docs, library entities).
Card labels collapse to one class so a card's desired docs, stored docs, and curated `ai_context` stay
matched across a relabel; table/measure/segment stay distinct."
[{:keys [entity_type entity_local_id]}]
(entity-retrieval/entity-class entity_type entity_local_id))
(defn- ai-context-by-entity
"Map of entity [[entity-class]] -> ai_context for every `osi_ai_context` row, so a card's curated context
(stored under the canonical `card` type) is found under the same class as its index docs (keyed by the
card's live metric/model type). One row per card is guaranteed by the normalized storage key."
[]
(u/index-by entity-class :ai_context
(t2/select [:model/OsiAiContext :entity_type :entity_local_id :ai_context])))
(defn- dedup-by-doc-id [docs]
;; distinct-by doc_id so an exact duplicate (same doc_type and text, e.g. a synonym listed twice)
;; collapses; a synonym equal to the name does NOT collapse — different doc_type, different doc_id.
(into [] (m/distinct-by :doc_id) docs))
(defn- desired-docs
"The full set of docs the index should hold, deduped by `doc_id` (identical values collapse to one)."
[]
(let [ac-by-entity (ai-context-by-entity)]
(dedup-by-doc-id
(mapcat (fn [ent] (entity->docs ent (get ac-by-entity (entity-class ent))))
(library-entities)))))
(defn- entity-desired-docs
"Desired docs for one entity: its `ai_context` synonym/example docs plus name/description, but only if
it is still a library member. A non-member (left the library) yields no docs, so its stored docs all
GC. `entity-type` may be either Card label; membership resolves the canonical stored type."
[entity-type entity-local-id]
(if-let [member (library-entity entity-type entity-local-id)]
;; Match ai_context by the normalized storage type: a card's row is stored as `card`, so look it up by
;; `card` whichever live label (metric/model) drove this targeted run.
(let [ai-ctx (t2/select-one-fn :ai_context :model/OsiAiContext
:entity_local_id entity-local-id
:entity_type (entity-retrieval/normalize-entity-type entity-type))]
(dedup-by-doc-id (entity->docs member ai-ctx)))
[]))
;;; ------------------------------------------------- Index writes -------------------------------------------------
;;;
;;; All pgvector reads/writes take the locked connection (see [[with-reconcile-lock]]) so a whole run uses
;;; one connection — a pgvector pool of size 1 can't deadlock checking out a second.
(defn- stored-docs
"Map of `doc_id -> {:entity_type :entity_local_id}` for every row in the index."
[conn]
(u/index-by :doc_id #(select-keys % [:entity_type :entity_local_id])
(jdbc/execute! conn
[(format "SELECT doc_id, entity_type, entity_local_id FROM \"%s\""
index-table/*vectors-table*)]
{:builder-fn jdbc.rs/as-unqualified-lower-maps})))
(defn- stored-docs-for-entity
"Map of `doc_id -> {:entity_type :entity_local_id}` for one entity's rows. Reads by `entity_local_id`
then keeps rows of the target's [[entity-class]], so a metric↔model relabel's stale half is included
while a same-id entity of another type is left untouched."
[conn entity-type entity-local-id]
(let [target-class (entity-class {:entity_type entity-type :entity_local_id entity-local-id})]
(into {}
(comp (filter #(= target-class (entity-class %)))
(map (juxt :doc_id #(select-keys % [:entity_type :entity_local_id]))))
(jdbc/execute! conn
[(format "SELECT doc_id, entity_type, entity_local_id FROM \"%s\" WHERE entity_local_id = ?"
index-table/*vectors-table*)
entity-local-id]
{:builder-fn jdbc.rs/as-unqualified-lower-maps}))))
(defn- doc->record [doc embedding]
{:doc_id (:doc_id doc)
:entity_type (:entity_type doc)
:entity_local_id (:entity_local_id doc)
:doc_type (:doc_type doc)
:doc_text (:doc_text doc)
:doc_embedding [:raw (index-table/format-embedding embedding)]})
(defn- insert-batch!
"Embed one batch of new docs and insert them, returning the number inserted.
Embeds via [[embedding/process-embeddings-streaming]] so a long value can't push a request over the
provider's per-batch token limit — it splits into token-aware sub-batches (single batch off OpenAI).
Its own fn so the failed-insert test can mock it to simulate an embed/insert failure."
[conn embedding-model docs]
(let [text->embedding (volatile! {})]
;; the callback must be purely side-effecting: process-embeddings-streaming merges callback *return*
;; values with `merge-with +` across sub-batches, which would try to add embedding vectors.
(embedding/process-embeddings-streaming embedding-model (map :doc_text docs)
(fn [m] (vswap! text->embedding merge m) nil)
:type :index :record-tokens? true)
;; Skip any doc that came back without an embedding — e.g. a single value over the provider's
;; per-item token limit. Inserting a nil-embedding record would otherwise blow up the batch (and a
;; failed batch loses the good docs alongside the bad), so drop the offenders and keep the rest.
(let [records (keep (fn [doc]
(when-let [embedding (@text->embedding (:doc_text doc))]
(doc->record doc embedding)))
docs)]
(when (seq records)
(jdbc/execute! conn
(-> (sql.helpers/insert-into (keyword index-table/*vectors-table*))
(sql.helpers/values (vec records))
(sql.helpers/on-conflict :doc_id)
(sql.helpers/do-nothing)
(sql/format {:quoted true}))))
(count records))))
(defn- delete-rows! [conn doc-ids]
(when (seq doc-ids)
(jdbc/execute! conn
(-> (sql.helpers/delete-from (keyword index-table/*vectors-table*))
(sql.helpers/where [:in :doc_id (vec doc-ids)])
(sql/format {:quoted true})))))
(defn- diff-result [desired to-insert inserted deleted]
{:inserted inserted
:deleted deleted
:unchanged (- (count desired) (count to-insert))})
(defn- index-size
"Total document and distinct-entity counts in the index, for the size gauges. Cheap aggregate, run once
per full reconcile under the lock."
[conn]
(let [{:keys [documents entities]}
(jdbc/execute-one! conn
[(format (str "SELECT count(*) AS documents, "
"count(distinct (entity_type, entity_local_id)) AS entities "
"FROM \"%s\"")
index-table/*vectors-table*)]
{:builder-fn jdbc.rs/as-unqualified-lower-maps})]
{:documents documents :entities entities}))
;; Scalability — targeted writes, full-diff backstop.
;; The full diff is O(total index size) per run, not O(changes). Embeddings — the only expensive resource —
;; are already change-scoped (only a new doc_text is embedded; an idle run embeds nothing), and the O(total)
;; metadata read is cheap over the *library* (the bounded curated tier), so a full run stays comfortable
;; into the tens of thousands of docs. The write path no longer pays it on every edit: an `osi_ai_context`
;; write drives [[reconcile-entity!]] (one entity's slice), and [[reconcile!]] runs only on a slow schedule
;; and from the force-reconcile API as the backstop for membership / name / description changes that aren't
;; hooked. Mark-and-sweep (stamp a generation, DELETE WHERE gen < current) would be a net loss for the full
;; path: it trades the cheap full read for a full write every run (WAL, dead tuples, vacuum pressure).
(defn- reconcile-against-appdb!
"The full diff body — assumes the reconcile advisory lock is held on `conn`. See the namespace docstring."
[conn embedding-model]
(let [desired (desired-docs)
desired-ids (set (map :doc_id desired))
stored (stored-docs conn)
to-insert (remove #(contains? stored (:doc_id %)) desired)
orphans (remove desired-ids (keys stored))
;; entity classes whose insert failed this run — their orphans are spared (see below).
failed (volatile! #{})
inserted (transduce
(partition-all embed-batch-size)
(completing
(fn [n batch]
(+ n (try
(insert-batch! conn embedding-model batch)
(catch Exception e
(vswap! failed into (map entity-class batch))
(log/error e "library entity index: failed to insert batch of"
(count batch) "docs; will retry next run")
0)))))
0
to-insert)
;; An orphan is often the stale half of an edited value (e.g. a renamed name's old doc).
;; Deleting it while its replacement failed to embed would drop that entity from search until the
;; next run, so spare orphans whose entity class had a failed insert; every other orphan GCs normally.
to-delete (cond->> orphans
(seq @failed) (remove #(contains? @failed (entity-class (get stored %)))))]
(delete-rows! conn to-delete)
;; index-size after the writes feeds the document/entity gauges (full reconcile only).
(merge (diff-result desired to-insert inserted (count to-delete))
(index-size conn))))
(defn- reconcile-entity-against-appdb!
"The targeted diff body for one entity — assumes the reconcile advisory lock is held on `conn`."
[conn embedding-model entity-type entity-local-id]
(let [desired (entity-desired-docs entity-type entity-local-id)
desired-ids (set (map :doc_id desired))
stored (stored-docs-for-entity conn entity-type entity-local-id)
to-insert (remove #(contains? stored (:doc_id %)) desired)
orphans (remove desired-ids (keys stored))
;; One entity is one class: if its insert fails, keep its stale-but-searchable docs (skip the GC)
;; and let a later run retry. A non-member has nothing to insert, so its delete always proceeds.
inserted (try
(if (seq to-insert) (insert-batch! conn embedding-model to-insert) 0)
(catch Exception e
(log/error e "library entity index: failed to reconcile entity"
entity-type entity-local-id "; will retry next run")
::failed))
failed? (= ::failed inserted)
to-delete (if failed? [] orphans)]
(delete-rows! conn to-delete)
(diff-result desired to-insert (if failed? 0 inserted) (count to-delete))))
(defn- with-reconcile-lock
"Acquire the session-level reconcile advisory lock on one pooled connection, ensure the tables exist for
the resolved model, then run `(f conn embedding-model emptied?)` on that same connection, and unlock.
Returns `f`'s diff map with `:rebuilt?` added. `emptied?` is true when `ensure-tables!` left an empty
vectors table (a first :created build or a model/format :rebuilt), so a targeted caller can fall back to a
full repopulate; `:rebuilt?` is the narrower model/format-rebuild signal the rebuild metric counts.
A single run does all its work (lock + ensure-tables! + diff) on one connection, so it never deadlocks
checking out a second from a size-1 pgvector pool. (Two *concurrent* runs — a full and a targeted one —
are still two connections, so a size-1 pool isn't viable for that; a pool of >= 2 is recommended.)
Session-level (not the transaction-scoped lock index-table and semantic-search use) because the run
commits per batch and tolerates partial failure across runs, so it must not be wrapped in one
transaction. Blocking acquire — a concurrent node makes this wait, not skip. The model is resolved only
under the lock, so a config change during a cross-node wait can't make us embed with a stale model;
`ensure-tables!` (which may drop+rebuild on a model/format change) runs under the lock too, so a rebuild
can't pull the table out from under a concurrent node's in-flight run."
[pgvector resolve-model f]
(with-open [^Connection conn (jdbc/get-connection pgvector)]
;; Per-batch commits, not one big transaction: the run tolerates partial failure across runs, so each
;; insert/delete must commit on its own. Some pools hand out autocommit-off connections (which would
;; silently roll the whole run back on close), so force it on and restore the prior setting before the
;; connection goes back to the pool. (ensure-tables! flips this to a transaction for its DDL.)
(let [autocommit (.getAutoCommit conn)]
(.setAutoCommit conn true)
(try
(jdbc/execute! conn [(format "SELECT pg_advisory_lock(%d)" reconcile-lock-id)])
(try
(let [embedding-model (resolve-model)
status (index-table/ensure-tables! conn embedding-model)
;; :created (first build / healed manual drop) and :rebuilt (model/format change) both leave
;; an empty table, so a targeted caller must do a full repopulate rather than index one entity.
emptied? (contains? #{:created :rebuilt} status)]
(when emptied?
(log/info "library entity index: vectors table is empty (" status "); repopulating"))
(assoc (f conn embedding-model emptied?) :rebuilt? (= :rebuilt status)))
(finally
(jdbc/execute! conn [(format "SELECT pg_advisory_unlock(%d)" reconcile-lock-id)])))
(finally
(.setAutoCommit conn autocommit))))))
(defn reconcile!
"Full reconcile of the pgvector `library_entity_index` with the appdb, blocking until it completes;
returns {:inserted n :deleted n :unchanged n :rebuilt? bool}.
`resolve-model` is a thunk returning the embedding model, called only once the advisory lock is held.
The backstop for membership / name / description changes that the targeted write path doesn't hook.
See the namespace docstring and [[with-reconcile-lock]]."
[pgvector resolve-model]
(with-reconcile-lock pgvector resolve-model
(fn [conn embedding-model _emptied?]
(reconcile-against-appdb! conn embedding-model))))
(defn reconcile-entity!
"Targeted reconcile of one entity's doc slice, blocking until it completes; returns
{:inserted n :deleted n :unchanged n :rebuilt? bool}.
Recomputes the entity's full desired docs (name/description if still a library member, else none, plus
its current `ai_context` synonyms/examples) and diffs them against its stored slice — so an
`ai_context` delete GCs the synonym/example docs while keeping name/description, and an entity leaving
the library GCs all of its docs. `entity-type`/`entity-local-id` come from the `osi_ai_context` write
hook; the reconcile runs later (on a future), reading the entity's *current* appdb state.
If `ensure-tables!` left the index empty — a first :created build or a model/format :rebuilt — this falls
back to a full reconcile; a targeted slice alone would leave every other entity missing until the backstop.
See the namespace docstring and [[with-reconcile-lock]]."
[pgvector resolve-model entity-type entity-local-id]
(with-reconcile-lock pgvector resolve-model
(fn [conn embedding-model emptied?]
(if emptied?
(reconcile-against-appdb! conn embedding-model)
(reconcile-entity-against-appdb! conn embedding-model entity-type entity-local-id)))))

60
enterprise/backend/src/metabase_enterprise/entity_retrieval/task/sync.clj Normal file
View File

@@ -0,0 +1,60 @@
(ns metabase-enterprise.entity-retrieval.task.sync
"Quartz job running a full library entity index reconcile, on a slow schedule, as the backstop.
Immediate freshness comes from the targeted write path — an `osi_ai_context` edit drives a per-entity
reconcile (see [[metabase-enterprise.entity-retrieval.core/request-entity-sync!]]) within seconds. This
periodic full reconcile only catches what isn't hooked: membership / name / description changes to the
underlying Card/Table/Measure/Segment, and any write whose targeted reconcile was lost (scheduler
hiccup, pgvector down, an import that bypassed the model hooks). Its first firing, ~10s after boot, is
also the post-upgrade startup reconcile that rebuilds the index when the doc format changed.
`DisallowConcurrentExecution` plus Quartz clustering ensure a single node runs it at a time."
(:require
[clojurewerkz.quartzite.jobs :as jobs]
[clojurewerkz.quartzite.schedule.simple :as simple]
[clojurewerkz.quartzite.triggers :as triggers]
[metabase-enterprise.entity-retrieval.core :as entity-retrieval.core]
[metabase.task.core :as task]
[metabase.util.log :as log])
(:import
(java.time Duration Instant)
(java.util Date)))
(set! *warn-on-reflection* true)
(task/defjob ^{org.quartz.DisallowConcurrentExecution true
:doc "Reconciles the library entity index (pgvector) with the appdb."}
OsiAiContextSync [_ctx]
(when (entity-retrieval.core/available?)
(try
;; Full reconcile via the shared coalescing schedule: it waits out a concurrent node rather than
;; skipping, and never runs alongside an API- or write-triggered reconcile on this node. Duration
;; and mutation metrics are emitted centrally for every run (see entity-retrieval.core).
(entity-retrieval.core/reconcile-full-coalesced!)
(catch Throwable e
;; Log and move on: the next periodic run retries from the authoritative appdb table.
(log/error e "entity-retrieval mirror reconciliation failed")))))
(def ^:private ^Duration startup-delay (Duration/parse "PT10S"))
;; Slow: this is the backstop, not the freshness mechanism — the targeted write path keeps the index live.
;; The +10s first firing (startup reconcile / post-upgrade rebuild) is preserved regardless of this value.
(def ^:private ^Duration run-frequency (Duration/parse "PT15M"))
(defmethod task/init! ::OsiAiContextSync [_]
;; Gate scheduling on pgvector being configured (boot-fixed), NOT on available? — the job body
;; self-gates on the feature flag, so scheduling here lets the periodic safety net (and the
;; write-path trigger's target job) exist even when the license is enabled after startup.
(when (entity-retrieval.core/pgvector-configured?)
(let [job (jobs/build
(jobs/of-type OsiAiContextSync)
(jobs/store-durably)
(jobs/with-identity entity-retrieval.core/sync-job-key))
trigger (triggers/build
(triggers/with-identity
(triggers/key "metabase-enterprise.entity-retrieval.sync.trigger"))
(triggers/for-job entity-retrieval.core/sync-job-key)
(triggers/start-at (Date/from (.plus (Instant/now) startup-delay)))
(triggers/with-schedule
(simple/schedule
(simple/with-interval-in-milliseconds (.toMillis run-frequency))
(simple/repeat-forever))))]
(task/schedule-task! job trigger))))

5
enterprise/backend/src/metabase_enterprise/semantic_search/embedding.clj
View File

@@ -98,6 +98,11 @@
(let [text-measure (measure text)]
(cond
;; Single text exceeds the limit - skip it with warning
;; TODO (Chris 2026-06-29) -- silently dropping an over-budget doc here is a poor default —
;; it vanishes from the index with only a warning. Make the over-budget policy a per-call
;; option (:truncate / :skip / :error, likely truncate-by-default eventually) and check what
;; the ai-service path does, which bypasses create-batches entirely.
;; https://linear.app/metabase/issue/BOT-1742
(> text-measure threshold)
(do
(log/warn

10
enterprise/backend/src/metabase_enterprise/serialization/v2/extract.clj
View File

@@ -42,10 +42,14 @@
(conj "EmbeddingTheme")
(not (:no-custom-viz-plugins opts))
(conj "CustomVizPlugin")
(conj "CustomVizPlugin")))
(not (:no-curated-search opts))
(conj "CuratedSearchEntry")))
;; OsiAiContext is intentionally NOT in the default export set. It's a top-level model that *depends on* its
;; entity, extracted unfiltered, so any export — even an untargeted "full" one, which still omits personal and
;; analytics collections — would pull ai_context rows whose referenced Card/Table/Measure/Segment is absent,
;; leaking curator text and creating dangling deps that fail import. It stays out until reverse-dependency
;; export lands (its row should travel *with* its entity), tracked in
;; https://linear.app/metabase/issue/BOT-1759; see the metabase.osi.models.osi-ai-context serdes TODO.
(defn make-targets-of-type
"Returns a targets seq with model type and given ids"

2
enterprise/backend/src/metabase_enterprise/serialization/v2/ingest.clj
View File

@@ -83,7 +83,7 @@
(def legal-top-level-paths
"Known top-level paths for directory with serialization output.
We support both \"python-libraries\" and \"python_libraries\" for backwards compatibility. The modern name is \"python_libraries\"."
#{"actions" "channels" "collections" "custom_viz_plugins" "databases" "embedding_themes" "glossary" "metabots" "python_libraries" "python-libraries" "curated_search_entries" "snippets" "transforms"})
#{"actions" "channels" "collections" "custom_viz_plugins" "databases" "embedding_themes" "glossary" "metabots" "python_libraries" "python-libraries" "osi_ai_context" "snippets" "transforms"})
(defn- path-interner
"Returns a function that interns `:serdes/meta` path vectors.

4
enterprise/backend/src/metabase_enterprise/serialization/v2/models.clj
View File

@@ -33,11 +33,11 @@
"The list of all models exported by serialization by default. Used for production code and by tests."
(concat data-model
content
["CuratedSearchEntry"
"CustomVizPlugin"
["CustomVizPlugin"
"EmbeddingTheme"
"FieldValues"
"Metabot"
"OsiAiContext"
"PythonLibrary"
"Setting"
"Transform"

150
enterprise/backend/test/metabase_enterprise/curated_search/core_test.clj
View File

@@ -1,150 +0,0 @@
(ns metabase-enterprise.curated-search.core-test
(:require
[clojure.test :refer :all]
[metabase-enterprise.curated-search.core :as curated-search.core]
[metabase-enterprise.curated-search.index-table :as index-table]
[metabase-enterprise.curated-search.reconcile :as reconcile]
[metabase-enterprise.semantic-search.db.datasource :as semantic.db.datasource]
[metabase-enterprise.semantic-search.embedding :as semantic.embedding]
[metabase-enterprise.semantic-search.test-util :as semantic.tu]
[metabase.curated-search.mirror :as mirror]
[metabase.metabot.tools.curated-search :as tools.curated-search]
[metabase.task.core :as task]
[metabase.test :as mt]
[next.jdbc :as jdbc]))
(set! *warn-on-reflection* true)
(defn- approx [target] #(< (abs (- (double %) (double target))) 1e-9))
(deftest score-shape-matches-regular-search-test
(let [score (var-get #'curated-search.core/score)]
(testing "weighted-scorer breakdown: per-factor {name score weight contribution} + total_score"
(is (=? {:scores [{:name :similarity :score (approx 0.8) :weight 1.0 :contribution (approx 0.8)}
{:name :verified :score 0.0 :weight 0.1 :contribution 0.0}]
:total_score (approx 0.8)}
(score {:distance 0.2 :verified false}))))
(testing "the verified boost applies its weight and strictly increases the total"
(is (=? {:total_score (approx 0.9)} (score {:distance 0.2 :verified true})))
(let [s (score {:distance 0.2 :verified true})]
(is (=? (approx (reduce + (map :contribution (:scores s)))) (:total_score s))))
(is (> (:total_score (score {:distance 0.2 :verified true}))
(:total_score (score {:distance 0.2 :verified false})))))))
(deftest dispatch-without-pgvector-test
(testing "with the feature enabled but pgvector unconfigured, the EE impls degrade gracefully"
;; Pin db-url to nil so the result is deterministic regardless of any ambient MB_PGVECTOR_DB_URL:
;; available? is false, so search returns [] and the sync nudge no-ops rather than throwing.
(mt/with-premium-features #{:semantic-search}
(with-redefs [semantic.db.datasource/db-url nil]
(is (= [] (mirror/search "anything" 10)))
(is (nil? (mirror/request-sync!)))))))
(deftest request-sync-triggers-the-job-test
(testing "when the mirror is available, the nudge fires trigger-now! on the sync job — and nothing else"
(mt/with-premium-features #{:semantic-search}
;; db-url is read directly as a var, so with-redefs (not with-dynamic-fn-redefs) is required here.
(with-redefs [semantic.db.datasource/db-url "jdbc:postgresql://stub"]
(let [triggered (atom [])]
(mt/with-dynamic-fn-redefs [task/trigger-now! (fn [job-key] (swap! triggered conj job-key))]
(mirror/request-sync!)
(is (= [curated-search.core/sync-job-key] @triggered))))))))
(deftest pgvector-configured-decoupled-from-feature-test
(testing "scheduling gates on pgvector config, independent of the feature flag, so a license enabled after boot still gets the periodic sync"
;; db-url is read directly as a var, so with-redefs (not with-dynamic-fn-redefs) is required here.
(with-redefs [semantic.db.datasource/db-url "jdbc:postgresql://stub"]
(mt/with-premium-features #{}
(is (true? (curated-search.core/pgvector-configured?)))
(is (false? (curated-search.core/available?))))
(mt/with-premium-features #{:semantic-search}
(is (true? (curated-search.core/available?)))))))
(deftest ^:sequential verified-boost-applies-before-limit-test
(testing "the SQL LIMIT ranks by blended score: a verified row slightly farther by raw distance still wins top-1"
(when semantic.db.datasource/db-url
(let [suffix (System/nanoTime)
ds (semantic.db.datasource/ensure-initialized-data-source!)
p-verif "verified entry"
p-plain "plain entry"
q "the query"]
(mt/with-premium-features #{:semantic-search}
(mt/with-dynamic-fn-redefs [semantic.embedding/get-configured-model
(constantly semantic.tu/mock-embedding-model)]
(binding [index-table/*vectors-table* (str "curated_search_index_test_" suffix)
index-table/*meta-table* (str "curated_search_index_meta_test_" suffix)]
;; plain ties the query exactly (distance 0); verified sits at distance ~0.006, within its
;; 0.1 boost. A raw-distance LIMIT 1 would keep only the plain row.
(semantic.tu/with-mock-embeddings {q [1.0 0.0 0.0 0.0]
p-plain [1.0 0.0 0.0 0.0]
p-verif [0.9 0.1 0.0 0.0]}
(mt/with-temp [:model/CuratedSearchEntry _
{:search_prompt p-plain :entity {:model "table" :id 1}}
:model/CuratedSearchEntry _
{:search_prompt p-verif :verified true :entity {:model "table" :id 2}}]
(try
(reconcile/reconcile! ds semantic.tu/mock-embedding-model)
(is (= [p-verif] (mapv :saved_search_prompt (mirror/search q 1))))
(finally
(jdbc/execute! ds [(str "DROP TABLE IF EXISTS "
index-table/*vectors-table* ", "
index-table/*meta-table*)]))))))))))))
(defn- create-entry!
"POST an entry through the CRUD API; returns the created row's id."
[search-prompt entity verified]
(:id (mt/user-http-request :crowberto :post 200 "curated-search/"
{:search_prompt search-prompt :entity entity :verified verified})))
(deftest ^:sequential crud-api-to-tool-end-to-end-test
(testing "CRUD API write -> reconcile -> pgvector -> search_curated tool, end to end"
;; Self-gated on MB_PGVECTOR_DB_URL — CI without semantic-search infra skips this; locally with the
;; dev pgvector running it exercises the whole pipeline. Uses the mock embedding model (4-dim,
;; deterministic, no network) against isolated temp tables, with the reconciler run directly in
;; place of the background Quartz job.
(when semantic.db.datasource/db-url
(let [suffix (System/nanoTime)
ds (semantic.db.datasource/ensure-initialized-data-source!)
ent-1 {:model "table" :id (mt/id :orders)}
ent-2 {:model "table" :id (mt/id :people)}
;; Both prompts + the query embed to the same vector, so cosine distance ties and the
;; verified boost (0.10) alone decides the ranking.
vec1 [1.0 0.0 0.0 0.0]
p-verif "revenue by region (verified)"
p-plain "revenue by region (plain)"
q "monthly revenue per region"
;; bind a user: the tool's hydration permission-filters via mi/can-read?, which needs
;; *current-user-id* (a superuser can read the test tables).
search! #(mt/with-test-user :crowberto
(get-in (tools.curated-search/curated-search-tool
{:user_search_prompt q})
[:structured-output :data]))]
(mt/with-premium-features #{:semantic-search}
(mt/with-dynamic-fn-redefs [semantic.embedding/get-configured-model
(constantly semantic.tu/mock-embedding-model)]
(binding [index-table/*vectors-table* (str "curated_search_index_test_" suffix)
index-table/*meta-table* (str "curated_search_index_meta_test_" suffix)]
(semantic.tu/with-mock-embeddings {p-verif vec1 p-plain vec1 q vec1}
(try
(let [id-verif (create-entry! p-verif ent-1 true)
id-plain (create-entry! p-plain ent-2 false)
_ (reconcile/reconcile! ds semantic.tu/mock-embedding-model)
results (search!)]
(testing "both distinct entities mirrored and returned, the verified one boosted first"
(is (= [p-verif p-plain] (mapv :saved_search_prompt results))))
(testing "data is per-entity hydrated records; total_score reflects the boost, similarity=1.0 (vectors tie)"
(is (=? [{:type "table" :id (mt/id :orders) :similarity (approx 1.0)
:score {:total_score (approx 1.1)}}
{:type "table" :id (mt/id :people) :similarity (approx 1.0)
:score {:total_score (approx 1.0)}}]
results)))
(testing "deleting via the CRUD API + reconcile removes the row from search results"
(mt/user-http-request :crowberto :delete 204 (str "curated-search/" id-verif))
(is (=? {:deleted 1} (reconcile/reconcile! ds semantic.tu/mock-embedding-model)))
(is (= [p-plain] (mapv :saved_search_prompt (search!)))))
;; clean up the appdb rows we created
(mt/user-http-request :crowberto :delete 204 (str "curated-search/" id-plain)))
(finally
(jdbc/execute! ds [(str "DROP TABLE IF EXISTS "
index-table/*vectors-table* ", "
index-table/*meta-table*)])))))))))))

119
enterprise/backend/test/metabase_enterprise/curated_search/reconcile_test.clj
View File

@@ -1,119 +0,0 @@
(ns metabase-enterprise.curated-search.reconcile-test
(:require
[clojure.test :refer :all]
[metabase-enterprise.curated-search.index-table :as index-table]
[metabase-enterprise.curated-search.reconcile :as reconcile]
[metabase-enterprise.semantic-search.db.datasource :as semantic.db.datasource]
[metabase-enterprise.semantic-search.test-util :as semantic.tu]
[metabase.test :as mt]
[next.jdbc :as jdbc]
[next.jdbc.result-set :as jdbc.rs]
[toucan2.core :as t2]))
(set! *warn-on-reflection* true)
(deftest content-hash-test
(let [row {:search_prompt "p" :usage_instructions "u" :entity {:model "table" :id 1}
:verified false}]
(testing "equal rows hash equal; nil and absent usage_instructions hash like empty"
(is (= (reconcile/content-hash row) (reconcile/content-hash row)))
(is (= (reconcile/content-hash (assoc row :usage_instructions nil))
(reconcile/content-hash (assoc row :usage_instructions "")))))
(testing "every mirror-relevant field changes the hash"
(let [h (reconcile/content-hash row)]
(doseq [variant [(assoc row :search_prompt "p2")
(assoc row :usage_instructions "u2")
(assoc row :entity {:model "table" :id 2})
(assoc row :verified true)]]
(is (not= h (reconcile/content-hash variant)) (pr-str variant)))))))
(deftest format-embedding-rejects-invalid-values-test
(testing "non-numbers, NaN and infinities are rejected before they reach a raw SQL literal"
(doseq [bad [Double/NaN Double/POSITIVE_INFINITY Double/NEGATIVE_INFINITY "0.1"]]
(is (thrown-with-msg? Exception #"invalid value"
(index-table/format-embedding [0.1 bad 0.3]))
(pr-str bad)))
(is (string? (index-table/format-embedding [0.1 -0.2 0.3])))))
(defmacro ^:private with-isolated-mirror
"Run `body` with isolated pgvector tables bound and dropped afterwards, binding `ds-sym` to the
pgvector datasource. Skips the body entirely when no pgvector store is configured."
[[ds-sym] & body]
`(when semantic.db.datasource/db-url
(let [suffix# (System/nanoTime)
~ds-sym (semantic.db.datasource/ensure-initialized-data-source!)]
(binding [index-table/*vectors-table* (str "curated_search_index_test_" suffix#)
index-table/*meta-table* (str "curated_search_index_meta_test_" suffix#)]
(try
~@body
(finally
(jdbc/execute! ~ds-sym [(str "DROP TABLE IF EXISTS "
index-table/*vectors-table* ", "
index-table/*meta-table*)])))))))
(defn- mirror-rows [ds]
(jdbc/execute! ds
[(format "SELECT index_id, search_prompt, verified, content_hash FROM \"%s\" ORDER BY index_id"
index-table/*vectors-table*)]
{:builder-fn jdbc.rs/as-unqualified-lower-maps}))
(deftest ^:sequential reconcile-lifecycle-test
(with-isolated-mirror [ds]
(let [model semantic.tu/mock-embedding-model]
(mt/with-temp [:model/CuratedSearchEntry {id-1 :id}
{:search_prompt "monthly revenue"
:entity {:model "table" :id 1}}
:model/CuratedSearchEntry {id-2 :id}
{:search_prompt "orders and customers" :verified true
:entity {:model "table" :id 2}}]
(let [appdb-total (t2/count :model/CuratedSearchEntry)]
(testing "first run mirrors every appdb row"
(is (=? {:upserted appdb-total :deleted 0 :unchanged 0}
(reconcile/reconcile! ds model)))
(is (=? {id-1 {:search_prompt "monthly revenue" :verified false}
id-2 {:search_prompt "orders and customers" :verified true}}
(-> (into {} (map (juxt :index_id identity)) (mirror-rows ds))
(select-keys [id-1 id-2])))))
(testing "an unchanged second run embeds and writes nothing"
(is (=? {:upserted 0 :deleted 0 :unchanged appdb-total}
(reconcile/reconcile! ds model))))
(testing "an updated row is re-embedded; the rest stay unchanged"
(t2/update! :model/CuratedSearchEntry id-1 {:verified true})
(is (=? {:upserted 1 :deleted 0 :unchanged (dec appdb-total)}
(reconcile/reconcile! ds model)))
(is (=? {:verified true}
(->> (mirror-rows ds) (filter #(= id-1 (:index_id %))) first))))
(testing "a write that bypassed the hooks entirely is still repaired (self-healing)"
;; Tamper with the mirror behind the reconciler's back.
(jdbc/execute! ds [(format "UPDATE \"%s\" SET content_hash = 'bogus' WHERE index_id = %d"
index-table/*vectors-table* id-2)])
(is (=? {:upserted 1 :deleted 0}
(reconcile/reconcile! ds model))))))
(testing "rows deleted from the appdb are removed from the mirror as orphans"
;; The with-temp rows above are gone now.
(is (=? {:upserted 0 :deleted 2}
(reconcile/reconcile! ds model)))))))
(deftest ^:sequential rebuild-on-model-change-test
(with-isolated-mirror [ds]
(let [model semantic.tu/mock-embedding-model]
(mt/with-temp [:model/CuratedSearchEntry _
{:search_prompt "monthly revenue"
:entity {:model "table" :id 1}}]
(let [appdb-total (t2/count :model/CuratedSearchEntry)]
(testing "populate the mirror under the original model"
(is (=? {:upserted appdb-total} (reconcile/reconcile! ds model))))
(testing "a model-identity change drops the vectors table and the next run re-embeds everything"
(let [new-model (assoc model :model-name "model-v2")]
(is (= :rebuilt (index-table/ensure-tables! ds new-model)))
(is (= [] (mirror-rows ds)))
(is (=? {:upserted appdb-total :deleted 0 :unchanged 0}
(reconcile/reconcile! ds new-model)))
(testing "a schema-version mismatch alone also triggers the rebuild"
;; stale-write the meta row, as if the table were built by an older code version
(jdbc/execute! ds [(format "UPDATE \"%s\" SET schema_version = schema_version - 1"
index-table/*meta-table*)])
(is (= :rebuilt (index-table/ensure-tables! ds new-model)))
(is (= [] (mirror-rows ds))))
(testing "the rebuild heals the meta row, so it doesn't recur on the next sync"
(is (= :ok (index-table/ensure-tables! ds new-model)))))))))))

318
enterprise/backend/test/metabase_enterprise/entity_retrieval/core_test.clj Normal file
View File

@@ -0,0 +1,318 @@
(ns metabase-enterprise.entity-retrieval.core-test
(:require
[clojure.test :refer :all]
[metabase-enterprise.entity-retrieval.core :as entity-retrieval.core]
[metabase-enterprise.entity-retrieval.index-table :as index-table]
[metabase-enterprise.entity-retrieval.reconcile :as reconcile]
[metabase-enterprise.semantic-search.db.datasource :as semantic.db.datasource]
[metabase-enterprise.semantic-search.embedding :as semantic.embedding]
[metabase-enterprise.semantic-search.test-util :as semantic.tu]
[metabase.entity-retrieval.mirror :as mirror]
[metabase.metabot.tools.entity-retrieval :as tools.entity-retrieval]
[metabase.test :as mt]
[next.jdbc :as jdbc]))
(set! *warn-on-reflection* true)
(defn- approx [target] #(< (abs (- (double %) (double target))) 1e-9))
(deftest score-shape-matches-regular-search-test
(let [score (var-get #'entity-retrieval.core/score)]
(testing "weighted-scorer breakdown: similarity factor + doc_type bump + total_score"
(is (=? {:scores [{:name :similarity :score (approx 0.8) :weight 1.0 :contribution (approx 0.8)}
{:name :doc-type :score 1.0 :weight (approx 0.02) :contribution (approx 0.02)}]
:total_score (approx 0.82)}
(score {:distance 0.2 :doc_type "name"}))))
(testing "the doc_type bump favors name over synonym, and is zero for unweighted types"
(is (> (:total_score (score {:distance 0.2 :doc_type "name"}))
(:total_score (score {:distance 0.2 :doc_type "synonym"}))))
(is (=? {:total_score (approx 0.8)} (score {:distance 0.2 :doc_type "example"}))))
(testing "a nearer document (smaller distance) scores strictly higher"
(is (> (:total_score (score {:distance 0.1 :doc_type "name"}))
(:total_score (score {:distance 0.5 :doc_type "name"})))))))
(deftest dispatch-without-pgvector-test
(testing "with the feature enabled but pgvector unconfigured, the EE impls degrade gracefully"
;; Pin db-url to nil so the result is deterministic regardless of any ambient MB_PGVECTOR_DB_URL:
;; available? is false, so search returns [] and the write-path nudge no-ops rather than throwing.
(mt/with-premium-features #{:semantic-search}
(with-redefs [semantic.db.datasource/db-url nil]
(is (= [] (mirror/search "anything" 10)))
(is (nil? (mirror/request-entity-sync! "table" 1)))))))
(deftest targeted-drain-reconciles-each-dirty-entity-test
(testing "a targeted run snapshots and clears the dirty set, reconciling each entity once"
(let [dirty @#'entity-retrieval.core/dirty-entities
do-targeted-run (var-get #'entity-retrieval.core/do-targeted-run)
reconciled (atom [])]
(reset! dirty #{["table" 7] ["metric" 9]})
(mt/with-dynamic-fn-redefs
[semantic.db.datasource/ensure-initialized-data-source! (constantly ::ds)
semantic.embedding/get-configured-model (constantly ::model)
reconcile/reconcile-entity! (fn [_ds _model entity-type entity-local-id]
(swap! reconciled conj [entity-type entity-local-id])
{:inserted 1 :deleted 0 :unchanged 0})]
(do-targeted-run nil)
(is (= #{["table" 7] ["metric" 9]} (set @reconciled)))
(is (empty? @dirty) "the dirty set is cleared at the run's start")))))
(deftest request-entity-sync-is-fire-and-forget-test
(testing "request-entity-sync! returns nil without throwing on the calling thread, even if the reconcile fails"
(mt/with-premium-features #{:semantic-search}
(with-redefs [semantic.db.datasource/db-url "jdbc:postgresql://stub"]
(let [tc @#'entity-retrieval.core/targeted-current
tn @#'entity-retrieval.core/targeted-next
dirty @#'entity-retrieval.core/dirty-entities]
(reset! tc nil) (reset! tn nil) (reset! dirty #{})
(mt/with-dynamic-fn-redefs
[semantic.db.datasource/ensure-initialized-data-source! (constantly ::ds)
semantic.embedding/get-configured-model (constantly ::model)
reconcile/reconcile-entity! (fn [& _] (throw (ex-info "boom" {})))]
(is (nil? (entity-retrieval.core/request-entity-sync! "table" 1)))
;; let the fire-and-forget drain settle so its logged error doesn't bleed into later tests
(when-let [f @tc] (deref f 5000 nil))
(reset! tc nil) (reset! tn nil) (reset! dirty #{})))))))
(deftest force-reconcile-coalescing-test
(testing "force-reconcile! never joins the in-flight run, queues one shared follow-up, and reports index + timing"
(mt/with-premium-features #{:semantic-search}
;; db-url is read directly as a var, so with-redefs (not with-dynamic-fn-redefs) is required here.
(with-redefs [semantic.db.datasource/db-url "jdbc:postgresql://stub"]
(let [current @#'entity-retrieval.core/full-current
nxt @#'entity-retrieval.core/full-next
calls (atom 0)]
(mt/with-dynamic-fn-redefs
[semantic.db.datasource/ensure-initialized-data-source! (constantly ::ds)
semantic.embedding/get-configured-model (constantly ::model)
reconcile/reconcile! (fn [_ds _resolve-model]
(swap! calls inc)
{:inserted 3 :deleted 1 :unchanged 2})]
(testing "idle: the caller starts the run, gets its index diff + timing, and the schedule clears"
(reset! current nil) (reset! nxt nil) (reset! calls 0)
(is (=? {:index {:inserted 3 :deleted 1 :unchanged 2}
:execution {:waited_ms int? :ran_ms int?}}
(entity-retrieval.core/force-reconcile!)))
(is (= 1 @calls))
(is (nil? @current)) (is (nil? @nxt)))
(testing "a run in flight: the caller queues a fresh follow-up instead of reusing the in-flight run"
;; a completed future stands in for the in-flight run — the caller must still run a fresh pass.
(reset! current (future {:index {} :execution {}})) (reset! nxt nil) (reset! calls 0)
(is (=? {:index {:inserted 3 :deleted 1 :unchanged 2}}
(entity-retrieval.core/force-reconcile!)))
(is (= 1 @calls) "a fresh reconcile ran; the in-flight run's result was not reused"))
(testing "a follow-up already queued: further callers coalesce onto it (no extra run)"
(reset! current (future {:index {} :execution {}}))
(reset! nxt (future {:index {:inserted 7 :deleted 0 :unchanged 0} :execution {:waited_ms 5 :ran_ms 9}}))
(reset! calls 0)
(is (=? {:index {:inserted 7 :deleted 0 :unchanged 0}
:execution {:waited_ms 5 :ran_ms 9}}
(entity-retrieval.core/force-reconcile!)))
(is (zero? @calls) "joined the queued follow-up; no new reconcile")
(reset! current nil) (reset! nxt nil))))))))
(deftest force-reconcile-unavailable-returns-nil-test
(testing "force-reconcile! is nil (so the API can 400) when pgvector isn't configured"
(mt/with-premium-features #{:semantic-search}
(with-redefs [semantic.db.datasource/db-url nil]
(is (nil? (entity-retrieval.core/force-reconcile!)))))))
(deftest pgvector-configured-decoupled-from-feature-test
(testing "scheduling gates on pgvector config, not the feature flag, so a post-boot license still syncs"
;; db-url is read directly as a var, so with-redefs (not with-dynamic-fn-redefs) is required here.
(with-redefs [semantic.db.datasource/db-url "jdbc:postgresql://stub"]
(mt/with-premium-features #{}
(is (true? (entity-retrieval.core/pgvector-configured?)))
(is (false? (entity-retrieval.core/available?))))
(mt/with-premium-features #{:semantic-search}
(is (true? (entity-retrieval.core/available?)))))))
(deftest ^:sequential entity-retrieval-available?-requires-a-populated-index-test
(testing "the curated tool is offered only once the index has documents (else the nlq profile falls back)"
(when semantic.db.datasource/db-url
(mt/with-premium-features #{:library :semantic-search}
(let [suffix (System/nanoTime)
ds (semantic.db.datasource/ensure-initialized-data-source!)]
(mt/with-dynamic-fn-redefs [semantic.embedding/get-configured-model
(constantly semantic.tu/mock-embedding-model)]
(binding [index-table/*vectors-table* (str "library_entity_index_test_" suffix)
index-table/*meta-table* (str "library_entity_index_meta_test_" suffix)]
(try
(testing "configured + licensed but no index table yet -> unavailable"
(is (true? (entity-retrieval.core/available?)))
(is (false? (entity-retrieval.core/entity-retrieval-available?))))
(mt/with-temp [:model/Collection {lib-id :id} {:type "library" :location "/"}
:model/Collection {data-id :id} {:type "library-data" :location (str "/" lib-id "/")}
:model/Database {db-id :id} {}
:model/Table _ {:db_id db-id :collection_id data-id :is_published true
:active true :name "t" :display_name "T"}]
(reconcile/reconcile! ds (constantly semantic.tu/mock-embedding-model))
(testing "after a reconcile populates the index -> available"
(is (true? (entity-retrieval.core/entity-retrieval-available?))))
(testing "but a configured model that no longer matches the built index -> unavailable"
;; e.g. an embedding-dimension change before the rebuild reconcile runs: querying the
;; stale index would degrade to empty, so the tool must not be offered.
(mt/with-dynamic-fn-redefs [semantic.embedding/get-configured-model
(constantly (assoc semantic.tu/mock-embedding-model
:vector-dimensions 8))]
(is (false? (entity-retrieval.core/entity-retrieval-available?))))))
(finally
(jdbc/execute! ds [(str "DROP TABLE IF EXISTS "
index-table/*vectors-table* ", "
index-table/*meta-table*)]))))))))))
(deftest ^:sequential ranks-by-similarity-test
(testing "search ranks library documents by cosine similarity, nearest first"
(when semantic.db.datasource/db-url
(mt/with-premium-features #{:library :semantic-search}
(let [suffix (System/nanoTime)
ds (semantic.db.datasource/ensure-initialized-data-source!)
q "the query"]
(mt/with-dynamic-fn-redefs [semantic.embedding/get-configured-model
(constantly semantic.tu/mock-embedding-model)]
(binding [index-table/*vectors-table* (str "library_entity_index_test_" suffix)
index-table/*meta-table* (str "library_entity_index_meta_test_" suffix)]
;; the two tables' name docs embed "near"/"far"; q ties "near" exactly and is orthogonal to "far".
(semantic.tu/with-mock-embeddings {q [1.0 0.0 0.0 0.0]
"near" [1.0 0.0 0.0 0.0]
"far" [0.0 1.0 0.0 0.0]}
(mt/with-temp [:model/Collection {lib-id :id} {:type "library" :location "/"}
:model/Collection {data-id :id} {:type "library-data" :location (str "/" lib-id "/")}
:model/Database {db-id :id} {}
:model/Table {near-id :id} {:db_id db-id :collection_id data-id :is_published true
:active true :name "near" :display_name "near"}
:model/Table {far-id :id} {:db_id db-id :collection_id data-id :is_published true
:active true :name "far" :display_name "far"}]
(try
(reconcile/reconcile! ds (constantly semantic.tu/mock-embedding-model))
(testing "the nearer table ranks first; each hit carries the entity ref + matched doc"
(is (=? [{:entity {:model "table" :id near-id} :doc_type "name" :doc_text "near"}
{:entity {:model "table" :id far-id}}]
(mirror/search q 10))))
(finally
(jdbc/execute! ds [(str "DROP TABLE IF EXISTS "
index-table/*vectors-table* ", "
index-table/*meta-table*)]))))))))))))
(defn- put-ai-context!
"PUT an ai_context entry through the CRUD API."
[entity-type entity-local-id ai-context]
(mt/user-http-request :crowberto :put 200 (format "osi/ai-context/%s/%d" entity-type entity-local-id)
{:ai_context ai-context}))
(deftest ^:sequential crud-api-to-tool-end-to-end-test
(testing "CRUD API write -> reconcile -> pgvector -> retrieve_library_entities tool, end to end"
;; Self-gated on MB_PGVECTOR_DB_URL — CI without semantic-search infra skips this; locally with the
;; dev pgvector running it exercises the whole pipeline. Uses the mock embedding model (4-dim,
;; deterministic, no network) against isolated temp tables, with the reconciler run directly in
;; place of the background Quartz job.
(when semantic.db.datasource/db-url
(mt/with-premium-features #{:library :semantic-search}
(let [suffix (System/nanoTime)
ds (semantic.db.datasource/ensure-initialized-data-source!)
q "monthly revenue per region"
synonym "revenue per region"
;; bind a user: the tool's hydration permission-filters via mi/can-read?, which needs
;; *current-user-id* (a superuser can read the test tables).
search! #(mt/with-test-user :crowberto
(get-in (tools.entity-retrieval/retrieve-library-entities-tool {:user_search_prompt q})
[:structured-output :data]))]
(mt/with-dynamic-fn-redefs [semantic.embedding/get-configured-model
(constantly semantic.tu/mock-embedding-model)]
(binding [index-table/*vectors-table* (str "library_entity_index_test_" suffix)
index-table/*meta-table* (str "library_entity_index_meta_test_" suffix)]
;; the curated synonym ties the query exactly; the table's own name "orders" is orthogonal.
(semantic.tu/with-mock-embeddings {q [1.0 0.0 0.0 0.0]
synonym [1.0 0.0 0.0 0.0]
"orders" [0.0 1.0 0.0 0.0]}
(mt/with-temp [:model/Collection {lib-id :id} {:type "library" :location "/"}
:model/Collection {data-id :id} {:type "library-data" :location (str "/" lib-id "/")}
:model/Database {db-id :id} {}
:model/Table {table-id :id} {:db_id db-id :collection_id data-id :is_published true
:active true :name "orders" :display_name "orders"}]
(try
(let [_ (put-ai-context! "table" table-id
{:instructions "Group by month." :synonyms [synonym]})]
(reconcile/reconcile! ds (constantly semantic.tu/mock-embedding-model))
(testing "the tool returns the table, matched on the curated synonym, with usage_instructions"
(is (=? [{:type "table" :id table-id :matched_doc_type "synonym" :matched_text synonym
:usage_instructions "Group by month." :similarity (approx 1.0)}]
(search!))))
(testing "an entity that has left the library is post-filtered out before the index catches up"
;; unpublish without reconciling: the index still holds the docs, but the tool drops
;; the hit because the table is no longer a current library member.
(mt/with-temp-vals-in-db :model/Table table-id {:is_published false}
(is (empty? (search!)))))
(testing "deleting the ai_context via the CRUD API + reconcile drops the synonym doc"
(mt/user-http-request :crowberto :delete 204 (format "osi/ai-context/table/%d" table-id))
(reconcile/reconcile! ds (constantly semantic.tu/mock-embedding-model))
(is (empty? (jdbc/execute!
ds
[(format (str "SELECT 1 FROM \"%s\" "
"WHERE doc_type = 'synonym' AND entity_local_id = %d")
index-table/*vectors-table* table-id)])))))
(finally
(jdbc/execute! ds [(str "DROP TABLE IF EXISTS "
index-table/*vectors-table* ", "
index-table/*meta-table*)]))))))))))))
(deftest ^:sequential doc-type-boost-breaks-ties-test
(testing "on a distance tie, a name match outranks a synonym match (blended ORDER BY, not raw NN)"
(when semantic.db.datasource/db-url
(mt/with-premium-features #{:library :semantic-search}
(let [suffix (System/nanoTime)
ds (semantic.db.datasource/ensure-initialized-data-source!)
q "the query"
synonym "a curated synonym"]
(mt/with-dynamic-fn-redefs [semantic.embedding/get-configured-model
(constantly semantic.tu/mock-embedding-model)]
(binding [index-table/*vectors-table* (str "library_entity_index_test_" suffix)
index-table/*meta-table* (str "library_entity_index_meta_test_" suffix)]
;; alpha's name and beta's curated synonym both tie the query exactly; beta's name is orthogonal.
(semantic.tu/with-mock-embeddings {q [1.0 0.0 0.0 0.0]
"alpha" [1.0 0.0 0.0 0.0]
synonym [1.0 0.0 0.0 0.0]
"beta" [0.0 1.0 0.0 0.0]}
(mt/with-temp [:model/Collection {lib-id :id} {:type "library" :location "/"}
:model/Collection {data-id :id} {:type "library-data" :location (str "/" lib-id "/")}
:model/Database {db-id :id} {}
:model/Table {alpha :id} {:db_id db-id :collection_id data-id :is_published true
:active true :name "alpha" :display_name "alpha"}
:model/Table {beta :id} {:db_id db-id :collection_id data-id :is_published true
:active true :name "beta" :display_name "beta"}]
(try
(put-ai-context! "table" beta {:synonyms [synonym]})
(reconcile/reconcile! ds (constantly semantic.tu/mock-embedding-model))
;; raw NN would tie both at distance 0; the 0.02 vs 0.01 doc_type bump puts alpha's name first.
(is (= [[alpha "name"] [beta "synonym"]]
(mapv (juxt (comp :id :entity) :doc_type) (take 2 (mirror/search q 10)))))
(finally
(jdbc/execute! ds [(str "DROP TABLE IF EXISTS "
index-table/*vectors-table* ", "
index-table/*meta-table*)]))))))))))))
(deftest ^:sequential search-degrades-on-dimension-mismatch-test
;; A post-upgrade embedding-dimension change leaves the query vector incompatible with the index column
;; until the next reconcile rebuilds it; search must degrade to [] rather than throw in that window.
(testing "an index/query vector dimension mismatch degrades search to [] instead of throwing"
(when semantic.db.datasource/db-url
(mt/with-premium-features #{:library :semantic-search}
(let [suffix (System/nanoTime)
ds (semantic.db.datasource/ensure-initialized-data-source!)]
(mt/with-dynamic-fn-redefs [semantic.embedding/get-configured-model
(constantly semantic.tu/mock-embedding-model)]
(binding [index-table/*vectors-table* (str "library_entity_index_test_" suffix)
index-table/*meta-table* (str "library_entity_index_meta_test_" suffix)]
(mt/with-temp [:model/Collection {lib-id :id} {:type "library" :location "/"}
:model/Collection {data-id :id} {:type "library-data" :location (str "/" lib-id "/")}
:model/Database {db-id :id} {}
:model/Table {_t :id} {:db_id db-id :collection_id data-id :is_published true
:active true :name "orders" :display_name "orders"}]
(try
(semantic.tu/with-mock-embeddings {"orders" [1.0 0.0 0.0 0.0]}
(reconcile/reconcile! ds (constantly semantic.tu/mock-embedding-model)))
;; a dim-1 query vector against the dim-4 index column -> pgvector SQLState 22000 -> degrade to []
(semantic.tu/with-mock-embeddings {"q" [1.0]}
(is (= [] (entity-retrieval.core/search "q" 10))))
(finally
(jdbc/execute! ds [(str "DROP TABLE IF EXISTS \"" index-table/*vectors-table*
"\", \"" index-table/*meta-table* "\"")])))))))))))

420
enterprise/backend/test/metabase_enterprise/entity_retrieval/reconcile_test.clj Normal file
View File

@@ -0,0 +1,420 @@
(ns metabase-enterprise.entity-retrieval.reconcile-test
(:require
[clojure.test :refer :all]
[metabase-enterprise.entity-retrieval.index-table :as index-table]
[metabase-enterprise.entity-retrieval.reconcile :as reconcile]
[metabase-enterprise.semantic-search.db.datasource :as semantic.db.datasource]
[metabase-enterprise.semantic-search.test-util :as semantic.tu]
[metabase.collections.test-utils :as collections.tu]
[metabase.entity-retrieval.mirror :as mirror]
[metabase.measures.test-util :as measures.tu]
[metabase.metabot.tools.search :as tools.search]
[metabase.test :as mt]
[next.jdbc :as jdbc]
[next.jdbc.result-set :as jdbc.rs]
[toucan2.core :as t2]))
(set! *warn-on-reflection* true)
;; OsiAiContext write hooks fire request-entity-sync!, which would spawn a background targeted reconcile
;; using the *real* configured embedding model and race these tests — they drive reconcile! / reconcile-entity!
;; with the mock model directly. Suppress the nudge here; the hook -> drain path is covered in core-test.
;; The ignore is for re-binding a bang fn in a fixture: validate-deftest flags it as a destructive call, but
;; we only redef it to a no-op, never invoke it.
#_{:clj-kondo/ignore [:metabase/validate-deftest]}
(use-fixtures :each
(fn [thunk]
(mt/with-dynamic-fn-redefs [mirror/request-entity-sync! (fn [& _] nil)]
(thunk))))
(deftest doc-id-test
(testing "equal (entity_type, entity_local_id, doc_type, doc_text) tuples hash equal"
(is (= (reconcile/doc-id "metric" 9 "name" "Revenue")
(reconcile/doc-id "metric" 9 "name" "Revenue"))))
(testing "each component changes the doc_id (instructions is deliberately not an input)"
(let [d (reconcile/doc-id "metric" 9 "name" "Revenue")]
(doseq [variant [(reconcile/doc-id "table" 9 "name" "Revenue")
(reconcile/doc-id "metric" 8 "name" "Revenue")
(reconcile/doc-id "metric" 9 "synonym" "Revenue")
(reconcile/doc-id "metric" 9 "name" "Sales")]]
(is (not= d variant) (pr-str variant))))))
(deftest entity-class-test
(let [entity-class (var-get #'reconcile/entity-class)]
(testing "a Card's question/metric/model labels collapse to one class, so a type flip still matches"
(is (= (entity-class {:entity_type "metric" :entity_local_id 5})
(entity-class {:entity_type "model" :entity_local_id 5})
(entity-class {:entity_type "question" :entity_local_id 5}))))
(testing "a same-id entity of another type is a distinct class (no false-positive sparing)"
(is (not= (entity-class {:entity_type "table" :entity_local_id 5})
(entity-class {:entity_type "metric" :entity_local_id 5})))
(is (not= (entity-class {:entity_type "measure" :entity_local_id 5})
(entity-class {:entity_type "segment" :entity_local_id 5}))))))
(deftest entity->docs-bounds-oversized-ai-context-test
(testing "list length and per-value text are bounded at index time, regardless of how the row was written"
(let [entity->docs (var-get #'reconcile/entity->docs)
ctx {:synonyms (mapv #(str "synonym-" %) (range 200))
:examples [(apply str (repeat 9000 \x))]}
docs (entity->docs {:entity_type "table" :entity_local_id 1 :name "Orders" :description "d"}
ctx)]
(testing "an unbounded synonym list is capped (bounds bloat from API-bypassing writes)"
(is (= 50 (count (filter #(= "synonym" (:doc_type %)) docs)))))
(testing "every doc's text is truncated to the char cap"
(is (every? #(<= (count (:doc_text %)) 8000) docs))
(is (= 8000 (count (:doc_text (first (filter #(= "example" (:doc_type %)) docs))))))))))
(deftest format-embedding-rejects-invalid-values-test
(testing "non-numbers, NaN and infinities are rejected before they reach a raw SQL literal"
(doseq [bad [Double/NaN Double/POSITIVE_INFINITY Double/NEGATIVE_INFINITY "0.1"]]
(is (thrown-with-msg? Exception #"invalid value"
(index-table/format-embedding [0.1 bad 0.3]))
(pr-str bad)))
(is (string? (index-table/format-embedding [0.1 -0.2 0.3])))))
(defmacro ^:private with-isolated-index
"Run `body` with isolated pgvector tables bound and dropped afterwards, binding `ds-sym` to the
pgvector datasource. Skips the body entirely when no pgvector store is configured."
[[ds-sym] & body]
`(when semantic.db.datasource/db-url
(let [suffix# (System/nanoTime)
~ds-sym (semantic.db.datasource/ensure-initialized-data-source!)]
(binding [index-table/*vectors-table* (str "library_entity_index_test_" suffix#)
index-table/*meta-table* (str "library_entity_index_meta_test_" suffix#)]
(try
~@body
(finally
(jdbc/execute! ~ds-sym [(str "DROP TABLE IF EXISTS "
index-table/*vectors-table* ", "
index-table/*meta-table*)])))))))
(defn- index-rows [ds]
(jdbc/execute! ds
[(format "SELECT doc_id, entity_type, entity_local_id, doc_type, doc_text FROM \"%s\""
index-table/*vectors-table*)]
{:builder-fn jdbc.rs/as-unqualified-lower-maps}))
(defn- docs-for
"Index rows describing one entity (robust to any other library content in the test appdb)."
[ds entity-type entity-local-id]
(filter #(and (= entity-type (:entity_type %)) (= entity-local-id (:entity_local_id %)))
(index-rows ds)))
(deftest ^:sequential reconcile-lifecycle-test
;; :library lets us publish a Table into a library-data collection; the reconcile enumerates the
;; library tree via collections/library-collection + descendant-ids, so we build a real library root.
(mt/with-premium-features #{:library :semantic-search}
(with-isolated-index [ds]
(let [model semantic.tu/mock-embedding-model]
(mt/with-temp [:model/Collection {lib-id :id} {:type "library" :location "/"}
:model/Collection {data-id :id} {:type "library-data" :location (str "/" lib-id "/")}
:model/Collection {metrics-id :id} {:type "library-metrics" :location (str "/" lib-id "/")}
:model/Database {db-id :id} {}
:model/Table {table-id :id} {:db_id db-id :collection_id data-id
:is_published true :active true
:name "orders" :display_name "Orders"
:description "All orders"}
:model/Card {metric-id :id} {:type "metric" :collection_id metrics-id
:name "Revenue" :description "Total revenue"
:database_id db-id}
:model/OsiAiContext _ {:entity_type "table" :entity_local_id table-id
:ai_context {:instructions "Group by month."
:synonyms ["sales" "revenue"]
:examples ["orders last month"]}}]
(testing "first run indexes name/description docs for every library entity, plus ai_context docs"
(reconcile/reconcile! ds (constantly model))
(testing "the published table: name + description + 2 synonyms + 1 example"
(let [docs (docs-for ds "table" table-id)]
(is (= {"name" 1 "description" 1 "synonym" 2 "example" 1}
(frequencies (map :doc_type docs))))
(is (= #{"Orders"} (set (map :doc_text (filter #(= "name" (:doc_type %)) docs)))))))
(testing "the library metric: name + description only (no ai_context)"
(let [docs (docs-for ds "metric" metric-id)]
(is (= {"name" 1 "description" 1} (frequencies (map :doc_type docs)))))))
(testing "an unchanged second run writes nothing"
(is (=? {:inserted 0 :deleted 0} (reconcile/reconcile! ds (constantly model)))))
(testing "editing instructions is a no-op for the index (instructions are read live, not stored)"
(let [before (set (map :doc_id (docs-for ds "table" table-id)))]
(t2/update! :model/OsiAiContext :entity_type "table" :entity_local_id table-id
{:ai_context {:instructions "Group by week."
:synonyms ["sales" "revenue"]
:examples ["orders last month"]}})
(is (=? {:inserted 0 :deleted 0} (reconcile/reconcile! ds (constantly model))))
(is (= before (set (map :doc_id (docs-for ds "table" table-id)))) "doc_ids unchanged")))
(testing "renaming the table mints a new name doc_id and GCs the old one"
(let [old-name-id (->> (docs-for ds "table" table-id)
(filter #(= "name" (:doc_type %))) first :doc_id)]
(t2/update! :model/Table table-id {:display_name "Sales Orders"})
(reconcile/reconcile! ds (constantly model))
(let [names (filter #(= "name" (:doc_type %)) (docs-for ds "table" table-id))]
(is (= 1 (count names)))
(is (= "Sales Orders" (:doc_text (first names))))
(is (not= old-name-id (:doc_id (first names)))))))
(testing "unpublishing the table removes it from the library, GCing all its docs"
(t2/update! :model/Table table-id {:is_published false})
(reconcile/reconcile! ds (constantly model))
(is (empty? (docs-for ds "table" table-id)))
(is (seq (docs-for ds "metric" metric-id)) "the metric is untouched")))))))
(deftest ^:sequential reconcile!-runs-to-completion-test
(testing "reconcile! blocks until the run completes and a second run is idempotent"
(mt/with-premium-features #{:library :semantic-search}
(with-isolated-index [ds]
(collections.tu/with-library [{data :data}]
(let [model semantic.tu/mock-embedding-model]
(mt/with-temp [:model/Database {db-id :id} {}
:model/Table {table-id :id} {:db_id db-id
:collection_id (:id data)
:is_published true
:active true
:name "orders"
:display_name "Orders"}]
(testing "the run returns the diff and populates the index"
(let [result (reconcile/reconcile! ds (constantly model))]
(is (pos? (:inserted result)))
(is (seq (docs-for ds "table" table-id)))))
(testing "a second run writes nothing"
(is (=? {:inserted 0 :deleted 0} (reconcile/reconcile! ds (constantly model))))))))))))
(deftest ^:sequential rebuild-on-model-change-test
(mt/with-premium-features #{:library :semantic-search}
(with-isolated-index [ds]
(let [model semantic.tu/mock-embedding-model]
(mt/with-temp [:model/Collection {lib-id :id} {:type "library" :location "/"}
:model/Collection {data-id :id} {:type "library-data" :location (str "/" lib-id "/")}
:model/Database {db-id :id} {}
:model/Table {table-id :id} {:db_id db-id :collection_id data-id
:is_published true :active true
:name "orders" :display_name "Orders"}]
(testing "populate the index under the original model"
(reconcile/reconcile! ds (constantly model))
(is (seq (docs-for ds "table" table-id))))
(testing "a model-identity change drops the vectors table and the next run re-embeds everything"
(let [new-model (assoc model :model-name "model-v2")]
(is (= :rebuilt (index-table/ensure-tables! ds new-model)))
(is (= [] (index-rows ds)))
(reconcile/reconcile! ds (constantly new-model))
(is (seq (docs-for ds "table" table-id)))
(testing "a schema-version mismatch alone also triggers the rebuild"
(jdbc/execute! ds [(format "UPDATE \"%s\" SET schema_version = schema_version - 1"
index-table/*meta-table*)])
(is (= :rebuilt (index-table/ensure-tables! ds new-model)))
(is (= [] (index-rows ds))))
(testing "the rebuild heals the meta row, so it doesn't recur on the next sync"
(is (= :ok (index-table/ensure-tables! ds new-model)))))))))))
(deftest ^:sequential measures-and-segments-indexed-and-hydrated-test
(testing "measures/segments on a published library table are indexed and hydrate with parent-table context"
(mt/with-premium-features #{:library :semantic-search}
(with-isolated-index [ds]
(let [model semantic.tu/mock-embedding-model
orders (mt/id :orders)
total (mt/id :orders :total)]
(mt/with-temp [:model/Collection {lib-id :id} {:type "library" :location "/"}
:model/Collection {data-id :id} {:type "library-data" :location (str "/" lib-id "/")}
:model/Measure {measure-id :id} {:name "Order Revenue" :description "sum of order totals"
:table_id orders :creator_id (mt/user->id :crowberto)
:definition (measures.tu/measure-definition orders total)}
:model/Segment {segment-id :id} {:name "Big Orders" :description "totals over 100"
:table_id orders
:definition (measures.tu/segment-definition orders total 100)}]
;; publish the (real, fielded) orders table into the library for the duration of the test
(mt/with-temp-vals-in-db :model/Table orders {:collection_id data-id :is_published true}
(reconcile/reconcile! ds (constantly model))
(testing "both are indexed with name + description docs"
(is (= {"name" 1 "description" 1}
(frequencies (map :doc_type (docs-for ds "measure" measure-id)))))
(is (= {"name" 1 "description" 1}
(frequencies (map :doc_type (docs-for ds "segment" segment-id))))))
(testing "the tool hydrates them as first-class results carrying parent-table context"
(mt/with-test-user :crowberto
(let [by-key (into {} (map (juxt (juxt :type :id) identity))
(tools.search/entity-refs->search-results
[{:model "measure" :id measure-id} {:model "segment" :id segment-id}]))]
;; portable_entity_id (the NanoID) lets the agent reference the measure/segment in a
;; [measure|segment, {}, <id>] clause, the way a metric carries its own.
(is (=? {:type "measure" :id measure-id :name "Order Revenue"
:database_id (mt/id) :base_table_id orders :portable_entity_id string?}
(get by-key ["measure" measure-id])))
(is (=? {:type "segment" :id segment-id :name "Big Orders"
:database_id (mt/id) :base_table_id orders :portable_entity_id string?}
(get by-key ["segment" segment-id])))))))))))))
(deftest ^:sequential failed-insert-spares-only-that-entitys-orphans-test
(testing "a failed insert spares that entity's orphans; an unrelated entity's orphans still GC"
(mt/with-premium-features #{:library :semantic-search}
(with-isolated-index [ds]
(let [model semantic.tu/mock-embedding-model]
(mt/with-temp [:model/Collection {lib-id :id} {:type "library" :location "/"}
:model/Collection {data-id :id} {:type "library-data" :location (str "/" lib-id "/")}
:model/Database {db-id :id} {}
:model/Table {edited :id} {:db_id db-id :collection_id data-id :is_published true
:active true :name "a" :display_name "Edited Orig"}
:model/Table {leaving :id} {:db_id db-id :collection_id data-id :is_published true
:active true :name "b" :display_name "Leaving"}]
(reconcile/reconcile! ds (constantly model))
;; `edited` is renamed (old name doc -> orphan, new name doc -> to-insert);
;; `leaving` is unpublished (its docs -> orphans, with nothing in to-insert).
(t2/update! :model/Table edited {:display_name "Edited New"})
(t2/update! :model/Table leaving {:is_published false})
;; force the insert of the only to-insert batch (edited's new name doc) to fail.
(mt/with-dynamic-fn-redefs [reconcile/insert-batch! (fn [& _] (throw (ex-info "boom" {})))]
(reconcile/reconcile! ds (constantly model)))
(testing "the edited entity's stale orphan is retained (its replacement insert failed)"
(is (= #{"Edited Orig"} (set (map :doc_text (docs-for ds "table" edited))))))
(testing "the entity that left the library is still GC'd (no failed insert of its own)"
(is (empty? (docs-for ds "table" leaving))))))))))
(deftest ^:sequential reconcile-entity!-targets-one-slice-test
(testing "reconcile-entity! reconciles only the given entity's docs, leaving other entities untouched"
(mt/with-premium-features #{:library :semantic-search}
(with-isolated-index [ds]
(collections.tu/with-library [{data :data}]
(let [model semantic.tu/mock-embedding-model]
(mt/with-temp [:model/Database {db-id :id} {}
:model/Table {a-id :id} {:db_id db-id :collection_id (:id data) :is_published true
:active true :name "a" :display_name "Orders"
:description "all orders"}
:model/Table {b-id :id} {:db_id db-id :collection_id (:id data) :is_published true
:active true :name "b" :display_name "Customers"}]
(reconcile/reconcile! ds (constantly model))
(let [b-before (set (map :doc_id (docs-for ds "table" b-id)))]
(testing "adding ai_context to A then reconciling only A inserts A's synonym, leaves B alone"
(mt/with-temp [:model/OsiAiContext _ {:entity_type "table" :entity_local_id a-id
:ai_context {:synonyms ["sales"]}}]
(is (=? {:inserted 1 :deleted 0}
(reconcile/reconcile-entity! ds (constantly model) "table" a-id)))
(is (= {"name" 1 "description" 1 "synonym" 1}
(frequencies (map :doc_type (docs-for ds "table" a-id)))))
(is (= b-before (set (map :doc_id (docs-for ds "table" b-id)))) "B untouched")))))))))))
(deftest ^:sequential reconcile-entity!-leaving-library-deletes-all-test
(testing "reconcile-entity! on an entity that has left the library GCs all of its docs"
(mt/with-premium-features #{:library :semantic-search}
(with-isolated-index [ds]
(collections.tu/with-library [{data :data}]
(let [model semantic.tu/mock-embedding-model]
(mt/with-temp [:model/Database {db-id :id} {}
:model/Table {a-id :id} {:db_id db-id :collection_id (:id data) :is_published true
:active true :name "a" :display_name "Orders"}
:model/Table {b-id :id} {:db_id db-id :collection_id (:id data) :is_published true
:active true :name "b" :display_name "Customers"}]
(reconcile/reconcile! ds (constantly model))
(is (seq (docs-for ds "table" a-id)))
(let [b-before (set (map :doc_id (docs-for ds "table" b-id)))]
(mt/with-temp-vals-in-db :model/Table a-id {:is_published false}
(is (=? {:inserted 0} (reconcile/reconcile-entity! ds (constantly model) "table" a-id)))
(is (empty? (docs-for ds "table" a-id)) "A (no longer a member) is GC'd")
(is (= b-before (set (map :doc_id (docs-for ds "table" b-id)))) "B untouched"))))))))))
(deftest ^:sequential reconcile-entity!-ai-context-removal-keeps-name-test
(testing "removing an entity's ai_context and reconciling it GCs synonym/example docs but keeps name/description"
(mt/with-premium-features #{:library :semantic-search}
(with-isolated-index [ds]
(collections.tu/with-library [{data :data}]
(let [model semantic.tu/mock-embedding-model]
(mt/with-temp [:model/Database {db-id :id} {}
:model/Table {a-id :id} {:db_id db-id :collection_id (:id data) :is_published true
:active true :name "a" :display_name "Orders"
:description "all orders"}
:model/OsiAiContext _ {:entity_type "table" :entity_local_id a-id
:ai_context {:synonyms ["sales" "revenue"]}}]
(reconcile/reconcile! ds (constantly model))
(is (= {"name" 1 "description" 1 "synonym" 2}
(frequencies (map :doc_type (docs-for ds "table" a-id)))))
(t2/delete! :model/OsiAiContext :entity_type "table" :entity_local_id a-id)
(is (=? {:inserted 0 :deleted 2} (reconcile/reconcile-entity! ds (constantly model) "table" a-id)))
(is (= {"name" 1 "description" 1} (frequencies (map :doc_type (docs-for ds "table" a-id))))
"the synonym docs are GC'd; name + description remain"))))))))
(deftest ^:sequential reconcile!-keeps-ai-context-across-a-card-type-flip-test
(testing "a full reconcile matches ai_context by entity class, so relabelling a card keeps its synonyms"
(mt/with-premium-features #{:library :semantic-search}
(with-isolated-index [ds]
(collections.tu/with-library [{metrics :metrics}]
(let [model semantic.tu/mock-embedding-model]
(mt/with-temp [:model/Database {db-id :id} {}
:model/Card {card-id :id} {:type "metric" :collection_id (:id metrics)
:name "Revenue" :database_id db-id}
;; ai_context curated while the card was a metric (stored under entity_type "metric")
:model/OsiAiContext _ {:entity_type "metric" :entity_local_id card-id
:ai_context {:synonyms ["sales" "turnover"]}}]
(reconcile/reconcile! ds (constantly model))
(is (= {"name" 1 "synonym" 2} (frequencies (map :doc_type (docs-for ds "metric" card-id))))
"indexed under metric with both curated synonyms")
(testing "relabelling the card a model re-keys its docs but the curated synonyms survive"
(mt/with-temp-vals-in-db :model/Card card-id {:type "model"}
(reconcile/reconcile! ds (constantly model))
(is (empty? (docs-for ds "metric" card-id)) "stale metric-typed docs are GC'd")
(is (= {"name" 1 "synonym" 2} (frequencies (map :doc_type (docs-for ds "model" card-id))))
"the ai_context (stored under metric) is matched by class and kept, re-keyed under model"))))))))))
(deftest ^:sequential library-entity-matches-library-entities-test
(testing "library-entity (point lookup) agrees with library-entities (full scan) for members and non-members"
(mt/with-premium-features #{:library :semantic-search}
(collections.tu/with-library [{data :data metrics :metrics}]
(mt/with-temp [:model/Database {db-id :id} {}
:model/Table {tbl :id} {:db_id db-id :collection_id (:id data) :is_published true
:active true :name "t" :display_name "T"}
:model/Card {metric :id} {:type "metric" :collection_id (:id metrics)
:name "M" :database_id db-id}
:model/Table {unpub :id} {:db_id db-id :collection_id (:id data) :is_published false
:active true :name "u"}
:model/Card {archived :id} {:type "metric" :collection_id (:id metrics) :archived true
:name "A" :database_id db-id}]
(let [by-key (into {} (map (juxt (juxt :entity_type :entity_local_id) identity))
(#'reconcile/library-entities))]
(testing "members resolve and match the full-scan entry"
(is (= (get by-key ["table" tbl]) (reconcile/library-entity "table" tbl)))
(is (= (get by-key ["metric" metric]) (reconcile/library-entity "metric" metric))))
(testing "non-members resolve to nil"
(is (nil? (reconcile/library-entity "table" unpub)))
(is (nil? (reconcile/library-entity "metric" archived))))))))))
(deftest library-root-collection-included-in-membership-test
(testing "the Library root id is in lib-ids, so an entity directly in the root would be indexed too"
;; Defensive: the appdb normally prevents leaf content directly in the Library root (it lives in the
;; Data/Metrics sub-collections), so this can't be exercised end-to-end — but membership covers the
;; root id, not just its descendants.
(mt/with-premium-features #{:library}
(collections.tu/with-library [{library :library}]
(is (contains? (set (#'reconcile/library-ids library)) (:id library)))))))
(deftest ^:sequential reconcile-entity!-on-first-build-repopulates-whole-library-test
(testing "a targeted reconcile that creates the index (first caller, empty table) repopulates the whole library"
(mt/with-premium-features #{:library :semantic-search}
(with-isolated-index [ds]
(collections.tu/with-library [{data :data}]
(let [model semantic.tu/mock-embedding-model]
(mt/with-temp [:model/Database {db-id :id} {}
:model/Table {a-id :id} {:db_id db-id :collection_id (:id data) :is_published true
:active true :name "a" :display_name "Orders"}
:model/Table {b-id :id} {:db_id db-id :collection_id (:id data) :is_published true
:active true :name "b" :display_name "Customers"}]
;; No prior reconcile! — reconcile-entity! is the first caller, so ensure-tables! returns :created
;; (empty table). It must repopulate the whole library, not index only the one dirty entity.
(reconcile/reconcile-entity! ds (constantly model) "table" a-id)
(is (seq (docs-for ds "table" a-id)) "A indexed")
(is (seq (docs-for ds "table" b-id))
"B indexed too: the empty-index build escalated to a full repopulate"))))))))
(deftest ^:sequential reconcile-entity!-on-rebuild-repopulates-whole-library-test
(testing "a targeted reconcile that triggers a model/format rebuild repopulates the whole library, not just the one entity"
(mt/with-premium-features #{:library :semantic-search}
(with-isolated-index [ds]
(collections.tu/with-library [{data :data}]
(let [model semantic.tu/mock-embedding-model]
(mt/with-temp [:model/Database {db-id :id} {}
:model/Table {a-id :id} {:db_id db-id :collection_id (:id data) :is_published true
:active true :name "a" :display_name "Orders"}
:model/Table {b-id :id} {:db_id db-id :collection_id (:id data) :is_published true
:active true :name "b" :display_name "Customers"}]
(reconcile/reconcile! ds (constantly model))
(is (seq (docs-for ds "table" a-id)))
(is (seq (docs-for ds "table" b-id)))
;; a targeted reconcile of A under a new model identity forces ensure-tables! to rebuild (empty);
;; it must repopulate B too, not leave it missing until the periodic backstop.
(let [new-model (assoc model :model-name "model-v2")]
(is (:rebuilt? (reconcile/reconcile-entity! ds (constantly new-model) "table" a-id)))
(is (seq (docs-for ds "table" a-id)) "A repopulated after the rebuild")
(is (seq (docs-for ds "table" b-id)) "B repopulated too (not dropped by the rebuild)")))))))))

3
enterprise/backend/test/metabase_enterprise/serialization/models/entity_id_test.clj
View File

@@ -24,6 +24,9 @@
:model/Field
:model/FieldValues
:model/FieldUserSettings
;; OsiAiContext is identified by the entity it describes (entity_type + the entity's portable ref); its
;; serdes path nests under that entity, so it has no generated entity_id.
:model/OsiAiContext
;; Settings have human-selected unique names.
:model/Setting
;; Glossary items have unique `term` key

136
enterprise/backend/test/metabase_enterprise/serialization/v2/curated_search_entry_test.clj
View File

@@ -1,136 +0,0 @@
(ns metabase-enterprise.serialization.v2.curated-search-entry-test
"Round-trip serialization tests for `:model/CuratedSearchEntry`, whose `entity` column holds a
polymorphic Card/Table reference that must survive export and import."
(:require
[clojure.java.io :as io]
[clojure.test :refer :all]
[metabase-enterprise.serialization.test-util :as ts]
[metabase-enterprise.serialization.v2.extract :as extract]
[metabase-enterprise.serialization.v2.ingest :as serdes.ingest]
[metabase-enterprise.serialization.v2.load :as serdes.load]
[metabase-enterprise.serialization.v2.storage :as storage]
[metabase-enterprise.serialization.v2.storage.files :as storage.files]
[metabase.models.serialization :as serdes]
[metabase.search.core :as search]
[metabase.test :as mt]
[metabase.warehouses.models.database :as models.database]
[toucan2.core :as t2]))
(set! *warn-on-reflection* true)
;; The references point at the H2 test-data Table, so keep the H2 guard off and re-enable H2 in extract.
#_{:clj-kondo/ignore [:metabase/validate-deftest]}
(use-fixtures :each (fn [thunk]
(mt/with-dynamic-fn-redefs [search/reindex! (constantly nil)
models.database/assert-not-h2! (constantly nil)]
(binding [models.database/*include-h2-in-extract?* true]
(thunk)))))
(defn- no-labels [path]
(mapv #(dissoc % :label) path))
(defn- ingestion-in-memory [extractions]
(let [mapped (into {} (for [entity (vec extractions)]
[(no-labels (serdes/path entity)) entity]))]
(reify serdes.ingest/Ingestable
(ingest-list [_] (keys mapped))
(ingest-one [_ path] (get mapped (no-labels path)))
(ingest-errors [_] []))))
(deftest card-backed-entity-round-trip-test
(testing "every card-backed entity model exports as the Card's entity_id and imports back"
(mt/with-temp [:model/Card {card-id :id card-eid :entity_id} {}]
(doseq [entity-model ["card" "model" "metric" "question"]]
(testing entity-model
(mt/with-temp [:model/CuratedSearchEntry {sli-id :id sli-eid :entity_id}
{:search_prompt (str "find " entity-model) :verified true
:usage_instructions "Use it directly."
:entity {:model entity-model :id card-id}}]
(let [extracted (ts/extract-one "CuratedSearchEntry" sli-id)]
(testing "the entity ref is exported as a portable reference"
(is (=? {:entity_id sli-eid
:search_prompt (str "find " entity-model)
:usage_instructions "Use it directly."
:verified true
:entity {:model entity-model :id card-eid}}
extracted)))
(testing "dependencies cover the referenced Card"
(is (= #{[{:model "Card" :id card-eid}]}
(serdes/dependencies extracted))))
(testing "importing resolves the ref back to the local id"
(t2/delete! :model/CuratedSearchEntry :id sli-id)
(serdes.load/load-metabase! (ingestion-in-memory [extracted]))
(is (= {:model entity-model :id card-id}
(t2/select-one-fn :entity :model/CuratedSearchEntry :entity_id sli-eid)))
;; load-metabase! re-inserts under a new id that with-temp can't reap; clean it up so the
;; row doesn't leak into other tests' view of the (small, fully-scanned) table.
(t2/delete! :model/CuratedSearchEntry :entity_id sli-eid)))))))))
(deftest unmapped-entity-model-passes-through-test
(testing "an unrecognized entity model exports as-is (raw id, no deps) instead of aborting the export"
;; The CRUD API rejects unknown models, but rows written before a model string was retired (or by
;; direct appdb writes) shouldn't take down a whole export.
(mt/with-temp [:model/CuratedSearchEntry {sli-id :id}
{:search_prompt "legacy" :entity {:model "garbage" :id 5}}]
(let [extracted (ts/extract-one "CuratedSearchEntry" sli-id)]
(is (=? {:entity {:model "garbage" :id 5}} extracted))
(is (= #{} (serdes/dependencies extracted)))))))
(deftest table-entity-round-trip-test
(let [table-id (mt/id :venues)]
(mt/with-temp [:model/CuratedSearchEntry {sli-id :id sli-eid :entity_id}
{:search_prompt "best venues" :entity {:model "table" :id table-id}}]
(let [extracted (ts/extract-one "CuratedSearchEntry" sli-id)
table-ref (:entity extracted)]
(testing "the table ref is exported as a [db schema table] path"
(is (=? {:entity_id sli-eid
:entity {:model "table" :id vector?}}
extracted))
(is (= ["VENUES"] (take-last 1 (:id table-ref)))))
(testing "dependencies cover the Table's Database (the Table is synthesized on import)"
(is (= #{[{:model "Database" :id (first (:id table-ref))}]}
(serdes/dependencies extracted))))
(testing "importing resolves the ref back to the local id"
(t2/delete! :model/CuratedSearchEntry :id sli-id)
(serdes.load/load-metabase! (ingestion-in-memory [extracted]))
(is (= {:model "table" :id table-id}
(t2/select-one-fn :entity :model/CuratedSearchEntry :entity_id sli-eid)))
;; load-metabase! re-inserts under a new id that with-temp can't reap; clean it up so the row
;; doesn't leak into other tests' view of the (small, fully-scanned) table.
(t2/delete! :model/CuratedSearchEntry :entity_id sli-eid))))))
(deftest disk-export-import-round-trip-test
(testing "CuratedSearchEntry survives a real on-disk export/import"
;; The in-memory tests above skip two layers that have silently dropped the model before:
;; `extract/model-set` (the default-export selection) and `ingest/legal-top-level-paths` (which
;; dirs the importer reads). This drives both through actual YAML on disk.
(let [serialized (atom nil)
sli-eid (atom nil)]
(ts/with-dbs [source-db dest-db]
(ts/with-db source-db
(let [db (ts/create! :model/Database :name "Prompt DB")
t1 (ts/create! :model/Table :name "ORDERS" :db_id (:id db) :schema "PUBLIC")
sli (ts/create! :model/CuratedSearchEntry
:search_prompt "orders" :verified true
:entity {:model "table" :id (:id t1)})]
(reset! sli-eid (:entity_id sli))
;; realize inside with-cache so the cached resolvers are actually used during extraction
(reset! serialized (serdes/with-cache (into [] (extract/extract {}))))))
(testing "default export selection includes it (guards extract/model-set)"
(is (= 1 (count (filter (fn [{[{:keys [model]}] :serdes/meta}]
(= model "CuratedSearchEntry"))
@serialized)))))
(ts/with-random-dump-dir [dump-dir "sli-serdes-"]
(storage/store! (seq @serialized) (storage.files/file-writer dump-dir))
(testing "stored under the curated_search_entries/ directory"
(is (.exists (io/file dump-dir "curated_search_entries"))))
(testing "ingest reads it back (guards ingest/legal-top-level-paths)"
(is (contains? (set (serdes.ingest/ingest-list (serdes.ingest/ingest-yaml dump-dir)))
[{:model "CuratedSearchEntry" :id @sli-eid}])))
(testing "loads into a fresh appdb with the table ref resolved"
(ts/with-db dest-db
(serdes/with-cache (serdes.load/load-metabase! (serdes.ingest/ingest-yaml dump-dir)))
(let [row (t2/select-one :model/CuratedSearchEntry :entity_id @sli-eid)]
(is (=? {:verified true} row))
(is (= "ORDERS"
(t2/select-one-fn :name :model/Table :id (get-in row [:entity :id]))))))))))))

160
enterprise/backend/test/metabase_enterprise/serialization/v2/osi_ai_context_test.clj Normal file
View File

@@ -0,0 +1,160 @@
(ns metabase-enterprise.serialization.v2.osi-ai-context-test
"Round-trip serialization tests for `:model/OsiAiContext`. Identity is the entity it describes: there is no
entity_id, so serdes nests each row under its entity's portable path (Card/Table/Measure/Segment) and
resolves the `(entity_type, entity_local_id)` key back from that path on import."
(:require
[clojure.test :refer :all]
[metabase-enterprise.serialization.test-util :as ts]
[metabase-enterprise.serialization.v2.extract :as extract]
[metabase-enterprise.serialization.v2.ingest :as serdes.ingest]
[metabase-enterprise.serialization.v2.load :as serdes.load]
[metabase-enterprise.serialization.v2.storage :as storage]
[metabase-enterprise.serialization.v2.storage.files :as storage.files]
[metabase.measures.test-util :as measures.tu]
[metabase.models.serialization :as serdes]
[metabase.search.core :as search]
[metabase.test :as mt]
[metabase.warehouses.models.database :as models.database]
[toucan2.core :as t2]))
(set! *warn-on-reflection* true)
;; The references point at the H2 test-data Table, so keep the H2 guard off and re-enable H2 in extract.
#_{:clj-kondo/ignore [:metabase/validate-deftest]}
(use-fixtures :each (fn [thunk]
(mt/with-dynamic-fn-redefs [search/reindex! (constantly nil)
models.database/assert-not-h2! (constantly nil)]
(binding [models.database/*include-h2-in-extract?* true]
(thunk)))))
(defn- no-labels [path]
(mapv #(dissoc % :label) path))
(defn- ingestion-in-memory [extractions]
(let [mapped (into {} (for [entity (vec extractions)]
[(no-labels (serdes/path entity)) entity]))]
(reify serdes.ingest/Ingestable
(ingest-list [_] (keys mapped))
(ingest-one [_ path] (get mapped (no-labels path)))
(ingest-errors [_] []))))
(defn- extract-for [entity-type entity-local-id]
(ts/extract-one "OsiAiContext" [:and [:= :entity_type entity-type] [:= :entity_local_id entity-local-id]]))
(defn- self-segment [extracted]
(last (serdes/path extracted)))
(deftest card-backed-entity-round-trip-test
(testing "every card-backed entity type stores as \"card\", nests under the Card's path, and imports back"
(doseq [entity-type ["card" "model" "metric" "question"]]
(testing entity-type
;; a fresh Card per flavor: they all normalize to the same `card` entity_type, so reusing one card
;; would collide on the (card, id) primary key across iterations.
(mt/with-temp [:model/Card {card-id :id card-eid :entity_id} {}
:model/OsiAiContext _ {:ai_context {:instructions "Use it directly." :synonyms [entity-type]}
:entity_type entity-type :entity_local_id card-id}]
(let [extracted (extract-for "card" card-id)]
(testing "path nests the row under the Card; ai_context copies verbatim; no entity_id/key fields"
(is (=? {:ai_context {:instructions "Use it directly." :synonyms [entity-type]}
:serdes/meta [{:model "Card" :id card-eid} {:model "OsiAiContext" :id "ai_context"}]}
extracted))
(is (not (contains? extracted :entity_id))))
(testing "depends on the referenced Card"
(is (= [[{:model "Card" :id card-eid}]] (serdes/dependencies extracted))))
(testing "importing resolves the path back to the local card row"
(t2/delete! :model/OsiAiContext :entity_type "card" :entity_local_id card-id)
(serdes.load/load-metabase! (ingestion-in-memory [extracted]))
(is (=? {:entity_type "card" :entity_local_id card-id}
(t2/select-one :model/OsiAiContext :entity_type "card" :entity_local_id card-id))))))))))
(deftest measure-segment-entity-round-trip-test
(testing "measure and segment refs nest under the entity's path and import back"
(let [orders (mt/id :orders)
total (mt/id :orders :total)]
(mt/with-temp [:model/Measure {measure-id :id measure-eid :entity_id}
{:name "M" :table_id orders :creator_id (mt/user->id :crowberto)
:definition (measures.tu/measure-definition orders total)}
:model/Segment {segment-id :id segment-eid :entity_id}
{:name "S" :table_id orders :definition (measures.tu/segment-definition orders total 100)}]
(doseq [[entity-type id eid serdes-model] [["measure" measure-id measure-eid "Measure"]
["segment" segment-id segment-eid "Segment"]]]
(testing entity-type
(mt/with-temp [:model/OsiAiContext _ {:ai_context {:synonyms [entity-type]}
:entity_type entity-type :entity_local_id id}]
(let [extracted (extract-for entity-type id)]
(testing "path nests under the entity"
(is (=? {:serdes/meta [{:model serdes-model :id eid} {:model "OsiAiContext" :id "ai_context"}]}
extracted)))
(testing "depends on the referenced entity"
(is (= [[{:model serdes-model :id eid}]] (serdes/dependencies extracted))))
(testing "importing resolves the path back to the local id"
(t2/delete! :model/OsiAiContext :entity_type entity-type :entity_local_id id)
(serdes.load/load-metabase! (ingestion-in-memory [extracted]))
(is (=? {:entity_type entity-type :entity_local_id id}
(t2/select-one :model/OsiAiContext :entity_type entity-type :entity_local_id id))))))))))))
(deftest table-entity-round-trip-test
(let [table-id (mt/id :venues)]
(mt/with-temp [:model/OsiAiContext _ {:ai_context {:instructions "best venues"}
:entity_type "table" :entity_local_id table-id}]
(let [extracted (extract-for "table" table-id)]
(testing "path nests under the Table's [Database Schema Table] path"
(is (= {:model "OsiAiContext" :id "ai_context"} (self-segment extracted)))
(is (= "Table" (:model (last (pop (vec (serdes/path extracted))))))))
(testing "depends on the Table"
(is (= [(vec (pop (vec (serdes/path extracted))))] (serdes/dependencies extracted))))
(testing "importing resolves the path back to the local id"
(t2/delete! :model/OsiAiContext :entity_type "table" :entity_local_id table-id)
(serdes.load/load-metabase! (ingestion-in-memory [extracted]))
(is (=? {:entity_type "table" :entity_local_id table-id}
(t2/select-one :model/OsiAiContext :entity_type "table" :entity_local_id table-id))))))))
(deftest import-over-existing-row-updates-by-compound-key-test
(testing "importing over an existing row updates it in place via the full compound key, not entity_type alone"
(mt/with-temp [:model/Card {card-id :id} {}
:model/OsiAiContext _ {:entity_type "metric" :entity_local_id card-id
:ai_context {:instructions "v1"}}]
(let [extracted (assoc-in (extract-for "card" card-id) [:ai_context :instructions] "v2")]
(serdes.load/load-metabase! (ingestion-in-memory [extracted]))
(is (= 1 (t2/count :model/OsiAiContext :entity_type "card" :entity_local_id card-id))
"still one row (updated, not duplicated)")
(is (=? {:ai_context {:instructions "v2"}}
(t2/select-one :model/OsiAiContext :entity_type "card" :entity_local_id card-id)))))))
(deftest disk-round-trip-card-context-test
(testing "a card-backed ai_context survives a real on-disk store/ingest/load (storage-path handles a non-table parent)"
(mt/with-temp [:model/Card {card-id :id} {}
:model/OsiAiContext _ {:entity_type "metric" :entity_local_id card-id
:ai_context {:instructions "card guidance"}}]
(let [extracted (serdes/with-cache (extract-for "card" card-id))]
(ts/with-random-dump-dir [dump-dir "osi-card-"]
;; store! exercises storage-path for a Card parent — the regression that used to throw.
(storage/store! [extracted] (storage.files/file-writer dump-dir))
(t2/delete! :model/OsiAiContext :entity_type "card" :entity_local_id card-id)
(serdes/with-cache (serdes.load/load-metabase! (serdes.ingest/ingest-yaml dump-dir)))
(is (=? {:ai_context {:instructions "card guidance"}}
(t2/select-one :model/OsiAiContext :entity_type "card" :entity_local_id card-id))))))))
(deftest osi-ai-context-excluded-from-default-export-test
(testing "OsiAiContext is kept out of the default export set"
;; it's a top-level model extracted unfiltered, so auto-exporting it orphans rows whose entity is in an
;; omitted (personal/analytics) collection — leaked curator text + dangling deps. Excluded until
;; reverse-dependency export lands (BOT-1759), so its row travels with its entity instead.
(is (not (contains? (#'extract/model-set {}) "OsiAiContext")))
(is (not (contains? (#'extract/model-set {:targets [["Collection" 1]]}) "OsiAiContext")))))
(deftest excluded-from-default-disk-export-test
(testing "a real default export omits OsiAiContext while still exporting its parent entity"
;; End-to-end guard through extract/model-set that we don't accidentally re-include it and leak orphans
;; (rows whose entity is in an omitted personal/analytics collection). Re-include via BOT-1759.
(ts/with-dbs [source-db]
(ts/with-db source-db
(let [db (ts/create! :model/Database :name "Prompt DB")
t1 (ts/create! :model/Table :name "ORDERS" :db_id (:id db) :schema "PUBLIC")]
(ts/create! :model/OsiAiContext
:ai_context {:instructions "orders guidance"}
:entity_type "table" :entity_local_id (:id t1))
(let [models (set (map (fn [e] (:model (last (:serdes/meta e))))
(serdes/with-cache (into [] (extract/extract {})))))]
(is (contains? models "Table") "the parent entity is exported")
(is (not (contains? models "OsiAiContext")) "but its ai_context is not")))))))

6
enterprise/backend/test/metabase_enterprise/serialization/v2/round_trip_test.clj
View File

@@ -133,9 +133,9 @@
(def ^:private covered-by-dedicated-round-trip-test?
"Models that have full export/import coverage in their own round-trip test and so don't need a
fixture in this shared baseline. CuratedSearchEntry is covered (in-memory + on-disk) by
metabase-enterprise.serialization.v2.curated-search-entry-test."
#{"CuratedSearchEntry"})
fixture in this shared baseline. OsiAiContext is covered (in-memory + on-disk) by
metabase-enterprise.serialization.v2.osi-ai-context-test."
#{"OsiAiContext"})
(defn add-to-baseline!
"Use this within v2.extract-test where relevant to add their fixtures to the baseline."

2
mage/src/mage/modules.clj
View File

@@ -104,7 +104,6 @@
collections
config
content-verification
curated-search
custom-viz-plugin
dashboards
documents
@@ -116,6 +115,7 @@
enterprise/sso
enterprise/transforms
enterprise/transforms-inspector
entity-retrieval
events
formatter
geojson

4
resources/metabot/prompts/llm_shape.selmer
View File

@@ -1,5 +1,5 @@
{% if is_measure %}
<measure measure_id="{{measure_id}}" name="{{measure_name}}" display_name="{{measure_display_name}}"{% if measure_portable_entity_id %} portable_entity_id="{{measure_portable_entity_id}}"{% endif %}>{% if measure_description %}
<measure measure_id="{{measure_id}}" name="{{measure_name}}" display_name="{{measure_display_name}}"{% if measure_database_name %} database_name="{{measure_database_name}}"{% endif %}{% if measure_base_table_fqn %} base_table_fully_qualified_name="{{measure_base_table_fqn}}"{% endif %}{% if measure_portable_entity_id %} portable_entity_id="{{measure_portable_entity_id}}"{% endif %}>{% if measure_description %}
{{measure_description}}{% endif %}{% if measure_definition %}
<definition>
{{measure_definition|safe}}
@@ -7,7 +7,7 @@
</measure>
{% endif %}
{% if is_segment %}
<segment id="{{segment_id}}" name="{{segment_name}}" display_name="{{segment_display_name}}"{% if segment_portable_entity_id %} portable_entity_id="{{segment_portable_entity_id}}"{% endif %}>{% if segment_description %}
<segment id="{{segment_id}}" name="{{segment_name}}" display_name="{{segment_display_name}}"{% if segment_database_name %} database_name="{{segment_database_name}}"{% endif %}{% if segment_base_table_fqn %} base_table_fully_qualified_name="{{segment_base_table_fqn}}"{% endif %}{% if segment_portable_entity_id %} portable_entity_id="{{segment_portable_entity_id}}"{% endif %}>{% if segment_description %}
{{segment_description}}{% endif %}{% if segment_definition %}
<definition>
{{segment_definition|safe}}

76
resources/metabot/prompts/system/natural-language-querying-fallback.selmer Normal file
View File

@@ -0,0 +1,76 @@
You are {{metabot_name}}, a query builder for Metabase. The user gives you a question in natural language; you produce a visualization that answers it against their actual data.
# What you're optimizing for
A single objective: **the query you build should return the data the user asked for on its first execution.**
You are judged on whether the produced query matches the user's intent against the actual schema and values in this Metabase instance — not on how many tool calls it took, how much you reasoned, or how thoroughly you verified. That means the important questions for you to ask yourself at each step are:
- Do I know enough about the data to build a query that will actually return rows?
- Am I building exactly what the user asked for — no more, no less?
- Am I using the best-available data source for this intent (a curated metric, a model, a saved question, or a raw table)?
{% include "metabot/prompts/shared/prompt_snippets/grounding.selmer" %}
{% include "metabot/prompts/shared/prompt_snippets/data-sources.selmer" %}
# Finding data: search for the right entity, then navigate in
You have two complementary ways to find the data behind a request:
- **`search`** ranks all matching data sources across the instance (tables, models, metrics, saved questions) by topic. Phrase the query as a full natural-language description of the data the user wants, not as keywords. There is no curated `usage_instructions` or confidence flag, so confirm a candidate fits with `read_resource` before building.
- **`read_resource`** navigates the instance by URI and returns *exactly* what exists within a scope you name — a database's schemas, a schema's tables, a collection's items, a table's fields and sample values. It never silently "misses" the way a ranked match can.
So: start from `search` to find the best entity for the user's intent, then drill into the chosen entity with `read_resource` to confirm its fields and sample values before building — don't re-search the same concept once you have a promising hit. When a request spans several distinct concepts, issue one search per concept (fire them in parallel) so one concept's results don't crowd out another's.
{% include "metabot/prompts/shared/prompt_snippets/mbql-shape.selmer" %}
{% include "metabot/prompts/shared/prompt_snippets/communication.selmer" %}
{% include "metabot/prompts/shared/prompt_snippets/field-references.selmer" %}
{% include "metabot/prompts/shared/prompt_snippets/personality.selmer" %}
# Scope
Your job is to answer analytical questions by producing visualizations. For other needs:
- Dashboards: direct the user to the dashboard builder.
- SQL queries: direct the user to the SQL editor.
- Finding existing saved content: direct the user to search/browse.
{% include "metabot/prompts/shared/skills.selmer" %}
{% if custom_instructions %}
# Custom Instructions
{{ custom_instructions|safe }}
{% endif %}
<<<METABOT_CACHE_BREAKPOINT>>>
# Important Context
Current date and time: {{current_time}}
In the system, the first day of the week (index 1) is: {{ first_day_of_week }}
{% if current_user_info %}
<user_context>
Here is some information about the user:
{{ current_user_info|safe }}
{% if viewing_context %}
{{ viewing_context|safe }}
{% endif %}
</user_context>
{% elif viewing_context %}
<user_context>
{{ viewing_context|safe }}
</user_context>
{% endif %}
{% if recent_views %}
{{ recent_views|safe }}
{% endif %}

9
resources/metabot/prompts/system/natural-language-querying-only.selmer
View File

@@ -14,7 +14,14 @@ You are judged on whether the produced query matches the user's intent against t
{% include "metabot/prompts/shared/prompt_snippets/data-sources.selmer" %}
{% include "metabot/prompts/shared/prompt_snippets/discovery.selmer" %}
# Finding data: search for the right entity, then navigate in
You have two complementary ways to find the data behind a request:
- **`retrieve_library_entities`** finds the best entity for the user's intent from a vetted *library* of curated entities (metrics, models, saved questions, tables), each mapped to a saved prompt by *meaning*. Phrase the query as a full natural-language description of the data the user wants, not as keywords. Each match carries curator `usage_instructions` and a confidence flag — a result flagged low-confidence (or a leading note) means nothing in the library clearly matches, so prefer asking the user to clarify or narrow the request rather than building on a weak guess.
- **`read_resource`** navigates the instance by URI and returns *exactly* what exists within a scope you name — a database's schemas, a schema's tables, a collection's items, a table's fields and sample values. It never silently "misses" the way a ranked match can.
So: start from `retrieve_library_entities` to find the best entity for the user's intent, then drill into the chosen entity with `read_resource` to confirm its fields and sample values before building — don't re-search the same concept once you have a promising hit. When a request spans several distinct concepts, issue one search per concept (fire them in parallel) so one concept's results don't crowd out another's.
{% include "metabot/prompts/shared/prompt_snippets/mbql-shape.selmer" %}

133
resources/migrations/063/20260623_osi_ai_context.yaml Normal file
View File

@@ -0,0 +1,133 @@
## osi_ai_context: one row per library entity holding OSI ai_context metadata
## ({instructions, synonyms[], examples[]}). Replaces curated_search_entries.
## The table is authoritative; an enterprise pgvector index (library_entity_index, many embedded
## docs per entity) is reconciled against it plus live library membership and backs the
## retrieve_library_entities Metabot tool.
##
## Identity is the logical (entity_type, entity_local_id) pair — that compound key is the primary key,
## so there is one row per entity and no surrogate id or serialization NanoID. Card flavors
## (question/metric/model) are stored under the canonical "card" entity_type (normalized on write), so a
## card keeps its one row across a type relabel.
databaseChangeLog:
- logicalFilePath: migrations/063_update_migrations.yaml
- changeSet:
id: v63.2026-06-23T00:00:01
author: christruter
comment: Create osi_ai_context table
changes:
- createTable:
tableName: osi_ai_context
remarks: OSI ai_context metadata for library entities, one row per entity
columns:
- column:
name: entity_type
type: varchar(32)
remarks: Type of entity this ai_context describes (table, card, measure, segment)
constraints:
nullable: false
- column:
name: entity_local_id
type: bigint
remarks: Local appdb id of the entity this ai_context describes
constraints:
nullable: false
- column:
name: ai_context
type: ${text.type}
remarks: OSI ai_context JSON blob ({instructions, synonyms[], examples[]})
constraints:
nullable: false
- column:
name: created_at
type: ${timestamp_type}
remarks: When the ai_context was created
defaultValueComputed: current_timestamp
constraints:
nullable: false
- column:
name: updated_at
type: ${timestamp_type}
remarks: When the ai_context was last updated
defaultValueComputed: current_timestamp
constraints:
nullable: false
rollback:
- dropTable:
tableName: osi_ai_context
- changeSet:
id: v63.2026-06-23T00:00:02
author: christruter
comment: Make (entity_type, entity_local_id) the primary key of osi_ai_context
changes:
- addPrimaryKey:
tableName: osi_ai_context
columnNames: entity_type, entity_local_id
constraintName: pk_osi_ai_context
rollback:
- dropPrimaryKey:
tableName: osi_ai_context
constraintName: pk_osi_ai_context
- changeSet:
id: v63.2026-06-23T00:00:03
author: christruter
comment: Drop the superseded curated_search_entries table
changes:
- dropTable:
tableName: curated_search_entries
## Faithful rollback recreates the old table so cross-version down-migrations don't error.
## (Data is not restored — the feature was unreleased, so there is none to preserve.)
rollback:
- createTable:
tableName: curated_search_entries
remarks: Curated search prompts mapped to the entities that answer them
columns:
- column:
name: id
type: int
autoIncrement: true
constraints:
primaryKey: true
nullable: false
- column:
name: search_prompt
type: ${text.type}
constraints:
nullable: false
- column:
name: usage_instructions
type: ${text.type}
constraints:
nullable: true
- column:
name: entity
type: ${text.type}
constraints:
nullable: false
- column:
name: verified
type: ${boolean.type}
defaultValueBoolean: false
constraints:
nullable: false
- column:
name: entity_id
type: char(21)
constraints:
nullable: true
unique: true
- column:
name: created_at
type: ${timestamp_type}
defaultValueComputed: current_timestamp
constraints:
nullable: false
- column:
name: updated_at
type: ${timestamp_type}
defaultValueComputed: current_timestamp
constraints:
nullable: false

19
src/metabase/analytics/prometheus.clj
View File

@@ -460,6 +460,25 @@
{:description "Number of successful semantic search DLQ retries"})
(prometheus/counter :metabase-search/semantic-indexer-dlq-failures
{:description "Number of failed semantic search DLQ retries"})
;; library entity index (entity-retrieval) reconciliation
(prometheus/histogram :metabase-entity-retrieval/reconcile-duration-ms
{:description "Duration (ms) of a library entity index reconcile run, labelled :scope full | targeted."
:labels [:scope]
;; One shared bucket set has to resolve both scopes, which differ by orders of
;; magnitude: a `targeted` slice reconcile is typically sub-second, a `full` run
;; seconds-to-minutes. Fine low-end buckets keep targeted from collapsing into one
;; bucket; the high end still covers a full run. (~1ms -> 10min)
:buckets [1 10 50 100 250 500 1000 5000 10000 30000 60000 120000 300000 600000]})
(prometheus/counter :metabase-entity-retrieval/docs-inserted
{:description "Number of documents embedded and inserted into the library entity index."})
(prometheus/counter :metabase-entity-retrieval/docs-deleted
{:description "Number of documents garbage-collected from the library entity index."})
(prometheus/counter :metabase-entity-retrieval/search-failed
{:description "Number of library entity index searches that errored and returned no results (e.g. a dimension mismatch before the next reconcile heals it). The nlq profile has no in-profile fallback once the curated tool is offered, so this is a real failure."})
(prometheus/gauge :metabase-entity-retrieval/index-documents
{:description "Number of documents in the library entity index, as of the last full reconcile."})
(prometheus/gauge :metabase-entity-retrieval/index-entities
{:description "Number of distinct entities in the library entity index, as of the last full reconcile."})
;; data-complexity-score timing
;; 1ms → 1min buckets; widen later if real-world runs push past a minute.
(prometheus/histogram :metabase-data-complexity/scoring-duration-ms

6
src/metabase/api_routes/routes.clj
View File

@@ -21,7 +21,6 @@
[metabase.collections-rest.api]
[metabase.comments.api]
[metabase.config.core :as config]
[metabase.curated-search.api]
[metabase.dashboards-rest.api]
[metabase.data-studio.api]
[metabase.documents.api]
@@ -45,6 +44,7 @@
[metabase.native-query-snippets.api]
[metabase.notification.api]
[metabase.oauth-server.api.admin]
[metabase.osi.ai-context.api]
[metabase.permissions-rest.api]
[metabase.premium-features.api]
[metabase.product-feedback.api]
@@ -89,7 +89,6 @@
metabase.cloud-migration.api/keep-me
metabase.comments.api/keep-me
metabase.collections-rest.api/keep-me
metabase.curated-search.api/keep-me
metabase.dashboards-rest.api/keep-me
metabase.data-studio.api/keep-me
metabase.documents.api/keep-me
@@ -103,6 +102,7 @@
metabase.mcp.api/keep-me
metabase.mcp.callback-api/keep-me
metabase.oauth-server.api.admin/keep-me
metabase.osi.ai-context.api/keep-me
metabase.measures.api/keep-me
metabase.metrics.api/keep-me
metabase.model-persistence.api/keep-me
@@ -180,7 +180,6 @@
"/cloud-migration" (+auth 'metabase.cloud-migration.api)
"/collection" (+auth 'metabase.collections-rest.api)
"/comment" (+auth metabase.comments.api/routes)
"/curated-search" (+auth 'metabase.curated-search.api)
"/dashboard" (+auth 'metabase.dashboards-rest.api)
"/data-studio" (+auth metabase.data-studio.api/routes)
"/database" (+auth 'metabase.warehouses-rest.api)
@@ -214,6 +213,7 @@
"/notification" metabase.notification.api/notification-routes
"/notify" (+static-apikey metabase.sync.api/notify-routes)
"/oauth" (+auth 'metabase.oauth-server.api.admin)
"/osi" {"/ai-context" (+auth 'metabase.osi.ai-context.api)}
"/permissions" (+auth 'metabase.permissions-rest.api)
"/persist" (+auth 'metabase.model-persistence.api)
"/premium-features" (+auth metabase.premium-features.api/routes)

59
src/metabase/app_db/connection.clj
View File

@@ -194,16 +194,60 @@
[]
(pos? *transaction-depth*))
;; Accumulates thunks to run after the outermost transaction commits. Bound to a fresh atom when the
;; outermost transaction starts (see [[do-with-transaction]]); nil outside any transaction.
(def ^:private ^:dynamic *after-commit-callbacks* nil)
(defn do-after-commit
"Run `thunk` after the current outermost transaction commits successfully — never on rollback.
Outside a transaction (autocommit), runs `thunk` immediately — the surrounding write already committed.
Use to *schedule* post-commit work — enqueue async work, fire a `future`, publish an event — that must
observe committed state (e.g. a reconcile that reads the row).
Do not do synchronous DB I/O in `thunk`: it runs while the transaction's connection is still checked out,
so a query here would hold a second connection and can deadlock a saturated pool. Hand DB work to the
async job you schedule, which acquires its own connection."
[thunk]
(if-let [callbacks *after-commit-callbacks*]
(do (swap! callbacks conj thunk) nil)
(thunk)))
(defn- run-after-commit-callbacks! [callbacks]
;; Bind the transaction connection and callback accumulator to nil so they are not conveyed into async work
;; (e.g. a reconcile `future`) a callback may start: that work must acquire its own connection rather than
;; reuse this transaction's connection after it returns to the pool, and a do-after-commit it makes must run
;; immediately rather than enqueue into this now-drained accumulator.
(binding [t2.conn/*current-connectable* nil
*transaction-depth* 0
*after-commit-callbacks* nil]
(doseq [thunk @callbacks]
;; the transaction already committed; a failing callback must not unwind it
(try (thunk) (catch Throwable t (log/error t "after-commit callback failed"))))
(reset! callbacks [])))
(defn- discard-after-commit-callbacks-after! [callback-count]
(when-let [callbacks *after-commit-callbacks*]
(swap! callbacks
(fn [callbacks]
;; copy rather than return the subvec view, which would retain the discarded callbacks (and their
;; captured closures) through the backing array until the outer transaction finishes
(into [] (subvec callbacks 0 (min callback-count (count callbacks))))))))
(defn- do-transaction [^java.sql.Connection connection f]
(letfn [(thunk []
(let [savepoint (.setSavepoint connection)]
(let [savepoint (.setSavepoint connection)
callback-count (some-> *after-commit-callbacks* deref count)]
(try
(let [result (f connection)]
(when (= *transaction-depth* 1)
;; top-level transaction, commit
;; top-level transaction; post-commit side effects run after the transaction bindings unwind
(.commit connection))
result)
(catch Throwable txn-e
;; the nested body failed, so its callbacks must never fire — discard them before attempting
;; rollback, otherwise a throwing .rollback would leave them in the shared accumulator to run
;; at outer-commit time for data that was rolled back
(when callback-count
(discard-after-commit-callbacks-after! callback-count))
(try
(.rollback connection savepoint)
(catch Exception rollback-e
@@ -253,8 +297,15 @@
{:options options}))
:else
(binding [*transaction-depth* (inc *transaction-depth*)]
(do-transaction connection f))))
(let [outermost? (zero? *transaction-depth*)
callbacks (if outermost? (atom []) *after-commit-callbacks*)
result (binding [*transaction-depth* (inc *transaction-depth*)
;; one accumulator for the whole tree, created when the outermost transaction starts
*after-commit-callbacks* callbacks]
(do-transaction connection f))]
(when outermost?
(run-after-commit-callbacks! callbacks))
result)))
(methodical/defmethod t2.pipeline/transduce-query :before :default
"Make sure application database calls are not done inside core.async dispatch pool threads. This is done relatively

1
src/metabase/app_db/core.clj
View File

@@ -36,6 +36,7 @@
application-db
data-source
db-type
do-after-commit
in-transaction?
quoting-style
unique-identifier]

12
src/metabase/app_db/query.clj
View File

@@ -217,9 +217,11 @@
[model select-map & [update-fn]]
(let [update-fn (or update-fn (constantly select-map))
select-kvs (mapcat identity select-map)
pks (t2/primary-keys model)
_ (assert (= 1 (count pks)) "This helper does not currently support compound keys")
pk-key (keyword (first pks))
pks (mapv keyword (t2/primary-keys model))
;; The "pk" used to address an update and as the return value: a scalar for a single-column key
;; (unchanged), a `[v1 v2 ...]` vector for a compound key — matching what `insert-returning-pk!`
;; returns and what `t2/update!` accepts for a composite key.
pk-of (fn [row] (if (= 1 (count pks)) ((first pks) row) (mapv #(% row) pks)))
update-fn (fn [existing]
(let [updated (update-fn existing)]
;; When called with an existing record, returning nil signals "no update needed".
@@ -231,11 +233,11 @@
(merge updated select-map))))]
(with-conflict-retry
(if-let [existing (apply t2/select-one model select-kvs)]
(let [pk (pk-key existing)]
(let [pk (pk-of existing)]
(if-let [updated (update-fn existing)]
(do (t2/update! model pk updated)
;; the primary key may have been changed by the update, and this is OK.
(pk-key updated pk))
(pk-of (merge existing updated)))
;; update-fn returned nil — no update needed, return existing pk.
pk))
;; Wrap INSERT in a savepoint so that a constraint violation inside an existing

5
src/metabase/cmd/copy.clj
View File

@@ -139,7 +139,7 @@
:model/MetabotSourceFeedback
:model/MetabotUsedTable
:model/MetabotPrompt
:model/CuratedSearchEntry]
:model/OsiAiContext]
(when config/ee-available?
[:model/MetabotPermissions
:model/MetabotGroupLimit
@@ -383,7 +383,8 @@
:model/FieldUserSettings
:model/QueryAction
:model/MetabotConversation
:model/ModelIndexValue})
:model/ModelIndexValue
:model/OsiAiContext})
(defmulti ^:private postgres-id-sequence-name
{:arglists '([model])}

117
src/metabase/curated_search/api.clj
View File

@@ -1,117 +0,0 @@
(ns metabase.curated-search.api
"Admin REST API for managing curated search entries: curated search prompts mapped to the entities
that answer them."
(:require
[metabase.api.common :as api]
[metabase.api.macros :as api.macros]
[metabase.request.core :as request]
[metabase.util.malli.schema :as ms]
[toucan2.core :as t2]))
(def ^:private EntityRef
"Entity ref as accepted on writes: the model must be a known agent-facing entity type (as used with
read_resource), plus plain \"card\"."
[:map
[:model [:enum "table" "card" "model" "metric" "question"]]
[:id ms/PositiveInt]
[:name {:optional true} :string]])
(def ^:private EntityRefOut
"Entity ref as returned on reads. `:model` is any string: rows can predate a model string's retirement
(serdes tolerates those too), and one legacy row must not fail response validation for a whole list."
[:map
[:model :string]
[:id :int]
[:name {:optional true} :string]])
(def ^:private Entry
[:map
[:id ms/PositiveInt]
[:search_prompt :string]
[:usage_instructions {:optional true} [:maybe :string]]
[:entity EntityRefOut]
[:verified :boolean]])
(def ^:private default-limit 50)
(def ^:private default-offset 0)
(api.macros/defendpoint :get "/"
:- [:map
[:data [:sequential Entry]]
[:total :int]
[:limit :int]
[:offset :int]]
"Get all curated search entries, paginated."
[_route-params
_query-params]
(api/check-superuser)
(let [limit (or (request/limit) default-limit)
offset (or (request/offset) default-offset)]
{:data (t2/select :model/CuratedSearchEntry
{:order-by [[:id :asc]]
:limit limit
:offset offset})
:total (t2/count :model/CuratedSearchEntry)
:limit limit
:offset offset}))
(api.macros/defendpoint :get "/:id"
:- Entry
"Get a curated search entry by ID."
[{:keys [id]} :- [:map [:id ms/PositiveInt]]
_query-params]
(api/check-superuser)
(api/check-404 (t2/select-one :model/CuratedSearchEntry :id id)))
(api.macros/defendpoint :post "/"
:- Entry
"Create a new curated search entry."
[_route-params
_query-params
{:keys [search_prompt usage_instructions entity verified]}
:- [:map
[:search_prompt ms/NonBlankString]
[:usage_instructions {:optional true} [:maybe :string]]
[:entity EntityRef]
[:verified {:optional true} [:maybe :boolean]]]]
(api/check-superuser)
(t2/insert-returning-instance! :model/CuratedSearchEntry
{:search_prompt search_prompt
:usage_instructions usage_instructions
:entity entity
:verified (boolean verified)}))
(api.macros/defendpoint :put "/:id"
:- Entry
"Update a curated search entry by ID. Only the provided fields are changed; sending a null
usage_instructions clears it."
[{:keys [id]} :- [:map [:id ms/PositiveInt]]
_query-params
;; patch semantics: presence of a key (not nil-ness of its value) decides what changes, so an explicit
;; null can clear the nullable usage_instructions. The non-nullable fields reject null in the schema.
{:keys [search_prompt usage_instructions entity verified] :as body}
:- [:map
[:search_prompt {:optional true} ms/NonBlankString]
[:usage_instructions {:optional true} [:maybe :string]]
[:entity {:optional true} EntityRef]
[:verified {:optional true} :boolean]]]
(api/check-superuser)
(api/check-404 (t2/select-one :model/CuratedSearchEntry :id id))
(let [changes (cond-> {}
(contains? body :search_prompt) (assoc :search_prompt search_prompt)
(contains? body :usage_instructions) (assoc :usage_instructions usage_instructions)
(contains? body :entity) (assoc :entity entity)
(contains? body :verified) (assoc :verified verified))]
(when (seq changes)
(t2/update! :model/CuratedSearchEntry id changes)))
(t2/select-one :model/CuratedSearchEntry :id id))
#_{:clj-kondo/ignore [:metabase/validate-defendpoint-has-response-schema]}
(api.macros/defendpoint :delete "/:id"
"Delete a curated search entry by ID."
[{:keys [id]} :- [:map [:id ms/PositiveInt]]
_query-params]
(api/check-superuser)
(api/check-404 (t2/select-one :model/CuratedSearchEntry :id id))
(t2/delete! :model/CuratedSearchEntry :id id)
api/generic-204-no-content)

16
src/metabase/curated_search/core.clj
View File

@@ -1,16 +0,0 @@
(ns metabase.curated-search.core
"Public API of the `curated-search` module.
The curated search index is a hand-curated library of saved search prompts, each mapped to the entity that
answers it (see [[metabase.curated-search.models.curated-search-entry]]).
[[search]] matches a natural-language request against it by vector similarity (enterprise; returns []
in OSS)."
(:require
[metabase.curated-search.mirror]
[potemkin :as p]))
(comment metabase.curated-search.mirror/keep-me)
(p/import-vars
[metabase.curated-search.mirror
search])

24
src/metabase/curated_search/mirror.clj
View File

@@ -1,24 +0,0 @@
(ns metabase.curated-search.mirror
"OSS entry points for the enterprise pgvector mirror of the `curated_search_entries` table.
The appdb table is authoritative; the mirror carries one embedding per row and serves the
`search_curated` Metabot tool's similarity search.
These `defenterprise` shims route to [[metabase-enterprise.curated-search.core]] when an
enterprise license with semantic search is present; otherwise they no-op or return empty."
(:require
[metabase.premium-features.core :refer [defenterprise]]))
(defenterprise request-sync!
"Ask the background sync to reconcile the pgvector mirror with the appdb table soon; no-op in OSS.
Fire-and-forget: never throws, does no embedding or pgvector work on the calling thread."
metabase-enterprise.curated-search.core
[]
nil)
(defenterprise search
"Similarity-search the curated search index for the saved prompts nearest `user-search-prompt`.
Returns up to `limit` matches shaped `{:saved_search_prompt :usage_instructions :entity :score}`,
best blended score first; [] in OSS or when no pgvector store is configured."
metabase-enterprise.curated-search.core
[_user-search-prompt _limit]
[])

102
src/metabase/curated_search/models/curated_search_entry.clj
View File

@@ -1,102 +0,0 @@
(ns metabase.curated-search.models.curated-search-entry
"The `curated_search_entries` appdb table: curated saved search prompts, each mapped to the single entity
that answers it, plus agent-facing `usage_instructions`.
This table is authoritative.
An enterprise pgvector mirror carries one embedding per row and serves the `search_curated`
Metabot tool's similarity search.
Writes here only nudge the mirror's background sync ([[mirror/request-sync!]]); they never touch
the embedding service or the pgvector store themselves."
(:require
[metabase.curated-search.mirror :as mirror]
[metabase.models.interface :as mi]
[metabase.models.serialization :as serdes]
[metabase.util :as u]
[methodical.core :as methodical]
[toucan2.core :as t2]))
(methodical/defmethod t2/table-name :model/CuratedSearchEntry [_model] :curated_search_entries)
(doto :model/CuratedSearchEntry
(derive :metabase/model)
(derive :hook/entity-id)
(derive :hook/timestamped?))
(t2/deftransforms :model/CuratedSearchEntry
{:entity mi/transform-json})
;;; Each write nudges the enterprise background sync, which reconciles the pgvector mirror against this
;;; table. The nudge is fire-and-forget (no-op in OSS, error-swallowing in EE), so appdb writes never
;;; fail or slow down because of the mirror.
(t2/define-after-insert :model/CuratedSearchEntry
[row]
(mirror/request-sync!)
row)
(t2/define-after-update :model/CuratedSearchEntry
[row]
(mirror/request-sync!)
row)
(t2/define-before-delete :model/CuratedSearchEntry
[row]
(mirror/request-sync!)
row)
;;; ------------------------------------------------- Serialization -------------------------------------------------
;;; `entity` is a polymorphic reference — `{:model "table"|"card"|"model"|"metric"|"question" :id <local-id>}`.
;;; On export we swap the local id for a portable reference (a Card's entity_id, or a Table's
;;; `[db schema table]` path) and reverse it on import, mirroring `serdes/export-viz-link-card`.
(def ^:private entity-model->toucan
"Toucan model for each entity ref `:model` string. Tables get the table-fk treatment instead.
All the card flavors are Cards (the `:model` string mirrors the `read_resource` resource type the
agent uses to fetch the entity, not the underlying table)."
{"card" :model/Card
"model" :model/Card
"metric" :model/Card
"question" :model/Card})
(defn- export-entity-ref [{:keys [model id] :as entity}]
(cond
(= model "table") (assoc entity :id (serdes/*export-table-fk* id))
(entity-model->toucan model) (assoc entity :id (serdes/*export-fk* id (entity-model->toucan model)))
;; Unmapped model string: leave the ref as-is rather than aborting the whole export
;; (same forgiving behavior as serdes/export-viz-link-card).
:else entity))
(defn- import-entity-ref [{:keys [model id] :as entity}]
(cond
(= model "table") (assoc entity :id (serdes/*import-table-fk* id))
(entity-model->toucan model) (assoc entity :id (serdes/*import-fk* id (entity-model->toucan model)))
:else entity))
(defmethod serdes/make-spec "CuratedSearchEntry" [_model-name _opts]
{:copy [:entity_id :search_prompt :usage_instructions :verified]
:transform {:created_at (serdes/date)
:updated_at (serdes/date)
:entity {:export export-entity-ref
:import import-entity-ref}}
:defaults {:verified false}})
(defmethod serdes/hash-fields :model/CuratedSearchEntry
[_model]
[:search_prompt])
(defmethod serdes/generate-path "CuratedSearchEntry" [_ entity]
(serdes/maybe-labeled "CuratedSearchEntry" entity :search_prompt))
(defmethod serdes/storage-path "CuratedSearchEntry" [entity _ctx]
[{:label "curated_search_entries"}
{:label (u/slugify (:search_prompt entity) {:unicode? true}) :key (:entity_id entity)}])
(defmethod serdes/dependencies "CuratedSearchEntry"
[{{:keys [model id]} :entity}]
;; A referenced Table is synthesized on import if missing, so (like link cards) we depend on its Database,
;; not the Table itself. Card refs depend on the Card directly.
(cond
(= model "table") #{[{:model "Database" :id (first id)}]}
(entity-model->toucan model) #{[{:model "Card" :id id}]}
:else #{}))

82
src/metabase/entity_retrieval/core.clj Normal file
View File

@@ -0,0 +1,82 @@
(ns metabase.entity-retrieval.core
"Public API of the `entity-retrieval` module.
The library entity index is a pgvector index of every library entity's name/description plus its OSI
`ai_context` synonyms/examples (see [[metabase.osi.models.osi-ai-context]]).
[[search]] matches a natural-language request against it by vector similarity (enterprise; returns []
in OSS). [[ai-context-instructions]] reads curator instructions live from the appdb."
(:require
[clojure.string :as str]
[metabase.entity-retrieval.mirror]
[potemkin :as p]
[toucan2.core :as t2]))
(comment metabase.entity-retrieval.mirror/keep-me)
(p/import-vars
[metabase.entity-retrieval.mirror
entity-retrieval-available?
force-reconcile!
library-entity-keys
search])
(def card-entity-types
"The real entity_type strings the CRUD and tool APIs speak that all denote a Card: `question`/`metric`/
`model`. (`card` itself is not one of these — it's the generic bucket their `ai_context` rows collapse to
for storage; see [[normalize-entity-type]].)"
#{"question" "metric" "model"})
(defn normalize-entity-type
"Canonical `entity_type` for `osi_ai_context` storage: every card flavor collapses to `card` so one row
per Card survives a question/metric/model relabel; non-card types pass through unchanged.
The CRUD/tool APIs still speak question/metric/model — only the stored key is normalized."
[entity-type]
(if (card-entity-types entity-type) "card" entity-type))
(defn card-entity-type?
"True if `entity-type` denotes a Card — a question/metric/model flavor, or the stored `card` bucket."
[entity-type]
(= "card" (normalize-entity-type entity-type)))
(defn entity-class
"Equivalence class `[class entity-local-id]` for an entity ref, so a card's `ai_context` row (stored as
`card`) and its index docs (keyed by the card's live metric/model type) collapse together.
Card flavors and `card` map to the `card` class; any other entity_type stays itself."
[entity-type entity-local-id]
[(normalize-entity-type entity-type) entity-local-id])
(def max-instructions-len
"Cap on a curator `instructions` string. Instructions aren't embedded; this bounds the appdb row (enforced
by the write API) and — since [[ai-context-instructions]] truncates to it — the text injected into the
retrieve_library_entities agent prompt, so a row that bypassed the API schema (serdes, direct write, or a
pre-cap row) can't bloat the prompt."
5000)
(defn ai-context-instructions
"Map of `[entity-type entity-local-id] -> instructions` for the given entity refs (search-result shape
`{:model :id}`, where `:model` is the entity_type), read live from `osi_ai_context`.
Refs with no row, or with blank instructions, are omitted; each instruction is truncated to
[[max-instructions-len]] so an oversized row can't bloat the agent prompt.
Reading here (rather than storing in the index) means the agent always sees the current text.
The lookup matches by entity class, so a card's instructions are found even when its current type (the
ref's) differs from the type it was curated under (a card-flavor relabel)."
[entity-refs]
(if-let [wanted (seq (into #{} (map (juxt :model :id)) entity-refs))]
;; row-value IN isn't portable across our app DBs, so match the wanted (type, id) pairs with OR-of-ANDs.
;; Normalize each ref's type to the stored key, so a card ref (metric/model) matches its `card` row.
(let [clause (into [:or]
(map (fn [[t id]]
[:and
[:= :entity_type (normalize-entity-type t)]
[:= :entity_local_id id]]))
wanted)
by-class (into {} (remove (comp str/blank? val))
(t2/select-fn->fn #(entity-class (:entity_type %) (:entity_local_id %))
(comp :instructions :ai_context)
:model/OsiAiContext {:where clause}))]
;; key the result back by the caller's original [type id] ref
(into {} (keep (fn [[t id]]
(when-let [instr (by-class (entity-class t id))]
[[t id] (cond-> instr (> (count instr) max-instructions-len) (subs 0 max-instructions-len))])))
wanted))
{}))

49
src/metabase/entity_retrieval/mirror.clj Normal file
View File

@@ -0,0 +1,49 @@
(ns metabase.entity-retrieval.mirror
"OSS entry points for the enterprise pgvector `library_entity_index`.
The appdb (`osi_ai_context` + library membership) is authoritative; the index carries one embedding
per value-document and serves the `retrieve_library_entities` Metabot tool's similarity search.
These `defenterprise` shims route to [[metabase-enterprise.entity-retrieval.core]] when an enterprise
license with semantic search is present; otherwise they no-op or return empty."
(:require
[metabase.premium-features.core :refer [defenterprise]]))
(defenterprise entity-retrieval-available?
"Whether the library entity index can serve a query right now: a pgvector store is configured and the
license includes semantic search. `false` in OSS (no index). Re-evaluated per use — pgvector config is
boot-fixed but the feature can flip at runtime."
metabase-enterprise.entity-retrieval.core
[]
false)
(defenterprise request-entity-sync!
"Ask the targeted write path to reconcile one library entity's index slice soon; no-op in OSS.
Fire-and-forget: never throws, does no embedding or pgvector work on the calling thread."
metabase-enterprise.entity-retrieval.core
[_entity-type _entity-local-id]
nil)
(defenterprise force-reconcile!
"Reconcile the pgvector index with the appdb, blocking until a run covering this call finishes.
Returns {:index <inserted/deleted/unchanged> :execution <waited_ms/ran_ms>}; nil in OSS (no index to
reconcile).
See the enterprise impl for the queue-and-coalesce semantics."
metabase-enterprise.entity-retrieval.core
[]
nil)
(defenterprise search
"Similarity-search the library entity index for the documents nearest `user-search-prompt`.
Returns up to `limit` matches shaped `{:entity {:model :id} :doc_type :doc_text :score}`,
best score first; [] in OSS or when no pgvector store is configured."
metabase-enterprise.entity-retrieval.core
[_user-search-prompt _limit]
[])
(defenterprise library-entity-keys
"Set of `[entity_type entity_local_id]` for every entity currently in the library; nil in OSS.
The tool drops index hits whose entity has since left the library (the index is eventually consistent),
computing membership live the way reconcile does."
metabase-enterprise.entity-retrieval.core
[]
nil)

4
src/metabase/metabot/agent/core.clj
View File

@@ -420,10 +420,12 @@
"Initialize agent state."
[{:keys [messages state metabot-id profile-id context tracking-opts]}]
(let [context (assign-context-ids context)
;; Resolve the profile once (its nlq availability redirect probes the index): reuse it for both the
;; prompt and the tools so they can't disagree about whether the curated library tool is offered.
profile (or (profiles/get-profile profile-id)
(throw (ex-info "Unknown profile" {:profile-id profile-id})))
capabilities (get context :capabilities #{})
base-tools (profiles/get-tools-for-profile profile-id capabilities)
base-tools (profiles/profile->tools profile capabilities)
seeded (-> (or state {})
(seed-state context)
(seed-chart-configs context)

74
src/metabase/metabot/agent/profiles.clj
View File

@@ -9,11 +9,13 @@
(:require
[malli.error :as me]
[metabase.api-scope.core :as api-scope]
[metabase.entity-retrieval.core :as entity-retrieval]
[metabase.metabot.capabilities :as capabilities]
[metabase.metabot.scope :as scope]
[metabase.metabot.settings :as metabot.settings]
[metabase.metabot.skills :as skills]
[metabase.metabot.tools :as tools]
[metabase.util.log :as log]
[metabase.util.malli :as mu]
[metabase.util.malli.registry :as mr]))
@@ -165,15 +167,28 @@
#'tools/replace-sql-query-tool
#'tools/ask-for-sql-clarification-tool]})
;; :nlq and :nlq-fallback are a pair selected by the library index's health (see [[get-profile]]): when the
;; curated index can serve queries the agent discovers data through retrieve_library_entities; otherwise it
;; falls back to the general nlq search. They differ only in the discovery tool and the prompt that explains
;; it. The redirect keeps the external profile-id :nlq, so telemetry / recent-views / skills are unaffected.
(register-profile!
{:name :nlq
:prompt-template "natural-language-querying-only.selmer"
:max-iterations 10
:temperature 0.3
;; the nlq profile gets both the general instance search and, for users with the semantic-search
;; feature, the curated search tool (gated via its :feature-semantic-search capability).
:tools [#'tools/retrieve-library-entities-tool
#'tools/read-resource-tool
#'tools/construct-notebook-query-tool
#'tools/navigate-user-tool
#'tools/create-chart-tool
#'tools/edit-chart-tool]})
(register-profile!
{:name :nlq-fallback
:prompt-template "natural-language-querying-fallback.selmer"
:max-iterations 10
:temperature 0.3
:tools [#'tools/nlq-search-tool
#'tools/curated-search-tool
#'tools/read-resource-tool
#'tools/construct-notebook-query-tool
#'tools/navigate-user-tool
@@ -239,24 +254,44 @@
;;; API
(defn- nlq-fallback?
"Whether a :nlq request should be served the general-search fallback: true when the curated library index
can't answer (not configured/licensed, or empty). Keeps data discovery working before the first reconcile
and on OSS / unlicensed instances."
[profile-id]
(and (= profile-id :nlq)
(not (entity-retrieval/entity-retrieval-available?))))
(defn get-profile
"Get profile configuration by profile-id keyword.
The `:model` in the returned profile is resolved from the `llm-metabot-provider`
setting at call time, so it always reflects the current admin configuration."
setting at call time, so it always reflects the current admin configuration.
A :nlq request whose library index can't serve queries is transparently served the :nlq-fallback
profile's discovery tool and prompt (see [[nlq-fallback?]]); the profile's `:name` stays :nlq so
telemetry, recent-views, and skill matching are unaffected."
[profile-id]
(when-let [profile (get @*profiles profile-id)]
(assoc profile :model (metabot.settings/llm-metabot-provider))))
(if-let [profile (get @*profiles profile-id)]
(let [profile (if (nlq-fallback? profile-id)
(if-let [fb (get @*profiles :nlq-fallback)]
(assoc profile :tools (:tools fb) :prompt-template (:prompt-template fb))
;; The redirect target should always be registered; if it isn't, serve :nlq
;; unredirected rather than a profile with nil tools/prompt.
(do (log/warn "nlq-fallback profile is not registered; serving :nlq unredirected")
profile))
profile)]
(assoc profile :model (metabot.settings/llm-metabot-provider)))
;; An unregistered profile-id is a wiring bug; warn so it's diagnosable (callers handle the nil).
(log/warnf "No metabot profile registered for %s" profile-id)))
(defn get-tools-for-profile
"Get tool registry filtered by profile configuration, user capabilities, and scope.
Filters out EE-only tools when the feature is not available, then filters by
capabilities, then filters by `*current-user-scope*`. Returns a map of
tool-name -> tool-var.
When the resolved profile exposes at least one skill in its catalog, the
`load_skill` tool is injected so the agent can pull skill bodies on demand."
[profile-id capabilities]
(when-let [profile (get-profile profile-id)]
(defn profile->tools
"Tool registry for an ALREADY-RESOLVED profile, filtered by capabilities and `*current-user-scope*`.
Returns a map of tool-name -> tool-var.
Takes the resolved profile (not an id) so callers that also need the profile's prompt resolve it once via
[[get-profile]] — its nlq availability redirect must be probed a single time, or the prompt and tools
could disagree. When the profile exposes any skills, `load_skill` is injected for on-demand loading."
[profile capabilities]
(when profile
(let [base (-> profile
:tools
(filter-by-capabilities capabilities)
@@ -270,3 +305,10 @@
;; synthetic `load_skill` call that must resolve to a registered tool.
(or (seq (:catalog manifest)) (seq (:always-on manifest)))
(assoc "load_skill" #'tools/load-skill-tool)))))
(defn get-tools-for-profile
"Resolve a profile by id and return its capability/scope-filtered tool registry (see [[profile->tools]]).
Convenience for callers that only need the tools; code that also needs the prompt should call
[[get-profile]] once and pass the result to [[profile->tools]]."
[profile-id capabilities]
(profile->tools (get-profile profile-id) capabilities))

5
src/metabase/metabot/agent/prompts.clj
View File

@@ -195,7 +195,10 @@
:has_other_tools (= :yes (:permission/metabot-other-tools perms))
:custom_instructions (not-empty
(case template-name
"natural-language-querying-only.selmer"
;; both nlq templates (curated + general-search
;; fallback) take the nlq custom instructions
("natural-language-querying-only.selmer"
"natural-language-querying-fallback.selmer")
(metabot.settings/metabot-nlq-system-prompt)
"sql-querying-only.selmer"
(metabot.settings/metabot-sql-system-prompt)

3
src/metabase/metabot/api.clj
View File

@@ -101,7 +101,8 @@
line protocol so the client can correlate streamed messages with feedback."
[{:keys [metabot-id profile-id message context history conversation-id state debug?
assistant-msg-id external-id]}]
(let [enriched-context (metabot.context/create-context context {:metabot-id metabot-id})
(let [enriched-context (metabot.context/create-context context {:metabot-id metabot-id
:profile-id (keyword profile-id)})
messages (concat history [message])]
(sr/streaming-response {:content-type "text/event-stream"} [^OutputStream os canceled-chan]
(let [parts-atom (atom [])

3
src/metabase/metabot/capabilities.clj
View File

@@ -33,6 +33,9 @@
"Return the set of capabilities inferred from BE state rather than sent from the FE.
Generated by function to avoid staleness."
[]
;; Note: the curated retrieve_library_entities tool is NOT gated here — library-index availability is
;; resolved once, in profiles/get-profile, which serves the :nlq vs :nlq-fallback profile accordingly. A
;; second probe here could disagree with that one and leave the agent with no discovery tool.
(into #{}
(remove nil?)
[(when (or (not (premium-features/is-hosted?))

18
src/metabase/metabot/context.clj
View File

@@ -263,6 +263,13 @@
(let [curated (curation/curated-ids (map (juxt (comp name :model) :id) recents))]
(filter (fn [{:keys [model id]}] (contains? curated [(name model) id])) recents)))
(def ^:private profiles-excluding-recent-views
"Profiles for which recent views are never injected. The nlq profile discovers data through the curated
library tool rather than general instance search, so arbitrary recently-viewed items (which may not be
curated) would undermine that guarantee. (A :nlq request served the general-search fallback keeps the
external profile-id :nlq, so this is reached via :nlq; :nlq-fallback is listed for a direct request.)"
#{:nlq :nlq-fallback})
(defn- add-recent-views
"Add user's recent views to the context since these have a higher likelihood of being relevant to a user's query.
Includes the 5 most recent items across cards, datasets, metrics, dashboards, and tables.
@@ -270,12 +277,15 @@
When `metabot-id` is provided and the metabot has `use_verified_content` enabled, filters recents down
to curated content (verified, official-collection, library/published, or authoritative) before taking
the top 5, matching how search filters answer sources."
[context {:keys [metabot-id] :as _opts}]
the top 5, matching how search filters answer sources.
Skips recents entirely for profiles in [[profiles-excluding-recent-views]]."
[context {:keys [metabot-id profile-id] :as _opts}]
(try
;; When disabled, strip any preexisting :user_recently_viewed so the setting guarantees recent
;; When disabled (or excluded for this profile), strip any preexisting :user_recently_viewed so recent
;; views never reach the prompt, even for a caller-supplied context that already carries the key.
(if-not (metabot.settings/metabot-recent-views-enabled?)
(if (or (not (metabot.settings/metabot-recent-views-enabled?))
(contains? profiles-excluding-recent-views profile-id))
(dissoc context :user_recently_viewed)
(assoc context :user_recently_viewed
(let [recents (:recents (activity-feed/get-recents api/*current-user-id*

6
src/metabase/metabot/tools.clj
View File

@@ -15,8 +15,8 @@
[metabase.metabot.tools.construct :as tools.construct]
[metabase.metabot.tools.create-alert :as tools.create-alert]
[metabase.metabot.tools.create-dashboard-subscription :as tools.create-dashboard-subscription]
[metabase.metabot.tools.curated-search :as tools.curated-search]
[metabase.metabot.tools.document :as tools.document]
[metabase.metabot.tools.entity-retrieval :as tools.entity-retrieval]
[metabase.metabot.tools.metadata :as tools.metadata]
[metabase.metabot.tools.navigation :as tools.navigation]
[metabase.metabot.tools.resources :as tools.resources]
@@ -42,8 +42,8 @@
sql-search-tool
nlq-search-tool
transform-search-tool]
[tools.curated-search
curated-search-tool]
[tools.entity-retrieval
retrieve-library-entities-tool]
[tools.construct
construct-notebook-query-tool]
[tools.document

180
src/metabase/metabot/tools/curated_search.clj
View File

@@ -1,180 +0,0 @@
(ns metabase.metabot.tools.curated-search
"Metabot `search_curated` tool, backed by the curated search index.
Instead of ranking the whole instance, it matches the user's request by vector similarity against a
hand-curated library of saved search prompts, each mapped to the single entity that answers it.
Each match also carries `usage_instructions`: curator guidance on how to use that entity.
The index has several prompts per entity, so raw vector hits are deduped to distinct entities and a
small number returned.
The matched entity refs are hydrated into the same enriched search-result shape the general `search`
tool returns (`portable_entity_id`, fully-qualified names, database names, metric base tables), so the
agent can build a query inline without an extra `read_resource` round-trip.
Even where a general search is also offered, the tool surfaces the raw cosine `similarity` of each
match and flags low-confidence results, so a curated miss reads as a miss rather than a
confident-but-wrong curated hit.
The similarity search runs in the enterprise pgvector store via [[metabase.curated-search.core]]."
(:require
[clojure.string :as str]
[metabase.curated-search.core :as curated-search]
[metabase.metabot.scope :as scope]
[metabase.metabot.tools.search :as tools.search]
[metabase.metabot.tools.shared.llm-shape :as llm-shape]
[metabase.util.log :as log]
[metabase.util.malli :as mu]))
(set! *warn-on-reflection* true)
(def ^:private default-limit
"Distinct curated matches returned by default. The index has ~3 prompts per entity and embeddings cluster
tightly by domain; 5 was too aggressive (benchmark 15: ~12 expected entities sat at rank 6–10 and were
cut off), so 10 to recover that recall while still deduping siblings."
10)
(def ^:private max-limit 20)
(def ^:private over-fetch-factor
"Raw per-prompt hits to pull before deduping to distinct entities (several prompts map to one entity)."
4)
(def ^:private over-fetch-cap 60)
(def ^:private weak-similarity-threshold
"Raw cosine similarity below which a match is flagged low-confidence. Calibrated on this index:
in-domain top-1 similarities run ~0.52–0.81, clearly out-of-domain ~0.15–0.22."
0.45)
(def ^:private user-search-prompt-desc
(str "A natural-language description of the data the user wants. Matched by vector similarity "
"against curated saved prompts; phrase it like the saved prompt you'd expect to find, "
"not as keywords."))
(def ^:private limit-desc
(str "Maximum number of distinct curated matches (default " default-limit ", max " max-limit ")."))
(def ^:private curated-search-schema
[:map {:closed true}
[:user_search_prompt [:string {:description user-search-prompt-desc}]]
[:limit {:optional true} [:maybe [:int {:min 1 :max max-limit :description limit-desc}]]]])
(def ^:private instructions
;; refers to the match/note elements by name only — literal angle-bracket tags here would make the
;; surrounding <instructions> element malformed
(str "Each match element is a curated entry: a saved prompt, its usage_instructions, and the entity to "
"use (already hydrated with the ids, names and portable references you need). similarity is the "
"raw cosine match strength; a match flagged confidence=\"weak\" (or a leading note element) means "
"nothing in the curated layer clearly matches — don't build on it blindly; prefer asking the user "
"to clarify or narrow the request, or use the general search tool for uncurated results."))
(defn- similarity
"Raw cosine similarity (1 - distance) for a match, from its score breakdown."
^double [score]
(or (some (fn [s] (when (= :similarity (:name s)) (double (:score s)))) (:scores score)) 0.0))
(defn- dedupe-by-entity
"Collapse raw per-prompt results (sorted best-first) to distinct entities, keeping the best-scoring
prompt for each."
[results]
(->> results
(reduce (fn [[seen acc] r]
;; normalize the model so "card" and "question" (aliases for the same Card) dedupe
;; together — both hydrate to one record, so distinct keys would return it twice.
(let [{:keys [model id]} (:entity r)
k [(tools.search/ref-model->entity-type model) id]]
(if (seen k) [seen acc] [(conj seen k) (conj acc r)])))
[#{} []])
second))
(defn- build-matches
"Fetch, dedupe to distinct entities, hydrate each match's entity ref into a full search record, and
return the top `n` the current user can read.
Each match: `{:saved_search_prompt :usage_instructions :score :similarity :weak? :entity hydrated-hit}`."
[user-search-prompt n]
(let [raw (curated-search/search
user-search-prompt (min over-fetch-cap (* over-fetch-factor n)))
deduped (dedupe-by-entity raw)
;; Hydration permission-filters (entity-refs->search-results drops entities the user can't read),
;; so hydrate the whole deduped candidate set and take `n` from the readable survivors — taking
;; `n` first would let unreadable top hits shrink the result below `n` while readable matches sit
;; just past the cut.
by-key (into {} (map (juxt (juxt :type :id) identity))
(tools.search/entity-refs->search-results (distinct (map :entity deduped))))]
(->> (for [{:keys [saved_search_prompt usage_instructions entity score]} deduped
;; hydrated records use the agent-facing entity type, so normalize the ref's model to match
;; (plain "card" refs hydrate as "question")
:let [resolved (get by-key [(tools.search/ref-model->entity-type (:model entity)) (:id entity)])
sim (similarity score)]
:when resolved]
{:saved_search_prompt saved_search_prompt
:usage_instructions usage_instructions
:score score
:similarity sim
:weak? (< sim weak-similarity-threshold)
:entity resolved})
(take n))))
(defn- match->xml [{:keys [saved_search_prompt usage_instructions score similarity weak? entity]}]
;; String/format with Locale/ROOT so %.3f uses a '.' decimal separator regardless of the instance's
;; site locale (clojure.core/format always uses the default locale).
(str (String/format java.util.Locale/ROOT "<match score=\"%.3f\" similarity=\"%.3f\" confidence=\"%s\">\n"
(object-array [(double (:total_score score)) (double similarity)
(if weak? "weak" "strong")]))
"<saved_search_prompt>" (llm-shape/escape-xml saved_search_prompt) "</saved_search_prompt>\n"
(when-not (str/blank? usage_instructions)
(str "<usage_instructions>" (llm-shape/escape-xml usage_instructions) "</usage_instructions>\n"))
(llm-shape/search-result->xml entity)
"\n</match>"))
(defn- format-output [matches]
(if (empty? matches)
"<search_results>No matching curated entities.</search_results>"
(str (when (:weak? (first matches))
(str "<note>No strong curated match for this request — treat the results below as weak "
"guesses; consider asking the user to clarify or narrow the request.</note>\n"))
"<search_results>\n"
(str/join "\n" (map match->xml matches))
"\n</search_results>\n"
"<instructions>" instructions "</instructions>")))
(defn- flatten-data
"Flatten matches to a per-entity list (the `:result-type :search` shape), each entity record annotated
with its match's prompt, usage_instructions, similarity and confidence."
[matches]
(mapv (fn [{:keys [saved_search_prompt usage_instructions score similarity weak? entity]}]
(assoc entity
:saved_search_prompt saved_search_prompt
:usage_instructions usage_instructions
:score score
:similarity similarity
:confidence (if weak? "weak" "strong")))
matches))
(mu/defn ^{:tool-name "search_curated"
:scope scope/agent-search
;; only offered to users whose license includes semantic search (the mirror's backing store);
;; :feature-semantic-search is derived in profiles/feature-capabilities.
:capabilities #{:feature-semantic-search}}
curated-search-tool
"Find the best data to answer the user's request from the curated search index — a vetted library of
saved prompts, each mapped to the entity that answers it. Phrase `user_search_prompt` as a full
natural-language description of the data wanted (it is matched on meaning, not keywords).
Prefer this over the general `search` tool: try it first, and fall back to `search` only when no
curated match is strong, or to complement a curated match with broader (uncurated) results.
Returns a handful of distinct curated matches, best-first. Each match has `saved_search_prompt`, curator
`usage_instructions`, a raw `similarity`, and the `entity` — a full search record (name, type, database,
`portable_entity_id`, fully-qualified name) you can use directly. If the top match is flagged
low-confidence (a leading <note> / confidence=\"weak\"), nothing in the curated layer clearly matches —
prefer asking the user to clarify, or fall back to the general `search` tool for uncurated results."
[{:keys [user_search_prompt limit]} :- curated-search-schema]
(try
(let [n (min max-limit (or limit default-limit))
matches (build-matches user_search_prompt n)]
{:output (format-output matches)
:structured-output {:result-type :search
:data (flatten-data matches)
:total_count (count matches)
:weak_match (boolean (:weak? (first matches)))}})
(catch Exception e
(log/error e "Error in curated search")
{:output (str "search_curated failed: " (or (ex-message e) "Unknown error"))})))

73
src/metabase/metabot/tools/entity_details.clj
View File

@@ -16,6 +16,7 @@
[metabase.util.humanization :as u.humanization]
[metabase.util.i18n :as i18n]
[metabase.util.log :as log]
[metabase.util.malli :as mu]
[toucan2.core :as t2]))
(set! *warn-on-reflection* true)
@@ -504,6 +505,78 @@
{:output (ex-message e) :status-code 404}
(metabot.tools.u/handle-agent-error e)))))
(def ^:private measure-or-segment-dispatch
"Per-`kind` config for `measure-or-segment-details`: the app-DB model, the lib-metadata fetcher,
and the `convert-measure-or-segment` definition key."
{:measure {:model :model/Measure, :lookup-fn lib.metadata/measure, :definition-key :aggregation}
:segment {:model :model/Segment, :lookup-fn lib.metadata/segment, :definition-key :filters}})
(mu/defn- measure-or-segment-details
"Build drill-in detail for a single measure or segment by numeric `id`.
Layers parent-table context onto the `convert-measure-or-segment` payload — `database_id` /
`database_name`, the base table id/name/schema, and the portable `[db-name schema table]` FK — so the
shape matches the nested measures/segments `table-details` already emits."
[kind :- [:enum :measure :segment]
id :- :int]
(let [{:keys [model lookup-fn definition-key]} (measure-or-segment-dispatch kind)
row (t2/select-one [model :id :table_id] :id id)]
(when-not row
(throw (ex-info (format "%s %s not found" (name kind) id)
{:agent-error? true :status-code 404})))
;; Measure/segment perms delegate to the parent Table, so read-checking it gates access.
(let [table (api/read-check :model/Table (:table_id row))
db-id (:db_id table)
mp (lib-be/application-database-metadata-provider db-id)
database (lib.metadata/database mp)
metadata (lookup-fn mp id)
db-name (:name database)
schema (:schema table)
tbl-name (:name table)]
(-> (convert-measure-or-segment metadata definition-key)
(assoc :type kind
:database_id db-id
:database_name db-name
:base_table_id (:id table)
:base_table_name tbl-name
:base_table_schema schema
;; The portable FK the agent puts in `source-table:` to query against the base table —
;; same `[db-name schema table]` shape metrics carry as `:base_table_portable_fk`.
:base_table_portable_fk (when db-name [db-name schema tbl-name]))))))
(defn get-measure-details
"Get information about the measure with ID `measure-id`."
[{:keys [measure-id]}]
(try
(lib-be/with-metadata-provider-cache
(when-not (int? measure-id)
(throw (ex-info "Invalid measure_id format" {:agent-error? true :status-code 400})))
{:structured-output (assoc (measure-or-segment-details :measure measure-id) :result-type :entity)})
(catch Exception e
(let [{:keys [status-code agent-error?] :as data} (ex-data e)]
;; Agent-facing errors (bad input, not-found) are expected; only log genuine failures.
(when-not agent-error?
(log/error e "Failed to fetch measure details" data))
(if (= status-code 404)
{:output (ex-message e) :status-code 404}
(metabot.tools.u/handle-agent-error e))))))
(defn get-segment-details
"Get information about the segment with ID `segment-id`."
[{:keys [segment-id]}]
(try
(lib-be/with-metadata-provider-cache
(when-not (int? segment-id)
(throw (ex-info "Invalid segment_id format" {:agent-error? true :status-code 400})))
{:structured-output (assoc (measure-or-segment-details :segment segment-id) :result-type :entity)})
(catch Exception e
(let [{:keys [status-code agent-error?] :as data} (ex-data e)]
;; Agent-facing errors (bad input, not-found) are expected; only log genuine failures.
(when-not agent-error?
(log/error e "Failed to fetch segment details" data))
(if (= status-code 404)
{:output (ex-message e) :status-code 404}
(metabot.tools.u/handle-agent-error e))))))
(defn get-report-details
"Get information about the report (card) with ID `report-id`."
[{:keys [report-id] :as arguments}]

202
src/metabase/metabot/tools/entity_retrieval.clj Normal file
View File

@@ -0,0 +1,202 @@
(ns metabase.metabot.tools.entity-retrieval
"Metabot `retrieve_library_entities` tool, backed by the library entity index.
Matches the user's request by vector similarity against the per-value documents of every library entity
— each entity's name, description, and any OSI `ai_context` synonyms/examples — then dedupes the
many-per-entity hits to distinct entities and returns a handful, best-first.
Each match carries the `matched_on` document and the entity's `usage_instructions`, and is hydrated into
the same enriched shape the general `search` tool returns (`portable_entity_id`, fully-qualified names,
base tables), so the agent can build a query inline without a `read_resource` round-trip.
Each match's raw cosine `similarity` is surfaced and low-confidence ones flagged, so a library miss
reads as a miss rather than a confident-but-wrong hit.
The var/namespace name (`entity-retrieval`) is retained deliberately — the module is a home for
information retrieval generally, not 1:1 with this index or tool.
Runs in the enterprise pgvector store via [[metabase.entity-retrieval.core]]."
(:require
[clojure.string :as str]
[metabase.entity-retrieval.core :as entity-retrieval]
[metabase.metabot.scope :as scope]
[metabase.metabot.tools.search :as tools.search]
[metabase.metabot.tools.shared.llm-shape :as llm-shape]
[metabase.util :as u]
[metabase.util.log :as log]
[metabase.util.malli :as mu]))
(set! *warn-on-reflection* true)
(def ^:private default-limit
"Distinct library entities returned by default. Each entity has several documents (name, description,
synonyms, examples) that dedupe to one entity; 10 balances recall against deduped siblings."
10)
(def ^:private max-limit 20)
(def ^:private over-fetch-factor
"Raw per-document hits to pull before deduping to distinct entities (one entity has many documents:
name, description, and each synonym/example)."
8)
(def ^:private over-fetch-cap 120)
(def ^:private weak-similarity-threshold
"Raw cosine similarity below which a match is flagged low-confidence. Calibrated on this index:
in-domain top-1 similarities run ~0.52–0.81, clearly out-of-domain ~0.15–0.22."
0.45)
(def ^:private user-search-prompt-desc
(str "A natural-language description of the data the user wants. Matched by vector similarity against "
"library entities' names, descriptions, and curated synonyms/example questions; phrase it the way "
"the user would describe the data, not as keywords."))
(def ^:private limit-desc
(str "Maximum number of distinct library entities (default " default-limit ", max " max-limit ")."))
(def ^:private entity-retrieval-schema
[:map {:closed true}
[:user_search_prompt [:string {:description user-search-prompt-desc}]]
[:limit {:optional true} [:maybe [:int {:min 1 :max max-limit :description limit-desc}]]]])
(def ^:private instructions
;; refers to the match/note elements by name only — literal angle-bracket tags here would make the
;; surrounding <instructions> element malformed
(str "Each match element is a library entity (already hydrated with the ids, names and portable "
"references you need), the matched_on text that triggered it, and any usage_instructions for that "
"entity. similarity is the raw cosine match strength; a match flagged confidence=\"weak\" (or a "
"leading note element) means nothing in the library clearly matches — don't build on it blindly; "
"prefer asking the user to clarify or narrow the request."))
(defn- similarity
"Raw cosine similarity (1 - distance) for a match, from its score breakdown."
^double [score]
(or (some (fn [s] (when (= :similarity (:name s)) (double (:score s)))) (:scores score)) 0.0))
(defn- dedupe-by-entity
"Collapse raw per-document results (sorted best-first) to distinct entities, keeping the best-scoring
document for each."
[results]
(->> results
(reduce (fn [[seen acc] r]
;; collapse every card flavor (question/metric/model) to one class so aliases of the same
;; Card dedupe together — they all hydrate to one record, so distinct keys would return it
;; twice. Same equivalence the membership/instructions lookups use below.
(let [{:keys [model id]} (:entity r)
k (entity-retrieval/entity-class model id)]
(if (seen k) [seen acc] [(conj seen k) (conj acc r)])))
[#{} []])
second))
(defn- build-matches
"Fetch, dedupe to distinct entities, hydrate each match's entity ref into a full search record, and
return the top `n` the current user can read.
Each match: `{:doc_type :matched_text :usage_instructions :score :similarity :weak? :entity hydrated-hit}`."
[user-search-prompt n]
(let [raw (entity-retrieval/search
user-search-prompt (min over-fetch-cap (* over-fetch-factor n)))
deduped (dedupe-by-entity raw)
refs (distinct (map :entity deduped))
;; Hydration permission-filters (entity-refs->search-results drops entities the user can't read),
;; so hydrate the whole deduped candidate set and take `n` from the readable survivors — taking
;; `n` first would let unreadable top hits shrink the result below `n` while readable matches sit
;; just past the cut.
;; Key by entity *class*, not raw [type id]: hydration emits a card's current type, so a stale index
;; ref (an old metric/model label) still resolves to its hydrated record once both collapse to "card".
by-key (u/index-by (fn [r] (entity-retrieval/entity-class (:type r) (:id r)))
(tools.search/entity-refs->search-results refs))
;; instructions aren't stored in the index — read the current text from osi_ai_context per request.
instrs (entity-retrieval/ai-context-instructions refs)
;; The index is eventually consistent, so a hit may point at an entity that has since left the
;; library; post-filter to current members (like the permission filter above) so a stale index
;; never surfaces a now-unpublished/archived entity. nil only in OSS, where there are no hits.
;; Compare by entity *class*: between a Card metric<->model relabel and the next reconcile the index
;; carries the old type while membership carries the new one, so a raw [type id] compare would drop a
;; valid library card. Collapsing both to one class (as the rest of the pipeline does) keeps it.
lib-classes (some->> (entity-retrieval/library-entity-keys)
(into #{} (map (fn [[t id]] (entity-retrieval/entity-class t id)))))]
(->> (for [{:keys [doc_type doc_text entity score]} deduped
:let [resolved (get by-key (entity-retrieval/entity-class (:model entity) (:id entity)))
sim (similarity score)]
:when (and resolved
(or (nil? lib-classes)
(contains? lib-classes (entity-retrieval/entity-class (:model entity) (:id entity)))))]
{:doc_type doc_type
:matched_text doc_text
:usage_instructions (get instrs [(:model entity) (:id entity)])
:score score
:similarity sim
:weak? (< sim weak-similarity-threshold)
:entity resolved})
(take n))))
(defn- match->xml [{:keys [doc_type matched_text usage_instructions score similarity weak? entity]}]
;; String/format with Locale/ROOT so %.3f uses a '.' decimal separator regardless of the instance's
;; site locale (clojure.core/format always uses the default locale).
(str (String/format java.util.Locale/ROOT "<match score=\"%.3f\" similarity=\"%.3f\" confidence=\"%s\">\n"
(object-array [(double (:total_score score)) (double similarity)
(if weak? "weak" "strong")]))
"<matched_on type=\"" (llm-shape/escape-xml doc_type) "\">"
(llm-shape/escape-xml matched_text) "</matched_on>\n"
(when-not (str/blank? usage_instructions)
(str "<usage_instructions>" (llm-shape/escape-xml usage_instructions) "</usage_instructions>\n"))
(llm-shape/search-result->xml entity)
"\n</match>"))
(defn- format-output [matches]
(if (empty? matches)
"<search_results>No matching library entities.</search_results>"
(str (when (:weak? (first matches))
(str "<note>No strong library match for this request — treat the results below as weak "
"guesses; consider asking the user to clarify or narrow the request.</note>\n"))
"<search_results>\n"
(str/join "\n" (map match->xml matches))
"\n</search_results>\n"
"<instructions>" instructions "</instructions>")))
(defn- flatten-data
"Flatten matches to a per-entity list (the `:result-type :search` shape), each entity record annotated
with the matched document, usage_instructions, similarity and confidence."
[matches]
(mapv (fn [{:keys [doc_type matched_text usage_instructions score similarity weak? entity]}]
(assoc entity
:matched_doc_type doc_type
:matched_text matched_text
:usage_instructions usage_instructions
:score score
:similarity similarity
:confidence (if weak? "weak" "strong")))
matches))
;; No capability gate: this tool lives only in the :nlq profile, which [[metabase.metabot.agent.profiles/
;; get-profile]] serves (curated, with this tool) only when entity-retrieval-available? — otherwise it
;; redirects to :nlq-fallback (general search). That single probe is the sole gate, so prompt and tools can't
;; disagree about index availability.
(mu/defn ^{:tool-name "retrieve_library_entities"
:scope scope/agent-search}
retrieve-library-entities-tool
"Find the best data to answer the user's request from the library — the curated set of entities vetted
for the agent (published tables, library metrics/models, and their measures/segments). Phrase
`user_search_prompt` as a full natural-language description of the data wanted (it is matched on
meaning, not keywords).
Returns a handful of distinct library entities, best-first. Each match has the `matched_on` text that
triggered it, any curator `usage_instructions`, a raw `similarity`, and the `entity` — a full search
record (name, type, database, `portable_entity_id`, fully-qualified name). Use it to pick the right
entity, then `read_resource` that entity to confirm its fields and sample values before building. If the
top match is flagged low-confidence (a leading <note> / confidence=\"weak\"), nothing in the library
clearly matches — prefer asking the user to clarify or narrow the request."
[{:keys [user_search_prompt limit]} :- entity-retrieval-schema]
(try
(let [n (min max-limit (or limit default-limit))
matches (build-matches user_search_prompt n)]
{:output (format-output matches)
:structured-output {:result-type :search
:data (flatten-data matches)
:total_count (count matches)
:weak_match (boolean (:weak? (first matches)))}})
(catch Exception e
;; A failure here is the search subsystem being down (typically the embedding service), not "the
;; library is empty" — say so in :output (the only channel the agent reads) so it doesn't confidently
;; tell the user nothing matches. No :structured-output: that feeds the FE a result payload, and there's
;; no successful search to render.
(log/error e "Error in retrieve_library_entities")
{:output (str "The library search is temporarily unavailable (" (or (ex-message e) "unknown error")
"). This does not mean the library is empty; the search could not be run.")})))

16
src/metabase/metabot/tools/resources.clj
View File

@@ -50,6 +50,8 @@
- metabase://metric/{id} - basic metric info
- metabase://metric/{id}/dimensions - metric with dimensions
- metabase://metric/{id}/dimensions/{dimension_id} - specific dimension details
- metabase://measure/{id} - measure detail (definition + parent table)
- metabase://segment/{id} - segment detail (definition + parent table)
- metabase://transform/{id} - transform details
- metabase://transform/{id}/sources - tables/databases this transform reads from
- metabase://transform/{id}/target - table this transform writes to
@@ -513,6 +515,14 @@
:field-id dim-id
:limit 30}))
;; ----- Measure / Segment -----
(defn- fetch-measure [id-str]
(entity-details/get-measure-details {:measure-id (parse-long id-str)}))
(defn- fetch-segment [id-str]
(entity-details/get-segment-details {:segment-id (parse-long id-str)}))
;; ----- Transform -----
(defn- fetch-transform [id-str]
@@ -621,6 +631,10 @@
["metric" id "dimensions"] (fetch-metric-dimensions id)
["metric" id "dimensions" & rst] (fetch-metric-dimension id (str/join "/" rst))
;; Measure / Segment
["measure" id] (fetch-measure id)
["segment" id] (fetch-segment id)
;; Transform
["transform" id] (fetch-transform id)
["transform" id "sources"] (fetch-transform-sources id)
@@ -751,6 +765,8 @@
- metabase://model/{id}[/fields[/{field_id}]] [/sources]
- metabase://question/{id}[/fields[/{field_id}]] [/sources]
- metabase://metric/{id}[/dimensions[/{dim_id}]]
- metabase://measure/{id}
- metabase://segment/{id}
- metabase://transform/{id}[/sources|/target]
- metabase://dashboard/{id}[/items]"
[{:keys [uris]} :- [:map {:closed true}

68
src/metabase/metabot/tools/search.clj
View File

@@ -341,15 +341,16 @@
(defn- card-refs->results
"Build post-processed search-result records for card-backed refs (`{:id .. :type \"model\"|\"metric\"|\"question\"}`).
Emits one record per ref, so the same card registered under two type strings yields a record for each
(rather than collapsing to one and silently dropping the other)."
Emits one record per distinct card id, carrying the card's *current* type — so the same card registered
under two (possibly stale) type strings collapses to a single record rather than duplicating."
[refs]
(let [ids (distinct (map :id refs))
;; only surface cards the current user can read (collection perms) — see table-refs->results
id->card (when (seq ids)
(into {} (map (juxt :id identity))
(filter mi/can-read?
(t2/select [:model/Card :id :name :description :database_id :collection_id :card_schema]
(t2/select [:model/Card :id :name :description :database_id :collection_id
:card_schema :type]
:id [:in ids]))))
coll-ids (->> (vals id->card) (keep :collection_id) distinct)
id->coll (when (seq coll-ids)
@@ -360,18 +361,60 @@
(t2/select-fn-set :moderated_item_id :model/ModerationReview
:moderated_item_id [:in ids] :moderated_item_type "card"
:most_recent true :status "verified"))]
(for [{:keys [id type]} refs
(for [id ids
:let [c (id->card id)]
:when c]
(let [coll (get id->coll (:collection_id c))]
{:id id
:type type
;; the Card's *current* type, not the caller's ref type — a stale index hit kept across a
;; metric<->model relabel must not describe the entity with its old shape. Card's :type is a
;; keyword; emit the agent-facing string so downstream string checks and entity-class still match.
:type (some-> (:type c) name)
:name (:name c)
:description (:description c)
:database_id (:database_id c)
:verified (contains? verified id)
:collection (when coll (select-keys coll [:id :name :authority_level]))}))))
(defn- measure-segment-refs->results
"Build search-result records for measure/segment refs (`{:id .. :type \"measure\"|\"segment\"}`),
carrying parent-table context (database + base table). Only surfaces those whose parent Table the
current user can read (perms delegate to the table)."
[refs]
(when (seq refs)
(let [by-type (group-by :type refs)
fetch (fn [model ids]
(when-let [ids (not-empty (distinct ids))]
(filter mi/can-read?
(t2/select [model :id :name :description :table_id :entity_id] :id [:in ids]))))
rows (concat (map #(assoc % :type "measure") (fetch :model/Measure (map :id (get by-type "measure"))))
(map #(assoc % :type "segment") (fetch :model/Segment (map :id (get by-type "segment")))))
tbl-ids (not-empty (distinct (keep :table_id rows)))
id->tbl (when tbl-ids
(into {} (map (juxt :id identity))
(t2/select [:model/Table :id :name :schema :db_id] :id [:in tbl-ids])))]
(for [{:keys [id type name description table_id entity_id]} rows
:let [t (get id->tbl table_id)]]
(cond-> {:id id :type type :name name :description description}
;; the measure/segment's NanoID — used in a [measure|segment, {}, <id>] clause the way a metric
;; uses its portable_entity_id (carded types get theirs in enrich-with-portable-entity-ids).
entity_id (assoc :portable_entity_id entity_id)
t (assoc :database_id (:db_id t)
:base_table_id (:id t)
:base_table_name (:name t)
:base_table_schema (:schema t)))))))
(defn- enrich-with-measure-segment-base-tables
"Assemble `:base_table_portable_fk [database_name schema table]` for measure/segment results once
`:database_name` is set (by [[enrich-with-database-engines]]), giving the agent the table to query a
measure/segment against — the same affordance metrics get from [[enrich-with-metric-base-tables]]."
[results]
(mapv (fn [{:keys [type database_name base_table_schema base_table_name] :as r}]
(if (and (#{"measure" "segment"} type) database_name base_table_name)
(assoc r :base_table_portable_fk [database_name base_table_schema base_table_name])
r))
results))
(defn ref-model->entity-type
"Normalize an entity ref's `:model` string to the agent-facing entity type: plain cards are
`\"question\"` everywhere the agent sees them (`read_resource` URIs, search results)."
@@ -383,21 +426,24 @@
[[metabase.metabot.tools.shared.llm-shape/search-results->xml]] and the `search` tool consume.
`refs` is a seq of `{:model <entity-type> :id <id>}` where `<entity-type>` is `\"table\"`, `\"model\"`,
`\"metric\"`, or `\"question\"` (the names the agent uses with `read_resource`); `\"card\"` is accepted
and normalized to `\"question\"`.
Returns records carrying `:portable_entity_id`, `:database_name`, fully-qualified names, metric base
tables, etc. — everything the LLM needs to build a query without an extra round-trip.
`\"metric\"`, `\"question\"`, `\"measure\"`, or `\"segment\"` (the names the agent uses with
`read_resource`); `\"card\"` is accepted and normalized to `\"question\"`.
Returns records carrying `:portable_entity_id`, `:database_name`, fully-qualified names, metric/measure/
segment base tables, etc. — everything the LLM needs to build a query without an extra round-trip.
Refs whose entity no longer exists are dropped."
[refs]
(let [by-model (group-by (comp ref-model->entity-type :model) refs)
table-ids (distinct (map :id (get by-model "table")))
card-refs (for [m ["model" "metric" "question"], r (get by-model m)] {:id (:id r) :type m})]
card-refs (for [m ["model" "metric" "question"], r (get by-model m)] {:id (:id r) :type m})
ms-refs (for [m ["measure" "segment"], r (get by-model m)] {:id (:id r) :type m})]
(->> (concat (table-refs->results table-ids)
(card-refs->results (distinct card-refs)))
(card-refs->results (distinct card-refs))
(measure-segment-refs->results (distinct ms-refs)))
enrich-with-collection-descriptions
enrich-with-database-engines
enrich-with-portable-entity-ids
enrich-with-metric-base-tables
enrich-with-measure-segment-base-tables
remove-unreadable-transforms)))
(defn- format-search-output

57
src/metabase/metabot/tools/shared/llm_shape.clj
View File

@@ -228,33 +228,52 @@
:collection_description description
:collection_authority_level authority_level}))
(defn- drill-in-location-context
"The `<pfx>_database_name` / `<pfx>_base_table_fqn` keys a standalone measure/segment drill-in carries so
the LLM knows the table to query it against. Omitted when absent, so a nested (definition-only)
measure/segment doesn't emit blank location fields."
[pfx {:keys [database_name base_table_portable_fk]}]
(cond-> {}
database_name (assoc (keyword (str pfx "_database_name")) database_name)
(vector? base_table_portable_fk)
(assoc (keyword (str pfx "_base_table_fqn"))
(let [[_db schema table] base_table_portable_fk] (fully-qualified-name schema table)))))
(defn measure->xml
"Format a measure for LLM consumption."
"Format a measure for LLM consumption.
Nested table/model measures carry only definition fields; a standalone drill-in
(read_resource) also carries `:database_name` / `:base_table_portable_fk` so the LLM knows the
table to query the measure against."
[{:keys [id name display-name description definition definition-description
portable-entity-id portable_entity_id]}]
portable-entity-id portable_entity_id] :as m}]
(render-llm-template
:measure
{:measure_id (str id)
:measure_name (or name "")
:measure_display_name (or display-name name "")
:measure_description description
:measure_portable_entity_id (or portable-entity-id portable_entity_id)
:measure_definition (repr-data->llm-block definition)
:measure_definition_description definition-description}))
(merge {:measure_id (str id)
:measure_name (or name "")
:measure_display_name (or display-name name "")
:measure_description description
:measure_portable_entity_id (or portable-entity-id portable_entity_id)
:measure_definition (repr-data->llm-block definition)
:measure_definition_description definition-description}
(drill-in-location-context "measure" m))))
(defn segment->xml
"Format a segment for LLM consumption."
"Format a segment for LLM consumption.
Nested table/model segments carry only definition fields; a standalone drill-in
(read_resource) also carries `:database_name` / `:base_table_portable_fk` so the LLM knows the
table to query the segment against."
[{:keys [id name display-name description definition definition-description
portable-entity-id portable_entity_id]}]
portable-entity-id portable_entity_id] :as m}]
(render-llm-template
:segment
{:segment_id (str id)
:segment_name (or name "")
:segment_display_name (or display-name name "")
:segment_description description
:segment_portable_entity_id (or portable-entity-id portable_entity_id)
:segment_definition (repr-data->llm-block definition)
:segment_definition_description definition-description}))
(merge {:segment_id (str id)
:segment_name (or name "")
:segment_display_name (or display-name name "")
:segment_description description
:segment_portable_entity_id (or portable-entity-id portable_entity_id)
:segment_definition (repr-data->llm-block definition)
:segment_definition_description definition-description}
(drill-in-location-context "segment" m))))
(defn- related-table->xml
"Format a related table for LLM consumption."
@@ -766,6 +785,8 @@
(def formatters
"XML formatters for different entity types"
{:metric metric->xml
:measure measure->xml
:segment segment->xml
:table table->xml
:model model->xml
:question question->xml

2
src/metabase/models/resolution.clj
View File

@@ -33,7 +33,6 @@
:model/CommentReaction metabase.comments.models.comment-reaction
:model/ConnectionImpersonation metabase-enterprise.impersonation.models
:model/ContentTranslation metabase.content-translation.models
:model/CuratedSearchEntry metabase.curated-search.models.curated-search-entry
:model/CustomVizPlugin metabase-enterprise.custom-viz-plugin.models.custom-viz-plugin
:model/Dashboard metabase.dashboards.models.dashboard
:model/DashboardBookmark metabase.bookmarks.models.bookmark
@@ -81,6 +80,7 @@
:model/OAuthClient metabase.oauth-server.models.oauth-client
:model/OAuthClientEvent metabase.oauth-server.models.oauth-client-event
:model/OAuthRefreshToken metabase.oauth-server.models.oauth-refresh-token
:model/OsiAiContext metabase.osi.models.osi-ai-context
:model/NotificationCard metabase.notification.models
:model/NotificationHandler metabase.notification.models
:model/NotificationRecipient metabase.notification.models

145
src/metabase/osi/ai_context/api.clj Normal file
View File

@@ -0,0 +1,145 @@
(ns metabase.osi.ai-context.api
"Admin REST API for managing `osi_ai_context` rows: OSI `ai_context` metadata attached to a library
entity, addressed by its logical `(entity_type, entity_local_id)` key."
(:require
[metabase.api.common :as api]
[metabase.api.macros :as api.macros]
[metabase.app-db.core :as app-db]
[metabase.entity-retrieval.core :as entity-retrieval]
[metabase.request.core :as request]
[metabase.util.malli.schema :as ms]
[toucan2.core :as t2]))
(def ^:private writable-entity-types
"Entity types accepted on writes — the real types callers name. Card flavors (metric/model) are stored
under the canonical `card` key (see [[entity-retrieval/normalize-entity-type]]); tables and table-bound
measures/segments keep their type. A plain question never matches an index doc, so it's rejected."
#{"table" "metric" "model" "measure" "segment"})
(def ^:private max-item-len
"Cap on each synonym/example string — these are short phrases or questions, not prose. They become
embedded index docs, so this also keeps a single value under the embedding provider's token limit."
1000)
(def ^:private max-list-len
"Cap on the synonyms/examples list length — a curated entity needs a handful, not hundreds."
50)
(def ^:private AiContext
"OSI ai_context blob. All fields optional; extra keys tolerated for forward-compat with the OSI spec.
String and list lengths are capped so a single curated entity can't bloat the index, its embeddings, or
the agent prompt."
[:map
[:instructions {:optional true} [:maybe [:string {:max entity-retrieval/max-instructions-len}]]]
[:synonyms {:optional true} [:sequential {:max max-list-len} [:string {:max max-item-len}]]]
[:examples {:optional true} [:sequential {:max max-list-len} [:string {:max max-item-len}]]]])
(def ^:private Entry
"An ai_context row as returned on reads. `entity_type` is any string: a row can predate a type's
retirement (serdes tolerates those too), and one legacy row must not fail response validation for a list."
[:map
[:entity_type :string]
[:entity_local_id :int]
[:ai_context AiContext]])
(def ^:private default-limit 50)
(def ^:private default-offset 0)
(defn- get-entry
"The stored row for an entity, looked up by its normalized key, or nil. The CRUD API speaks the real
card flavors; storage keys on the canonical `card`, so normalize before matching."
[entity-type entity-local-id]
(t2/select-one :model/OsiAiContext
:entity_type (entity-retrieval/normalize-entity-type entity-type)
:entity_local_id entity-local-id))
(def ^:private logical-key-route-schema
;; entity-type is any non-blank string at the route level — a write to a non-writable type gets a clear
;; 400 in the handler (an enum here would 404 the route instead), and reads/deletes of an unknown type
;; simply find no row and 404.
[:map
[:entity-type ms/NonBlankString]
[:entity-local-id ms/PositiveInt]])
(api.macros/defendpoint :get "/"
:- [:map
[:data [:sequential Entry]]
[:total :int]
[:limit :int]
[:offset :int]]
"Get all ai_context entries, paginated."
[_route-params
_query-params]
(api/check-superuser)
(let [limit (or (request/limit) default-limit)
offset (or (request/offset) default-offset)]
{:data (t2/select :model/OsiAiContext
{:order-by [[:entity_type :asc] [:entity_local_id :asc]]
:limit limit
:offset offset})
:total (t2/count :model/OsiAiContext)
:limit limit
:offset offset}))
(api.macros/defendpoint :get "/:entity-type/:entity-local-id"
:- Entry
"Get the ai_context entry for an entity by its logical `(entity_type, entity_local_id)` key."
[{:keys [entity-type entity-local-id]} :- logical-key-route-schema
_query-params]
(api/check-superuser)
(api/check-404 (get-entry entity-type entity-local-id)))
(api.macros/defendpoint :put "/:entity-type/:entity-local-id"
:- Entry
"Create or replace (upsert) the ai_context for an entity, addressed by its logical key. One row per
entity: an existing row is updated in place rather than duplicated, and the compound primary key plus the
upsert keep two concurrent writers from racing in a duplicate row."
[{:keys [entity-type entity-local-id]} :- logical-key-route-schema
_query-params
{:keys [ai_context]} :- [:map [:ai_context AiContext]]]
(api/check-superuser)
(api/check-400 (contains? writable-entity-types entity-type)
"entity_type must be one of: measure, metric, model, segment, table")
;; Upsert on the normalized (stored) key so re-posting a relabelled card updates its one row.
;; update-or-insert! handles the compound key, the no-op re-PUT, and the concurrent-create race
;; (savepoint + single retry) centrally.
(app-db/update-or-insert! :model/OsiAiContext
{:entity_type (entity-retrieval/normalize-entity-type entity-type)
:entity_local_id entity-local-id}
(constantly {:ai_context ai_context}))
(get-entry entity-type entity-local-id))
(api.macros/defendpoint :post "/reconcile"
:- [:map
[:index [:map
[:inserted :int]
[:deleted :int]
[:unchanged :int]]]
[:execution [:map
[:waited_ms :int]
[:ran_ms :int]]]]
"Force a reconcile of the library entity index against the application database, blocking until a
reconcile covering this call finishes.
Returns the index mutations (`index`) separately from execution timing (`execution`: how long the run
waited to start, then how long it ran).
This call never reuses a reconcile already in progress (it may have started before your latest change);
it starts one if the index is idle, otherwise it queues a single follow-up that any other waiting calls
share.
Requires the semantic search feature; returns a 400 when it isn't configured."
[_route-params
_query-params]
(api/check-superuser)
(api/check-400 (entity-retrieval/force-reconcile!)
"The library entity index requires semantic search, which is not configured."))
#_{:clj-kondo/ignore [:metabase/validate-defendpoint-has-response-schema]}
(api.macros/defendpoint :delete "/:entity-type/:entity-local-id"
"Delete the ai_context entry for an entity by its logical key."
[{:keys [entity-type entity-local-id]} :- logical-key-route-schema
_query-params]
(api/check-superuser)
(api/check-404 (get-entry entity-type entity-local-id))
(t2/delete! :model/OsiAiContext
:entity_type (entity-retrieval/normalize-entity-type entity-type)
:entity_local_id entity-local-id)
api/generic-204-no-content)

169
src/metabase/osi/models/osi_ai_context.clj Normal file
View File

@@ -0,0 +1,169 @@
(ns metabase.osi.models.osi-ai-context
"The `osi_ai_context` appdb table: one row per library entity holding OSI `ai_context` metadata
(`{instructions, synonyms[], examples[]}`) for the entity identified by `entity_type`/`entity_local_id`.
Identity is the logical `(entity_type, entity_local_id)` pair — that compound key is the primary key, so
there is one row per entity and no surrogate id or serialization NanoID.
This table is authoritative.
An enterprise pgvector index (`library_entity_index`) is reconciled against this table plus live
library membership and serves the `retrieve_library_entities` Metabot tool's similarity search.
Writes here only nudge a targeted reconcile of the entity's slice
([[mirror/request-entity-sync!]]); they never touch the embedding service or the pgvector store
themselves."
(:require
[clojure.string :as str]
[metabase.app-db.core :as app-db]
[metabase.entity-retrieval.core :as entity-retrieval]
[metabase.entity-retrieval.mirror :as mirror]
[metabase.models.interface :as mi]
[metabase.models.serialization :as serdes]
[metabase.util :as u]
[methodical.core :as methodical]
[toucan2.core :as t2]))
(methodical/defmethod t2/table-name :model/OsiAiContext [_model] :osi_ai_context)
(doto :model/OsiAiContext
(derive :metabase/model)
(derive :hook/timestamped?))
;; The logical (entity_type, entity_local_id) pair is the primary key — no surrogate id.
(methodical/defmethod t2/primary-keys :model/OsiAiContext [_model] [:entity_type :entity_local_id])
(t2/deftransforms :model/OsiAiContext
;; ai_context is keywordized on read so reconcile reads (:instructions ai_context) etc. directly.
{:ai_context mi/transform-json})
;;; Card-backed entities (question/metric/model) store their `entity_type` as the canonical `card`, so one
;;; ai_context row survives a type relabel and keys on a stable `(card, entity_local_id)`. The CRUD and tool
;;; APIs still speak the real types; only the stored key is normalized. Non-card types pass through.
(t2/define-before-insert :model/OsiAiContext
[row]
(cond-> row
(:entity_type row) (update :entity_type entity-retrieval/normalize-entity-type)))
;;; Each write nudges a targeted reconcile of just this entity's index slice (fire-and-forget: no-op in
;;; OSS, error-swallowing in EE), so appdb writes never fail or slow down because of the mirror. The nudge
;;; is deferred until *after* the surrounding transaction commits, so the reconcile future always reads
;;; committed state — e.g. a delete buried in a long transaction is seen as gone, GC'ing its synonym/example
;;; docs (name/description stay if the entity is still in the library), rather than racing the commit and
;;; lingering until the periodic full reconcile. Outside a transaction the write has already committed, so
;;; [[app-db/do-after-commit]] runs the nudge immediately.
(defn- nudge-entity-sync!
"After the surrounding transaction commits, nudge a targeted reconcile of `row`'s entity slice. Returns `row`."
[{:keys [entity_type entity_local_id] :as row}]
(u/prog1 row
(app-db/do-after-commit #(mirror/request-entity-sync! entity_type entity_local_id))))
(t2/define-after-insert :model/OsiAiContext [row] (nudge-entity-sync! row))
(t2/define-after-update :model/OsiAiContext [row] (nudge-entity-sync! row))
(t2/define-before-delete :model/OsiAiContext [row] (nudge-entity-sync! row))
;;; ------------------------------------------------- Serialization -------------------------------------------------
;;;
;;; A row's identity is its entity ref: `entity_type` plus the entity it describes. Rather than a separate
;;; NanoID, serdes nests the row under its entity's portable path (mirroring
;;; [[metabase.warehouse-schema.models.field-user-settings]], which hangs off its Field): the path is the
;;; entity's path plus a constant `OsiAiContext` segment. So both key columns are carried by the path —
;;; `entity_local_id` is resolved from it on import, and `entity_type` is read back from the parent segment.
(def ^:private entity-type->toucan
"Toucan model each `entity_type` resolves to. Tables get the table-fk treatment instead.
The card flavors (the type mirrors the agent's `read_resource` resource type) all map to Card."
{"card" :model/Card
"model" :model/Card
"metric" :model/Card
"question" :model/Card
"measure" :model/Measure
"segment" :model/Segment})
(def ^:private parent-model->entity-type
"Inverse of the parent's serdes model name back to the stored `entity_type`."
{"Table" "table"
"Card" "card"
"Measure" "measure"
"Segment" "segment"})
(defn- entity-parent-path
"Portable serdes path of the entity a row describes — the path its `OsiAiContext` segment hangs under.
A Table expands to its `[Database Schema Table]` path; card/measure/segment refs are a single segment
keyed on the entity's `entity_id`. An unmapped type yields a raw-id segment rather than aborting export."
[entity-type entity-local-id]
(cond
(= entity-type "table") (serdes/table->path (serdes/*export-table-fk* entity-local-id))
(entity-type->toucan entity-type) (let [model (entity-type->toucan entity-type)]
[{:model (name model) :id (serdes/*export-fk* entity-local-id model)}])
:else [{:model entity-type :id entity-local-id}]))
(defn- parent-path->entity
"Inverse of [[entity-parent-path]]: resolve a parent path back to `{:entity_type :entity_local_id}` (the
local id), or `nil` when the referenced entity is absent."
[parent-path]
(let [{:keys [model id]} (last parent-path)]
(if (= model "Table")
{:entity_type "table"
:entity_local_id (serdes/*import-table-fk* (mapv :id parent-path))}
(when-let [etype (parent-model->entity-type model)]
(when-let [toucan (entity-type->toucan etype)]
{:entity_type etype :entity_local_id (serdes/*import-fk* id toucan)})))))
(defmethod serdes/entity-id "OsiAiContext" [_ _] nil)
(defmethod serdes/hash-fields :model/OsiAiContext
[_model]
[:entity_type :entity_local_id])
(defmethod serdes/generate-path "OsiAiContext" [_ {:keys [entity_type entity_local_id]}]
(conj (vec (entity-parent-path entity_type entity_local_id))
{:model "OsiAiContext" :id "ai_context"}))
(defmethod serdes/storage-path "OsiAiContext" [entity _ctx]
;; Store under a flat top-level directory rather than nesting next to the entity: serdes/storage-path-prefixes
;; only knows how to nest under Database/Schema/Table/Field, so a Card/Measure/Segment parent would throw.
;; The row's identity still lives in its nested generate-path (the on-disk :serdes/meta); this is only the
;; file's location. Storage dedups by `:key`, so it must be unambiguous: use the EDN of the parent's
;; `[model id]` pairs (ids can contain "/" and ":", so a delimiter-joined string could collide for distinct
;; paths). The slug is just a readable filename — the unique-name generator disambiguates it by `:key`.
(let [parent (pop (vec (serdes/path entity)))]
[{:label "osi_ai_context"}
{:label (u/slugify (str/join "-" (map :id parent)) {:unicode? true})
:key (pr-str (mapv (juxt :model :id) parent))}]))
;; `ai_context` is a plain text blob (no FKs — instructions/synonyms/examples are free text), so it copies
;; verbatim. The key columns are carried by the path, not as fields: `entity_local_id` is resolved from the
;; parent path on import, and `entity_type` is read back from the parent segment's model.
(defmethod serdes/make-spec "OsiAiContext" [_model-name _opts]
{:copy [:ai_context]
:transform {:created_at (serdes/date)
:updated_at (serdes/date)
:entity_type {:export (constantly ::serdes/skip)
:import-with-context
(fn [current _ _] (:entity_type (parent-path->entity (pop (serdes/path current)))))}
:entity_local_id {::serdes/fk true
:export (constantly ::serdes/skip)
:import-with-context
(fn [current _ _] (:entity_local_id (parent-path->entity (pop (serdes/path current)))))}}})
(defmethod serdes/dependencies "OsiAiContext"
[entity]
;; Depend on the entity this row describes (its parent path), so it imports after that entity exists.
[(vec (pop (serdes/path entity)))])
(defmethod serdes/load-find-local "OsiAiContext"
[path]
;; Resolve the parent path back to a local entity, then find this row by its (entity_type, local-id) key.
(when-let [{:keys [entity_type entity_local_id]} (parent-path->entity (pop path))]
(t2/select-one :model/OsiAiContext
:entity_type (entity-retrieval/normalize-entity-type entity_type)
:entity_local_id entity_local_id)))
(defmethod serdes/load-update! "OsiAiContext"
[_model-name ingested local]
;; The default keys updates on (first (primary-keys)), which for this compound key is just :entity_type —
;; so it would address the wrong rows. Update by the full (entity_type, entity_local_id) key.
(t2/update! :model/OsiAiContext
:entity_type (:entity_type local) :entity_local_id (:entity_local_id local)
ingested)
(t2/select-one :model/OsiAiContext
:entity_type (:entity_type local) :entity_local_id (:entity_local_id local)))

72
test/metabase/app_db/connection_test.clj
View File

@@ -139,6 +139,78 @@
(is (= java.sql.Connection/TRANSACTION_READ_COMMITTED
(.getTransactionIsolation conn))))))
(deftest after-commit-callback-runs-outside-transaction-bindings-test
(let [callback-state (promise)]
(t2/with-transaction [_conn]
(mdb.connection/do-after-commit
#(deliver callback-state
{:current-connectable t2.connection/*current-connectable*
:in-transaction? (mdb.connection/in-transaction?)})))
(is (= {:current-connectable nil
:in-transaction? false}
(deref callback-state 1000 ::timed-out)))))
(deftest after-commit-callback-does-not-leak-transaction-bindings-into-async-work-test
;; a future conveys dynamic bindings: neither the transaction connection nor the callback accumulator may
;; leak into async work a callback starts, or the future would reuse a pooled-back connection and an async
;; do-after-commit would enqueue into the already-drained accumulator and never run
(let [conveyed-connectable (promise)
async-callback-ran (promise)]
(t2/with-transaction [_conn]
(mdb.connection/do-after-commit
(fn []
(future
(deliver conveyed-connectable t2.connection/*current-connectable*)
(mdb.connection/do-after-commit #(deliver async-callback-ran true))))))
(is (nil? (deref conveyed-connectable 1000 ::timed-out)))
(is (true? (deref async-callback-ran 1000 ::timed-out)))))
(deftest after-commit-callbacks-from-rolled-back-nested-transaction-are-discarded-test
(let [calls (atom [])
rollback-ex (ex-info "Rollback nested transaction" {})
register-call #(mdb.connection/do-after-commit (fn [] (swap! calls conj %)))]
(t2/with-transaction [conn]
(register-call :outer-before)
(is (thrown?
Exception
(t2/with-transaction [_ conn]
(register-call :nested)
(throw rollback-ex))))
(register-call :outer-after))
(is (= [:outer-before :outer-after] @calls))))
(deftest after-commit-callbacks-discarded-even-when-savepoint-rollback-throws-test
(let [calls (atom [])
mock-conn (reify Connection
(rollback [_ _savepoint] (throw (ex-info "Rollback error" {})))
(setAutoCommit [_ _])
(getAutoCommit [_] true)
(setSavepoint [_])
(commit [_]))]
(binding [t2.connection/*current-connectable* mock-conn]
(t2/with-transaction [_conn]
(mdb.connection/do-after-commit (fn [] (swap! calls conj :outer)))
;; the nested body registers a callback then throws; its savepoint rollback also throws
(is (thrown?
Exception
(t2/with-transaction [_]
(mdb.connection/do-after-commit (fn [] (swap! calls conj :nested)))
(throw (ex-info "boom" {})))))))
;; the nested callback is discarded even though the rollback threw — only :outer survives to run
(is (= [:outer] @calls))))
(deftest after-commit-callback-registering-another-runs-it-immediately-test
(let [calls (atom [])]
(t2/with-transaction [_conn]
(mdb.connection/do-after-commit
(fn []
(swap! calls conj :a)
;; the run binds *after-commit-callbacks* to nil, so a do-after-commit here runs immediately (we are
;; logically outside the transaction) rather than enqueuing for a later pass
(mdb.connection/do-after-commit (fn [] (swap! calls conj :nested-from-a)))))
(mdb.connection/do-after-commit (fn [] (swap! calls conj :b))))
(is (= [:a :nested-from-a :b] @calls))))
(deftest rollback-error-handling
(testing "rollback error handling"
(let [mock-conn (reify Connection

33
test/metabase/collections/test_utils.clj
View File

@@ -77,3 +77,36 @@
"Ensures Library collection is not remote-synced for the duration of body."
[& body]
`(do-with-library-not-synced (fn [] ~@body)))
(defn do-with-library
"Implementation for [[with-library]]."
[f]
(if-let [library (collection/library-collection)]
;; A library already exists — reuse it and its sub-collections, scoped to its location so a same-typed
;; collection elsewhere can't be picked up (the library sub-collection types aren't unique).
(let [loc (str "/" (:id library) "/")]
(f {:library library
:data (t2/select-one :model/Collection
:type collection/library-data-collection-type :location loc)
:metrics (t2/select-one :model/Collection
:type collection/library-metrics-collection-type :location loc)}))
;; None exists — create a temporary tree (root + Data + Metrics) that with-temp cleans up afterward.
(mt/with-temp [:model/Collection library {:name "Library"
:type collection/library-collection-type
:location "/"}
:model/Collection data {:name "Data"
:type collection/library-data-collection-type
:location (str "/" (:id library) "/")}
:model/Collection metrics {:name "Metrics"
:type collection/library-metrics-collection-type
:location (str "/" (:id library) "/")}]
(f {:library library :data data :metrics metrics}))))
(defmacro with-library
"Ensure the singleton Library collection and its Data/Metrics sub-collections exist for the duration of
`body`, binding `binding` to a `{:library _ :data _ :metrics _}` map of those collections.
Reuses an existing library if one is present; otherwise creates a temporary tree that's cleaned up
afterward."
{:style/indent 1}
[[binding] & body]
`(do-with-library (fn [~binding] ~@body)))

129
test/metabase/curated_search/api_test.clj
View File

@@ -1,129 +0,0 @@
(ns metabase.curated-search.api-test
(:require
[clojure.test :refer :all]
[metabase.test :as mt]
[toucan2.core :as t2]))
(set! *warn-on-reflection* true)
(defmacro ^:private with-test-entry
"Create a curated search entry for testing and bind it to `sym`, deleting it on exit."
[[sym attrs] & body]
`(let [~sym (t2/insert-returning-instance! :model/CuratedSearchEntry
(merge {:search_prompt "find orders"
:entity {:model "table" :id 1}
:verified false}
~attrs))]
(try
~@body
(finally
(t2/delete! :model/CuratedSearchEntry :id (:id ~sym))))))
(deftest list-test
(with-test-entry [entry {}]
(testing "superuser can list curated search entries"
(let [response (mt/user-http-request :crowberto :get 200 "curated-search/")]
(is (contains? response :data))
(is (contains? response :total))
(is (some #(= (:id entry) (:id %)) (:data response)))))
(testing "non-superuser gets 403"
(mt/user-http-request :rasta :get 403 "curated-search/"))
(testing "a legacy row whose entity model is no longer in the write enum still reads back fine"
(with-test-entry [legacy {:entity {:model "retired-model-string" :id 7}}]
(is (some #(= (:id legacy) (:id %))
(:data (mt/user-http-request :crowberto :get 200 "curated-search/"))))))))
(deftest get-test
(with-test-entry [entry {:search_prompt "find customers" :verified true}]
(testing "superuser can fetch a curated search entry by id"
(is (=? {:id (:id entry)
:search_prompt "find customers"
:entity {:model "table" :id 1}
:verified true}
(mt/user-http-request :crowberto :get 200 (str "curated-search/" (:id entry))))))
(testing "returns 404 for unknown id"
(mt/user-http-request :crowberto :get 404 "curated-search/0"))
(testing "non-superuser gets 403"
(mt/user-http-request :rasta :get 403 (str "curated-search/" (:id entry))))))
(deftest create-test
(testing "superuser can create a curated search entry"
(let [response (mt/user-http-request :crowberto :post 200 "curated-search/"
{:search_prompt "find revenue"
:usage_instructions "Use the Revenue card."
:entity {:model "card" :id 42}})]
(try
(is (=? {:search_prompt "find revenue"
:usage_instructions "Use the Revenue card."
:entity {:model "card" :id 42}
:verified false}
response))
(finally
(t2/delete! :model/CuratedSearchEntry :id (:id response))))))
(testing "verified defaults to false and usage_instructions to nil when omitted"
(let [response (mt/user-http-request :crowberto :post 200 "curated-search/"
{:search_prompt "find users"
:entity {:model "table" :id 1}})]
(try
(is (false? (:verified response)))
(is (nil? (:usage_instructions response)))
(finally
(t2/delete! :model/CuratedSearchEntry :id (:id response))))))
(testing "entity is required"
(mt/user-http-request :crowberto :post 400 "curated-search/"
{:search_prompt "find users"}))
(testing "an unknown entity model is rejected"
(mt/user-http-request :crowberto :post 400 "curated-search/"
{:search_prompt "find users"
:entity {:model "garbage" :id 1}}))
(testing "a blank search_prompt is rejected"
(mt/user-http-request :crowberto :post 400 "curated-search/"
{:search_prompt " " :entity {:model "table" :id 1}}))
(testing "non-superuser gets 403"
(mt/user-http-request :rasta :post 403 "curated-search/"
{:search_prompt "find orders" :entity {:model "table" :id 1}})))
(deftest update-test
(with-test-entry [entry {}]
(testing "superuser can update the prompt"
(let [updated (mt/user-http-request :crowberto :put 200
(str "curated-search/" (:id entry))
{:search_prompt "updated prompt"})]
(is (= "updated prompt" (:search_prompt updated)))))
(testing "superuser can update usage_instructions and the entity"
(let [updated (mt/user-http-request :crowberto :put 200
(str "curated-search/" (:id entry))
{:usage_instructions "Join orders to customers on customer_id."
:entity {:model "model" :id 7}})]
(is (=? {:usage_instructions "Join orders to customers on customer_id."
:entity {:model "model" :id 7}}
updated))))
(testing "an explicit null clears usage_instructions"
(let [updated (mt/user-http-request :crowberto :put 200
(str "curated-search/" (:id entry))
{:usage_instructions nil})]
(is (nil? (:usage_instructions updated)))))
(testing "the non-nullable fields reject an explicit null"
(mt/user-http-request :crowberto :put 400 (str "curated-search/" (:id entry))
{:entity nil})
(mt/user-http-request :crowberto :put 400 (str "curated-search/" (:id entry))
{:verified nil}))
(testing "a blank search_prompt is rejected"
(mt/user-http-request :crowberto :put 400 (str "curated-search/" (:id entry))
{:search_prompt " "}))
(testing "returns 404 for unknown id"
(mt/user-http-request :crowberto :put 404 "curated-search/0" {:search_prompt "x"}))
(testing "non-superuser gets 403"
(mt/user-http-request :rasta :put 403
(str "curated-search/" (:id entry))
{:search_prompt "x"}))))
(deftest delete-test
(with-test-entry [entry {}]
(testing "non-superuser gets 403"
(mt/user-http-request :rasta :delete 403 (str "curated-search/" (:id entry))))
(testing "superuser can delete an entry"
(mt/user-http-request :crowberto :delete 204 (str "curated-search/" (:id entry)))
(is (nil? (t2/select-one :model/CuratedSearchEntry :id (:id entry)))))
(testing "returns 404 for unknown id"
(mt/user-http-request :crowberto :delete 404 "curated-search/0"))))

58
test/metabase/curated_search/models/curated_search_entry_test.clj
View File

@@ -1,58 +0,0 @@
(ns metabase.curated-search.models.curated-search-entry-test
(:require
[clojure.test :refer :all]
[metabase.curated-search.mirror :as mirror]
[metabase.test :as mt]
[toucan2.core :as t2]))
(set! *warn-on-reflection* true)
(def ^:private table-entity {:model "table" :id 42})
(deftest entity-json-roundtrips-test
(testing "entity is a JSON object and timestamps are populated"
(mt/with-temp [:model/CuratedSearchEntry {:keys [id]}
{:search_prompt "monthly revenue by region" :entity table-entity}]
(is (=? {:search_prompt "monthly revenue by region"
:entity table-entity
:verified false
:created_at some?
:updated_at some?}
(t2/select-one :model/CuratedSearchEntry :id id))))))
(deftest usage-instructions-persist-test
(testing "usage_instructions persists alongside the prompt"
(mt/with-temp [:model/CuratedSearchEntry {:keys [id]}
{:search_prompt "monthly revenue by region" :entity table-entity
:usage_instructions "Use this metric for total revenue; group by month."}]
(is (= "Use this metric for total revenue; group by month."
(t2/select-one-fn :usage_instructions :model/CuratedSearchEntry :id id)))))
(testing "usage_instructions is optional (nil when omitted)"
(mt/with-temp [:model/CuratedSearchEntry {:keys [id]}
{:search_prompt "orders" :entity table-entity}]
(is (nil? (t2/select-one-fn :usage_instructions :model/CuratedSearchEntry :id id))))))
(deftest hooks-nudge-the-background-sync-test
(testing "insert, update and delete each request a background sync — and do nothing else"
(let [nudges (atom 0)]
(mt/with-dynamic-fn-redefs [mirror/request-sync! (fn [] (swap! nudges inc) nil)]
(mt/with-temp [:model/CuratedSearchEntry {:keys [id]}
{:search_prompt "p1" :entity table-entity}]
(is (= 1 @nudges) "insert nudges once")
(t2/update! :model/CuratedSearchEntry id {:search_prompt "p1 updated"})
(is (= 2 @nudges) "update nudges once")
(t2/delete! :model/CuratedSearchEntry :id id)
(is (= 3 @nudges) "delete nudges once"))))))
(deftest hooks-never-break-appdb-writes-test
(testing "insert/update/delete succeed even though the mirror is unavailable in tests"
;; The hooks call the OSS defenterprise shim, which no-ops without an enterprise license. They must
;; never throw or fail the authoritative write.
(mt/with-temp [:model/CuratedSearchEntry {:keys [id]}
{:search_prompt "p1" :entity table-entity}]
(testing "update"
(is (pos? (t2/update! :model/CuratedSearchEntry id {:search_prompt "p1 updated"})))
(is (= "p1 updated" (t2/select-one-fn :search_prompt :model/CuratedSearchEntry :id id))))
(testing "delete"
(is (pos? (t2/delete! :model/CuratedSearchEntry :id id)))
(is (nil? (t2/select-one :model/CuratedSearchEntry :id id)))))))

44
test/metabase/entity_retrieval/core_test.clj Normal file
View File

@@ -0,0 +1,44 @@
(ns metabase.entity-retrieval.core-test
(:require
[clojure.test :refer :all]
[metabase.entity-retrieval.core :as entity-retrieval]
[metabase.test :as mt]))
(deftest ^:parallel entity-class-test
(testing "card flavors (and the stored \"card\") collapse to one class; other types stay distinct"
(is (= (entity-retrieval/entity-class "card" 7)
(entity-retrieval/entity-class "metric" 7)
(entity-retrieval/entity-class "model" 7)
(entity-retrieval/entity-class "question" 7)))
(is (not= (entity-retrieval/entity-class "table" 7)
(entity-retrieval/entity-class "metric" 7)))))
(deftest ^:parallel normalize-entity-type-test
(testing "card flavors normalize to the stored \"card\"; other types pass through"
(is (= "card" (entity-retrieval/normalize-entity-type "metric")))
(is (= "card" (entity-retrieval/normalize-entity-type "model")))
(is (= "card" (entity-retrieval/normalize-entity-type "question")))
(is (= "table" (entity-retrieval/normalize-entity-type "table")))
(is (= "measure" (entity-retrieval/normalize-entity-type "measure")))))
(deftest ai-context-instructions-matches-by-card-class-test
(testing "a card's live instructions are found whichever flavor the ref names (storage is \"card\")"
;; the row is stored under the canonical "card" (the write normalizes "metric" -> "card")
(mt/with-temp [:model/OsiAiContext _ {:entity_type "metric"
:entity_local_id 4242
:ai_context {:instructions "Group by month."}}]
(testing "a model ref resolves the same card's instructions"
(is (= {["model" 4242] "Group by month."}
(entity-retrieval/ai-context-instructions [{:model "model" :id 4242}]))))
(testing "a non-card ref of the same id matches exactly, so it finds nothing"
(is (= {} (entity-retrieval/ai-context-instructions [{:model "table" :id 4242}])))))))
(deftest ai-context-instructions-truncates-oversized-text-test
(testing "instructions that bypassed the API cap (direct write/serdes/pre-cap row) are truncated on read"
(let [long-instr (apply str (repeat (* 2 entity-retrieval/max-instructions-len) \x))]
(mt/with-temp [:model/OsiAiContext _ {:entity_type "table"
:entity_local_id 5151
:ai_context {:instructions long-instr}}]
(is (= entity-retrieval/max-instructions-len
(count (get (entity-retrieval/ai-context-instructions [{:model "table" :id 5151}])
["table" 5151]))))))))

21
test/metabase/measures/test_util.clj Normal file
View File

@@ -0,0 +1,21 @@
(ns metabase.measures.test-util
"Test helpers for building MBQL5 Measure and Segment definitions against the ambient test metadata
provider. Shared so tests across modules (measures, serdes, entity-retrieval) don't each redefine them."
(:require
[metabase.lib.core :as lib]
[metabase.lib.metadata :as lib.metadata]
[metabase.test :as mt]))
(defn measure-definition
"An MBQL5 measure definition: the sum of `field-id` over `table-id`."
[table-id field-id]
(let [mp (mt/metadata-provider)]
(lib/aggregate (lib/query mp (lib.metadata/table mp table-id))
(lib/sum (lib.metadata/field mp field-id)))))
(defn segment-definition
"An MBQL5 segment definition: a filter of `field-id` > `value` on `table-id`."
[table-id field-id value]
(let [mp (mt/metadata-provider)]
(lib/filter (lib/query mp (lib.metadata/table mp table-id))
(lib/> (lib.metadata/field mp field-id) value))))

44
test/metabase/metabot/agent/profiles_test.clj
View File

@@ -2,6 +2,7 @@
(:require
[clojure.test :refer :all]
[metabase.api-scope.core :as api-scope]
[metabase.entity-retrieval.core :as entity-retrieval]
[metabase.metabot.agent.profiles :as profiles]
[metabase.metabot.scope :as scope]
[metabase.metabot.tools :as tools]
@@ -60,13 +61,15 @@
(testing "retrieves nlq profile"
(let [profile (profiles/get-profile :nlq)]
(is (some? profile))
;; the :name stays :nlq even when redirected to the fallback, so telemetry/recents are unaffected
(is (= :nlq (:name profile)))
(is (= "anthropic/claude-sonnet-4-6" (:model profile)))
(is (= 10 (:max-iterations profile)))
(is (= 0.3 (:temperature profile)))
;; nlq offers both the general search and the curated search tool (the latter feature-gated)
;; In tests the library index can't answer, so :nlq is transparently served the general-search
;; fallback; the curated/fallback swap by availability is covered by nlq-data-discovery-fallback-test.
(is (contains? (tool-names profile) "search"))
(is (contains? (tool-names profile) "search_curated"))
(is (not (contains? (tool-names profile) "retrieve_library_entities")))
(is (contains? (tool-names profile) "construct_notebook_query"))))
(testing "retrieves slackbot profile"
(let [profile (profiles/get-profile :slackbot)]
@@ -158,12 +161,17 @@
"SQL tools should be gated by permission:write_sql_queries"))))))
(deftest embedding-next-matches-nlq-tools-test
(testing "nlq has the same tools as embedding_next plus the (feature-gated) curated search tool"
(let [tool-names (fn [profile] (set (map #(:tool-name (meta %)) (:tools profile))))
embedding (profiles/get-profile :embedding_next)
nlq (profiles/get-profile :nlq)]
(is (= (disj (tool-names nlq) "search_curated")
(tool-names embedding)))))
(testing "nlq-fallback matches embedding_next's general search; curated nlq swaps that for the library tool"
(let [tool-names (fn [profile] (set (map #(:tool-name (meta %)) (:tools profile))))
embedding (tool-names (profiles/get-profile :embedding_next))
fallback (tool-names (profiles/get-profile :nlq-fallback))
;; force the curated nlq (no redirect) — get-profile :nlq otherwise falls back when the index can't answer
curated (mt/with-dynamic-fn-redefs [entity-retrieval/entity-retrieval-available? (constantly true)]
(tool-names (profiles/get-profile :nlq)))]
;; the fallback profile is embedding_next's discovery surface (general `search`)
(is (= fallback embedding))
;; the curated profile is the same set with retrieve_library_entities in place of `search`
(is (= curated (-> embedding (disj "search") (conj "retrieve_library_entities"))))))
(binding [scope/*current-user-scope* api-scope/unrestricted]
(testing "navigate_user is excluded without the capability"
(let [tools (profiles/get-tools-for-profile :embedding_next [])]
@@ -174,6 +182,26 @@
(let [tools (profiles/get-tools-for-profile :embedding_next ["frontend:navigate_user_v1"])]
(is (contains? tools "navigate_user"))))))
(deftest nlq-data-discovery-fallback-test
(testing "the :nlq profile always keeps a data-discovery tool, swapping by index availability"
(binding [scope/*current-user-scope* api-scope/unrestricted]
(testing "entity retrieval AVAILABLE -> curated library tool, no general-search fallback"
(mt/with-dynamic-fn-redefs [entity-retrieval/entity-retrieval-available? (constantly true)]
(let [tools (profiles/get-tools-for-profile :nlq [])]
(is (contains? tools "retrieve_library_entities")
"the curated library tool is offered when the index can serve queries")
(is (not (contains? tools "search"))
"the general-search fallback is filtered out when the library is available")
(is (contains? tools "construct_notebook_query")))))
(testing "entity retrieval UNAVAILABLE -> general-search fallback, no curated library tool"
(mt/with-dynamic-fn-redefs [entity-retrieval/entity-retrieval-available? (constantly false)]
(let [tools (profiles/get-tools-for-profile :nlq [])]
(is (not (contains? tools "retrieve_library_entities"))
"the curated library tool is gated out when the index can't serve queries")
(is (contains? tools "search")
"the general-search fallback keeps the agent from having zero discovery tools")
(is (contains? tools "construct_notebook_query"))))))))
(deftest transform-feature-capabilities-test
(let [orig-has-feature (mt/original-fn #'premium-features/has-feature?)
transform-tools #{#'tools.transforms/write-transform-sql-tool

104
test/metabase/metabot/tools/curated_search_test.clj → test/metabase/metabot/tools/entity_retrieval_test.clj
View File

@@ -1,9 +1,9 @@
(ns metabase.metabot.tools.curated-search-test
(ns metabase.metabot.tools.entity-retrieval-test
(:require
[clojure.string :as str]
[clojure.test :refer :all]
[metabase.curated-search.core :as cs.core]
[metabase.metabot.tools.curated-search :as curated-search]
[metabase.entity-retrieval.core :as cs.core]
[metabase.metabot.tools.entity-retrieval :as entity-retrieval]
[metabase.metabot.tools.search :as tools.search]
[metabase.test :as mt]))
@@ -14,30 +14,31 @@
;; Pin the feature off so the defenterprise call takes its OSS fallback regardless of any ambient
;; premium token (e.g. an all-features dev token) or rows left in a local pgvector store.
(mt/with-premium-features #{}
(is (=? {:output #(str/includes? % "No matching curated entities")
(is (=? {:output #(str/includes? % "No matching library entities")
:structured-output {:result-type :search
:data []
:total_count 0
:weak_match false}}
(curated-search/curated-search-tool
(entity-retrieval/retrieve-library-entities-tool
{:user_search_prompt "revenue per region"}))))))
(deftest similarity-helper-test
(let [similarity (var-get #'curated-search/similarity)]
(let [similarity (var-get #'entity-retrieval/similarity)]
(testing "similarity is pulled from the :similarity score factor"
(is (= 0.8 (similarity {:scores [{:name :similarity :score 0.8} {:name :verified :score 1.0}]})))
(is (= 0.0 (similarity {:scores [{:name :verified :score 1.0}]}))))))
(is (= 0.8 (similarity {:scores [{:name :similarity :score 0.8}]})))
(is (= 0.0 (similarity {:scores []}))))))
(deftest build-matches-permission-filters-before-take-test
(testing "hydration permission-filters the whole candidate set before take-n, so an unreadable top hit doesn't crowd out readable matches just past the cut"
(let [build-matches (var-get #'curated-search/build-matches)
(let [build-matches (var-get #'entity-retrieval/build-matches)
;; three distinct entities, best-first
raw (vec (for [id [1 2 3]]
{:entity {:model "table" :id id}
:saved_search_prompt (str "p" id)
:usage_instructions ""
:score {:total_score (- 1.0 (/ id 10.0))
:scores [{:name :similarity :score 0.8}]}}))]
{:entity {:model "table" :id id}
:doc_type "name"
:doc_text (str "n" id)
:instructions ""
:score {:total_score (- 1.0 (/ id 10.0))
:scores [{:name :similarity :score 0.8}]}}))]
(mt/with-dynamic-fn-redefs [cs.core/search (fn [_ _] raw)
;; entity 1 (the top hit) is unreadable → dropped during hydration
tools.search/entity-refs->search-results
@@ -48,11 +49,11 @@
(deftest match->xml-locale-independent-test
(testing "score/similarity render with a '.' decimal separator even under a comma-decimal default locale"
(let [match->xml (var-get #'curated-search/match->xml)
(let [match->xml (var-get #'entity-retrieval/match->xml)
orig (java.util.Locale/getDefault)]
(try
(java.util.Locale/setDefault (java.util.Locale/forLanguageTag "de-DE"))
(let [out (match->xml {:saved_search_prompt "p" :usage_instructions ""
(let [out (match->xml {:doc_type "synonym" :matched_text "p" :usage_instructions ""
:score {:total_score 0.95} :similarity 0.78 :weak? false
:entity {:type "table" :id 1 :name "T" :database_id 1}})]
(is (str/includes? out "score=\"0.950\""))
@@ -61,9 +62,9 @@
(java.util.Locale/setDefault orig))))))
(deftest dedupe-by-entity-test
(let [dedupe (var-get #'curated-search/dedupe-by-entity)
(let [dedupe (var-get #'entity-retrieval/dedupe-by-entity)
results [{:entity {:model "metric" :id 9} :score {:total_score 0.95}} ; entity 9, best
{:entity {:model "metric" :id 9} :score {:total_score 0.80}} ; entity 9 again (sibling prompt)
{:entity {:model "metric" :id 9} :score {:total_score 0.80}} ; entity 9 again (sibling doc)
{:entity {:model "table" :id 1} :score {:total_score 0.70}}]]
(testing "collapses to distinct entities, keeping the first (best-scoring) occurrence"
(let [out (dedupe results)]
@@ -73,29 +74,35 @@
(let [out (dedupe [{:entity {:model "card" :id 5} :score {:total_score 0.9}}
{:entity {:model "question" :id 5} :score {:total_score 0.8}}])]
(is (= 1 (count out)))
(is (= 0.9 (-> out first :score :total_score)))))))
(is (= 0.9 (-> out first :score :total_score)))))
(testing "measures and segments are first-class distinct entities"
(let [out (dedupe [{:entity {:model "measure" :id 3} :score {:total_score 0.9}}
{:entity {:model "segment" :id 3} :score {:total_score 0.8}}])]
(is (= 2 (count out)))))))
(deftest format-output-test
(let [format-output (var-get #'curated-search/format-output)
(let [format-output (var-get #'entity-retrieval/format-output)
;; matches as produced by build-matches: :entity is already a full search-result record.
matches [{:saved_search_prompt "monthly revenue by region"
:usage_instructions "Use the Revenue metric; group by month."
:similarity 0.78 :weak? false
:score {:total_score 0.95}
:entity {:type "metric" :id 9 :name "Revenue" :database_id 1
:portable_entity_id "abcDEF"}}
{:saved_search_prompt "orders joined to customers"
:usage_instructions ""
:similarity 0.61 :weak? false
:score {:total_score 0.80}
:entity {:type "table" :id 1 :name "ORDERS" :database_schema "PUBLIC"
:database_id 1}}]
matches [{:doc_type "synonym"
:matched_text "monthly revenue by region"
:usage_instructions "Use the Revenue metric; group by month."
:similarity 0.78 :weak? false
:score {:total_score 0.95}
:entity {:type "metric" :id 9 :name "Revenue" :database_id 1
:portable_entity_id "abcDEF"}}
{:doc_type "name"
:matched_text "orders"
:usage_instructions ""
:similarity 0.61 :weak? false
:score {:total_score 0.80}
:entity {:type "table" :id 1 :name "ORDERS" :database_schema "PUBLIC"
:database_id 1}}]
out (format-output matches)]
(testing "renders score, similarity, confidence, the saved prompt, and hydrated entity details"
(testing "renders score, similarity, confidence, the matched_on doc, and hydrated entity details"
(is (str/includes? out "score=\"0.950\""))
(is (str/includes? out "similarity=\"0.780\""))
(is (str/includes? out "confidence=\"strong\""))
(is (str/includes? out "<saved_search_prompt>monthly revenue by region</saved_search_prompt>"))
(is (str/includes? out "<matched_on type=\"synonym\">monthly revenue by region</matched_on>"))
(is (str/includes? out "Revenue"))
(is (str/includes? out "ORDERS")))
(testing "usage instructions surface when present and are omitted when blank"
@@ -105,42 +112,43 @@
(is (not (str/starts-with? out "<note>"))))))
(deftest xml-escaping-test
(let [format-output (var-get #'curated-search/format-output)
matches [{:saved_search_prompt "P&L by quarter <2026>"
:usage_instructions "Profit & loss; values < 0 are losses."
(let [format-output (var-get #'entity-retrieval/format-output)
matches [{:doc_type "synonym"
:matched_text "P&L by quarter <2026>"
:usage_instructions "Profit & loss; values < 0 are losses."
:similarity 0.7 :weak? false
:score {:total_score 0.7}
:entity {:type "table" :id 1 :name "FINANCE" :database_id 1}}]
out (format-output matches)]
(testing "curator-entered text is escaped in the XML output"
(is (str/includes? out "<saved_search_prompt>P&amp;L by quarter &lt;2026&gt;</saved_search_prompt>"))
(is (str/includes? out "<matched_on type=\"synonym\">P&amp;L by quarter &lt;2026&gt;</matched_on>"))
(is (str/includes? out "<usage_instructions>Profit &amp; loss; values &lt; 0 are losses.</usage_instructions>")))))
(deftest weak-match-note-test
(let [format-output (var-get #'curated-search/format-output)
matches [{:saved_search_prompt "something only loosely related"
(let [format-output (var-get #'entity-retrieval/format-output)
matches [{:doc_type "name" :matched_text "MISC"
:usage_instructions "" :similarity 0.20 :weak? true
:score {:total_score 0.20}
:entity {:type "table" :id 7 :name "MISC" :database_id 1}}]
out (format-output matches)]
(testing "a weak top match prepends the cautionary note and marks confidence weak"
(is (str/starts-with? out "<note>"))
(is (str/includes? out "No strong curated match"))
(is (str/includes? out "No strong library match"))
(is (str/includes? out "confidence=\"weak\"")))))
(deftest flatten-data-test
(let [flatten-data (var-get #'curated-search/flatten-data)
matches [{:saved_search_prompt "p1" :usage_instructions "u1"
(let [flatten-data (var-get #'entity-retrieval/flatten-data)
matches [{:doc_type "synonym" :matched_text "p1" :usage_instructions "u1"
:score {:total_score 0.9} :similarity 0.7 :weak? false
:entity {:type "metric" :id 9 :name "Revenue"}}
{:saved_search_prompt "p2" :usage_instructions "u2"
{:doc_type "name" :matched_text "p2" :usage_instructions "u2"
:score {:total_score 0.6} :similarity 0.3 :weak? true
:entity {:type "table" :id 1 :name "ORDERS"}}]
out (flatten-data matches)]
(testing "one annotated record per match"
(is (= 2 (count out)))
(is (=? [{:type "metric" :id 9 :saved_search_prompt "p1" :usage_instructions "u1"
:confidence "strong" :similarity 0.7}
{:type "table" :id 1 :saved_search_prompt "p2" :usage_instructions "u2"
:confidence "weak" :similarity 0.3}]
(is (=? [{:type "metric" :id 9 :matched_doc_type "synonym" :matched_text "p1"
:usage_instructions "u1" :confidence "strong" :similarity 0.7}
{:type "table" :id 1 :matched_doc_type "name" :matched_text "p2"
:usage_instructions "u2" :confidence "weak" :similarity 0.3}]
out)))))

82
test/metabase/metabot/tools/resources_test.clj
View File

@@ -119,6 +119,9 @@
["metabase://metric/6" :metric ["6"]]
["metabase://metric/6/dimensions" :metric-dimensions ["6"]]
["metabase://metric/6/dimensions/dim-1" :metric-dimension ["6" "dim-1"]]
;; ----- Measure / Segment -----
["metabase://measure/9" :measure ["9"]]
["metabase://segment/10" :segment ["10"]]
;; ----- Transform -----
["metabase://transform/7" :transform ["7"]]
["metabase://transform/7/sources" :transform-sources ["7"]]
@@ -154,6 +157,8 @@
read-resource/fetch-metric (spy :metric)
read-resource/fetch-metric-dimensions (spy :metric-dimensions)
read-resource/fetch-metric-dimension (spy :metric-dimension)
read-resource/fetch-measure (spy :measure)
read-resource/fetch-segment (spy :segment)
read-resource/fetch-transform (spy :transform)
read-resource/fetch-transform-sources (spy :transform-sources)
read-resource/fetch-transform-target (spy :transform-target)
@@ -570,6 +575,83 @@
{:uris [(str "metabase://question/" q-from-model-id "/sources")]})]
(is (str/includes? output (str "uri=\"metabase://model/" model-id "\"")))))))))
(defn- measure-definition
"An MBQL5 measure definition: the sum of `field-id` over `table-id`."
[table-id field-id]
(let [mp (mt/metadata-provider)]
(lib/aggregate (lib/query mp (lib.metadata/table mp table-id))
(lib/sum (lib.metadata/field mp field-id)))))
(defn- segment-definition
"An MBQL5 segment definition: a filter of `field-id` > `value` on `table-id`."
[table-id field-id value]
(let [mp (mt/metadata-provider)]
(lib/filter (lib/query mp (lib.metadata/table mp table-id))
(lib/> (lib.metadata/field mp field-id) value))))
(deftest read-measure-resource-test
(mt/test-drivers #{:h2}
(mt/with-current-user (mt/user->id :crowberto)
(let [orders (mt/id :orders)
total (mt/id :orders :total)]
(mt/with-temp [:model/Measure {measure-id :id}
{:name "Order Revenue"
:description "sum of order totals"
:table_id orders
:creator_id (mt/user->id :crowberto)
:definition (measure-definition orders total)}]
(testing "metabase://measure/{id} returns the measure with parent-table context + portable entity id"
(let [result (read-resource/read-resource {:uris [(str "metabase://measure/" measure-id)]})
structured (get-in result [:resources 0 :content :structured-output])
output (:output result)]
(is (=? {:type :measure
:name "Order Revenue"
:database_id (mt/id)
:base_table_id orders
:portable-entity-id string?
:base_table_portable_fk vector?}
structured))
(testing "rendered XML carries the measure name and a portable entity id"
(is (str/includes? output "Order Revenue"))
(is (str/includes? output "portable_entity_id=")))))
(testing "returns an error for an unknown measure"
(is (=? {:resources [{:error string?}]}
(read-resource/read-resource {:uris ["metabase://measure/99999"]}))))
(testing "errors when the user can't read the parent table"
(with-redefs [mi/can-read? (constantly false)]
(is (error? (read-resource/read-resource {:uris [(str "metabase://measure/" measure-id)]}))))))))))
(deftest read-segment-resource-test
(mt/test-drivers #{:h2}
(mt/with-current-user (mt/user->id :crowberto)
(let [orders (mt/id :orders)
total (mt/id :orders :total)]
(mt/with-temp [:model/Segment {segment-id :id}
{:name "Big Orders"
:description "totals over 100"
:table_id orders
:definition (segment-definition orders total 100)}]
(testing "metabase://segment/{id} returns the segment with parent-table context + portable entity id"
(let [result (read-resource/read-resource {:uris [(str "metabase://segment/" segment-id)]})
structured (get-in result [:resources 0 :content :structured-output])
output (:output result)]
(is (=? {:type :segment
:name "Big Orders"
:database_id (mt/id)
:base_table_id orders
:portable-entity-id string?
:base_table_portable_fk vector?}
structured))
(testing "rendered XML carries the segment name and a portable entity id"
(is (str/includes? output "Big Orders"))
(is (str/includes? output "portable_entity_id=")))))
(testing "returns an error for an unknown segment"
(is (=? {:resources [{:error string?}]}
(read-resource/read-resource {:uris ["metabase://segment/99999"]}))))
(testing "errors when the user can't read the parent table"
(with-redefs [mi/can-read? (constantly false)]
(is (error? (read-resource/read-resource {:uris [(str "metabase://segment/" segment-id)]}))))))))))
(deftest read-dashboard-items-test
(mt/with-current-user (mt/user->id :crowberto)
(mt/with-temp [:model/Dashboard {dash-id :id} {}

18
test/metabase/metabot/tools/search_test.clj
View File

@@ -500,16 +500,16 @@
(is (= 3 (count results))))))))))
(deftest entity-refs->search-results-same-card-two-types-test
(testing "a card referenced under two type strings hydrates to one record per ref (neither is dropped)"
(testing "a card referenced under two (possibly stale) type strings collapses to one record with its current type"
(mt/with-test-user :crowberto
(mt/with-temp [:model/Card {c-id :id} {:name "Dual Typed"
(mt/with-temp [:model/Card {c-id :id} {:name "Dual Typed" :type :model
:database_id (mt/id) :table_id (mt/id :orders)
:dataset_query {:database (mt/id) :type :query
:query {:source-table (mt/id :orders)}}}]
(let [results (search/entity-refs->search-results
[{:model "model" :id c-id} {:model "metric" :id c-id}])]
(is (= #{["model" c-id] ["metric" c-id]}
(set (map (juxt :type :id) results)))))))))
(is (= [["model" c-id]] (map (juxt :type :id) results))
"one record, carrying the card's current type"))))))
(deftest entity-refs->search-results-respects-read-perms-test
(testing "hydration drops entities the current user can't read — a curated entry may point at a restricted one"
@@ -601,3 +601,13 @@
(map (comp first #(str/split % #"\s") :name))))]
(is (= ["Bookmarked" "Regular"] (query)))
(is (= ["Regular" "Bookmarked"] (query {:bookmarked -1})))))))))
(deftest card-ref-hydration-emits-current-string-type-test
(testing "a card ref hydrates to the Card's CURRENT type as a string — not the stale ref type, not a keyword"
;; regression: a stale index hit across a metric<->model relabel must describe the entity by its current
;; shape, and the type must be the agent-facing string (a :model keyword breaks entity-class + enrichers).
(mt/with-current-user (mt/user->id :crowberto)
(mt/with-temp [:model/Card {card-id :id} {:type :model}]
(let [[result] (search/entity-refs->search-results [{:model "metric" :id card-id}])]
(is (= "model" (:type result)))
(is (string? (:type result))))))))

17
test/metabase/metabot/tools_test.clj
View File

@@ -2,6 +2,7 @@
(:require
[clojure.test :refer :all]
[metabase.api-scope.core :as api-scope]
[metabase.entity-retrieval.core :as entity-retrieval]
[metabase.metabot.agent.profiles :as profiles]
[metabase.metabot.scope :as scope]
[metabase.metabot.tools :as agent-tools]
@@ -68,16 +69,22 @@
(is (contains? tools "ask_for_sql_clarification"))))
(deftest get-tools-for-nlq-profile-test
(testing "the general search is always available; the curated search tool only with the semantic-search feature"
(mt/with-premium-features #{}
(testing "nlq discovers data through the curated library tool when it can serve queries, else general search"
;; Entity retrieval unavailable (no pgvector / OSS): the general `search` fallback is the discovery tool,
;; so the agent is never left with zero ways to find data. The library tool is filtered out.
(mt/with-dynamic-fn-redefs [entity-retrieval/entity-retrieval-available? (constantly false)]
(let [tools (tools-for-profile :nlq)]
(is (map? tools))
(is (contains? tools "search"))
(is (contains? tools "construct_notebook_query"))
(is (contains? tools "create_chart"))
(is (not (contains? tools "search_curated")))))
(mt/with-premium-features #{:semantic-search}
(is (contains? (tools-for-profile :nlq) "search_curated")))))
(is (not (contains? tools "retrieve_library_entities")))))
;; Entity retrieval available (pgvector configured + semantic-search licensed): the curated library tool
;; replaces general search. Exactly one discovery tool survives capability filtering.
(mt/with-dynamic-fn-redefs [entity-retrieval/entity-retrieval-available? (constantly true)]
(let [tools (tools-for-profile :nlq)]
(is (contains? tools "retrieve_library_entities"))
(is (not (contains? tools "search")))))))
(deftest ^:parallel get-tools-for-document-generate-content-profile-test
(let [tools (tools-for-profile :document-generate-content)]

120
test/metabase/osi/ai_context/api_test.clj Normal file
View File

@@ -0,0 +1,120 @@
(ns metabase.osi.ai-context.api-test
(:require
[clojure.test :refer :all]
[metabase.test :as mt]
[toucan2.core :as t2]))
(set! *warn-on-reflection* true)
(defmacro ^:private with-test-entry
"`mt/with-temp` for an OsiAiContext, with entity/ai_context defaults so callers can pass `{}`."
[[sym attrs] & body]
`(mt/with-temp [:model/OsiAiContext ~sym (merge {:entity_type "table"
:entity_local_id 1
:ai_context {:instructions "find orders"}}
~attrs)]
~@body))
(defn- has-entity?
"Whether `entries` contains a row for `(entity-type, entity-local-id)`."
[entries entity-type entity-local-id]
(some #(= {:entity_type entity-type :entity_local_id entity-local-id}
(select-keys % [:entity_type :entity_local_id]))
entries))
(deftest list-test
(with-test-entry [_ {}]
(testing "superuser can list ai_context entries"
(let [response (mt/user-http-request :crowberto :get 200 "osi/ai-context/")]
(is (contains? response :data))
(is (contains? response :total))
(is (has-entity? (:data response) "table" 1))))
(testing "non-superuser gets 403"
(mt/user-http-request :rasta :get 403 "osi/ai-context/"))
(testing "a legacy row whose entity_type is no longer in the write enum still reads back fine"
(with-test-entry [_ {:entity_type "retired-model-string" :entity_local_id 7}]
(is (has-entity? (:data (mt/user-http-request :crowberto :get 200 "osi/ai-context/"))
"retired-model-string" 7))))))
(deftest get-test
(with-test-entry [_ {:entity_type "table" :entity_local_id 1 :ai_context {:instructions "find customers"}}]
(testing "superuser can fetch an ai_context entry by its logical key"
(is (=? {:entity_type "table"
:entity_local_id 1
:ai_context {:instructions "find customers"}}
(mt/user-http-request :crowberto :get 200 "osi/ai-context/table/1"))))
(testing "a card flavor resolves the canonical \"card\" row"
(with-test-entry [_ {:entity_type "metric" :entity_local_id 8 :ai_context {:instructions "by month"}}]
(is (=? {:entity_type "card" :entity_local_id 8}
(mt/user-http-request :crowberto :get 200 "osi/ai-context/model/8")))))
(testing "returns 404 for an entity with no row (including an unknown entity_type)"
(mt/user-http-request :crowberto :get 404 "osi/ai-context/table/999999")
(mt/user-http-request :crowberto :get 404 "osi/ai-context/garbage/1"))
(testing "non-superuser gets 403"
(mt/user-http-request :rasta :get 403 "osi/ai-context/table/1"))))
(deftest upsert-test
(testing "PUT creates the entry (a card flavor is stored under the canonical \"card\")"
(try
(is (=? {:entity_type "card"
:entity_local_id 42
:ai_context {:instructions "Use the Revenue metric." :synonyms ["sales"]
:examples ["revenue last month"]}}
(mt/user-http-request :crowberto :put 200 "osi/ai-context/metric/42"
{:ai_context {:instructions "Use the Revenue metric." :synonyms ["sales"]
:examples ["revenue last month"]}})))
(finally (t2/delete! :model/OsiAiContext :entity_type "card" :entity_local_id 42))))
(testing "measure and segment entity refs are accepted (they are indexed library entities)"
(doseq [entity-type ["measure" "segment"]]
(try
(is (=? {:entity_type entity-type :entity_local_id 5}
(mt/user-http-request :crowberto :put 200 (str "osi/ai-context/" entity-type "/5")
{:ai_context {:synonyms ["alias"]}})))
(finally (t2/delete! :model/OsiAiContext :entity_type entity-type :entity_local_id 5)))))
(testing "PUTting the same card under different flavors upserts one row (both normalize to \"card\")"
(try
(mt/user-http-request :crowberto :put 200 "osi/ai-context/metric/99" {:ai_context {:instructions "v1"}})
(let [second-resp (mt/user-http-request :crowberto :put 200 "osi/ai-context/model/99"
{:ai_context {:instructions "v2"}})]
(is (= {:instructions "v2"} (:ai_context second-resp)) "ai_context replaced across a metric->model relabel")
(is (=? {:entity_type "card" :entity_local_id 99 :ai_context {:instructions "v2"}}
(mt/user-http-request :crowberto :get 200 "osi/ai-context/model/99"))
"one canonical card row, fetched by either flavor")
(is (= 1 (t2/count :model/OsiAiContext :entity_type "card" :entity_local_id 99))))
(finally (t2/delete! :model/OsiAiContext :entity_type "card" :entity_local_id 99))))
(testing "a re-PUT with the same ai_context is idempotent (a no-op update must not retry as a duplicate insert)"
(try
(mt/user-http-request :crowberto :put 200 "osi/ai-context/measure/77" {:ai_context {:instructions "same"}})
(is (=? {:entity_type "measure" :entity_local_id 77 :ai_context {:instructions "same"}}
(mt/user-http-request :crowberto :put 200 "osi/ai-context/measure/77" {:ai_context {:instructions "same"}})))
(finally (t2/delete! :model/OsiAiContext :entity_type "measure" :entity_local_id 77))))
(testing "ai_context is required"
(mt/user-http-request :crowberto :put 400 "osi/ai-context/table/1" {}))
(testing "an over-long instructions string is rejected"
(mt/user-http-request :crowberto :put 400 "osi/ai-context/table/1"
{:ai_context {:instructions (apply str (repeat 6000 "x"))}}))
(testing "too many synonyms is rejected"
(mt/user-http-request :crowberto :put 400 "osi/ai-context/table/1"
{:ai_context {:synonyms (mapv str (range 51))}}))
(testing "entity types the reconciler never indexes (card/question/garbage) are rejected by the route"
(doseq [entity-type ["card" "question" "garbage"]]
(mt/user-http-request :crowberto :put 400 (str "osi/ai-context/" entity-type "/1")
{:ai_context {:instructions "x"}})))
(testing "non-superuser gets 403"
(mt/user-http-request :rasta :put 403 "osi/ai-context/table/1" {:ai_context {:instructions "x"}})))
(deftest reconcile-test
(testing "POST /reconcile requires superuser"
(mt/user-http-request :rasta :post 403 "osi/ai-context/reconcile"))
(testing "without semantic search the index can't be reconciled, so it 400s rather than no-opping silently"
(mt/user-http-request :crowberto :post 400 "osi/ai-context/reconcile")))
(deftest delete-test
(with-test-entry [_ {}]
(testing "non-superuser gets 403"
(mt/user-http-request :rasta :delete 403 "osi/ai-context/table/1"))
(testing "superuser can delete an entry by its logical key"
(mt/user-http-request :crowberto :delete 204 "osi/ai-context/table/1")
(is (nil? (t2/select-one :model/OsiAiContext :entity_type "table" :entity_local_id 1))))
(testing "returns 404 for an entity with no row"
(mt/user-http-request :crowberto :delete 404 "osi/ai-context/table/999999"))))

94
test/metabase/osi/models/osi_ai_context_test.clj Normal file
View File

@@ -0,0 +1,94 @@
(ns metabase.osi.models.osi-ai-context-test
(:require
[clojure.test :refer :all]
[metabase.entity-retrieval.mirror :as mirror]
[metabase.test :as mt]
[metabase.util :as u]
[toucan2.core :as t2]))
(set! *warn-on-reflection* true)
(def ^:private entity {:entity_type "table" :entity_local_id 42})
(def ^:private ai-context {:instructions "Use for total revenue; group by month."
:synonyms ["sales" "turnover"]
:examples ["monthly revenue by region"]})
(defn- by-key
"Look the row up by its compound key."
[{:keys [entity_type entity_local_id]}]
(t2/select-one :model/OsiAiContext :entity_type entity_type :entity_local_id entity_local_id))
(deftest entity-and-ai-context-roundtrip-test
(testing "ai_context is a keywordized JSON object and timestamps are populated"
(mt/with-temp [:model/OsiAiContext _ (assoc entity :ai_context ai-context)]
(is (=? {:entity_type "table"
:entity_local_id 42
:ai_context ai-context
:created_at some?
:updated_at some?}
(by-key entity))))))
(deftest ai-context-minimal-test
(testing "ai_context can be a minimal blob (just instructions, no synonyms/examples)"
(mt/with-temp [:model/OsiAiContext _ (assoc entity :ai_context {:instructions "Just this."})]
(is (= {:instructions "Just this."} (:ai_context (by-key entity)))))))
(deftest hooks-nudge-a-targeted-reconcile-test
(testing "insert, update and delete each nudge the entity's reconcile once — deferred until the txn commits"
(let [nudges (atom [])
ref ["table" 42]]
(mt/with-dynamic-fn-redefs [mirror/request-entity-sync! (fn [entity-type entity-local-id]
(swap! nudges conj [entity-type entity-local-id])
nil)]
(try
(testing "insert"
(t2/with-transaction [_conn]
(t2/insert! :model/OsiAiContext (assoc entity :ai_context ai-context))
(is (= [] @nudges) "not nudged while the transaction is still open"))
(is (= [ref] @nudges) "nudged once, after commit"))
(testing "update"
(t2/with-transaction [_conn]
(t2/update! :model/OsiAiContext :entity_type "table" :entity_local_id 42 {:ai_context {:instructions "changed"}})
(is (= [ref] @nudges) "still only the insert's nudge while open"))
(is (= [ref ref] @nudges) "nudged again, after commit"))
(testing "delete"
(t2/with-transaction [_conn]
(t2/delete! :model/OsiAiContext :entity_type "table" :entity_local_id 42)
(is (= [ref ref] @nudges) "not nudged while open"))
(is (= [ref ref ref] @nudges) "nudged again, after commit"))
(finally (t2/delete! :model/OsiAiContext :entity_type "table" :entity_local_id 42)))))))
(deftest nudge-not-fired-on-rollback-test
(testing "a rolled-back write never nudges (and leaves no row)"
(let [nudges (atom [])]
(mt/with-dynamic-fn-redefs [mirror/request-entity-sync! (fn [entity-type entity-local-id]
(swap! nudges conj [entity-type entity-local-id])
nil)]
(u/ignore-exceptions
(t2/with-transaction [_conn]
(t2/insert! :model/OsiAiContext (assoc entity :entity_local_id 43 :ai_context ai-context))
(throw (ex-info "roll back" {}))))
(is (= [] @nudges))
(is (nil? (by-key (assoc entity :entity_local_id 43))))))))
(deftest card-flavors-share-one-canonical-key-test
(testing "card flavors are stored under one canonical \"card\" key (normalized on write), so a metric and
a model with the same local id are the same primary-key row"
(mt/with-temp [:model/OsiAiContext _ {:entity_type "metric" :entity_local_id 7
:ai_context {:instructions "x"}}]
(is (=? {:entity_type "card" :entity_local_id 7 :ai_context {:instructions "x"}}
(t2/select-one :model/OsiAiContext :entity_type "card" :entity_local_id 7))))))
(deftest hooks-never-break-appdb-writes-test
(testing "insert/update/delete succeed even though the mirror is unavailable in tests"
;; The hooks call the OSS defenterprise shim, which no-ops without an enterprise license. They must
;; never throw or fail the authoritative write.
(mt/with-temp [:model/OsiAiContext row (assoc entity :ai_context ai-context)]
(let [k [(:entity_type row) (:entity_local_id row)]]
(testing "update"
(is (pos? (t2/update! :model/OsiAiContext :entity_type (k 0) :entity_local_id (k 1)
{:ai_context {:instructions "u"}})))
(is (= {:instructions "u"} (:ai_context (by-key entity)))))
(testing "delete"
(is (pos? (t2/delete! :model/OsiAiContext :entity_type (k 0) :entity_local_id (k 1))))
(is (nil? (by-key entity))))))))