Compare commits

...

46 Commits

Author SHA1 Message Date
Sloan Sparger
6e4e20090d TEMP COMMIT 2025-07-30 15:51:14 -05:00
Tyler Plude
89bdd982d1 Fix serialization baseline test 2025-07-30 12:53:16 -07:00
Sloan Sparger
a24b6daabd fix type issue 2025-07-30 13:33:38 -05:00
Sloan Sparger
fd7059b2de add basic setting switch for enabling semantic search, upsell component, logic for falling back to default search in settings, and basic e2e test 2025-07-30 13:19:13 -05:00
Dan Stone
15eed35a49 BOT-255 Index metadata system (#61379)
* Add index-metadata system

Adds the necessary indirection to support more advanced index table
management.

- Allows multiple indexes to exist
- Indexes are tracked by a new index_metadata table
- A single active pointer determines which index to use for the current
search engine backend
- It is assumed table == logical index and there is no other
partitioning, hence the unique constraint
- Some small simplifications to index.clj (drop index-name, assume names
are prefixed with the table name)
- The top level orchestration code happens in a new pgvector-api
namespace to avoid bloating the enterprise hooks (which assume a global
environment), and allow the api level to be properly parameterized.

* Kondo

* Appease kondo with require for next.jdbc in comment block

---------

Co-authored-by: Mike Appleby <86076+appleby@users.noreply.github.com>
2025-07-28 10:25:53 +01:00
appleby
e05b159dbb Add prometheus metrics for semantic search backend (#61309)
* Add :metabase-search/semantic-index-size prometheus metric

And update it on upsert-index! / delete-from-index!

This metric mirrors the corresponding :metabase-search/appdb-index-size

* Add :metabase-search/semantic-embedding-tokens prometheus metric

* Make log levels / messages more consistent

* Remove vestigial populate! reference

* Record actual token usage as reported by OpenAI API response
2025-07-25 16:23:53 -07:00
appleby
92a080aa61 Remove indexed-entities from the semantic search model-set and documents reducible (#61304)
* Remove indexed-entities from the semantic search model-set

* Add comment about filter-read-permitted and indexed-entity docs
2025-07-25 09:16:02 -07:00
Tyler Plude
61c413fd8e Merge branch 'master' into nm-semantic-search-in-metabase 2025-07-24 16:24:34 -07:00
Sloan Sparger
9496e659f4 [Metabot] Hybrid search w/ RRF (#61353)
* rrf based hybrid search implemented

* keyword results weighted more heavily by name, corrected some issues with the rrf calc

* fix final result ranking

* move tsvector-related util functions into search.util and expose in the search interface to fix kondo errors

* Taking a stab at fixing SQL injection risks following existing postgres tsvector querying pattern

* Move postgres tsvector util tests

* gracefully handle nil names

* switch back to raw embedding literals but validate their contents

* add back perm checks (whoops)

* add back str/replace for ts-search-expr, bump up cosine distance cutoff

* tweak cutoff again, fix create index test

* remove str/replace usage bc it seems unneeded

* fix rrf calc

---------

Co-authored-by: Noah Moss <noahbmoss@gmail.com>
2025-07-24 15:35:56 -05:00
Noah Moss
1f81d2f35d Quick follow-up optimization to prior PR: don't iterate over all of filtered-documents in each batch 2025-07-22 17:36:06 -04:00
Noah Moss
b307dc5c8f Fix bug in associating batched embeddings with original documents (#61347)
* Fix bug in associating batched embeddings with original documents

* tweak

* add log
2025-07-22 17:09:20 -04:00
Dan Stone
4fbefd6dcf Parameterization refactor (#61328)
Adds the `index` and `embedding-model` parameters to functions in `index` and `embedding` respectively. This allows a degree more flexibility and less coupling to global variables which restrict how indexing functions can be used. 

The goal is this PR supports changes regarding active/pending and the embedding model being something you can play with at runtime.

* Add embedding-model parameter

* Restore previous top-level transaction behaviour on init/reindex

* Reduced redundancy

* Embedding test fixes

* Fixup comments

* Fix rebase goofs

* Fix tests again

* Joy

* Add `index` parameter, composed of table-name and embedding model

* Keep index name in the index structure

* English is a hard language

* Remove `:redef` we no longer want

* Another merge goof fix
2025-07-22 17:09:20 -04:00
Noah Moss
29bf9a91ad Better batching for semantic search (#61299)
* streaming embedding results to the index-updating code

* rework

* suppress kondo
2025-07-22 17:09:19 -04:00
appleby
c61c604145 Only create HNSW embeddings index :if-not-exists (#61305)
Since init! no longer wipes the index table if it already exists, we also need to only create the embeddings HNSW
index if it does not exist.
2025-07-22 17:09:19 -04:00
Noah Moss
a6d39b60de fix document filtering 2025-07-22 17:09:19 -04:00
Noah Moss
a45843559d alter create-index-table! to not wipe the index if it already exists, and use the upsert API during index population (stop-gap changes to try to get the staging instance working) 2025-07-22 17:09:19 -04:00
appleby
a856712245 Create HNSW index on semantic search embeddings column (#61238)
* Create HNSW index on semantic search embeddings column

* Add test for HNSW index creation

* Ensure identifiers are quoted when formatting honey sql in semantic search index

The nano-ids used in temp table names when testing (and various other places) might contain uppercase chars or hyphens which need to be quoted.

* Fix batch-delete! log message
2025-07-22 17:09:19 -04:00
Noah Moss
084d8b9ace Filter indexed entities from semantic search index 2025-07-22 17:09:19 -04:00
appleby
7c647de9b0 Remove transaction from batched semantic search index updates (#61229)
* Drop uneccessary ! from table-exists-in-db?!

* Don't wrap batched index updates in a transaction

https://github.com/metabase/metabase/pull/61117#discussion_r2214828302
2025-07-22 17:09:19 -04:00
Noah Moss
8f70a19dca Lower batch size, improve logging, and make it configurable via env var (#61274)
* Lower batch size, improve logging, and make it configurable via env var

* fix whitespace
2025-07-22 17:09:19 -04:00
Sloan Sparger
e21d011708 fix fe tests 2025-07-22 17:09:19 -04:00
Noah Moss
be980c252f Add more DB existance checks for safety 2025-07-22 17:09:19 -04:00
Noah Moss
68790b70ce Init db connection in reindex! if needed 2025-07-22 17:09:19 -04:00
Noah Moss
855b60a016 Batched requests to OpenAI for fetching vector embeddings (#61217)
* Batch calls to openai according to token counts

* refactor

* mock the new batch embedding function

* mock both batched and unbatched embedding functions in tests

* a little code consolidation

* address comments
2025-07-22 17:09:19 -04:00
Dan Stone
be1c82dcb8 Fix accidental blocking search index init (#61200)
* kick ci

* Fix accidental eager index initialisation

---------

Co-authored-by: Noah Moss <noahbmoss@gmail.com>
2025-07-22 17:09:19 -04:00
appleby
848192ccfb Batched index updates for semantic search (#61117)
* kick ci

* Add comment to trigger reindex job

* Fix docstring typo

* Batch index updates for semantic search

* Wrap {create,drop}-index-table! in mt/as-admin in with-temp-index-table!

* Move binding for *vector-dimensions* into with-mock-embeddings

* Create metabase-enterprise.semantic-search.test-util and move once-fixture there

* Move other test helpes and macros into the semantic.test-util ns

* Factor out with-indexable-documents! from with-index!

* Move binding of search.ingestion/*force-sync* into with-temp-index-table!

Questionable, but realistically can't see tests wanting to use a temp index table without forcing sync reindex.

* Add table-exists-in-db?! to semantic.test-util

* Add some basic tests for funcs in the semantic.index namespace

* Add tests for batched index updates

* Disable flaky tests

* Fix up issues post-rebase

---------

Co-authored-by: Noah Moss <noahbmoss@gmail.com>
2025-07-22 17:09:19 -04:00
Noah Moss
778e210968 Swappable embedding models + semantic search kill switch (#61176)
* first stab at swappable embedding models + setting to disable semantic search globally

* fixups

* fixups

* fix up error in tests

* add semantic-search to stats ping to quiet test

* quiet another test

* fix whitespace

* fix kondo warnings

* fix up and simplify embedding tests

* fix core init test

* fix more tests

* tweak setting defaults + comments, remove debugging code

* fix compilation

* change oss version of supported? to return false instead of throwing

* fix lint issues
2025-07-22 17:09:19 -04:00
Dan Stone
6ed6ffadb7 Test collection based permission checks apply to search results (#61152)
* Fix t2-instance derivation from search docs

Lookup in database as the search document model is lossy

* Test collection based permission checks apply to search results
2025-07-22 17:09:19 -04:00
Sloan Sparger
0820d6af6e FE - default to app db search, enable semantic in command palette and search page, add disable semantic search on search page 2025-07-22 17:09:19 -04:00
Noah Moss
98d3cb58fa Register semantic-search premium feature 2025-07-22 17:09:19 -04:00
Noah Moss
a7b719807f Use EE feature for semantic search defenterprises (#61139)
* switch defenterprise features from :none to :semantic-search

* update tests
2025-07-22 17:09:19 -04:00
appleby
5a21df8919 Add nm-semantic-search-in-metabase to uberjar workflow and remove metabot-v3-main (#61142)
In preparation to deploy the nm-semantic-search-in-metabase branch to stats.
2025-07-22 17:09:18 -04:00
Noah Moss
5ffb202684 Semantic search filtering, basic querying tests, connection pooling, etc (#61098)
* fix results postprocessing, add verified field to search index

* scaffold a basic query test

* index filtering tests

* quicky and dirty implementation of the requested ui changes

* Add basic connection pooling + tests

* Fix test indentation

* add test for verified content, refactor tests a bit

* address comments

---------

Co-authored-by: Sloan Sparger <sloansparger@gmail.com>
2025-07-22 17:08:50 -04:00
Dan Stone
beaf957d2b Add semantic search test workflow (#61065)
* Add new workflow for semantic search

* Add a quick example test to see if things work

* Have the new tests only run if db-url is defined

* Sparks of joy

* Add a delay to appease the kondo gods

Not sure this rule should apply to once-fixture, but for now lets fix
it.
2025-07-22 17:08:50 -04:00
Noah Moss
3896229518 fix remaining clj-kondo issues 2025-07-22 17:08:50 -04:00
Noah Moss
47387d8302 remove a blank line 2025-07-22 17:08:50 -04:00
Noah Moss
e351fd7d82 fix up syntax in metabase.search.impl and whitespace 2025-07-22 17:08:50 -04:00
Noah Moss
27dac611ea sort kondo modules 2025-07-22 17:08:50 -04:00
Dan Stone
0635595e7c Adds mi/can-read? checks to search results (#61060)
* Add a couple of helpful comments for manual testing

* Add ability to change db-url at the repl

* Add can-read? checks for search results

* Semantic search engine should not be used if the db is not setup
2025-07-22 17:08:50 -04:00
Noah Moss
3a62376cff remove ^:const to try to fix compilation errorin tests 2025-07-22 17:08:50 -04:00
Noah Moss
d730468b24 kick ci 2025-07-22 17:08:50 -04:00
Noah Moss
205b88f6bb fix whitespace 2025-07-22 17:08:50 -04:00
Noah Moss
75dd41a0d6 fix up reindex! call 2025-07-22 17:08:50 -04:00
Noah Moss
894f808c0d WIP getting the integration working 2025-07-22 17:08:50 -04:00
Noah Moss
11b6380431 copy docker compose, add defenterprise integration 2025-07-22 17:08:50 -04:00
Noah Moss
d309f2a7f9 copy over relevant code from semantic search service 2025-07-22 17:08:50 -04:00
63 changed files with 3431 additions and 246 deletions

View File

@@ -1500,6 +1500,7 @@
premium-features
enterprise/advanced-config
enterprise/scim
enterprise/semantic-search
enterprise/sso}}
enterprise/api
@@ -1615,7 +1616,6 @@
warehouse-schema
warehouses}}
enterprise/email
{:team "Admin Webapp"
:api #{metabase-enterprise.email.api}
@@ -1667,7 +1667,8 @@
enterprise/llm
{:team "Metabot"
:api #{metabase-enterprise.llm.api
metabase-enterprise.llm.init}
metabase-enterprise.llm.init
metabase-enterprise.llm.settings}
:uses #{analytics
analyze
api
@@ -1765,6 +1766,19 @@
:uses #{premium-features
search}}
enterprise/semantic-search
{:team "Metabot"
:api #{metabase-enterprise.semantic-search.core
metabase-enterprise.semantic-search.init}
:uses #{analytics
connection-pool
models
premium-features
search
settings
util
enterprise/llm}}
enterprise/serialization
{:team "Workflows"
:api #{metabase-enterprise.serialization.api

View File

@@ -74,6 +74,13 @@ jobs:
with:
skip: ${{ needs.files-changed.outputs.backend_all != 'true' && needs.static-viz-files-changed.outputs.static_viz != 'true' }}
semantic-search-tests:
needs: files-changed
uses: ./.github/workflows/semantic-search.yml
secrets: inherit
with:
skip: ${{ needs.files-changed.outputs.backend_all != 'true' }}
app-db-tests:
needs: files-changed
uses: ./.github/workflows/app-db.yml

145
.github/workflows/semantic-search.yml vendored Normal file
View File

@@ -0,0 +1,145 @@
name: Semantic Search Tests
on:
workflow_call:
inputs:
skip:
type: boolean
default: false
concurrency:
group: ${{ github.workflow }}-${{ github.head_ref && github.ref || github.run_id }}-semantic-search
cancel-in-progress: true
jobs:
semantic-search-tests-pgvector:
if: ${{ !inputs.skip }}
runs-on: ubuntu-22.04
timeout-minutes: 60
strategy:
fail-fast: false
matrix:
version:
- name: pgvector pg16
junit-name: semantic-search-tests-pgvector-16
image: pgvector/pgvector:pg16
env:
enable-ssl-tests: 'false'
job:
- name: Semantic Search Tests
build-static-viz: false
test-args: >-
:only '"enterprise/backend/test/metabase_enterprise/semantic_search"'
exclude-tag: ':mb/driver-tests'
services:
pgvector:
image: ${{ matrix.version.image }}
ports:
- "5432:5432"
env:
POSTGRES_DB: mb_semantic_search
POSTGRES_USER: postgres
POSTGRES_PASSWORD: postgres
env:
# Only run old migrations tests on pushes to master or release branches. All other branches should skip tests
# with the tag `mb/old-migrations-test`. `__ADDITIONAL_EXCLUDED_TAG__` is not used anywhere outside of splicing
# it in to the command below.
__ADDITIONAL_EXCLUDED_TAG__: >-
${{
(
github.event_name == 'push' &&
(
github.ref == 'master' ||
startsWith(github.ref, 'release-')
) &&
''
) ||
':mb/old-migrations-test'
}}
MB_PGVECTOR_DB_URL: 'jdbc:postgres://localhost:5432/mb_semantic_search?user=postgres&password=postgres'
CI: 'true'
name: "${{ matrix.version.name }} ${{ matrix.job.name }}"
steps:
- uses: actions/checkout@v4
- name: Test ${{ matrix.version.name }}
uses: ./.github/actions/test-driver
with:
build-static-viz: ${{ matrix.job.build-static-viz }}
junit-name: ${{ matrix.version.junit-name }}
test-args: >-
${{ matrix.job.test-args }}
:exclude-tags '[ ${{ matrix.job.exclude-tag }} ${{ env.__ADDITIONAL_EXCLUDED_TAG__ }}]'
- name: Upload Test Results
uses: ./.github/actions/upload-test-results
if: always()
with:
input-path: ./target/junit/
output-name: ${{ github.job }}
bucket: ${{ vars.AWS_S3_TEST_RESULTS_BUCKET }}
aws-access-key-id: ${{ secrets.AWS_TEST_RESULTS_ACCESS_KEY_ID }}
aws-secret-access-key: ${{ secrets.AWS_TEST_RESULTS_SECRET_ACCESS_KEY }}
aws-region: ${{ vars.AWS_REGION }}
trunk-api-token: ${{ secrets.TRUNK_API_TOKEN }}
semantic-search-tests-result:
needs:
- semantic-search-tests-pgvector
runs-on: ubuntu-latest
timeout-minutes: 5
name: semantic-search-tests-result
if: always() && !cancelled()
steps:
- name: Check job status
uses: actions/github-script@v7
env:
needs: ${{ toJson(needs) }}
with:
script: | # js
const needs = JSON.parse(process.env.needs);
const jobs = Object.entries(needs).map(
([jobName, jobValues]) => ({
name: jobName,
result: jobValues.result
}));
// are all jobs skipped or successful?
if (jobs.every(job => (job.result === 'skipped' || job.result === 'success'))) {
console.log("");
console.log(" _------. ");
console.log(" / , \_ ");
console.log(" / / /{}\ |o\_ ");
console.log(" / \ `--' /-' \ ");
console.log(" | \ \ | ");
console.log(" | |`-, | ");
console.log(" / /__/)/ ");
console.log(" | | ");
console.log("");
console.log("All semantic search tests have passed (or have been skipped). Cam is very proud of you.");
process.exit(0);
}
// otherwise, something failed
console.log("");
console.log(" .::::::::::. .::::::::::. ");
console.log(" .::::''''''::::. .::::''''''::::. ");
console.log(" .:::' `::::.... ....::::' `:::. ");
console.log(" .::' `:::::::| |:::::::' `::. ");
console.log(" .::| |::::::|_ ___ __|::::::| |::. ");
console.log(" `--' |::::::|_()__()_|::::::| `--' ");
console.log(" ::: |::-o::| |::o-::| ::: ");
console.log(" `::. .|::::::| |::::::|. .::' ");
console.log(" `:::. .::\-----' `-----/::. .:::' ");
console.log(" `::::......::::' `::::......::::' ");
console.log(" `::::::::::' `::::::::::' ");
console.log("");
console.log("Semantic search tests have failed. You have been sentenced to MetaJail.");
jobs.forEach((job) => {
if (job.result !== 'success') {
console.log(`${job.name} - ${job.result}`);
}
});
process.exit(1);

View File

@@ -7,8 +7,9 @@ on:
push:
branches:
- master
- metabot-v3-main
- feature-visualizer
- internal-tools
- nm-semantic-search-in-metabase
paths-ignore:
# config files
- ".**"

View File

@@ -56,6 +56,7 @@
com.h2database/h2 ^:antq/exclude ; https://metaboat.slack.com/archives/CKZEMT1MJ/p1727191010259979
{:mvn/version "2.1.214"} ; embedded SQL database
com.ibm.icu/icu4j {:mvn/version "77.1"} ; Used to transliterate Unicode characters into Latin characters
com.knuddels/jtokkit {:mvn/version "1.1.0"} ; Java tokenizer library for OpenAI models, used for token counting in semantic search
com.mchange/c3p0 {:mvn/version "0.10.1"} ; DB connection pool
com.snowplowanalytics/snowplow-java-tracker ^:antq/exclude ; 2+ depend on apache httpclient 5.x, but saml and clj-http depend on 4.x
{:mvn/version "1.0.1"} ; Snowplow analytics

View File

@@ -24,7 +24,11 @@ export const commandPaletteInput = () =>
cy.findByPlaceholderText("Search for anything or jump somewhere…");
export const commandPaletteSearch = (query, viewAll = true) => {
cy.intercept("GET", "/api/search?q=*").as("paletteSearch");
cy.intercept({
method: "GET",
pathname: "/api/search",
query: { q: "*" },
}).as("paletteSearch");
commandPaletteButton().click();
commandPaletteInput().type(query);
cy.wait("@paletteSearch");

View File

@@ -18,7 +18,11 @@ describe("scenarios > metrics > search", () => {
H.restore();
cy.signInAsNormalUser();
cy.intercept("POST", "/api/dataset").as("dataset");
cy.intercept("GET", "/api/search?q=*").as("search");
cy.intercept({
method: "GET",
pathname: "/api/search",
query: { q: "*" },
}).as("search");
});
it("should be able to search for metrics in global search", () => {

View File

@@ -34,9 +34,14 @@ describe("scenarios > question > new", () => {
});
it("new question data picker search should work for both saved questions and database tables", () => {
cy.intercept("GET", "/api/search?q=*", cy.spy().as("searchQuery")).as(
"search",
);
cy.intercept(
{
method: "GET",
pathname: "/api/search",
query: { q: "*" },
},
cy.spy().as("searchQuery"),
).as("search");
H.startNewQuestion();

View File

@@ -80,7 +80,11 @@ const TEST_CREATED_AT_FILTERS = [
describe("scenarios > search", () => {
beforeEach(() => {
H.restore();
cy.intercept("GET", "/api/search?q=*").as("search");
cy.intercept({
method: "GET",
pathname: "/api/search",
query: { q: "*" },
}).as("search");
cy.signInAsAdmin();
});

View File

@@ -20,7 +20,11 @@ const visitEmbeddingWithSearch = (url = "/") => {
describe("scenarios > search", () => {
beforeEach(() => {
H.restore();
cy.intercept("GET", "/api/search?q=*").as("search");
cy.intercept({
method: "GET",
pathname: "/api/search",
query: { q: "*" },
}).as("search");
cy.signInAsAdmin();
});

View File

@@ -0,0 +1,57 @@
const { H } = cy;
const semanticSearch = () => cy.findByTestId("search-engine-setting");
const semanticSearchInput = () => semanticSearch().findByRole("switch");
describe("scenarios > search > semantic search", () => {
beforeEach(() => {
H.restore();
cy.signInAsAdmin();
});
describe("settings", () => {
it("should not show setting to self-hosted customers", () => {
H.activateToken("pro-self-hosted");
cy.visit("/admin/settings");
cy.findByTestId("site-url-setting").should("exist");
cy.findByTestId("search-engine-setting").should("not.exist");
});
it("should upsell hosted starter customers", () => {
H.mockSessionProperty("is-hosted?", true);
H.activateToken("starter"); // TODO still "pro-self-hosted" somehow...
cy.visit("/admin/settings");
cy.findByTestId("search-engine-setting").within(() => {
cy.findByText("Get this with Pro. Try for free.");
});
});
it("should be able to enable/disable the setting", () => {
H.mockSessionProperty("is-hosted?", true);
H.activateToken("pro-cloud"); // still "pro-self-hosted" somehow...
cy.visit("/admin/settings");
cy.intercept("PUT", "/api/setting/search-engine").as("search-engine");
semanticSearchInput().should("not.be.checked");
semanticSearch()
.contains(/Advanced semantic search/)
.click();
cy.wait("@search-engine").should(({ response }) => {
expect(response?.statusCode).to.eq(204);
});
semanticSearchInput().should("be.checked");
semanticSearch()
.contains(/Advanced semantic search/)
.click();
cy.wait("@search-engine").should(({ response }) => {
expect(response?.statusCode).to.eq(204);
});
semanticSearchInput().should("not.be.checked");
});
});
});

View File

@@ -191,7 +191,7 @@ describe("scenarios > visualizations > table", () => {
it("should preserve set widths after reordering (VIZ-439)", () => {
cy.intercept(
"GET",
"/api/search?models=dataset&models=table&table_db_id=*",
"/api/search?search_engine=appdb&models=dataset&models=table&table_db_id=*",
).as("getSearchResults");
H.startNewNativeQuestion({
query: 'select 1 "first_column", 2 "second_column"',

View File

@@ -2,6 +2,7 @@
(:require
[metabase-enterprise.advanced-config.settings :as advanced-config.settings]
[metabase-enterprise.scim.core :as scim]
[metabase-enterprise.semantic-search.core :as semantic-search]
[metabase-enterprise.sso.settings :as sso-settings]
[metabase.driver :as driver]
[metabase.premium-features.core :as premium-features :refer [defenterprise]]
@@ -27,4 +28,7 @@
:enabled (t2/exists? :model/GroupTableAccessPolicy)}
{:name :email-allow-list
:available (premium-features/enable-email-allow-list?)
:enabled (boolean (some? (advanced-config.settings/subscription-allowed-domains)))}])
:enabled (boolean (some? (advanced-config.settings/subscription-allowed-domains)))}
{:name :semantic-search
:available (premium-features/enable-semantic-search?)
:enabled (semantic-search/supported?)}])

View File

@@ -14,5 +14,6 @@
[metabase-enterprise.llm.init]
[metabase-enterprise.metabot-v3.init]
[metabase-enterprise.scim.init]
[metabase-enterprise.semantic-search.init]
[metabase-enterprise.sso.init]
[metabase-enterprise.stale.init]))

View File

@@ -0,0 +1,98 @@
(ns metabase-enterprise.semantic-search.core
"Enterprise implementations of semantic search core functions using defenterprise."
(:require
[metabase-enterprise.semantic-search.db :as semantic.db]
[metabase-enterprise.semantic-search.embedding :as semantic.embedding]
[metabase-enterprise.semantic-search.index-metadata :as semantic.index-metadata]
[metabase-enterprise.semantic-search.pgvector-api :as semantic.pgvector-api]
[metabase-enterprise.semantic-search.settings :as semantic.settings]
[metabase.premium-features.core :refer [defenterprise]]
[next.jdbc :as jdbc]))
(defn- get-pgvector-datasource! []
(or @semantic.db/data-source (semantic.db/init-db!)))
(defn- get-configured-embedding-model []
(semantic.embedding/get-configured-model))
(defn- get-index-metadata []
semantic.index-metadata/default-index-metadata)
(defenterprise supported?
"Enterprise implementation of semantic search engine support check."
:feature :semantic-search
[]
(and
(some? semantic.db/db-url)
(semantic.settings/semantic-search-enabled)))
(defenterprise results
"Enterprise implementation of semantic search results."
:feature :semantic-search
[search-ctx]
(semantic.pgvector-api/query
(get-pgvector-datasource!)
(get-index-metadata)
search-ctx))
(defenterprise update-index!
"Enterprise implementation of semantic index updating."
:feature :semantic-search
[document-reducible]
(semantic.pgvector-api/index-documents!
(get-pgvector-datasource!)
(get-index-metadata)
(vec document-reducible)))
(defenterprise delete-from-index!
"Enterprise implementation of semantic index deletion."
:feature :semantic-search
[model ids]
(semantic.pgvector-api/delete-documents!
(get-pgvector-datasource!)
(get-index-metadata)
model
(vec ids)))
;; TODO: add reindexing/table-swapping logic when index is detected as stale
(defenterprise init!
"Initialize the semantic search table and populate it with initial data."
:feature :semantic-search
[searchable-documents _opts]
(let [pgvector (get-pgvector-datasource!)
index-metadata (get-index-metadata)
embedding-model (get-configured-embedding-model)]
(jdbc/with-transaction [tx pgvector]
(semantic.pgvector-api/init-semantic-search! tx index-metadata embedding-model))
(semantic.pgvector-api/index-documents! pgvector index-metadata (vec searchable-documents))))
(defenterprise reindex!
"Reindex the semantic search index."
:feature :semantic-search
[searchable-documents _opts]
(let [pgvector (get-pgvector-datasource!)
index-metadata (get-index-metadata)
embedding-model (get-configured-embedding-model)]
;; todo force a new index
(jdbc/with-transaction [tx pgvector]
(semantic.pgvector-api/init-semantic-search! tx index-metadata embedding-model))
(semantic.pgvector-api/index-documents! pgvector index-metadata (vec searchable-documents))))
;; TODO: implement
(defenterprise reset-tracking!
"Enterprise implementation of semantic search tracking reset."
:feature :semantic-search
[]
nil)
(comment
(update-index! [{:model "card"
:id "1"
:searchable_text "This is a test card"}
{:model "card"
:id "2"
:searchable_text "This is a test card too"}
{:model "dashboard"
:id "3"
:searchable_text "This is a test dashboard"}])
(delete-from-index! "card" ["1" "2"]))

View File

@@ -0,0 +1,68 @@
(ns metabase-enterprise.semantic-search.db
(:require
[environ.core :refer [env]]
[metabase.connection-pool :as connection-pool]
[metabase.util.log :as log]
[next.jdbc :as jdbc])
(:import
(com.mchange.v2.c3p0 DataSources)))
(set! *warn-on-reflection* true)
(def data-source
"Atom to hold the pooled JDBC data source for the semantic search database."
(atom nil))
(def db-url
"The database URL used to connect to pgvector"
(env :mb-pgvector-db-url))
;; See metabase.app-db.connection-pool-setup for more details on these properties
;; TODO: not sure if we need the MetabaseConnectionCustomizer like in the app DB connection pool seutp
(def ^:private semantic-search-connection-pool-props
"Connection pool properties for semantic search database using c3p0."
{"idleConnectionTestPeriod" 60
"maxIdleTimeExcessConnections" (* 10 60) ; 5 minutes
"maxConnectionAge" (* 30 60) ; 30 minutes
"maxPoolSize" 5 ; Small pool to start
"minPoolSize" 1
"initialPoolSize" 1
"dataSourceName" "metabase-semantic-search-db"})
(defn- build-db-config
"Build database configuration from environment variables"
[]
(if db-url
{:jdbcUrl db-url}
(throw (ex-info "MB_PGVECTOR_DB_URL environment variable is required for semantic search" {}))))
(defn init-db!
"Initialize c3p0 connection pool for semantic search database.
Requires MB_PGVECTOR_DB_URL environment variable."
[]
(let [db-config (build-db-config)
unpooled-ds (jdbc/get-datasource db-config)
pooled-ds (DataSources/pooledDataSource
^javax.sql.DataSource unpooled-ds
(connection-pool/map->properties semantic-search-connection-pool-props))]
(log/info "Initializing semantic search connection pool with properties:" semantic-search-connection-pool-props)
(reset! data-source pooled-ds)))
(defn test-connection!
"Test database connectivity"
[]
(if @data-source
(try
(let [result (jdbc/execute-one! @data-source ["SELECT 1 as test"])]
(log/info "Semantic search database connection successful:" result)
result)
(catch Exception e
(log/error "Semantic search database connection failed:" (.getMessage e))
(throw e)))
(throw (ex-info "Semantic search connection pool is not initialized. Call init-db! first." {}))))
(comment
;; docker-compose.yml
(.doReset #'db-url "jdbc:postgres://localhost:55432/mb_semantic_search?user=postgres&password=postgres")
(init-db!)
(test-connection!))

View File

@@ -0,0 +1,34 @@
version: '3.8'
services:
postgres:
image: pgvector/pgvector:pg16
environment:
POSTGRES_DB: mb_semantic_search
POSTGRES_USER: postgres
POSTGRES_PASSWORD: postgres
ports:
- "55432:5432"
volumes:
- postgres_data:/var/lib/postgresql/data
healthcheck:
test: ["CMD-SHELL", "pg_isready -U postgres"]
interval: 5s
timeout: 5s
retries: 5
ollama:
image: ollama/ollama:latest
ports:
- "11434:11434"
volumes:
- ollama_data:/root/.ollama
healthcheck:
test: ["CMD", "curl", "-f", "http://localhost:11434/api/version"]
interval: 30s
timeout: 10s
retries: 3
volumes:
postgres_data:
ollama_data:

View File

@@ -0,0 +1,250 @@
(ns metabase-enterprise.semantic-search.embedding
(:require
[clj-http.client :as http]
[metabase-enterprise.semantic-search.settings :as semantic-settings]
[metabase.analytics.core :as analytics]
[metabase.util.json :as json]
[metabase.util.log :as log])
(:import
[com.knuddels.jtokkit Encodings]
[com.knuddels.jtokkit.api Encoding EncodingType]))
(set! *warn-on-reflection* true)
(def ^:private ollama-supported-models
"Map of supported Ollama models to their embedding dimensions."
{"mxbai-embed-large" 1024
"snowflake-arctic-embed2:568m" 1024
"nomic-embed-text" 768
"all-minilm" 384})
(def ^:private openai-supported-models
"Map of supported OpenAI models to their embedding dimensions."
{"text-embedding-3-small" 1536
"Snowflake/snowflake-arctic-embed-l-v2.0" 1024
"text-embedding-ada-002" 1536})
;;; Token Counting for OpenAI Models
(def ^:private ^Encoding openai-encoding
"OpenAI tokenizer encoding for cl100k_base (used by text-embedding models)."
(delay (.getEncoding (Encodings/newDefaultEncodingRegistry) EncodingType/CL100K_BASE)))
(defn- count-tokens
"Count the number of tokens in a text string using OpenAI's cl100k_base encoding."
[^String text]
(when text
(let [^Encoding encoding @openai-encoding]
(.size (.encode encoding text)))))
(defn- count-tokens-batch
"Count the total number of tokens across multiple text strings."
[texts]
(reduce + 0 (map count-tokens texts)))
;;; Batching Logic
(defn- create-batches
"Split texts into batches that don't exceed the threshold.
Returns a vector of batches, where each batch is a vector of texts.
- threshold: Maximum allowed measure per batch
- measure: Function to measure each text (e.g., count-tokens)
- texts: Collection of texts to batch"
[threshold measure texts]
(let [step (fn [{:keys [current-batch current-measure batches] :as acc} text]
(let [text-measure (measure text)]
(cond
;; Single text exceeds the limit - skip it with warning
(> text-measure threshold)
(do
(log/warn
(format "Skipping text that exceeds maximum measure per batch: %s"
(str (subs text 0 10) "..."))
{:measure text-measure :threshold threshold})
acc)
;; Adding this text would exceed the limit - start new batch
(> (+ current-measure text-measure) threshold)
(assoc acc
:current-batch [text]
:current-measure text-measure
:batches (conj batches current-batch))
;; Add to current batch
:else
(assoc acc
:current-batch (conj current-batch text)
:current-measure (+ current-measure text-measure)))))
{:keys [batches current-batch]}
(reduce step
{:current-batch [] :current-measure 0 :batches []}
texts)]
(if (seq current-batch)
(conj batches current-batch)
batches)))
(def default-model-for-provider
"Get the default model for a given provider."
{"openai" "text-embedding-3-small"
"ollama" "mxbai-embed-large"})
(def supported-models-for-provider
"Get all supported models for a give provider."
{"openai" openai-supported-models
"ollama" ollama-supported-models})
;;;; Provider SPI
(defn- dispatch-provider [embedding-model & _] (:provider embedding-model))
(defmulti get-embedding
"Returns a single embedding vector for the given text"
{:arglists '([embedding-model text])} dispatch-provider)
(defmulti get-embeddings-batch
"Returns a sequential collection of embedding vectors, in the same order as the input texts."
{:arglists '([embedding-model texts])} dispatch-provider)
(defmulti pull-model
"If a model needs to be downloaded (which is the case for ollama), downloads it."
{:arglists '([embedding-model])} dispatch-provider)
;;;; Ollama impl
(defn- ollama-get-embedding [model-name text]
(try
;; TODO count ollama tokens into :metabase-search/semantic-embedding-tokens?
(log/debug "Generating Ollama embedding for text of length:" (count text))
(-> (http/post "http://localhost:11434/api/embeddings"
{:headers {"Content-Type" "application/json"}
:body (json/encode {:model model-name
:prompt text})})
:body
(json/decode true)
:embedding)
(catch Exception e
(log/error e "Failed to generate Ollama embedding for text of length:" (count text))
(throw e))))
(defn- ollama-get-embeddings-batch [model-name texts]
;; Ollama doesn't have a native batch API, so we fall back to individual calls
;; No special batching needed for Ollama - just process all texts
(log/debug "Generating" (count texts) "Ollama embeddings (using individual calls)")
(mapv #(ollama-get-embedding model-name %) texts))
(defn- ollama-pull-model [model-name]
(try
(log/debug "Pulling embedding model from Ollama...")
(http/post "http://localhost:11434/api/pull"
{:headers {"Content-Type" "application/json"}
:body (json/encode {:model model-name})})
(catch Exception e
(log/error e "Failed to pull embedding model")
(throw e))))
(defmethod get-embedding "ollama" [{:keys [model-name]} text] (ollama-get-embedding model-name text))
(defmethod get-embeddings-batch "ollama" [{:keys [model-name]} text] (ollama-get-embeddings-batch model-name text))
(defmethod pull-model "ollama" [{:keys [model-name]}] (ollama-pull-model model-name))
;;;; OpenAI impl
(defn- openai-get-embeddings-batch [model-name texts]
(let [api-key (semantic-settings/openai-api-key)
endpoint "https://vllm.coredev.metabase.com/v1"]
;;endpoint (str (semantic-settings/openai-api-host) "/v1/embeddings")]
(when-not api-key
(throw (ex-info "OpenAI API key not configured" {:setting "ee-openai-api-key"})))
(try
(log/debug "Generating" (count texts) "OpenAI embeddings in batch")
(let [response (-> (http/post endpoint
{:headers {"Content-Type" "application/json"
"Authorization" (str "Bearer " api-key)}
:body (json/encode {:model model-name
:input texts})})
:body
(json/decode true))]
(analytics/inc! :metabase-search/semantic-embedding-tokens
{:provider "openai", :model model-name}
(->> response :usage :total_tokens))
(->> response
:data
(map :embedding)
vec))
(catch Exception e
(log/error e "Failed to generate OpenAI embeddings for batch of" (count texts) "texts"
"with token count:" (count-tokens-batch texts))
(throw e)))))
(defn- openai-get-embedding [model-name text]
(try
(log/debug "Generating OpenAI embedding for text of length:" (count text))
(first (openai-get-embeddings-batch model-name text))
(catch Exception e
(log/error e "Failed to generate OpenAI embedding for text of length:" (count text))
(throw e))))
(defmethod get-embedding "openai" [{:keys [model-name]} text] (openai-get-embedding model-name text))
(defmethod get-embeddings-batch "openai" [{:keys [model-name]} text] (openai-get-embeddings-batch model-name text))
(defmethod pull-model "openai" [_] (log/debug "OpenAI provider does not require pulling a model"))
;;;; Global embedding model
(defn get-configured-model
"Get the environments default embedding model according to the ee-embedding-provider / ee-embedding-model settings.
Requires the model dimensions to be defined in ollama-supported-models or openai-supported-models (throws if not)."
[]
(let [provider (semantic-settings/ee-embedding-provider)
models (case provider
"ollama" ollama-supported-models
"openai" openai-supported-models
(throw (ex-info (format "Unknown embedding provider: %s" provider) {:provider provider})))
model-name (or (semantic-settings/ee-embedding-model) (default-model-for-provider provider))
vector-dimensions (or (get models model-name)
(throw (ex-info (format "Not a supported model: %s" model-name) {:provider provider, :model-name model-name})))]
{:provider provider
:model-name model-name
:vector-dimensions vector-dimensions}))
(defn process-embeddings-streaming
"Process texts in provider-appropriate batches, calling process-fn for each batch. process-fn will be called with
a map from text to embedding for each batch."
[embedding-model texts process-fn]
(when (seq texts)
(let [{:keys [model-name provider]} embedding-model]
(if (= "openai" provider)
;; For OpenAI, use token-aware batching and stream processing
(let [batches (create-batches (semantic-settings/openai-max-tokens-per-batch) count-tokens texts)]
(run! (fn [batch-texts]
(let [embeddings (openai-get-embeddings-batch model-name batch-texts)
text-embedding-map (zipmap batch-texts embeddings)]
(process-fn text-embedding-map)))
batches))
;; For other providers, process all at once (existing behavior)
(let [embeddings (get-embeddings-batch embedding-model texts)
text-embedding-map (zipmap texts embeddings)]
(process-fn text-embedding-map))))))
(comment
;; This gets loaded after metabase.analytics.prometheus/setup-metrics! so the known labels don't get initialized. If
;; we care about this, we need to ensure this namespace gets loaded before setup-metrics! is called, e.g. by requiring
;; this namespace in the semantic search module's init file. See https://github.com/metabase/metabase/pull/52834
(defmethod analytics/known-labels :metabase-search/semantic-embedding-tokens
[_]
(for [provider (keys supported-models-for-provider)
model (keys (supported-models-for-provider provider))]
{:provider provider :model model})))
(comment
;; Configuration:
;; MB_EE_EMBEDDING_PROVIDER: "openai" or "ollama" (default)
;; MB_EE_EMBEDDING_MODEL: optional override (leave empty for provider defaults)
;; - OpenAI default: "text-embedding-3-small"
;; - Ollama default: "mxbai-embed-large"
;; MB_EE_OPENAI_API_KEY your OpenAI API key
(def embedding-model (get-configured-model))
embedding-model
(pull-model embedding-model)
(get-embedding embedding-model "hello"))

View File

@@ -0,0 +1,455 @@
(ns metabase-enterprise.semantic-search.index
(:require
[clojure.string :as str]
[honey.sql :as sql]
[honey.sql.helpers :as sql.helpers]
[metabase-enterprise.semantic-search.embedding :as embedding]
[metabase.analytics.core :as analytics]
[metabase.models.interface :as mi]
[metabase.search.core :as search]
[metabase.util :as u]
[metabase.util.json :as json]
[metabase.util.log :as log]
[nano-id.core :as nano-id]
[next.jdbc :as jdbc]
[toucan2.core :as t2])
(:import
[java.time LocalDate]
[org.postgresql.util PGobject]))
(set! *warn-on-reflection* true)
(comment
((requiring-resolve 'metabase-enterprise.semantic-search.db/init-db!))
(def db @@(requiring-resolve 'metabase-enterprise.semantic-search.db/data-source)))
(defn sql-format-quoted
"Call [[sql/format]] with {:quoted true}.
Ensures identifiers are quoted since the nano-ids used in temp table names when testing (and various other places)
might contain uppercase chars or hyphens which need to be quoted."
[honey-sql & {:as opts}]
(sql/format honey-sql (merge opts {:quoted true})))
(def ^:dynamic *batch-size*
"The number of documents to process per batch when updating the index."
150)
(defn- index-table-schema
"Schema for the index table."
[vector-dimensions]
;; Generate unique constraint name to avoid index name conflicts between multiple test tables
(let [unique-constraint-name (keyword (str "unique_constraint_" (nano-id/nano-id) "_model_model_id"))]
[[:id :bigint [:primary-key] [:raw "GENERATED BY DEFAULT AS IDENTITY"]]
[:model :text :not-null]
[:model_id :text :not-null]
[:created_at :timestamp-with-time-zone [:default [:raw "CURRENT_TIMESTAMP"]] :not-null]
[:creator_id :int]
[:model_created_at :timestamp-with-time-zone]
[:last_editor_id :int]
[:model_updated_at :timestamp-with-time-zone]
[:archived :boolean [:default false]]
[:verified :boolean]
[:official_collection :boolean]
[:legacy_input :jsonb]
[:embedding [:raw (format "vector(%d)" vector-dimensions)] :not-null]
[:text_search_vector :tsvector :not-null]
[:content :text :not-null]
[:metadata :jsonb]
[[:constraint unique-constraint-name]
[:unique [:composite :model :model_id]]]]))
(defn- format-embedding
"Formats and validates the embedding vector for SQL insertion."
[embedding]
;; Validate all values are numeric as a safe-guard against malformed/malicious embeddings
(doseq [value embedding]
(when-not (number? value)
(throw (ex-info "Embedding contains non-numeric value"
{:invalid-value value
:embedding embedding}))))
(str "'[" (str/join ", " embedding) "]'::vector"))
(defn- doc->db-record
"Convert a document to a database record with a provided embedding."
[embedding-vec {:keys [model id searchable_text created_at creator_id updated_at
last_editor_id archived verified official_collection legacy_input] :as doc}]
{:model model
:model_id id
:creator_id creator_id
:model_created_at created_at
:last_editor_id last_editor_id
:model_updated_at updated_at
:archived archived
:verified verified
:official_collection official_collection
:embedding [:raw (format-embedding embedding-vec)]
:content searchable_text
:text_search_vector (if (:name doc)
[:||
(search/weighted-tsvector "A" (:name doc))
(search/weighted-tsvector "B" (:searchable_text doc ""))]
(search/weighted-tsvector "A" (:searchable_text doc "")))
:legacy_input [:cast (json/encode legacy_input) :jsonb]
:metadata [:cast (json/encode doc) :jsonb]})
(defn- analytics-set-index-size!
"Set the semantic-index-size metric to the number of rows in the index table."
[connectable table-name]
(try
(->> (jdbc/execute-one! connectable
(-> (sql.helpers/select [:%count.* :count])
(sql.helpers/from (keyword table-name))
(sql.helpers/limit 1)
sql-format-quoted))
:count
(analytics/set! :metabase-search/semantic-index-size))
(catch Exception e
(log/warn e "Failed to set :metabase-search/semantic-index-size metric"))))
(defn- batch-update!
[connectable table-name records->sql documents embeddings]
(when (seq documents)
(u/prog1 (transduce (comp (map (fn [[doc embedding]] (doc->db-record embedding doc)))
(partition-all *batch-size*)
(map (fn [db-records]
(jdbc/execute! connectable (records->sql db-records))
;; TODO should this return (or at least log) the number of docs actually
;; updated, not just the number in the batch?
(u/prog1 (->> db-records (map :model) frequencies)
(log/debug "semantic search processed a batch of" (count db-records)
"documents with frequencies" <>)))))
(partial merge-with +)
(map vector documents embeddings))
(log/info "semantic search processed" (count documents) "total documents with frequencies" <>)
(analytics-set-index-size! connectable table-name))))
(defn- batch-delete-ids!
[connectable table-name model ids->sql ids]
(when (seq ids)
(u/prog1 (->> (transduce (comp (partition-all *batch-size*)
(map (fn [ids]
(jdbc/execute! connectable (ids->sql ids))
(u/prog1 (count ids)
(log/debug "semantic search deleted a batch of" <>
"documents with model type" model)))))
+
ids)
(array-map model))
(log/info "semantic search deleted" (get <> model) "total documents with model type" model)
(analytics-set-index-size! connectable table-name))))
(defn- db-records->update-set
[db-records]
(let [update-keys (-> db-records first (dissoc :id :model :model_id) keys)
excluded-kw (fn [column] (keyword (str "excluded." (name column))))]
(zipmap update-keys (map excluded-kw update-keys))))
(defn default-index
"Returns the default index spec for a model."
[embedding-model]
(let [{:keys [model-name provider vector-dimensions]} embedding-model
sanitized-model-name (clojure.string/replace model-name #"[/:.]" "_")
table-name (str "index_table__" provider "__" sanitized-model-name "__" vector-dimensions)]
{:embedding-model embedding-model
:table-name table-name
:version 0}))
(defn upsert-index!
"Inserts or updates documents in the index table. If a document with the same
model + model_id already exists, it will be replaced."
[connectable index documents]
(log/info "Populating semantic search index with" (count documents) "documents")
(when (seq documents)
(let [text->docs (group-by :searchable_text documents)
searchable-texts (map :searchable_text documents)]
(embedding/process-embeddings-streaming
(:embedding-model index)
searchable-texts
(fn [text->embedding]
(let [batch-documents
(mapcat (fn [text]
(if-let [embedding (text->embedding text)]
(map #(assoc % :embedding embedding) (get text->docs text))
(when-let [docs (get text->docs text)]
(log/warn "No embedding found for" (count docs) "documents with searchable text:"
{:searchable_text text
:document_count (count docs)}))))
(keys text->embedding))]
(batch-update!
connectable
(:table-name index)
(fn [db-records]
(-> (sql.helpers/insert-into (keyword (:table-name index)))
(sql.helpers/values db-records)
(sql.helpers/on-conflict :model :model_id)
(sql.helpers/do-update-set (db-records->update-set db-records))
sql-format-quoted))
batch-documents
(map :embedding batch-documents))))))))
(defn- drop-index-table-sql
[{:keys [table-name]}]
(-> (sql.helpers/drop-table :if-exists (keyword table-name))
sql-format-quoted))
(defn drop-index-table!
"Drops the index table for the given embedding model if it exists."
[connectable index]
(jdbc/execute! connectable (drop-index-table-sql index)))
(defn hnsw-index-name
"Returns the name for a HNSW database index for the given semantic search index configuration."
[index]
(str (:table-name index) "__hnsw_idx"))
(defn fts-index-name
"Returns the name for a full-text database index for the given semantic search index configuration."
[index]
(str (:table-name index) "__gin_idx"))
(defn create-index-table-if-not-exists!
"Ensure that the index table exists and is ready to be populated. If
force-reset? is true, drops and recreates the table if it exists."
[connectable index & {:keys [force-reset?] :or {force-reset? false}}]
(try
(let [{:keys [embedding-model table-name]} index
_ (assert (not (re-find #"\s" table-name))
(format "whitespace in the table name (%s) is not currently supported" table-name))
{:keys [vector-dimensions]} embedding-model]
(log/info "Creating index table" table-name)
(jdbc/execute! connectable (sql/format (sql.helpers/create-extension :vector :if-not-exists)))
(when force-reset? (drop-index-table! connectable index))
(jdbc/execute!
connectable
(-> (sql.helpers/create-table (keyword table-name) :if-not-exists)
(sql.helpers/with-columns (index-table-schema vector-dimensions))
sql-format-quoted))
(jdbc/execute!
connectable
(-> (sql.helpers/create-index
[(keyword (hnsw-index-name index)) :if-not-exists]
[(keyword table-name) :using-hnsw [:raw "embedding vector_cosine_ops"]])
sql-format-quoted))
(jdbc/execute!
connectable
(-> (sql.helpers/create-index
[(keyword (fts-index-name index)) :if-not-exists]
[(keyword table-name) :using-gin [:raw "text_search_vector"]])
sql-format-quoted)))
(catch Exception e
(throw (ex-info "Failed to create index table" {} e)))))
(comment
(def embedding-model {:provider "ollama"
:model-name "mxbai-embed-large"
:vector-dimensions 1024})
(def index (default-index embedding-model))
(drop-index-table! db index)
(create-index-table-if-not-exists! db index)
(jdbc/execute! db ["select table_name from INFORMATION_SCHEMA.tables where table_name like 'index_table__%'"]))
(defn- search-filters
"Generate WHERE conditions based on search context filters."
[{:keys [archived? verified models created-at created-by last-edited-at last-edited-by]}]
(let [conditions (filter some?
[(when (some? archived?)
[:= :archived archived?])
(when (some? verified)
[:= :verified verified])
(when (seq models)
[:in :model models])
(when (seq created-by)
[:in :creator_id created-by])
(when (seq last-edited-by)
[:in :last_editor_id last-edited-by])
(when (and created-at (:start created-at) (:end created-at))
[:between :model_created_at
(LocalDate/parse (:start created-at))
(LocalDate/parse (:end created-at))])
(when (and last-edited-at (:start last-edited-at) (:end last-edited-at))
[:between :model_updated_at
(LocalDate/parse (:start last-edited-at))
(LocalDate/parse (:end last-edited-at))])])]
(when (seq conditions)
(into [:and] conditions))))
(defn- keyword-search-query [index search-context]
(let [filters (search-filters search-context)
ts-search-expr (search/to-tsquery-expr (:search-string search-context))
tsv-lang (search/tsv-language)]
{:select [[:id :id]
[:model_id :model_id]
[:model :model]
[:content :content]
[:verified :verified]
[:metadata :metadata]
[[:raw "row_number() OVER (ORDER BY ts_rank_cd(text_search_vector, query) DESC)"] :keyword_rank]]
:from [(keyword (:table-name index))]
;; Using a join allows us to share the query expression between our SELECT and WHERE clauses.
;; This follows the same secure pattern as metabase.search.appdb.specialization.postgres/base-query
:join [[[:raw "to_tsquery('" tsv-lang "', " [:lift ts-search-expr] ")"]
:query] [:= 1 1]]
:where (if (seq filters)
(into [:and [:raw "text_search_vector @@ query"]] [filters])
[:raw "text_search_vector @@ query"])
:order-by [[:keyword_rank :asc]]
:limit 100}))
(defn- semantic-search-query
"Build a semantic search query using vector similarity."
[index embedding search-context]
(let [filters (search-filters search-context)
embedding-literal (format-embedding embedding)
max-cosine-distance 0.7
base-query {:select [[:id :id]
[:model_id :model_id]
[:model :model]
[:content :content]
[:verified :verified]
[:metadata :metadata]
[[:raw (str "embedding <=> " embedding-literal)] :semantic_score]
[[:raw (str "row_number() OVER (ORDER BY embedding <=> " embedding-literal " ASC)")] :semantic_rank]]
:from [(keyword (:table-name index))]
:where [:<= [:raw (str "embedding <=> " embedding-literal)] max-cosine-distance]
:order-by [[:semantic_rank :asc]]
:limit 100}]
(if filters
(update base-query :where #(into [:and] [% filters]))
base-query)))
(defn- hybrid-search-query
"Build a hybrid search query using vector + keyword based searches and reranking with RRF"
[index embedding search-context]
(let [semantic-results (semantic-search-query index embedding search-context)
keyword-results (keyword-search-query index search-context)
k 60
keyword-weight 0.51
semantic-weight 0.49
full-query {:with [[:vector_results semantic-results]
[:text_results keyword-results]]
:select [[[:coalesce :v.id :t.id] :id]
[[:coalesce :v.model_id :t.model_id] :model_id]
[[:coalesce :v.model :t.model] :model]
[[:coalesce :v.content :t.content] :content]
[[:coalesce :v.verified :t.verified] :verified]
[[:coalesce :v.metadata :t.metadata] :metadata]
[[:coalesce :v.semantic_rank -1] :semantic_rank]
[[:coalesce :t.keyword_rank -1] :keyword_rank]
[[:+
[:* [:cast semantic-weight :float]
[:coalesce [:/ 1.0 [:+ k [:. :v :semantic_rank]]] 0]]
[:* [:cast keyword-weight :float]
[:coalesce [:/ 1.0 [:+ k [:. :t :keyword_rank]]] 0]]]
:rrf_rank]]
:from [[:vector_results :v]]
:full-join [[:text_results :t] [:= :v.id :t.id]]
:order-by [[:rrf_rank :desc]]
:limit 100}]
full-query))
(defn- legacy-input-with-score
"Fetches the legacy_input field from a result's metadata and attaches a score based on the
embedding distance."
[row]
(-> (get-in row [:metadata :legacy_input])
(assoc :score (:rrf_rank row 1.0))))
;; TODO: can the query return unqualified keys directly?
(defn- unqualify-keys
"Remove table namespace from namespaced keywords in a result row."
[row]
(into {} (map (fn [[k v]] [(keyword (name k)) v]) row)))
(defn- decode-pgobject
"Decode a PGObject (returned from a jsonb field) into a Clojure map."
[^PGobject obj]
(json/decode (.getValue ^PGobject obj) true))
(defn- decode-metadata
"Decode `row`s `:metadata`."
[row]
(update row :metadata decode-pgobject))
(defn- filter-read-permitted
"Returns only those documents in `docs` whose corresponding t2 instances pass an mi/can-read? check for the bound api user."
[docs]
(let [doc->t2-model (fn [doc] (:model (search/spec (:model doc))))
t2-instances (for [[t2-model docs] (group-by doc->t2-model docs)
;; NOTE an "indexed-entity" (:model/ModelIndexValue) does not have an :id column. For now
;; we are filtering indexed-entities out and they should not appear here.
t2-instance (t2/select t2-model :id [:in (map :id docs)])]
t2-instance)
doc->t2 (comp (u/index-by (juxt :id t2/model) t2-instances)
(fn [doc]
[(:id doc) (doc->t2-model doc)]))]
(filterv (fn [doc] (some-> doc doc->t2 mi/can-read?)) docs)))
(defn query-index
"Query the index for documents similar to the search string."
[db index search-context]
(let [{:keys [embedding-model]} index
search-string (:search-string search-context)]
(when-not (str/blank? search-string)
(let [embedding (embedding/get-embedding embedding-model search-string)
query (hybrid-search-query index embedding search-context)
xform (comp (map unqualify-keys)
(map decode-metadata)
(map legacy-input-with-score))
reducible (jdbc/plan db (sql-format-quoted query))]
(comment
(jdbc/execute! db (sql-format-quoted query)))
(-> (transduce xform conj [] reducible)
filter-read-permitted)))))
(comment
(def embedding-model (embedding/get-configured-model))
(def index (default-index embedding-model))
(def search-ctx {:search-string "pasta"})
(def embed (embedding/get-embedding embedding-model (:search-string search-ctx)))
(keyword-search-query index search-ctx)
(sql-format-quoted (keyword-search-query index search-ctx))
(jdbc/execute! db (sql-format-quoted (keyword-search-query index search-ctx)))
(semantic-search-query index embed search-ctx)
(sql/format (semantic-search-query index embed search-ctx))
(jdbc/execute! db (sql-format-quoted (semantic-search-query index embed search-ctx)))
(hybrid-search-query index embed search-ctx)
(sql-format-quoted (hybrid-search-query index embed search-ctx))
(jdbc/execute! db (sql-format-quoted (hybrid-search-query index embed search-ctx)))
(query-index db index search-ctx))
(defn delete-from-index!
"Deletes documents from the index table based on model and model_ids."
[db index model model-ids]
(batch-delete-ids!
db
(:table-name index)
model
(fn [batch-ids]
(-> (sql.helpers/delete-from (keyword (:table-name index)))
(sql.helpers/where [:and
[:= :model model]
[:in :model_id batch-ids]])
sql-format-quoted))
model-ids))
(comment
(def embedding-model (embedding/get-configured-model))
(def index (default-index embedding-model))
(create-index-table-if-not-exists! index {:force-reset? true})
(upsert-index! db index [{:model "card"
:id "1"
:searchable_text "This is a test card"}])
(delete-from-index! db index "card" ["1"])
(delete-from-index! db index "dashboard" ["13"])
;; no user
(query-index db index {:search-string "Copper knife"})
#_:clj-kondo/ignore
(require '[metabase.test :as mt])
(mt/with-test-user :crowberto
(doall (query-index db index {:search-string "Copper knife"}))))

View File

@@ -0,0 +1,283 @@
(ns metabase-enterprise.semantic-search.index-metadata
"Multi-index management for pgvector semantic search.
Tracks multiple embedding model indexes through metadata tables, supporting model switching
without data loss. Each embedding model (provider/name/dimensions) gets its own index table,
with a control table managing which index is currently active."
(:require
[honey.sql :as sql]
[honey.sql.helpers :as sql.helpers]
[metabase-enterprise.semantic-search.index :as semantic.index]
[metabase.util :as u]
[metabase.util.log :as log]
[next.jdbc :as jdbc]
[next.jdbc.result-set :as jdbc.rs]))
;; metadata and control tables allow multiple indexes to coexist
;; while maintaining a single active pointer.
;; This enables model switching without data loss - old indexes remain available and can still be targeted with index functions.
;; Garbage / pending / switchover policy beyond 'one is active' TBD, right now we want the active index to reflect the global
;; embedding configuration in metabase settings / envars.
(def ^:private schema-0
"The version 0 schema for metadata / control tables"
{:metadata
[[:id :bigint [:primary-key] [:raw "GENERATED BY DEFAULT AS IDENTITY"]]
[:provider :text :not-null]
[:model_name :text :not-null]
[:vector_dimensions :int :not-null]
[:table_name :text :not-null :unique]
[:index_version :int :not-null]
[:index_created_at :timestamp-with-time-zone :not-null]]
:control
[[:id :bigint [:primary-key]] ;; not auto-inc, only one row - still useful to ensure only one row when inserting.
[:version :text :not-null]
[:active_id :int :null]
[:active_updated_at :timestamp-with-time-zone :null]]})
(defn qualify-index
"Qualifies an index-map, returning a new index-map that might have additional disambiguating prefixes applied
to table names.
This ensures tables/constraints created as part of operating a meta index can be isolated within
the same database."
[index index-metadata]
;; if created databases for namespacing in tests we could probably remove this whole idea.
(let [{:keys [index-table-qualifier]} index-metadata
qualify (partial format index-table-qualifier)]
(-> index
(update :table-name qualify))))
(def default-index-metadata
"The default index metadata configuration that will be used for the search engine integration."
{:version "0"
:metadata-table-name "index_metadata"
:control-table-name "index_control"
:index-table-qualifier "%s"})
(defn- create-index-metadata-table-if-not-exists-sql [index-metadata]
(let [{:keys [metadata-table-name]} index-metadata
schema (:metadata schema-0)]
(-> (sql.helpers/create-table (keyword metadata-table-name) :if-not-exists)
(sql.helpers/with-columns schema)
(sql/format :quoted true))))
(defn- create-control-table-if-not-exists-sql [index-metadata]
(let [{:keys [control-table-name]} index-metadata
schema (:control schema-0)]
(-> (sql.helpers/create-table (keyword control-table-name) :if-not-exists)
(sql.helpers/with-columns schema)
(sql/format :quoted true))))
(comment
(create-index-metadata-table-if-not-exists-sql default-index-metadata)
(create-control-table-if-not-exists-sql default-index-metadata))
(defn ensure-control-row-exists!
"Ensures the singleton control row exists in the control table.
Idempotent - uses ON CONFLICT DO NOTHING to avoid duplicate key errors."
[pgvector index-metadata]
(let [{:keys [control-table-name version]} index-metadata]
(jdbc/execute!
pgvector
(-> (sql.helpers/insert-into (keyword control-table-name))
;; Control table uses hardcoded id=0 to enforce a singleton row.
;; This ensures exactly one control row exists, simplifying active index lookups
;; and avoiding race conditions when multiple nodes initialize simultaneously.
(sql.helpers/values [{:id 0
:version version
:active_id nil
:active_updated_at nil}])
(sql.helpers/on-conflict :id)
(sql.helpers/do-nothing)
(sql/format :quoted true)))
nil))
(defn create-tables-if-not-exists!
"Creates the metadata and control tables if they do not already exist"
[pgvector index-metadata]
(let [{:keys [metadata-table-name control-table-name version]} index-metadata]
;; **note** we do not currently deal with version mismatches as there is only one version.
;; metadata
(log/infof "Creating metadata table if not exists %s (%s)" metadata-table-name version)
(jdbc/execute!
pgvector
(create-index-metadata-table-if-not-exists-sql index-metadata))
;; control
(log/infof "Creating control table if not exists %s (%s)" control-table-name version)
(jdbc/execute!
pgvector
(create-control-table-if-not-exists-sql index-metadata))
nil))
(defn drop-tables-if-exists!
"Drops both metadata and control tables if they exist.
Used for cleanup in tests and development - destructive operation."
[pgvector index-metadata]
;; we could move this to test.util as it is currently redundant in prod. I like it being here for locality and later
;; we will probably use it trigger 'start fresh', even on deployed nodes.
(let [{:keys [metadata-table-name
control-table-name]} index-metadata]
(jdbc/execute!
pgvector
(-> (sql.helpers/drop-table :if-exists (keyword metadata-table-name))
(sql/format :quoted true)))
(jdbc/execute!
pgvector
(-> (sql.helpers/drop-table :if-exists (keyword control-table-name))
(sql/format :quoted true)))
nil))
(defn- row->index [{:keys [provider
model_name
vector_dimensions
table_name
index_version]}]
{:embedding-model {:provider provider
:model-name model_name
:vector-dimensions vector_dimensions}
:table-name table_name
:version index_version})
(defn get-active-index-state
"Returns the currently active index configuration, or nil if none active.
Expect a map with:
- :index
The indexes parameters, table-name, model etc, that you can use with functions taking `index` in the semantic search package)
- :metadata-row
The metadata record for this index, as returned by next.jdbc. Keys are unqualified. "
[pgvector index-metadata]
(let [{:keys [metadata-table-name
control-table-name]}
index-metadata
;; Returns nil if no active index is set (active_id is null).
active-row-sql (-> {:select [:m.*]
:from [[(keyword control-table-name) :c]]
:join [[(keyword metadata-table-name) :m] [:= :m.id :c.active_id]]}
(sql/format :quoted true))
active-row (jdbc/execute-one! pgvector active-row-sql {:builder-fn jdbc.rs/as-unqualified-lower-maps})]
(when active-row
{:index (row->index active-row)
:metadata-row active-row})))
(defn- index-table-exists? [pgvector index]
(-> (jdbc/execute-one! pgvector
["SELECT exists (select 1 FROM information_schema.tables WHERE table_name = ?) table_exists"
(:table-name index)]
{:builder-fn jdbc.rs/as-unqualified-lower-maps})
(:table_exists false)))
(defn activate-index!
"Sets the specified index as active by updating the control table. The index-id is the `id` column value for
the indexes metadata row."
[pgvector index-metadata index-id]
(assert index-id)
(assert (nat-int? index-id) (format "expected an integer id (hint do not pass the index map!)"))
(let [{:keys [control-table-name]} index-metadata
;; only one row, no traditional WHERE for now
activation-sql (-> {:update (keyword control-table-name)
:set {:active_id index-id
:active_updated_at [:now]}}
(sql/format :quoted true))]
(jdbc/execute! pgvector activation-sql)
nil))
(defn find-best-index!
"Locates the optimal index to be used for the given embedding model.
Returns information that helps you decide what to do next in order to start using the index.
Returns map with
- :index the index parameters, table name, version etc.
- :metadata-row the unqualified next.jdbc row, nil if it is a new index or if the metadata row has somehow been deleted
- :index-table-exists whether the table named by the `:index` actually exists in the database"
[pgvector index-metadata embedding-model]
(let [{:keys [metadata-table-name
control-table-name]}
index-metadata
{:keys [provider
model-name
vector-dimensions]}
embedding-model
model-rows-sql (-> {:select [:m.* [[:coalesce [:= :c.active_id :m.id] false] :is_active]]
:from [[(keyword control-table-name) :c]
[(keyword metadata-table-name) :m]]
:where [:and
[:= :provider provider]
[:= :model_name model-name]
[:= :vector_dimensions vector-dimensions]]}
(sql/format :quoted true))
model-rows (jdbc/execute! pgvector model-rows-sql {:builder-fn jdbc.rs/as-unqualified-lower-maps})
{[active] true
inactive false}
(u/group-by :is_active #(dissoc % :is_active) model-rows)]
;; Priority order: active matching index > inactive matching index > new
;; This preserves existing data when possible and avoids unnecessary index creation.
(cond
;; already active index matching model, return it.
active
(let [index (row->index active)]
{:index index
:index-table-exists (index-table-exists? pgvector index)
:metadata-row active
:active true})
;; has an inactive index
(seq inactive)
(let [best-index-row (first inactive) ; 'best' is undefined right now, not sure if it matters
best-index (row->index best-index-row)]
{:index best-index
:index-table-exists (index-table-exists? pgvector best-index)
:metadata-row best-index-row
:active false})
;; new or untracked index
:else
(let [index (qualify-index (semantic.index/default-index embedding-model) index-metadata)]
{:index index
:index-table-exists (index-table-exists? pgvector index)
:active false}))))
(defn record-new-index-table!
"Records an index in the metadata table and returns its assigned ID.
The indexes :table-name must be unique, or you will receive a constraint violation."
[pgvector index-metadata index]
(let [{:keys [metadata-table-name]}
index-metadata
{index-table-name :table-name
index-version :version
:keys [embedding-model]}
index
{:keys [provider
model-name
vector-dimensions]}
embedding-model
insert-sql (-> (sql.helpers/insert-into (keyword metadata-table-name))
(sql.helpers/values [{:provider provider
:model_name model-name
:vector_dimensions vector-dimensions
:table_name index-table-name
:index_version index-version
:index_created_at [:now]}])
(sql.helpers/returning :id)
(sql/format :quoted true))
{:keys [id]} (jdbc/execute-one! pgvector insert-sql {:builder-fn jdbc.rs/as-unqualified-lower-maps})]
id))
(comment
(def pgvector ((requiring-resolve 'metabase-enterprise.semantic-search.core/get-pgvector-datasource!)))
(def index-metadata ((requiring-resolve 'metabase-enterprise.semantic-search.core/get-index-metadata)))
(def embedding-model ((requiring-resolve 'metabase-enterprise.semantic-search.core/get-configured-embedding-model)))
(create-tables-if-not-exists! pgvector index-metadata)
(ensure-control-row-exists! pgvector index-metadata)
(drop-tables-if-exists! pgvector index-metadata)
(get-active-index-state pgvector index-metadata)
(find-best-index! pgvector index-metadata embedding-model)
(record-new-index-table! pgvector index-metadata (:index (find-best-index! pgvector index-metadata embedding-model)))
(activate-index! pgvector index-metadata (:id (:metadata-row (find-best-index! pgvector index-metadata embedding-model)))))

View File

@@ -0,0 +1,3 @@
(ns metabase-enterprise.semantic-search.init
(:require
[metabase-enterprise.semantic-search.settings]))

View File

@@ -0,0 +1,111 @@
(ns metabase-enterprise.semantic-search.pgvector-api
"High-level API for pgvector based semantic search over metabase.search documents.
Provides the primary interface for search, indexing, and document management.
Determines which index to target in the database using the index-metadata lib.
Important: The pgvector database must be setup for metadata by calling (init-semantic-search!)
After this, the document management and query functions will work as long as you pass the same index-metadata configuration."
(:require
[metabase-enterprise.semantic-search.index :as semantic.index]
[metabase-enterprise.semantic-search.index-metadata :as semantic.index-metadata]
[metabase.util :as u]
[metabase.util.log :as log]))
#_{:clj-kondo/ignore [unresolved-require]}
(comment
(require '[metabase-enterprise.semantic-search.db :as semantic.db])
(def pgvector (or @semantic.db/data-source (semantic.db/init-db!)))
(def index-metadata semantic.index-metadata/default-index-metadata)
(require '[metabase-enterprise.semantic-search.embedding :as semantic.embedding])
(def embedding-model (semantic.embedding/get-configured-model)))
(defn init-semantic-search!
"Initialises a pgvector database for semantic search if it does not exist and creates an index for the provided
embedding model (if it does not exist).
Returns the index that you can use with semantic.search.index functions to operate on the index.
Designed to be called once at application startup (or in tests)."
[pgvector index-metadata embedding-model]
;; called async on application startup see task/search_index.clj for the job entrypoint.
;; called under cluster lock, should not race across nodes.
;; each node does call this function independently so we should avoid redundant work
;; each node _should_ have the same environmental settings, and therefore the same global model
(let [_ (semantic.index-metadata/create-tables-if-not-exists! pgvector index-metadata)
_ (semantic.index-metadata/ensure-control-row-exists! pgvector index-metadata)
active-index-state (semantic.index-metadata/get-active-index-state pgvector index-metadata)
active-index (:index active-index-state)
active-model (:embedding-model active-index)
;; Model switching: compare configured embedding-model vs currently active model.
;; If different, find/create appropriate index and activate it. This handles
;; environment changes (model config updates) without losing existing indexes.
;; nil active-model (no active index) is treated as model change so that a new index is created and made active
model-changed (not= embedding-model active-model)
model-switching (and active-model model-changed)]
(when model-switching
(log/infof "Configured model does not match active index, switching. Previous active: %s" (u/pprint-to-str active-index)))
(if model-changed
(let [{:keys [index
index-table-exists
metadata-row]}
(semantic.index-metadata/find-best-index! pgvector index-metadata embedding-model)]
;; Metadata might exist without table (deleted manually) or table without metadata
;; (created outside this system). Both cases are handled gracefully.
;; We might delete some of this fancyness later once schema / setup etc solidifies
(when-not index-table-exists
(semantic.index/create-index-table-if-not-exists! pgvector index))
(let [index-id (or (:id metadata-row) (semantic.index-metadata/record-new-index-table! pgvector index-metadata index))]
(semantic.index-metadata/activate-index! pgvector index-metadata index-id))
index)
active-index)))
;; query/index-mgmt require an active index to be established first.
;; init-semantic-search! must be called on startup
(defn- ensure-active-index-state [pgvector index-metadata]
(or (semantic.index-metadata/get-active-index-state pgvector index-metadata)
(throw (ex-info "No active semantic search index found" {:index-metadata index-metadata}))))
(defn query
"Executes a semantic search query against the active index.
Requires init-semantic-search! to have been called first to establish active index, otherwise an exception will be thrown."
[pgvector index-metadata search-context]
(let [{:keys [index]} (ensure-active-index-state pgvector index-metadata)]
(semantic.index/query-index pgvector index search-context)))
(defn index-documents!
"Indexes documents into the active semantic search index.
Documents are upserted - existing documents with same id are replaced.
`documents` is a logical collection, but can be reducible to save memory usage."
[pgvector index-metadata documents]
(let [{:keys [index]} (ensure-active-index-state pgvector index-metadata)]
(semantic.index/upsert-index! pgvector index documents)))
(defn delete-documents!
"Removes documents from the active index by model (search.spec model string) and ids.
`ids` is a logical collection, but can be reducible to save memory usage."
[pgvector index-metadata model ids]
(let [{:keys [index]} (ensure-active-index-state pgvector index-metadata)]
(semantic.index/delete-from-index! pgvector index model ids)))
#_{:clj-kondo/ignore [:unresolved-require :metabase/modules]}
(comment
(init-semantic-search! pgvector index-metadata embedding-model)
(index-documents! pgvector index-metadata [{:model "card"
:id "1"
:searchable_text "This is a test card"}])
(require '[metabase.test :as mt])
(mt/as-admin (query pgvector index-metadata {:search-string "test card"}))
(def index-state (ensure-active-index-state pgvector index-metadata))
(def index (:index index-state))
;; warning: deletes everything
(require 'next.jdbc)
(next.jdbc/execute! pgvector ((requiring-resolve 'honey.sql/format) {:delete-from (keyword (:table-name index))} :quoted true))
(def searchable-documents (vec ((requiring-resolve 'metabase.search.ingestion/searchable-documents))))
(index-documents! pgvector index-metadata searchable-documents)
(mt/as-admin (query pgvector index-metadata {:search-string "sharp objects"})))

View File

@@ -0,0 +1,57 @@
(ns metabase-enterprise.semantic-search.settings
(:require
[metabase-enterprise.llm.settings :as llm-settings]
[metabase.premium-features.core :as premium-features]
[metabase.settings.core :as setting :refer [defsetting]]
[metabase.util.i18n :refer [deferred-tru]]))
(defsetting ee-embedding-provider
(deferred-tru "The embedding provider to use (:openai or :ollama)")
:encryption :no
:visibility :settings-manager
:default "ollama"
:type :string
:export? false
:doc "This feature is experimental.")
(defsetting ee-embedding-model
(deferred-tru "Override the default embedding model for the selected provider")
:encryption :no
:visibility :settings-manager
:default nil
:export? false
:doc "This feature is experimental. Leave empty to use provider defaults.")
(defsetting openai-api-host
(deferred-tru "Override the default OpenAI host with a compitable API")
:encryption :no
:visibility :settings-manager
:default "https://vllm.coredev.metabase.com/v1" ;; "https://api.openai.com"
:export? false
:doc "This feature is experimental.")
(defn openai-api-key
"Get the OpenAI API key from the existing LLM settings."
[]
(llm-settings/ee-openai-api-key))
(defsetting semantic-search-enabled
(deferred-tru "Enable the semantic search engine? Intended as a kill switch for the semantic search feature while dogfooding.")
:visibility :internal
:export? false
:encryption :no
:default true
:getter (fn []
(and (setting/get-value-of-type :boolean :semantic-search-enabled)
(premium-features/enable-semantic-search?)))
:type :boolean)
(defsetting openai-max-tokens-per-batch
(deferred-tru "The maximum number of tokens sent in a single OpenAI embedding API call.")
:type :integer
:default 4000
:encryption :no
:export? false
:visibility :internal
:description "Maximum number of texts per batch for embedding.")

View File

@@ -35,6 +35,7 @@
:query-reference-validation
:sandboxes
:scim
:semantic-search
:serialization
:session-timeout-config
:snippet-collections
@@ -74,6 +75,7 @@
:query_reference_validation true
:sandboxes true
:scim true
:semantic_search true
:serialization true
:session_timeout_config true
:snippet_collections true

View File

@@ -0,0 +1,49 @@
(ns metabase-enterprise.semantic-search.db-test
(:require
[clojure.test :refer :all]
[metabase-enterprise.semantic-search.db :as semantic.db])
(:import
(com.mchange.v2.c3p0 PoolBackedDataSource)))
(set! *warn-on-reflection* true)
(deftest connection-pool-initialization-test
(testing "Connection pool can be initialized and tested"
(when semantic.db/db-url
;; Reset the data source to ensure clean test
(reset! semantic.db/data-source nil)
(testing "Data source is nil before initialization"
(is (nil? @semantic.db/data-source)))
(testing "init-db! creates a pooled data source"
(semantic.db/init-db!)
(is (some? @semantic.db/data-source))
(is (instance? PoolBackedDataSource @semantic.db/data-source)))
(testing "test-connection! works with pooled connection"
(let [result (semantic.db/test-connection!)]
(is (= {:test 1} result))))
(testing "Connection pool properties are configured correctly"
(when (instance? PoolBackedDataSource @semantic.db/data-source)
(let [pool ^PoolBackedDataSource @semantic.db/data-source]
;; Test that pool is properly configured
(is (<= (.getNumConnections pool) 5))
(is (>= (.getNumConnections pool) 0))))))))
(deftest db-url-validation-test
(testing "init-db! throws exception when DB URL is missing"
(with-redefs [semantic.db/db-url nil]
(is (thrown-with-msg?
clojure.lang.ExceptionInfo
#"MB_PGVECTOR_DB_URL environment variable is required"
(semantic.db/init-db!))))))
(deftest test-connection-before-init-test
(testing "test-connection! throws exception when pool not initialized"
(reset! semantic.db/data-source nil)
(is (thrown-with-msg?
clojure.lang.ExceptionInfo
#"Semantic search connection pool is not initialized"
(semantic.db/test-connection!)))))

View File

@@ -0,0 +1,117 @@
(ns metabase-enterprise.semantic-search.embedding-test
(:require
[clojure.test :refer :all]
[metabase-enterprise.semantic-search.embedding :as embedding]
[metabase.test :as mt]))
(set! *warn-on-reflection* true)
(deftest test-get-provider
(testing "get-active-model returns based on setting"
(mt/with-temporary-setting-values [ee-embedding-provider "ollama"]
(is (= {:provider "ollama"
:model-name "mxbai-embed-large"
:vector-dimensions 1024}
(embedding/get-configured-model))))
(mt/with-temporary-setting-values [ee-embedding-provider "openai"]
(is (= {:provider "openai"
:model-name "text-embedding-3-small"
:vector-dimensions 1536}
(embedding/get-configured-model)))))
(testing "get-provider throws on unknown provider"
(mt/with-temporary-setting-values [ee-embedding-provider "unknown"]
(is (thrown-with-msg?
clojure.lang.ExceptionInfo
#"Unknown embedding provider: unknown"
(embedding/get-configured-model))))))
(deftest test-model-dimensions-with-settings
(testing "model-dimensions uses provider defaults when override is nil"
(mt/with-temporary-setting-values [ee-embedding-provider "ollama"
ee-embedding-model nil]
(is (= 1024 (:vector-dimensions (embedding/get-configured-model)))))
(mt/with-temporary-setting-values [ee-embedding-provider "openai"
ee-embedding-model nil]
(is (= 1536 (:vector-dimensions (embedding/get-configured-model))))))
(testing "model-dimensions uses override when specified"
(mt/with-temporary-setting-values [ee-embedding-provider "ollama"
ee-embedding-model "nomic-embed-text"]
(is (= 768 (:vector-dimensions (embedding/get-configured-model)))))
(mt/with-temporary-setting-values [ee-embedding-provider "openai"
ee-embedding-model "text-embedding-3-large"]
(is (= 3072 (:vector-dimensions (embedding/get-configured-model)))))))
(deftest test-default-models
(testing "Provider defaults are used when no override is set"
(is (= "text-embedding-3-small" (#'embedding/default-model-for-provider "openai")))
(is (= "mxbai-embed-large" (#'embedding/default-model-for-provider "ollama")))
(is (nil? (#'embedding/default-model-for-provider "unknown"))))
(testing "get-model uses defaults when override is nil or empty"
(mt/with-temporary-setting-values [ee-embedding-provider "openai"
ee-embedding-model nil]
(is (= "text-embedding-3-small" (:model-name (embedding/get-configured-model)))))
(mt/with-temporary-setting-values [ee-embedding-provider "openai"
ee-embedding-model ""]
(is (= "text-embedding-3-small" (:model-name (embedding/get-configured-model)))))
(mt/with-temporary-setting-values [ee-embedding-provider "ollama"
ee-embedding-model nil]
(is (= "mxbai-embed-large" (:model-name (embedding/get-configured-model))))))
(testing "get-model uses override when specified"
(mt/with-temporary-setting-values [ee-embedding-provider "openai"
ee-embedding-model "text-embedding-3-large"]
(is (= "text-embedding-3-large" (:model-name (embedding/get-configured-model)))))))
(deftest test-openai-provider-validation
(testing "OpenAIProvider throws when API key not configured"
(let [embedding-model {:provider "openai"
:model-name "text-embedding-3-small"
:vector-dimensions 1536}]
(mt/with-temporary-setting-values [ee-openai-api-key nil]
(is (thrown-with-msg?
clojure.lang.ExceptionInfo
#"OpenAI API key not configured"
(embedding/get-embedding embedding-model "test text")))
(is (thrown-with-msg?
clojure.lang.ExceptionInfo
#"OpenAI API key not configured"
(embedding/get-embeddings-batch embedding-model ["test text"])))))))
(deftest test-token-counting
(testing "count-tokens returns reasonable counts for text"
(is (= 2 (#'embedding/count-tokens "Hello world")))
(is (= 9 (#'embedding/count-tokens "This is a longer sentence with more tokens.")))
(is (zero? (#'embedding/count-tokens "")))
(is (nil? (#'embedding/count-tokens nil)))))
(deftest test-batching-logic
(testing "create-batches handles empty input"
(is (empty? (#'embedding/create-batches 10 count [])))
(is (empty? (#'embedding/create-batches 10 count nil))))
(testing "create-batches with single short text"
(let [texts ["Short text"]
batches (#'embedding/create-batches 10 count texts)]
(is (= 1 (count batches)))
(is (= texts (first batches)))))
(testing "create-batches splits texts appropriately with smaller token limit"
;; Use a smaller token limit to make testing more predictable
(let [texts ["This is document 1" "This is document 2" "This is document 3"]
batches (#'embedding/create-batches 10 #'embedding/count-tokens texts)]
(is (= [["This is document 1" "This is document 2"] ["This is document 3"]]
batches))))
(testing "create-batches skips texts that exceed token limit"
(let [texts ["Short" "This is a much longer text that exceeds the limit" "Also short"]
batches (#'embedding/create-batches 5 #'embedding/count-tokens texts)]
;; Should skip the long text and batch the short ones
(is (= [["Short" "Also short"]] batches)))))

View File

@@ -0,0 +1,212 @@
(ns metabase-enterprise.semantic-search.index-metadata-test
(:require
[clojure.test :refer :all]
[metabase-enterprise.semantic-search.index :as semantic.index]
[metabase-enterprise.semantic-search.index-metadata :as semantic.index-metadata]
[metabase-enterprise.semantic-search.test-util :as semantic.tu]
[metabase.util :as u]))
(use-fixtures :once #'semantic.tu/once-fixture)
;; NOTE: isolation tests are absent, in prod there is only one index-metadata
(deftest qualify-index-test
(let [sut semantic.index-metadata/qualify-index]
(testing "default qualifier sanity"
(is (= "%s" (:index-table-qualifier semantic.index-metadata/default-index-metadata))))
(testing "qualify-index applies table qualifier to index map"
(is (= {:table-name "foo_bar_baz"} (sut {:table-name "bar"} {:index-table-qualifier "foo_%s_baz"})))
(is (= {:table-name "bar"} (sut {:table-name "bar"} {:index-table-qualifier "%s"})))
(testing "qualifier applies to derived index names"
(is (= "foo_bar__hnsw_idx" (-> (sut {:table-name "bar"} {:index-table-qualifier "foo_%s"})
semantic.index/hnsw-index-name)))))))
(deftest create-tables-if-not-exists!-test
(let [pgvector semantic.tu/db
index-metadata (semantic.tu/unique-index-metadata)
;; could use open-tables here but I chose to be explicit
sut #(semantic.tu/closeable
(semantic.index-metadata/create-tables-if-not-exists! %1 %2)
(fn [_] (semantic.index-metadata/drop-tables-if-exists! %1 %2)))]
(testing "creates metadata and control tables when they don't exist"
(with-open [_ (sut pgvector index-metadata)]
(is (semantic.tu/table-exists-in-db? (:metadata-table-name index-metadata)))
(is (semantic.tu/table-exists-in-db? (:control-table-name index-metadata)))))
(testing "is idempotent when tables already exist"
(with-open [_ (sut pgvector index-metadata)]
(let [table-names-snap (semantic.tu/get-table-names pgvector)
_ (sut pgvector index-metadata)]
(is (= table-names-snap (semantic.tu/get-table-names pgvector))))))))
(defn- open-tables! [pgvector index-metadata]
(semantic.tu/closeable
(semantic.index-metadata/create-tables-if-not-exists! pgvector index-metadata)
(fn [_] (semantic.tu/cleanup-index-metadata! pgvector index-metadata))))
(deftest drop-tables-if-exists!-test
(let [pgvector semantic.tu/db
index-metadata (semantic.tu/unique-index-metadata)
sut semantic.index-metadata/drop-tables-if-exists!]
(with-open [_ (open-tables! pgvector index-metadata)]
(let [table-names-snap (semantic.tu/get-table-names pgvector)
_ (sut pgvector index-metadata)
now-expected (remove (hash-set (:control-table-name index-metadata)
(:metadata-table-name index-metadata))
table-names-snap)]
(is (= now-expected (semantic.tu/get-table-names pgvector)))))))
(deftest ensure-control-row-exists!-test
(let [pgvector semantic.tu/db
index-metadata (semantic.tu/unique-index-metadata)
sut semantic.index-metadata/ensure-control-row-exists!]
(with-open [_ (open-tables! pgvector index-metadata)]
(testing "creates singleton control row with id=0"
(is (= [] (semantic.tu/get-control-rows pgvector index-metadata)) "should be empty at start")
(sut pgvector index-metadata)
(is (= [{:id 0, :version (:version index-metadata), :active_id nil, :active_updated_at nil}]
(semantic.tu/get-control-rows pgvector index-metadata))))
(testing "is idempotent - doesn't create duplicate rows"
(let [control-snap (semantic.tu/get-control-rows pgvector index-metadata)]
(sut pgvector index-metadata)
(is (= control-snap (semantic.tu/get-control-rows pgvector index-metadata))))))))
(deftest activate-index!-test
(let [pgvector semantic.tu/db
index-metadata (semantic.tu/unique-index-metadata)
un-index (semantic.index/default-index semantic.tu/mock-embedding-model)
index1 (semantic.index-metadata/qualify-index (assoc un-index :table-name "i1") index-metadata)
index2 (semantic.index-metadata/qualify-index (assoc un-index :table-name "i2") index-metadata)
sut semantic.index-metadata/activate-index!]
(with-open [_ (open-tables! pgvector index-metadata)]
(let [index-id1 (semantic.index-metadata/record-new-index-table! pgvector index-metadata index1)
index-id2 (semantic.index-metadata/record-new-index-table! pgvector index-metadata index2)]
(testing "does nothing if no control row"
;; note: I might recommend this to throw if no control row, this behaviour is ok for now
(sut pgvector index-metadata index-id1)
(is (= [] (semantic.tu/get-control-rows pgvector index-metadata))))
(semantic.index-metadata/ensure-control-row-exists! pgvector index-metadata)
(testing "sets specified index as active"
(sut pgvector index-metadata index-id1)
(is (=? [{:active_id index-id1}] (semantic.tu/get-control-rows pgvector index-metadata)))
(testing "switch again"
(sut pgvector index-metadata index-id2)
(is (=? [{:active_id index-id2}] (semantic.tu/get-control-rows pgvector index-metadata)))))))))
(deftest get-active-index-state-test
(let [pgvector semantic.tu/db
index-metadata (semantic.tu/unique-index-metadata)
un-index (semantic.index/default-index semantic.tu/mock-embedding-model)
index (semantic.index-metadata/qualify-index un-index index-metadata)
sut semantic.index-metadata/get-active-index-state]
(with-open [_ (open-tables! pgvector index-metadata)]
(semantic.index-metadata/ensure-control-row-exists! pgvector index-metadata)
(testing "by default there is no active index"
(is (nil? (sut pgvector index-metadata))))
(testing "returns active index configuration when one is set"
(let [index-id (semantic.index-metadata/record-new-index-table! pgvector index-metadata index)]
(semantic.index-metadata/activate-index! pgvector index-metadata index-id)
(is (=? [{:active_id index-id}] (semantic.tu/get-control-rows pgvector index-metadata)))
(is (=? {:index index
:metadata-row (first (semantic.tu/get-metadata-rows pgvector index-metadata))}
(sut pgvector index-metadata))))))))
(defn- default-index [embedding-model index-metadata]
(-> (semantic.index/default-index embedding-model)
(semantic.index-metadata/qualify-index index-metadata)))
(defn- add-index! [pgvector index-metadata embedding-model]
(let [index (default-index embedding-model index-metadata)]
(semantic.index/create-index-table-if-not-exists! pgvector index)
{:index index
:index-id (semantic.index-metadata/record-new-index-table! pgvector index-metadata index)}))
(defn- setup-scenario! [pgvector index-metadata {:keys [active inactive]}]
(semantic.index-metadata/ensure-control-row-exists! pgvector index-metadata)
(when active
(let [{:keys [index-id]} (add-index! pgvector index-metadata active)]
(semantic.index-metadata/activate-index! pgvector index-metadata index-id)))
(run! #(add-index! pgvector index-metadata %) inactive))
(deftest find-best-index!-test
(let [pgvector semantic.tu/db
embedding-model1 semantic.tu/mock-embedding-model
embedding-model2 (assoc semantic.tu/mock-embedding-model
:model-name "mock-harder")
sut semantic.index-metadata/find-best-index!
;; warning: the setup-scenario can currently only set up a happy path
;; other variables include:
;; - table not existing
;; - metadata not existing
;; - duplicate indexes with different table names
;; in practice early on these are not likely to be major concerns, so letting them slide for now
scenarios [{:desc "no metadata"}
{:desc "single active model, happy path"
:active embedding-model1}
{:desc "single active model, happy path, model 2"
:active embedding-model2}
{:desc "active + inactive models"
:active embedding-model2
:inactive [embedding-model1]}
{:desc "multiple inactive models"
:inactive [embedding-model1
embedding-model2]}]
all-models [embedding-model1 embedding-model2]]
(doseq [{:keys [desc active inactive] :as scenario} scenarios
:let [index-metadata (semantic.tu/unique-index-metadata)
index' #(default-index % index-metadata)
sut' #(sut pgvector index-metadata %)]]
(testing (str "scenario: " desc)
(with-open [_ (open-tables! pgvector index-metadata)]
(setup-scenario! pgvector index-metadata scenario)
(let [metadata-rows (semantic.tu/get-metadata-rows pgvector index-metadata)
model-name->row (u/index-by :model_name metadata-rows)
model-row (comp model-name->row :model-name)]
(doseq [model all-models
:let [is-active (= model active)
is-inactive (some #{model} inactive)]]
(cond
is-active
(testing "is already active"
(is (= {:index (index' model)
:index-table-exists true
:metadata-row (model-row model)
:active true}
(sut' model))))
is-inactive
(testing "is inactive"
(is (= {:index (index' model)
:index-table-exists true
:metadata-row (model-row model)
:active false}
(sut' model))))
:else
(testing "no metadata"
(is (= {:index (index' model)
:index-table-exists false
:active false}
(sut' model))))))))))))
(deftest record-new-index-table!-test
(let [pgvector semantic.tu/db
index-metadata (semantic.tu/unique-index-metadata)
embedding-model semantic.tu/mock-embedding-model
index (default-index embedding-model index-metadata)
sut semantic.index-metadata/record-new-index-table!]
(with-open [_ (open-tables! pgvector index-metadata)]
(testing "empty to start"
(is (= [] (semantic.tu/get-metadata-rows pgvector index-metadata))))
(testing "records index metadata and returns assigned ID"
(let [index-id (sut pgvector index-metadata index)]
(is (int? index-id))
(is (=? [{:id index-id
:provider (:provider embedding-model)
:model_name (:model-name embedding-model)
:vector_dimensions (:vector-dimensions embedding-model)
:table_name (:table-name index)
:index_version (:version index)}]
(semantic.tu/get-metadata-rows pgvector index-metadata)))))
(testing "enforces unique table-name constraint"
(is (thrown-with-msg? Exception #"duplicate key" (sut pgvector index-metadata index)))))))

View File

@@ -0,0 +1,234 @@
(ns metabase-enterprise.semantic-search.index-test
(:require
[clojure.test :refer :all]
[honey.sql.helpers :as sql.helpers]
[metabase-enterprise.semantic-search.index :as semantic.index]
[metabase-enterprise.semantic-search.test-util :as semantic.tu]
[metabase.test :as mt]
[next.jdbc :as jdbc]))
(use-fixtures :once #'semantic.tu/once-fixture)
(deftest create-index-table!-test
(mt/with-premium-features #{:semantic-search}
(with-open [index-ref (semantic.tu/open-temp-index!)]
;; open-temp-index-table! creates the temp table, so drop it in order to test create!.
(semantic.index/drop-index-table! semantic.tu/db semantic.tu/mock-index)
(testing "index table is not present before create!"
(is (not (semantic.tu/table-exists-in-db? (:table-name @index-ref))))
(is (not (semantic.tu/table-has-index? (:table-name @index-ref) (semantic.index/hnsw-index-name @index-ref))))
(is (not (semantic.tu/table-has-index? (:table-name @index-ref) (semantic.index/fts-index-name @index-ref)))))
(testing "index table is present after create!"
(semantic.index/create-index-table-if-not-exists! semantic.tu/db semantic.tu/mock-index {:force-reset? false})
(is (semantic.tu/table-exists-in-db? (:table-name @index-ref)))
(is (semantic.tu/table-has-index? (:table-name @index-ref) (semantic.index/hnsw-index-name @index-ref)))
(is (semantic.tu/table-has-index? (:table-name @index-ref) (semantic.index/fts-index-name @index-ref)))))))
(deftest drop-index-table!-test
(mt/with-premium-features #{:semantic-search}
(with-open [index-ref (semantic.tu/open-temp-index!)]
(testing "index table is present before drop!"
(is (semantic.tu/table-exists-in-db? (:table-name @index-ref))))
(testing "index table is not present after drop!"
(semantic.index/drop-index-table! semantic.tu/db semantic.tu/mock-index)
(is (not (semantic.tu/table-exists-in-db? (:table-name @index-ref))))))))
(defn- decode-embedding
"Decode `row`s `:embedding`."
[row]
(update row :embedding #'semantic.index/decode-pgobject))
#_:clj-kondo/ignore
(defn- full-index
"Query the full index table and return all documents with decoded embeddings.
Not used in tests, but useful for debugging."
[]
(->> (jdbc/execute! semantic.tu/db
(-> (sql.helpers/select :model :model_id :content :creator_id :embedding)
(sql.helpers/from (keyword (:table-name semantic.tu/mock-index)))
semantic.index/sql-format-quoted))
(map #'semantic.index/unqualify-keys)
(map decode-embedding)))
(defn- query-embeddings
[{:keys [model model_id]}]
(->> (jdbc/execute! semantic.tu/db
(-> (sql.helpers/select :model :model_id :content :creator_id :embedding)
(sql.helpers/from (keyword (:table-name semantic.tu/mock-index)))
(sql.helpers/where :and
[:= :model model]
[:= :model_id model_id])
semantic.index/sql-format-quoted))
(map #'semantic.index/unqualify-keys)
(mapv decode-embedding)))
(defn- check-index-has-no-mock-card []
(testing "no mock card present"
(is (= []
(query-embeddings {:model "card"
:model_id "123"})))))
(defn- check-index-has-no-mock-dashboard []
(testing "no mock dashboard present"
(is (= []
(query-embeddings {:model "dashboard"
:model_id "456"})))))
(defn- check-index-has-no-mock-docs []
(let [{:keys [table-name]} semantic.tu/mock-index
table-exists-sql "select exists(select * from information_schema.tables where table_name = ?) table_exists"
[{:keys [table_exists]}] (jdbc/execute! semantic.tu/db [table-exists-sql table-name])]
(when table_exists
(check-index-has-no-mock-card)
(check-index-has-no-mock-dashboard))))
(defn- check-index-has-mock-card []
(is (= [{:model "card"
:model_id "123"
:creator_id 1
:content "Dog Training Guide"
:embedding (semantic.tu/get-mock-embedding "Dog Training Guide")}]
(query-embeddings {:model "card"
:model_id "123"}))))
(defn- check-index-has-mock-dashboard []
(is (= [{:model "dashboard"
:model_id "456"
:creator_id 2
:content "Elephant Migration"
:embedding (semantic.tu/get-mock-embedding "Elephant Migration")}]
(query-embeddings {:model "dashboard"
:model_id "456"}))))
(defn- check-index-has-mock-docs []
(check-index-has-mock-card)
(check-index-has-mock-dashboard))
(deftest upsert-index!-test
(mt/with-premium-features #{:semantic-search}
(with-open [_ (semantic.tu/open-temp-index!)]
(check-index-has-no-mock-docs)
(testing "upsert-index! returns nil if you pass it an empty collection"
(is (nil? (semantic.tu/upsert-index! [])))
(check-index-has-no-mock-docs))
(testing "upsert-index! works on a fresh index"
(is (= {"card" 1, "dashboard" 1}
(semantic.tu/upsert-index! semantic.tu/mock-documents)))
(check-index-has-mock-docs))
(testing "upsert-index! works with duplicate documents"
(is (= {"card" 1, "dashboard" 1}
(semantic.tu/upsert-index! semantic.tu/mock-documents)))
(check-index-has-mock-docs)))))
(deftest delete-from-index!-test
(mt/with-premium-features #{:semantic-search}
(with-open [_ (semantic.tu/open-temp-index!)]
(check-index-has-no-mock-docs)
(testing "upsert-index! before delete!"
(is (= {"card" 1, "dashboard" 1}
(semantic.tu/upsert-index! semantic.tu/mock-documents)))
(check-index-has-mock-docs))
(testing "delete-from-index! returns nil if you pass it an empty collection"
(is (nil? (semantic.tu/delete-from-index! "card" [])))
(check-index-has-mock-docs))
(testing "delete-from-index! works for cards"
(is (= {"card" 1}
(semantic.tu/delete-from-index! "card" ["123"])))
(check-index-has-no-mock-card)
(check-index-has-mock-dashboard))
(testing "delete-from-index! works for dashboards"
(is (= {"dashboard" 1}
(semantic.tu/delete-from-index! "dashboard" ["456"])))
(check-index-has-no-mock-docs))
(testing "delete-from-index! doesn't complain if you delete a document that doesn't exist"
(is (= {"card" 1}
(semantic.tu/delete-from-index! "card" ["123"])))
(check-index-has-no-mock-docs)))))
(deftest batch-process-mock-docs!-test
(mt/with-premium-features #{:semantic-search}
(with-open [_ (semantic.tu/open-temp-index!)]
(binding [semantic.index/*batch-size* 1]
(let [extra-ids (->> (range 1337 1347) (map str))
extra-docs (map (fn [id doc]
(assoc doc :id id))
extra-ids
(flatten (repeat semantic.tu/mock-documents)))
mock-docs (into semantic.tu/mock-documents extra-docs)]
(testing "ensure upsert! and delete! work when batch size is exceeded"
(check-index-has-no-mock-docs)
(testing "upsert-index! with batch processing"
(is (= {"card" 6, "dashboard" 6}
(semantic.tu/upsert-index! mock-docs)))
(check-index-has-mock-docs))
(testing "delete-from-index! with batch processing"
(testing "delete just the card"
(is (= {"card" 11}
(semantic.tu/delete-from-index! "card" (into ["123"] extra-ids))))
(check-index-has-no-mock-card)
(check-index-has-mock-dashboard)))
(testing "delete the dashboard"
(is (= {"dashboard" 11}
(semantic.tu/delete-from-index! "dashboard" (into ["456"] extra-ids))))
(check-index-has-no-mock-docs))))))))
(deftest upsert-index-batched-embeddings-pairing-test
(mt/with-premium-features #{:semantic-search}
(with-open [_ (semantic.tu/open-temp-index!)]
(testing "Documents with different searchable texts get associated with their correct embeddings"
(let [test-documents [{:model "card"
:id "1"
:searchable_text "Dog Training Guide"
:creator_id 1
:legacy_input {:model "card" :id "1"}
:metadata {}}
{:model "card"
:id "2"
:searchable_text "Elephant Migration"
:creator_id 2
:legacy_input {:model "card" :id "2"}
:metadata {}}
{:model "card"
:id "3"
:searchable_text "Tiger Conservation"
:creator_id 3
:legacy_input {:model "card" :id "3"}
:metadata {}}]]
;; Each embedding should be processed separately
(mt/with-temporary-setting-values [openai-max-tokens-per-batch 3]
(binding [semantic.index/*batch-size* 1]
(semantic.tu/upsert-index! test-documents)
;; Verify each document has its own correct embedding
(is (= [{:model "card" :model_id "1" :creator_id 1
:content "Dog Training Guide"
:embedding (semantic.tu/get-mock-embedding "Dog Training Guide")}]
(query-embeddings {:model "card" :model_id "1"})))
(is (= [{:model "card" :model_id "2" :creator_id 2
:content "Elephant Migration"
:embedding (semantic.tu/get-mock-embedding "Elephant Migration")}]
(query-embeddings {:model "card" :model_id "2"})))
(is (= [{:model "card" :model_id "3" :creator_id 3
:content "Tiger Conservation"
:embedding (semantic.tu/get-mock-embedding "Tiger Conservation")}]
(query-embeddings {:model "card" :model_id "3"}))))))))))
(deftest prometheus-metrics-test
(mt/with-premium-features #{:semantic-search}
(with-open [_ (semantic.tu/open-temp-index!)]
(mt/with-prometheus-system! [_ system]
(testing "semantic-index-size starts at zero"
(is (= 0.0 (mt/metric-value system :metabase-search/semantic-index-size))))
(testing "semantic-index-size is updated after upsert-index! on empty db"
(is (= {"card" 1, "dashboard" 1}
(semantic.tu/upsert-index! semantic.tu/mock-documents)))
(is (= 2.0 (mt/metric-value system :metabase-search/semantic-index-size))))
(testing "semantic-index-size is updated after delete-from-index!"
(is (= {"card" 1}
(semantic.tu/delete-from-index! "card" ["123"])))
(is (= 1.0 (mt/metric-value system :metabase-search/semantic-index-size))))
(testing "semantic-index-size is updated after upsert-index! on populated db"
(is (= {"card" 1, "dashboard" 1}
(semantic.tu/upsert-index! semantic.tu/mock-documents)))
(is (= 2.0 (mt/metric-value system :metabase-search/semantic-index-size))))))))

View File

@@ -0,0 +1,190 @@
(ns metabase-enterprise.semantic-search.pgvector-api-test
(:require
[clojure.test :refer :all]
[honey.sql :as sql]
[metabase-enterprise.semantic-search.index :as semantic.index]
[metabase-enterprise.semantic-search.index-metadata :as semantic.index-metadata]
[metabase-enterprise.semantic-search.pgvector-api :as semantic.pgvector-api]
[metabase-enterprise.semantic-search.test-util :as semantic.tu]
[metabase.search.ingestion :as search.ingestion]
[metabase.test :as mt]
[metabase.util :as u]
[next.jdbc :as jdbc]))
(use-fixtures :once #'semantic.tu/once-fixture)
;; NOTE: isolation tests are absent, in prod there is only one index-metadata
(deftest init-semantic-search!-test
(let [pgvector semantic.tu/db
index-metadata (semantic.tu/unique-index-metadata)
model1 semantic.tu/mock-embedding-model
model2 (assoc semantic.tu/mock-embedding-model :model-name "embed-harder")
sut* semantic.pgvector-api/init-semantic-search!
cleanup (fn [_] (semantic.tu/cleanup-index-metadata! pgvector index-metadata))
sut #(semantic.tu/closeable (apply sut* %&) cleanup)]
(with-open [index-ref (sut pgvector index-metadata model1)]
(testing "sets up active index for model1"
(let [active-state (semantic.index-metadata/get-active-index-state pgvector index-metadata)]
(is (= @index-ref (:index active-state)))
(is (= model1 (:embedding-model (:index active-state))))
(is (= (:model-name model1) (:model_name (:metadata-row active-state))))
(testing "idempotent"
(let [new-index @(sut pgvector index-metadata model1)]
(is (= @index-ref new-index))
(is (= active-state (semantic.index-metadata/get-active-index-state pgvector index-metadata)))))))
(testing "switch to model1"
(let [new-index @(sut pgvector index-metadata model2)
active-state (semantic.index-metadata/get-active-index-state pgvector index-metadata)]
(is (= model2 (:embedding-model new-index)))
(is (= new-index (:index active-state))))
(testing "model1 index still exists"
(is (=? {:index {:embedding-model model1}
:active false
:index-table-exists true}
(semantic.index-metadata/find-best-index! pgvector index-metadata model1))))
(testing "switch back!"
(let [new-index @(sut pgvector index-metadata model1)
active-state (semantic.index-metadata/get-active-index-state pgvector index-metadata)]
(is (= model1 (:embedding-model new-index)))
(is (= new-index (:index active-state)))
(testing "model2 index still exists"
(is (=? {:index {:embedding-model model2}
:active false
:index-table-exists true}
(semantic.index-metadata/find-best-index! pgvector index-metadata model2))))))))))
(defn- open-semantic-search! [pgvector index-metadata embedding-model]
(semantic.tu/closeable
(semantic.pgvector-api/init-semantic-search! pgvector index-metadata embedding-model)
(fn [_]
(semantic.tu/cleanup-index-metadata! pgvector index-metadata))))
(defn- test-not-initialized [sut & args]
(testing "throws if not initialized"
;; not defined specific exception behaviour here yet, but best not to return empty or something silly like that
(is (thrown? Exception (apply sut args)))))
(defn- spy [f]
(let [calls (atom [])]
{:calls calls
:proxy (fn [& args]
(let [ret (apply f args)]
(swap! calls conj {:args args, :ret ret})
ret))}))
(deftest index-documents!-test
(let [pgvector semantic.tu/db
index-metadata (semantic.tu/unique-index-metadata)
model1 semantic.tu/mock-embedding-model
model2 (assoc semantic.tu/mock-embedding-model :model-name "embedagain")
sut semantic.pgvector-api/index-documents!]
(test-not-initialized sut pgvector index-metadata model1)
(with-open [index-ref (open-semantic-search! pgvector index-metadata model1)]
;; no specific behaviour, only proxies the active index to the index search
(testing "is only a proxy for the active index call"
(let [{:keys [proxy calls]} (spy semantic.index/upsert-index!)
documents semantic.tu/mock-documents]
(with-redefs [semantic.index/upsert-index! proxy]
(testing "check proxies correct args and ret is untouched"
(let [ret (sut pgvector index-metadata documents)]
(is (= [{:args [pgvector @index-ref documents]
:ret ret}]
@calls))))
(testing "check proxies after switch"
(reset! calls [])
(let [new-index (semantic.pgvector-api/init-semantic-search! pgvector index-metadata model2)
ret (sut pgvector index-metadata documents)]
(is (= [{:args [pgvector new-index documents]
:ret ret}]
@calls))))))))))
(deftest delete-documents!-test
(let [pgvector semantic.tu/db
index-metadata (semantic.tu/unique-index-metadata)
dash "dashboard"
card "card"
model1 semantic.tu/mock-embedding-model
model2 (assoc semantic.tu/mock-embedding-model :model-name "embedagain")
sut semantic.pgvector-api/delete-documents!
documents semantic.tu/mock-documents
{card-ids "card", dash-ids "dashboard"} (u/group-by :model :id documents)]
(test-not-initialized sut pgvector index-metadata model1)
(with-open [index-ref (open-semantic-search! pgvector index-metadata model1)]
;; no specific behaviour, only proxies the active index to the index search
(testing "is only a proxy for the active index call"
(semantic.pgvector-api/index-documents! pgvector index-metadata documents)
(let [{:keys [calls proxy]} (spy semantic.index/delete-from-index!)]
(with-redefs [semantic.index/delete-from-index! proxy]
(testing "check proxies correct args and ret is untouched"
(let [ret1 (sut pgvector index-metadata dash dash-ids)
ret2 (sut pgvector index-metadata card card-ids)
ret3 (sut pgvector index-metadata card dash-ids)
ret4 (sut pgvector index-metadata card [])]
(is (= [{:args [pgvector @index-ref dash dash-ids]
:ret ret1}
{:args [pgvector @index-ref card card-ids]
:ret ret2}
{:args [pgvector @index-ref card dash-ids]
:ret ret3}
{:args [pgvector @index-ref card []]
:ret ret4}]
@calls))))
(testing "check proxies after switch"
(reset! calls [])
(let [new-index (semantic.pgvector-api/init-semantic-search! pgvector index-metadata model2)
ret (sut pgvector index-metadata dash dash-ids)]
(is (= [{:args [pgvector new-index dash dash-ids]
:ret ret}]
@calls))))))))))
(deftest query-test
(semantic.tu/with-indexable-documents!
(let [pgvector semantic.tu/db
index-metadata (semantic.tu/unique-index-metadata)
model1 semantic.tu/mock-embedding-model
model2 (assoc semantic.tu/mock-embedding-model :model-name "judge-embedd")
sut* semantic.pgvector-api/query
sut #(mt/as-admin (apply sut* %&)) ; see notes below about perms
search-string "insect patterns" ; specifics of search will be handled under index tests
_ (assert semantic.tu/mock-embeddings "search string should have test embedding")
search {:search-string search-string}]
(test-not-initialized sut pgvector index-metadata search)
(with-open [index-ref (open-semantic-search! pgvector index-metadata model1)]
(semantic.pgvector-api/index-documents! pgvector index-metadata (vec (search.ingestion/searchable-documents)))
(testing "search results are the same as direct index query"
(let [index-results (mt/as-admin (semantic.index/query-index pgvector @index-ref search))
api-results (sut pgvector index-metadata search)]
(testing "sanity check the test is setup correctly" (is (seq api-results)))
(is (= api-results index-results))))
;; it would be better to test permissions at a higher level than we currently do, but for now to save time...
(testing "assumption: permissions are applied at the index query level, check the expected fn is called"
(let [query-index semantic.index/query-index
called (atom false)]
(with-redefs [semantic.index/query-index (fn [& args] (reset! called true) (apply query-index args))]
(sut pgvector index-metadata search)
(is @called))))
(testing "same results after reinit"
(let [current-results (sut pgvector index-metadata search)]
(semantic.pgvector-api/init-semantic-search! pgvector index-metadata model1)
(is (= current-results (sut pgvector index-metadata search)))))
(testing "switching to an empty index"
(let [current-results (sut pgvector index-metadata search)]
(semantic.pgvector-api/init-semantic-search! pgvector index-metadata model2)
;; sanity
(is (not= @index-ref (:index (semantic.index-metadata/get-active-index-state pgvector index-metadata))))
(testing "empty to start but works"
(is (= [] (sut pgvector index-metadata search))))
(testing "after indexing returns again"
(semantic.pgvector-api/index-documents! pgvector index-metadata (vec (search.ingestion/searchable-documents)))
;; todo test the results with different mock embeddings to show the embeddings
;; can actually change
(is (seq (sut pgvector index-metadata search))))
(testing "querying previous index directly still works"
(is (= current-results (mt/as-admin (semantic.index/query-index pgvector @index-ref search)))))))
(testing "throws exception when no active index exists"
;; corrupt the control table
(jdbc/execute! pgvector (sql/format {:delete-from (keyword (:control-table-name index-metadata))}
:quoted true))
(is (thrown-with-msg? Exception #"No active semantic search index" (sut pgvector index-metadata search))))))))

View File

@@ -0,0 +1,189 @@
(ns metabase-enterprise.semantic-search.query-test
(:require
[clojure.test :refer :all]
[metabase-enterprise.semantic-search.db :as semantic.db]
[metabase-enterprise.semantic-search.index :as semantic.index]
[metabase-enterprise.semantic-search.test-util :as semantic.tu]
[metabase.permissions.models.data-permissions :as data-perms]
[metabase.permissions.models.permissions :as perms]
[metabase.permissions.models.permissions-group :as perms-group]
[metabase.test :as mt]
[toucan2.core :as t2]))
(use-fixtures :once #'semantic.tu/once-fixture)
(deftest database-initialised-test
(is (some? @semantic.db/data-source))
(is (= {:test 1} (semantic.db/test-connection!))))
(deftest basic-query-test
(testing "Simple queries with no filters"
(mt/with-premium-features #{:semantic-search}
(mt/as-admin
(semantic.tu/with-index!
(testing "Dog-related query finds dog content"
(let [results (semantic.tu/query-index {:search-string "puppy"})]
(is (= "Dog Training Guide" (-> results first :name)))))
(testing "Bird-related query finds bird content"
(let [results (semantic.tu/query-index {:search-string "avian"})]
(is (= "Bird Watching Tips" (-> results first :name))))))))))
(deftest model-filtering-test
(testing "Filter results by model type"
(mt/with-premium-features #{:semantic-search}
(mt/as-admin
(semantic.tu/with-index!
(testing "Filter for cards only"
(let [card-results (semantic.tu/query-index {:search-string "avian" :models ["card"]})
filtered-results (semantic.tu/filter-for-mock-embeddings card-results)]
(is (pos? (count filtered-results)))
(is (every? #(= "card" (:model %)) card-results))))
(testing "Filter for dashboards only"
(let [dashboard-results (semantic.tu/query-index {:search-string "marine mammal" :models ["dashboard"]})
filtered-results (semantic.tu/filter-for-mock-embeddings dashboard-results)]
(is (pos? (count filtered-results)))
(is (every? #(= "dashboard" (:model %)) dashboard-results))))
(testing "Filter for multiple model types"
(let [mixed-results (semantic.tu/query-index {:search-string "predator" :models ["card" "dashboard"]})
filtered-results (semantic.tu/filter-for-mock-embeddings mixed-results)]
(is (pos? (count filtered-results)))
(is (every? #(contains? #{"card" "dashboard"} (:model %)) mixed-results)))))))))
(deftest archived-filtering-test
(testing "Filter results by archived status"
(mt/with-premium-features #{:semantic-search}
(mt/as-admin
(semantic.tu/with-index!
(testing "Include only non-archived items"
(let [active-results (semantic.tu/query-index {:search-string "feline" :archived? false})]
(is (every? #(not (:archived %)) active-results))))
(testing "Include only archived items"
(let [archived-results (semantic.tu/query-index {:search-string "feline" :archived? true})
filtered-results (semantic.tu/filter-for-mock-embeddings archived-results)]
(is (pos? (count filtered-results)))
(is (every? :archived archived-results))))
(testing "Include all items regardless of archived status"
(let [all-results (semantic.tu/query-index {:search-string "feline"})
filtered-results (semantic.tu/filter-for-mock-embeddings all-results)]
(is (pos? (count filtered-results))))))))))
(deftest creator-filtering-test
(testing "Filter results by creator"
(mt/with-premium-features #{:semantic-search}
(mt/as-admin
(semantic.tu/with-index!
(testing "Filter by single creator"
(let [crowberto-results (semantic.tu/query-index {:search-string "aquatic" :created-by [(mt/user->id :crowberto)]})
filtered-results (semantic.tu/filter-for-mock-embeddings crowberto-results)]
(is (pos? (count filtered-results)))
(is (every? #(= (mt/user->id :crowberto) (:creator_id %)) crowberto-results))))
(testing "Filter by multiple creators"
(let [multi-creator-results (semantic.tu/query-index {:search-string "endangered species"
:created-by [(mt/user->id :crowberto) (mt/user->id :rasta)]})]
(is (pos? (count multi-creator-results)))
(is (every? #(contains? #{(mt/user->id :crowberto) (mt/user->id :rasta)} (:creator_id %)) multi-creator-results)))))))))
(deftest verified-filtering-test
(testing "Filter results by verified status"
(mt/with-premium-features #{:semantic-search}
(mt/as-admin
(semantic.tu/with-index!
(testing "Include only verified items"
(let [verified-results (semantic.tu/query-index {:search-string "puppy" :verified true})
filtered-results (semantic.tu/filter-for-mock-embeddings verified-results)]
(is (pos? (count filtered-results)))
(is (every? :verified verified-results))))
(testing "Include all items regardless of verified status when filter not specified"
(let [all-results (semantic.tu/query-index {:search-string "puppy"})
filtered-results (semantic.tu/filter-for-mock-embeddings all-results)]
(is (pos? (count filtered-results))))))))))
(deftest complex-filtering-test
(testing "Complex filters with multiple criteria"
(mt/with-premium-features #{:semantic-search}
(mt/as-admin
(semantic.tu/with-index!
(testing "Non-archived cards by specific creator"
(let [results (semantic.tu/query-index {:search-string "equine"
:models ["card"]
:archived? false
:created-by [(mt/user->id :rasta)]})
filtered-results (semantic.tu/filter-for-mock-embeddings results)]
(is (pos? (count filtered-results)))
(is (every? #(and (= "card" (:model %))
(not (:archived %))
(= (mt/user->id :rasta) (:creator_id %))) results))))
(testing "Archived dashboards by multiple creators"
(let [results (semantic.tu/query-index {:search-string "Antarctic wildlife"
:models ["dashboard"]
:archived? true
:created-by [(mt/user->id :crowberto) (mt/user->id :rasta)]})
filtered-results (semantic.tu/filter-for-mock-embeddings results)]
(is (pos? (count filtered-results)))
(is (every? #(and (= "dashboard" (:model %))
(:archived %)
(contains? #{(mt/user->id :crowberto) (mt/user->id :rasta)} (:creator_id %))) results))))
(testing "Verified items with additional filters"
(let [results (semantic.tu/query-index {:search-string "marine mammal"
:models ["dashboard"]
:verified true
:archived? false})]
(is (every? #(and (= "dashboard" (:model %))
(:verified %)
(not (:archived %))) results)))))))))
(deftest permissions-test
(mt/with-premium-features #{:semantic-search}
(semantic.tu/with-index!
(let [monsters-table (t2/select-one-pk :model/Table :name "Monsters Table")
q (fn [model s] (map :name (semantic.tu/query-index {:search-string s, :models [model]})))
all-users (perms-group/all-users)
unrestrict-table (fn [table-id]
(data-perms/set-table-permission! all-users table-id :perms/create-queries :query-builder)
(data-perms/set-table-permission! all-users table-id :perms/view-data :unrestricted))
restrict-table (fn [table-id]
(data-perms/set-table-permission! all-users table-id :perms/create-queries :no)
(data-perms/set-table-permission! all-users table-id :perms/view-data :blocked))]
(perms/revoke-collection-permissions! (perms-group/all-users) (t2/select-one :model/Collection :name "Cryptozoology"))
(restrict-table monsters-table)
(testing "admin"
(mt/with-test-user :crowberto
(testing "collection permissions"
(is (some #{"Loch Ness Stuff"} (q "dashboard" "prehistoric monsters")))
(is (some #{"Bigfoot Sightings"} (q "card" "spooky video evidence"))))
(testing "data permissions"
(is (some #{"Monsters Table"} (q "table" "monster facts"))))))
(testing "all-users"
(mt/with-test-user :rasta
(testing "collection permissions"
(is (not-any? #{"Loch Ness Stuff"} (q "dashboard" "prehistoric monsters")))
(is (not-any? #{"Bigfoot Sightings"} (q "card" "spooky video evidence"))))
(testing "data permissions"
(is (not-any? #{"Monsters Table"} (q "table" "monster facts"))))
(testing "give data permissions"
(unrestrict-table monsters-table)
(data-perms/disable-perms-cache
(is (some #{"Monsters Table"} (q "table" "monster facts")))))))))))
(deftest basic-batched-query-test
(testing "Small-batch reindexing"
(mt/with-premium-features #{:semantic-search}
(mt/as-admin
(binding [semantic.index/*batch-size* 1]
(semantic.tu/with-index!
(testing "Horse-related query finds horse content"
(let [results (semantic.tu/query-index {:search-string "equine"})]
(is (= "Horse Racing Analysis" (-> results first :name)))))
(testing "Tiger-related query finds tiger content"
(let [results (semantic.tu/query-index {:search-string "endangered species"})]
(is (= "Tiger Conservation" (-> results first :name)))))))))))

View File

@@ -0,0 +1,269 @@
(ns metabase-enterprise.semantic-search.test-util
(:require
[honey.sql :as sql]
[metabase-enterprise.semantic-search.core]
[metabase-enterprise.semantic-search.db :as semantic.db]
[metabase-enterprise.semantic-search.embedding :as semantic.embedding]
[metabase-enterprise.semantic-search.index :as semantic.index]
[metabase-enterprise.semantic-search.index-metadata :as semantic.index-metadata]
[metabase.search.core :as search.core]
[metabase.search.ingestion :as search.ingestion]
[metabase.test :as mt]
[metabase.util :as u]
[metabase.util.log :as log]
[next.jdbc :as jdbc]
[next.jdbc.protocols :as jdbc.protocols]
[next.jdbc.result-set :as jdbc.rs])
(:import (clojure.lang IDeref)
(java.io Closeable)))
(set! *warn-on-reflection* true)
(def ^:private init-delay
(delay
(when-not @semantic.db/data-source
(semantic.db/init-db!))))
(defn once-fixture [f]
(when semantic.db/db-url
@init-delay
(f)))
(def db
"Proxies the semantic.db/data-source, avoids the deref and prettifies a little"
;; proxy because semantic.db/data-source is not initialised until the fixture runs
(reify jdbc.protocols/Sourceable
(get-datasource [_] (jdbc.protocols/get-datasource @semantic.db/data-source))))
(comment
(jdbc/execute! db ["select 1"]))
(def mock-embeddings
"Static mapping from strings to (made-up) 4-dimensional embedding vectors for testing. Each pair of strings represents a
document and a search query that should be most semantically similar to it, according to the embeddings."
{"Dog Training Guide" [0.12 -0.34 0.56 -0.78]
"puppy" [0.13 -0.33 0.57 -0.77]
"Bird Watching Tips" [0.23 0.45 -0.67 0.89]
"avian" [0.24 0.46 -0.66 0.88]
"Cat Behavior Study" [0.11 -0.22 0.33 -0.44]
"feline" [0.12 -0.21 0.34 -0.43]
"Horse Racing Analysis" [0.55 0.66 -0.77 0.88]
"equine" [0.56 0.67 -0.76 0.87]
"Fish Tank Setup" [0.10 0.20 -0.30 0.40]
"aquatic" [0.11 0.21 -0.29 0.39]
"Elephant Migration" [0.15 -0.25 0.35 -0.45]
"pachyderm" [0.16 -0.24 0.36 -0.44]
"Lion Pride Dynamics" [0.31 0.42 -0.53 0.64]
"predator" [0.32 0.43 -0.52 0.63]
"Penguin Colony Study" [0.75 -0.86 0.97 -0.18]
"Antarctic wildlife" [0.76 -0.85 0.96 -0.17]
"Whale Communication" [0.29 0.38 -0.47 0.56]
"marine mammal" [0.30 0.39 -0.46 0.55]
"Tiger Conservation" [0.65 -0.74 0.83 -0.92]
"endangered species" [0.66 -0.73 0.84 -0.91]
"Butterfly Migration" [0.17 0.28 -0.39 0.50]
"insect patterns" [0.18 0.29 -0.38 0.49]})
(defn get-mock-embedding
"Lookup the embedding for `text` in [[mock-embeddings]]."
[text]
(get mock-embeddings text [0.01 0.02 0.03 0.04]))
(defn get-mock-embeddings-batch
"Lookup embeddings for multiple texts in [[mock-embeddings]]."
[texts]
(mapv get-mock-embedding texts))
;;;; mock provider
(def mock-embedding-model
{:provider "mock"
:model-name "model"
:vector-dimensions 4})
(def mock-index-metadata
"An index metadata to qualify and isolate mock indexes"
{:version "0"
:metadata-table-name "mock_index_metadata"
:control-table-name "mock_index_control"
:index-table-qualifier "mock_%s"})
(defn unique-index-metadata
[]
(let [nano-id (u/generate-nano-id)
fmt (str "mock_%s_" nano-id "_")]
{:version "0"
:metadata-table-name (format fmt "index_metadata")
:control-table-name (format fmt "index_control")
:index-table-qualifier fmt}))
(def mock-index
"A mock index for testing low level indexing functions.
Coincides with what the index-metadata system would create for the mock-embedding-model."
(-> (semantic.index/default-index mock-embedding-model)
(semantic.index-metadata/qualify-index mock-index-metadata)))
(defmethod semantic.embedding/get-embedding "mock" [_ text] (get-mock-embedding text))
(defmethod semantic.embedding/get-embeddings-batch "mock" [_ texts] (get-mock-embeddings-batch texts))
(defmethod semantic.embedding/pull-model "mock" [_])
(defn query-index [search-context]
(semantic.index/query-index db mock-index search-context))
(defn upsert-index! [documents]
(semantic.index/upsert-index! db mock-index documents))
(defn delete-from-index! [model ids]
(semantic.index/delete-from-index! db mock-index model ids))
(def mock-documents
[{:model "card"
:id "123"
:searchable_text "Dog Training Guide"
:created_at #t "2025-01-01T12:00:00Z"
:creator_id 1
:archived false
:legacy_input {:model "card" :id "123"}
:metadata {:title "Dog Training Guide" :description "How to teach an old dog new tricks"}}
{:model "dashboard"
:id "456"
:searchable_text "Elephant Migration"
:created_at #t "2025-02-01T12:00:00Z"
:creator_id 2
:archived true
:legacy_input {:model "dashboard" :id "456"}
:metadata {:title "Elephant Migration" :description "How do elephants deal with schema upgrades?"}}])
(defn filter-for-mock-embeddings
"Filter results to only include items whose names are keys in mock-embeddings map."
[results]
(filter #(contains? mock-embeddings (:name %)) results))
(defn closeable [o close-fn]
(reify
IDeref
(deref [_] o)
Closeable
(close [_] (close-fn o))))
(defn open-temp-index! []
(closeable
(do (semantic.index/create-index-table-if-not-exists! db mock-index {:force-reset? true})
mock-index)
(fn cleanup-temp-index-table! [{:keys [table-name]}]
(try
(semantic.index/drop-index-table! db mock-index)
(catch Exception e
(log/error e "Warning: failed to clean up test table" table-name))))))
(defmacro with-indexable-documents!
"Add a collection of test documents to that can be indexed to the appdb."
[& body]
`(mt/dataset ~(symbol "test-data")
(mt/with-temp [:model/Collection {col1# :id} {:name "Wildlife Collection" :archived false}
:model/Collection {col2# :id} {:name "Archived Animals" :archived true}
:model/Collection {col3# :id} {:name "Cryptozoology", :archived false}
:model/Card {card1# :id} {:name "Dog Training Guide" :collection_id col1# :creator_id (mt/user->id :crowberto) :archived false}
:model/Card {} {:name "Bird Watching Tips" :collection_id col1# :creator_id (mt/user->id :rasta) :archived false}
:model/Card {} {:name "Cat Behavior Study" :collection_id col2# :creator_id (mt/user->id :crowberto) :archived true}
:model/Card {} {:name "Horse Racing Analysis" :collection_id col1# :creator_id (mt/user->id :rasta) :archived false}
:model/Card {} {:name "Fish Tank Setup" :collection_id col2# :creator_id (mt/user->id :crowberto) :archived true}
:model/Card {} {:name "Bigfoot Sightings" :collection_id col3# :creator_id (mt/user->id :crowberto), :archived false}
:model/ModerationReview {} {:moderated_item_type "card"
:moderated_item_id card1#
:moderator_id (mt/user->id :crowberto)
:status "verified"
:most_recent true}
:model/Dashboard {} {:name "Elephant Migration" :collection_id col1# :creator_id (mt/user->id :rasta) :archived false}
:model/Dashboard {} {:name "Lion Pride Dynamics" :collection_id col1# :creator_id (mt/user->id :crowberto) :archived false}
:model/Dashboard {} {:name "Penguin Colony Study" :collection_id col2# :creator_id (mt/user->id :rasta) :archived true}
:model/Dashboard {} {:name "Whale Communication" :collection_id col1# :creator_id (mt/user->id :crowberto) :archived false}
:model/Dashboard {} {:name "Tiger Conservation" :collection_id col2# :creator_id (mt/user->id :rasta) :archived true}
:model/Dashboard {} {:name "Loch Ness Stuff" :collection_id col3# :creator_id (mt/user->id :crowberto), :archived false}
:model/Database {db-id# :id} {:name "Animal Database"}
:model/Table {} {:name "Species Table", :db_id db-id#}
:model/Table {} {:name "Monsters Table", :db_id db-id#, :active true}]
~@body)))
(defmacro with-index!
"Ensure a clean, small index for testing populated with a few collections, cards, and dashboards."
[& body]
`(with-indexable-documents!
(with-redefs [semantic.embedding/get-configured-model (fn [] mock-embedding-model)
semantic.index-metadata/default-index-metadata mock-index-metadata]
(with-open [_# (open-temp-index!)]
(binding [search.ingestion/*force-sync* true]
(search.core/reindex! :search.engine/semantic {:force-reset true})
~@body)))))
(defn table-exists-in-db?
"Check if a table actually exists in the database"
[table-name]
(when table-name
(try
(let [result (jdbc/execute! db
["SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name = ?)"
(name table-name)])]
(-> result first vals first))
(catch Exception _ false))))
(defn table-has-index?
[table-name index-name]
(when table-name
(try
(let [result (jdbc/execute! db
["SELECT EXISTS (SELECT 1 FROM pg_indexes WHERE tablename = ? AND indexname = ?)"
(name table-name)
(name index-name)])]
(-> result first vals first))
(catch Exception _ false))))
(defn get-metadata-rows [pgvector index-metadata]
(jdbc/execute! pgvector
(-> {:select [:*]
:from [(keyword (:metadata-table-name index-metadata))]}
(sql/format :quoted true))
{:builder-fn jdbc.rs/as-unqualified-lower-maps}))
(defn get-control-rows [pgvector index-metadata]
(jdbc/execute! pgvector
(-> {:select [:*]
:from [(keyword (:control-table-name index-metadata))]}
(sql/format :quoted true))
{:builder-fn jdbc.rs/as-unqualified-lower-maps}))
(defn cleanup-index-metadata!
"A-la-carte drop everything related to the index-metdata root (including index tables themselves)"
[pgvector index-metadata]
(doseq [{:keys [table_name]}
(when (table-exists-in-db? (:metadata-table-name index-metadata))
(get-metadata-rows pgvector index-metadata))]
(semantic.index/drop-index-table! pgvector {:table-name table_name}))
(semantic.index-metadata/drop-tables-if-exists! pgvector index-metadata))
(defn get-table-names [pgvector]
(->> ["select table_name from information_schema.tables"]
(jdbc/execute! pgvector)
(mapv :tables/table_name)))

View File

@@ -134,6 +134,7 @@ export const createMockTokenFeatures = (
ai_sql_generation: false,
database_routing: false,
development_mode: false,
semantic_search: false,
...opts,
});
@@ -329,6 +330,7 @@ export const createMockSettings = (
"scim-enabled": false,
"scim-base-url": "http://localhost:3000/api/ee/scim/v2/",
"snowplow-url": "",
"search-engine": "appdb",
"search-typeahead-enabled": true,
"setup-token": null,
"session-cookies": null,

View File

@@ -263,6 +263,7 @@ export const tokenFeatures = [
"development_mode",
"etl_connections",
"etl_connections_pg",
"semantic_search",
] as const;
export type TokenFeature = (typeof tokenFeatures)[number];
@@ -552,6 +553,8 @@ export type ColorSettings = Record<string, string>;
export type IllustrationSettingValue = "default" | "none" | "custom";
export type TimeoutValue = { amount: number; unit: string };
export type SearchEngineSettingValue = "semantic" | "appdb" | "in-place";
export type DatabaseReplicationConnections = Record<
DatabaseId,
{ connection_id: string }
@@ -573,6 +576,7 @@ export interface EnterpriseSettings extends Settings {
"ee-openai-api-key"?: string;
"ee-openai-model"?: string;
"session-timeout": TimeoutValue | null;
"search-engine": SearchEngineSettingValue | null;
"scim-enabled"?: boolean | null;
"scim-base-url"?: string;
"send-new-sso-user-admin-email?"?: boolean;

View File

@@ -14,6 +14,7 @@ import { AdminSettingInput } from "../widgets/AdminSettingInput";
import { AnonymousTrackingInput } from "../widgets/AnonymousTrackingInput";
import { CustomHomepageDashboardSetting } from "../widgets/CustomHomepageDashboardSetting";
import { HttpsOnlyWidget } from "../widgets/HttpsOnlyWidget";
import { SearchSettingsWidget } from "../widgets/SearchSettingsWidget";
import { SiteUrlWidget } from "../widgets/SiteUrlWidget";
export function GeneralSettingsPage() {
@@ -36,6 +37,8 @@ export function GeneralSettingsPage() {
<HttpsOnlyWidget />
<SearchSettingsWidget />
<CustomHomepageDashboardSetting />
<PLUGIN_LANDING_PAGE.LandingPageWidget />

View File

@@ -31,6 +31,7 @@ const generalSettings = {
"humanization-strategy": "simple",
"enable-xrays": false,
"allowed-iframe-hosts": "https://cooldashboards.limo",
"search-engine": "appdb",
} as const;
const setup = async () => {

View File

@@ -0,0 +1,64 @@
import { t } from "ttag";
import { UpsellSemanticSearchPill } from "metabase/admin/upsells/UpsellSemanticSearch";
import { useAdminSetting } from "metabase/api/utils";
import { useHasTokenFeature } from "metabase/common/hooks";
import { getPlan, isProPlan } from "metabase/common/utils/plan";
import { useSelector } from "metabase/lib/redux";
import { getSetting } from "metabase/selectors/settings";
import { getIsHosted } from "metabase/setup";
import { Stack } from "metabase/ui";
import { SettingHeader } from "../SettingHeader";
import { BasicAdminSettingInput } from "../widgets/AdminSettingInput";
export function SearchSettingsWidget() {
const hasFeature = useHasTokenFeature("semantic_search");
const isHosted = useSelector(getIsHosted);
const plan = useSelector((state) =>
getPlan(getSetting(state, "token-features")),
);
const shouldUpsell = !hasFeature && !isProPlan(plan);
const { value, updateSetting } = useAdminSetting("search-engine");
const handleChange = async (enabled: boolean) => {
await updateSetting({
key: "search-engine",
// null will pick the appropriate default for the app-db
value: enabled ? "semantic" : null,
});
};
const isPartOfLimitedRollout = hasFeature && isHosted;
if (!isPartOfLimitedRollout) {
return null;
}
return (
<Stack data-testid="search-engine-setting">
<Stack gap="0">
<SettingHeader
id="search-engine"
title={t`Advanced semantic search`}
description={t`Provides more relevant search results.`}
/>
{shouldUpsell && (
<div>
<UpsellSemanticSearchPill source="settings-general" />
</div>
)}
</Stack>
{!shouldUpsell && (
<BasicAdminSettingInput
name="search-engine"
inputType="boolean"
value={value === "semantic"}
onChange={(newValue) => handleChange(Boolean(newValue))}
/>
)}
</Stack>
);
}

View File

@@ -0,0 +1,12 @@
import { t } from "ttag";
import { UpsellPill } from "./components";
import { UPGRADE_URL } from "./constants";
export const UpsellSemanticSearchPill = ({ source }: { source: string }) => {
return (
<UpsellPill campaign="semantic-search" link={UPGRADE_URL} source={source}>
{t`Get this with Pro. Try for free.`}
</UpsellPill>
);
};

View File

@@ -10,7 +10,12 @@ export const searchApi = Api.injectEndpoints({
query: (params) => ({
method: "GET",
url: "/api/search",
params,
params: {
// default FE to using app db for search to so it
// can be enabled only in specific features
search_engine: "appdb",
...params,
},
}),
providesTags: (response, error, { models }) =>
provideSearchItemListTags(response?.data ?? [], models),

View File

@@ -1,13 +1,12 @@
import type { TokenFeatures } from "metabase-types/api";
import { tokenFeatures } from "metabase-types/api";
export type Plan =
| "oss"
| "starter"
| "starter-with-dwh"
| "pro-cloud"
| "pro-cloud-with-dwh"
| "pro-self-hosted";
export type ProPlan = "pro-cloud" | "pro-cloud-with-dwh" | "pro-self-hosted";
export type Plan = "oss" | "starter" | "starter-with-dwh" | ProPlan;
export const isProPlan = (plan: Plan): plan is ProPlan =>
plan.startsWith("pro-");
export const getPlan = (features?: TokenFeatures | null): Plan => {
if (features) {

View File

@@ -201,6 +201,7 @@ describe("LinkedEntityPicker", () => {
expect(urlSearchParamsToObject(urlObject.searchParams)).toEqual({
context: "entity-picker",
models: "dashboard",
search_engine: "appdb",
q: typedText,
filter_items_in_personal_collection: "exclude",
});
@@ -262,6 +263,7 @@ describe("LinkedEntityPicker", () => {
expect(urlSearchParamsToObject(urlObject.searchParams)).toEqual({
context: "entity-picker",
models: "dashboard",
search_engine: "appdb",
q: typedText,
});
});
@@ -347,6 +349,7 @@ describe("LinkedEntityPicker", () => {
models: ["card", "dataset"],
q: typedText,
filter_items_in_personal_collection: "exclude",
search_engine: "appdb",
});
});
});
@@ -404,6 +407,7 @@ describe("LinkedEntityPicker", () => {
expect(urlSearchParamsToObject(urlObject.searchParams)).toEqual({
context: "entity-picker",
models: ["card", "dataset"],
search_engine: "appdb",
q: typedText,
});
});

View File

@@ -78,6 +78,7 @@ export const useCommandPalette = ({
context: "command-palette",
include_dashboard_questions: true,
limit: 20,
search_engine: "semantic",
},
{
skip: !debouncedSearchText || !isSearchTypeaheadEnabled,

View File

@@ -555,6 +555,7 @@ describe("AddToDashSelectDashModal", () => {
expect(Object.fromEntries(urlObject.searchParams.entries())).toEqual({
context: "entity-picker",
models: "dashboard",
search_engine: "appdb",
q: typedText,
filter_items_in_personal_collection: "only",
});
@@ -585,6 +586,7 @@ describe("AddToDashSelectDashModal", () => {
expect(Object.fromEntries(urlObject.searchParams.entries())).toEqual({
context: "entity-picker",
models: "dashboard",
search_engine: "appdb",
q: typedText,
});
});

View File

@@ -1,3 +1,4 @@
import { t } from "ttag";
import _ from "underscore";
import { PLUGIN_CONTENT_VERIFICATION } from "metabase/plugins";
@@ -14,6 +15,7 @@ import { SearchFilterKeys } from "metabase/search/constants";
import type {
FilterTypeKeys,
SearchFilterComponent,
SearchFilterToggle,
SearchQueryParamValue,
URLSearchFilterQueryParams,
} from "metabase/search/types";
@@ -25,6 +27,14 @@ type SearchSidebarProps = {
};
export const SearchSidebar = ({ value, onChange }: SearchSidebarProps) => {
// TODO: don't ship this to prod
const DisableSemanticSearchFilter: SearchFilterToggle = {
label: () => t`Disable semantic search`,
type: "toggle",
fromUrl: (value) => value === "true",
toUrl: (value: boolean) => (value ? "true" : null),
};
const filterMap: Record<FilterTypeKeys, SearchFilterComponent> = {
[SearchFilterKeys.Type]: TypeFilter,
[SearchFilterKeys.CreatedBy]: CreatedByFilter,
@@ -32,17 +42,28 @@ export const SearchSidebar = ({ value, onChange }: SearchSidebarProps) => {
[SearchFilterKeys.LastEditedBy]: LastEditedByFilter,
[SearchFilterKeys.LastEditedAt]: LastEditedAtFilter,
[SearchFilterKeys.Verified]: PLUGIN_CONTENT_VERIFICATION.VerifiedFilter,
[SearchFilterKeys.DisableSemanticSearch]: DisableSemanticSearchFilter,
[SearchFilterKeys.NativeQuery]: NativeQueryFilter,
[SearchFilterKeys.SearchTrashedItems]: SearchTrashedItemsFilter,
};
const onOutputChange = (key: FilterTypeKeys, val?: SearchQueryParamValue) => {
if (!val) {
onChange(_.omit(value, key));
const omitKeys = _.compact([
key,
// when enabling semantic search, prevent native queries from being searched
key === SearchFilterKeys.DisableSemanticSearch &&
SearchFilterKeys.NativeQuery,
]);
onChange(_.omit(value, omitKeys));
} else {
onChange({
...value,
[key]: val,
// when enabling native queries search, prevent semantic search from being used
...(key === SearchFilterKeys.NativeQuery && val
? { [SearchFilterKeys.DisableSemanticSearch]: "true" }
: {}),
});
}
};
@@ -90,6 +111,7 @@ export const SearchSidebar = ({ value, onChange }: SearchSidebarProps) => {
{getFilter(SearchFilterKeys.LastEditedAt)}
</Stack>
{getFilter(SearchFilterKeys.Verified)}
{getFilter(SearchFilterKeys.DisableSemanticSearch)}
{getFilter(SearchFilterKeys.NativeQuery)}
{getFilter(SearchFilterKeys.SearchTrashedItems)}
</Stack>

View File

@@ -9,6 +9,7 @@ export const SearchFilterKeys = {
LastEditedAt: "last_edited_at",
NativeQuery: "search_native_query",
SearchTrashedItems: "archived",
DisableSemanticSearch: "disable_semantic_search",
} as const;
export const enabledSearchTypes: EnabledSearchModel[] = [

View File

@@ -49,12 +49,18 @@ function SearchApp({ location }) {
const query = {
q: searchText,
..._.omit(searchFilters, SearchFilterKeys.Type),
..._.omit(searchFilters, [
SearchFilterKeys.Type,
SearchFilterKeys.DisableSemanticSearch,
]),
models: models && (Array.isArray(models) ? models : [models]),
limit: PAGE_SIZE,
offset: PAGE_SIZE * page,
context: SearchContextTypes.SEARCH_APP,
include_dashboard_questions: true,
search_engine: searchFilters[SearchFilterKeys.DisableSemanticSearch]
? "appdb"
: "semantic",
};
const onChangeLocation = useCallback(

View File

@@ -29,6 +29,8 @@ export type VerifiedFilterProps = true | null;
export type NativeQueryFilterProps = true | null;
export type SearchTrashedItemsFilterProps = true | undefined;
export type SemanticSearchFilterProps = true | undefined;
export type SearchFilterPropTypes = {
[SearchFilterKeys.Type]: TypeFilterProps;
[SearchFilterKeys.Verified]: VerifiedFilterProps;
@@ -38,6 +40,7 @@ export type SearchFilterPropTypes = {
[SearchFilterKeys.LastEditedAt]: LastEditedAtFilterProps;
[SearchFilterKeys.NativeQuery]: NativeQueryFilterProps;
[SearchFilterKeys.SearchTrashedItems]: SearchTrashedItemsFilterProps;
[SearchFilterKeys.DisableSemanticSearch]: SemanticSearchFilterProps;
};
export type FilterTypeKeys = keyof SearchFilterPropTypes;

View File

@@ -276,7 +276,9 @@
;; 1ms -> 10minutes
:buckets [1 500 1000 5000 10000 30000 60000 120000 300000 600000]})
(prometheus/gauge :metabase-search/appdb-index-size
{:description "Number of rows in the active index table."})
{:description "Number of rows in the active appdb index table."})
(prometheus/gauge :metabase-search/semantic-index-size
{:description "Number of rows in the active semantic index table."})
(prometheus/gauge :metabase-search/queue-size
{:description "Number of updates on the search indexing queue."})
(prometheus/counter :metabase-search/response-ok
@@ -289,6 +291,10 @@
(prometheus/gauge :metabase-search/engine-active
{:description "Whether a given engine is active. This does NOT mean that it is the default."
:labels [:engine]})
(prometheus/counter :metabase-search/semantic-embedding-tokens
{:description (str "Number of tokens consumed by the given embedding model and provider. "
"Not all providers track token use.")
:labels [:model :provider]})
;; notification metrics
(prometheus/counter :metabase-notification/send-ok
{:description "Number of successful notification sends."

View File

@@ -735,7 +735,7 @@
(defn- ee-snowplow-features-data'
[]
(let [features [:sso-jwt :sso-saml :scim :sandboxes :email-allow-list]]
(let [features [:sso-jwt :sso-saml :scim :sandboxes :email-allow-list :semantic-search]]
(map
(fn [feature]
{:name feature

View File

@@ -61,6 +61,7 @@
enable-query-reference-validation?
enable-sandboxes?
enable-scim?
enable-semantic-search?
enable-serialization?
enable-session-timeout-config?
enable-snippet-collections?

View File

@@ -270,6 +270,10 @@
"Does the Metabase Cloud instance have ETL connections with PG?"
:etl-connections-pg)
(define-premium-feature ^{:added "0.56.0"} enable-semantic-search?
"Should we enable the semantic search backend?"
:semantic-search)
(defn- -token-features []
{:advanced_permissions (enable-advanced-permissions?)
:ai_sql_fixer (enable-ai-sql-fixer?)
@@ -303,6 +307,7 @@
:query_reference_validation (enable-query-reference-validation?)
:sandboxes (enable-sandboxes?)
:scim (enable-scim?)
:semantic_search (enable-semantic-search?)
:serialization (enable-serialization?)
:session_timeout_config (enable-session-timeout-config?)
:snippet_collections (enable-snippet-collections?)

View File

@@ -1,52 +1,10 @@
(ns metabase.search.appdb.specialization.postgres
(:require
[clojure.string :as str]
[metabase.app-db.core :as mdb]
[metabase.search.appdb.specialization.api :as specialization]
[metabase.search.ingestion :as search.ingestion]
[metabase.search.settings :as search.settings]
[metabase.util :as u]
[metabase.util.i18n :as i18n]
[metabase.util.string :as u.str]
[metabase.search.util :as search.util]
[toucan2.core :as t2]))
(def ^:private available-tsv-languages
"Mapping of our available locals to the names of the postgres tsvector languages.
Queries the pg_ts_config table to find out which languages are actually available."
(if (= :postgres (mdb/db-type))
(let [default-mapping {:ar :arabic
:ar_SA :arabic
:ca :catalan
:da :danish
:en :english
:fi :finnish
:fr :french
:de :german
:hu :hungarian
:id :indonesian
:it :italian
:nb :norwegian
:pt_BR :portuguese
:ru :russian
:sr :serbian
:es :spanish
:sv :swedish
:tr :turkish}
available-languages (->> (t2/query {:select [:cfgname]
:from [:pg_ts_config]})
(map :cfgname)
(map keyword)
set)]
(into {} (filter (comp available-languages val) default-mapping)))
{}))
(defn- tsv-language []
(if-let [custom-language (search.settings/search-language)]
custom-language
(if-let [lang ((keyword (i18n/site-locale-string)) available-tsv-languages)]
(name lang)
"simple")))
(defmethod specialization/table-schema :postgres [base-schema]
(into [[:id :bigint [:primary-key] [:raw "GENERATED BY DEFAULT AS IDENTITY"]]
[:search_vector :tsvector :not-null]
@@ -76,75 +34,14 @@
:on-conflict [:model :model_id]
:do-update-set (zipmap update-keys (map excluded-kw update-keys))}))))
(defn- quote* [s]
(str "'" (str/replace s "'" "''") "'"))
(defn- process-phrase [word-or-phrase]
;; a phrase is quoted even if the closing quotation mark has not been typed yet
(cond
;; trailing quotation mark
(= word-or-phrase "\"") nil
;; quoted phrases must be matched sequentially
(str/starts-with? word-or-phrase "\"")
(as-> word-or-phrase <>
;; remove the quote mark(s)
(str/replace <> #"^\"|\"$" "")
(str/trim <>)
(str/split <> #"\s+")
(map quote* <>)
(str/join " <-> " <>))
;; negation
(re-find #"^-\w" word-or-phrase)
(str "!" (quote* (subs word-or-phrase 1)))
;; just a regular word
:else
(quote* word-or-phrase)))
(defn- split-preserving-quotes
"Break up the words in the search input, preserving quoted and partially quoted segments."
[s]
(re-seq #"\"[^\"]*(?:\"|$)|[^\s\"]+|\s+" (u/lower-case-en s)))
(defn- process-clause [words-and-phrases]
(->> words-and-phrases
(remove #{"and"})
(map process-phrase)
(remove str/blank?)
(str/join " & ")))
(defn- complete-last-word
"Add wildcards at the end of the final word, so that we match ts completions."
[expression]
(str/replace expression #"(\S+)(?=\s*$)" "$1:*"))
(defn- to-tsquery-expr
"Given the user input, construct a query in the Postgres tsvector query language."
[input]
(str
(when input
(let [trimmed (str/trim input)
complete? (not (str/ends-with? trimmed "\""))
;; TODO also only complete if the :context is appropriate
maybe-complete (if complete? complete-last-word identity)]
(->> (str/replace trimmed "\\" "\\\\")
split-preserving-quotes
(remove str/blank?)
(partition-by #{"or"})
(remove #(= (first %) "or"))
(map process-clause)
(str/join " | ")
maybe-complete)))))
(defmethod specialization/base-query :postgres
[active-table search-term search-ctx select-items]
{:select select-items
:from [[active-table :search_index]]
;; Using a join allows us to share the query expression between our SELECT and WHERE clauses.
:join [[[:raw "to_tsquery('"
(tsv-language) "', "
[:lift (to-tsquery-expr search-term)] ")"]
(search.util/tsv-language) "', "
[:lift (search.util/to-tsquery-expr search-term)] ")"]
:query] [:= 1 1]]
:where (if (str/blank? search-term)
[:= [:inline 1] [:inline 1]]
@@ -155,20 +52,16 @@
"search_vector")
" @@ query")])})
(defn- weighted-tsvector [weight text]
;; tsvector has a max value size of 1048575 bytes, limit to less than that because the multiple values get concatenated together
[:setweight [:to_tsvector [:inline (tsv-language)] [:cast (u.str/limit-bytes text search.ingestion/max-searchable-value-length) :text]] [:inline weight]])
(defmethod specialization/extra-entry-fields :postgres [entity]
{:search_vector
[:||
(weighted-tsvector "A" (:name entity))
(weighted-tsvector "B" (:searchable_text entity ""))]
(search.util/weighted-tsvector "A" (:name entity))
(search.util/weighted-tsvector "B" (:searchable_text entity ""))]
:with_native_query_vector
[:||
(weighted-tsvector "A" (:name entity))
(weighted-tsvector "B" (str/join " " (keep entity [:searchable_text :native_query])))]})
(search.util/weighted-tsvector "A" (:name entity))
(search.util/weighted-tsvector "B" (str/join " " (keep entity [:searchable_text :native_query])))]})
;; See https://www.postgresql.org/docs/current/textsearch-controls.html#TEXTSEARCH-RANKING
;; 0 (the default) ignores the document length

View File

@@ -9,6 +9,7 @@
[metabase.search.impl :as search.impl]
[metabase.search.in-place.legacy :as search.legacy]
[metabase.search.ingestion :as search.ingestion]
[metabase.search.semantic.core :as search.engines.semantic]
[metabase.search.spec :as search.spec]
[metabase.search.util :as search.util]
[metabase.util :as u]
@@ -21,6 +22,7 @@
;; TODO -- maybe engine loading should be moved to [[metabase.search.init]] instead
search.engine/keep-me
search.engines.appdb/keep-me
search.engines.semantic/keep-me
search.legacy/keep-me
search.config/keep-me
@@ -43,7 +45,13 @@
searchable-value-trim-sql]
[search.spec
define-spec])
spec
define-spec]
[search.util
tsv-language
to-tsquery-expr
weighted-tsvector])
(defmethod analytics/known-labels :metabase-search/index-updates
[_]

View File

@@ -44,7 +44,7 @@
search-engine))
(defmulti init!
"Ensure that the search index exists, an is ready to take search queries."
"Ensure that the search index exists, and is ready to take search queries."
{:arglists '([engine opts])}
(fn [engine _opts]
engine))

View File

@@ -210,6 +210,7 @@
:search.engine/in-place))
(first (filter search.engine/supported-engine?
[:search.engine/appdb
:search.engine/semantic
:search.engine/in-place]))))
(defn- parse-engine [value]
@@ -393,17 +394,16 @@
;; We get to do this slicing and dicing with the result data because
;; the pagination of search is for UI improvement, not for performance.
;; We intend for the cardinality of the search results to be below the default max before this slicing occurs
(cond->
{:data (cond->> total-results
(some? (:offset-int search-ctx)) (drop (:offset-int search-ctx))
(some? (:limit-int search-ctx)) (take (:limit-int search-ctx))
true (map add-perms-for-col))
:limit (:limit-int search-ctx)
:models (:models search-ctx)
:offset (:offset-int search-ctx)
:table_db_id (:table-db-id search-ctx)
:engine (:search-engine search-ctx)
:total (count total-results)}
(cond-> {:data (cond->> total-results
(some? (:offset-int search-ctx)) (drop (:offset-int search-ctx))
(some? (:limit-int search-ctx)) (take (:limit-int search-ctx))
true (map add-perms-for-col))
:limit (:limit-int search-ctx)
:models (:models search-ctx)
:offset (:offset-int search-ctx)
:table_db_id (:table-db-id search-ctx)
:engine (:search-engine search-ctx)
:total (count total-results)}
(:calculate-available-models? search-ctx)
(assoc :available_models (model-set-fn search-ctx)))))

View File

@@ -0,0 +1,141 @@
(ns metabase.search.semantic.core
(:require
[metabase.premium-features.core :refer [defenterprise]]
[metabase.search.config :as search.config]
[metabase.search.engine :as search.engine]
[metabase.search.ingestion :as search.ingestion]
[metabase.util.log :as log]))
(defn- oss-semantic-search-error
"Helper function to throw semantic search enterprise feature error."
[]
(throw (ex-info "Semantic search is a premium feature and requires an enterprise license."
{:type :semantic-search-not-available})))
;; defenterprise functions that map directly to search.engine methods
(defenterprise supported?
"Check if semantic search is supported on this instance."
metabase-enterprise.semantic-search.core
[]
false)
(defenterprise results
"Execute a semantic search query and return results."
metabase-enterprise.semantic-search.core
[_search-ctx]
(oss-semantic-search-error))
(defenterprise update-index!
"Update the semantic search index with the given documents."
metabase-enterprise.semantic-search.core
[_document-reducible]
(oss-semantic-search-error))
(defenterprise delete-from-index!
"Delete documents from the semantic search index."
metabase-enterprise.semantic-search.core
[_model _ids]
(oss-semantic-search-error))
(defenterprise init!
"Initialize the semantic search index."
metabase-enterprise.semantic-search.core
[_searchable-documents _opts]
(oss-semantic-search-error))
(defenterprise reindex!
"Perform a full reindex of the semantic search engine."
metabase-enterprise.semantic-search.core
[_searchable-documents _opts]
(oss-semantic-search-error))
(defenterprise reset-tracking!
"Reset tracking for the semantic search engine."
metabase-enterprise.semantic-search.core
[]
(oss-semantic-search-error))
;; Search engine method implementations
(defmethod search.engine/supported-engine? :search.engine/semantic [_]
(try
(supported?)
(catch Exception e
(log/warn e "Semantic search engine not supported")
false)))
(defmethod search.engine/results :search.engine/semantic
[search-ctx]
(try
(results search-ctx)
(catch Exception e
(log/error e "Error executing semantic search")
[])))
(defmethod search.engine/model-set :search.engine/semantic
[_search-ctx]
(disj search.config/all-models "indexed-entity"))
;; NOTE if you remove this you need to also modify filter-read-permitted to handle indexed-entity docs.
;; See the corresponding note in metabase-enterprise.semantic-search.index/filter-read-permitted.
(defn- remove-indexed-entities
"Remove indexed-entities from `document-reducible`"
[document-reducible]
(eduction (remove (comp #{"indexed-entity"} :model))
document-reducible))
(defmethod search.engine/update! :search.engine/semantic
[_ document-reducible]
(try
(log/info "Updating semantic search engine")
(update-index! (-> document-reducible
remove-indexed-entities))
(catch Exception e
(log/error e "Error updating semantic search engine")
{})))
(defmethod search.engine/delete! :search.engine/semantic
[_ model ids]
(try
(log/info "Deleting from semantic search engine")
(delete-from-index! model ids)
(catch Exception e
(log/error e "Error deleting from semantic search engine")
{})))
(defmethod search.engine/init! :search.engine/semantic
[_ opts]
(try
(log/info "Initializing semantic search engine")
(init! (-> (search.ingestion/searchable-documents)
remove-indexed-entities)
opts)
(catch Exception e
(log/error e "Error initializing semantic search engine")
(throw e))))
(defmethod search.engine/reindex! :search.engine/semantic
[_ opts]
(try
(log/info "Reindexing semantic search engine")
(reindex! (-> (search.ingestion/searchable-documents)
remove-indexed-entities)
opts)
(catch Exception e
(log/error e "Error reindexing semantic search engine")
(throw e))))
(comment
(def docs (->> (search.ingestion/searchable-documents)
remove-indexed-entities
vec))
(into #{} (map :model docs))
(init! docs {})
(reindex! docs {}))
(defmethod search.engine/reset-tracking! :search.engine/semantic [_]
(try
(reset-tracking!)
(catch Exception e
(log/debug e "Error resetting tracking for semantic search engine"))))

View File

@@ -1,6 +1,8 @@
(ns metabase.search.settings
(:require
[metabase.app-db.core :as app-db]
[metabase.appearance.core :as appearance]
[metabase.premium-features.core :as premium-features]
[metabase.settings.core :as setting :refer [defsetting]]
[metabase.util.i18n :as i18n]))
@@ -14,11 +16,20 @@
:audit :getter)
(defsetting search-engine
(i18n/deferred-tru "Which engine to use when performing search. Supported values are :in-place and :appdb")
:visibility :internal
:export? false
(i18n/deferred-tru "Which engine to use by default when performing search. Supported values are :in-place, :appdb, and :semantic")
:visibility :authenticated
:export? true
:default :appdb
:type :keyword)
:type :keyword
:setter (fn [new-value]
(let [value (or new-value
(if (= :postgres (app-db/db-type)) :appdb :in-place))]
(when (some? value)
(assert (#{:in-place :appdb :semantic} (keyword value))
(i18n/tru "Invalid search-engine, only values of in-place, appdb, and semantic are allowed."))
(when (= :semantic (keyword value))
(premium-features/assert-has-feature :semantic-search (i18n/tru "Semantic search"))))
(setting/set-value-of-type! :keyword :search-engine value))))
(defsetting experimental-search-weight-overrides
(i18n/deferred-tru "Used to override weights used for search ranking")

View File

@@ -43,7 +43,7 @@
(search/reindex!))
(defmethod startup/def-startup-logic! ::SearchIndexInit [_]
(quick-task/submit-task! (init!)))
(quick-task/submit-task! init!))
(defmethod task/init! ::SearchIndexReindex [_]
(let [job (jobs/build
@@ -65,4 +65,5 @@
(ingestion/start-listener!))
(comment
(task/job-exists? reindex-job-key))
(task/job-exists? reindex-job-key)
(task/trigger-now! reindex-job-key))

View File

@@ -1,4 +1,13 @@
(ns metabase.search.util)
(ns metabase.search.util
(:require
[clojure.string :as str]
[metabase.app-db.core :as mdb]
[metabase.search.ingestion :as search.ingestion]
[metabase.search.settings :as search.settings]
[metabase.util :as u]
[metabase.util.i18n :as i18n]
[metabase.util.string :as u.str]
[toucan2.core :as t2]))
(defn impossible-condition?
"An (incomplete) check where queries will definitely return nothing, to help avoid spurious index update queries."
@@ -12,3 +21,113 @@
:and (boolean (some impossible-condition? (rest where)))
:or (every? impossible-condition? (rest where))
false)))
;;; ============================================================================
;;; Postgres-specific utilities
;;; ============================================================================
(def ^:private available-tsv-languages
"Mapping of our available locals to the names of the postgres tsvector languages.
Queries the pg_ts_config table to find out which languages are actually available."
(if (= :postgres (mdb/db-type))
(let [default-mapping {:ar :arabic
:ar_SA :arabic
:ca :catalan
:da :danish
:en :english
:fi :finnish
:fr :french
:de :german
:hu :hungarian
:id :indonesian
:it :italian
:nb :norwegian
:pt_BR :portuguese
:ru :russian
:sr :serbian
:es :spanish
:sv :swedish
:tr :turkish}
available-languages (->> (t2/query {:select [:cfgname]
:from [:pg_ts_config]})
(map :cfgname)
(map keyword)
set)]
(into {} (filter (comp available-languages val) default-mapping)))
{}))
(defn tsv-language
"Get the appropriate text search configuration language for Postgres tsvector operations."
[]
(if-let [custom-language (search.settings/search-language)]
custom-language
(if-let [lang ((keyword (i18n/site-locale-string)) available-tsv-languages)]
(name lang)
"simple")))
(defn- quote* [s]
(str "'" (str/replace s "'" "''") "'"))
(defn- process-phrase [word-or-phrase]
;; a phrase is quoted even if the closing quotation mark has not been typed yet
(cond
;; trailing quotation mark
(= word-or-phrase "\"") nil
;; quoted phrases must be matched sequentially
(str/starts-with? word-or-phrase "\"")
(as-> word-or-phrase <>
;; remove the quote mark(s)
(str/replace <> #"^\"|\"$" "")
(str/trim <>)
(str/split <> #"\s+")
(map quote* <>)
(str/join " <-> " <>))
;; negation
(re-find #"^-\w" word-or-phrase)
(str "!" (quote* (subs word-or-phrase 1)))
;; just a regular word
:else
(quote* word-or-phrase)))
(defn- split-preserving-quotes
"Break up the words in the search input, preserving quoted and partially quoted segments."
[s]
(re-seq #"\"[^\"]*(?:\"|$)|[^\s\"]+|\s+" (u/lower-case-en s)))
(defn- process-clause [words-and-phrases]
(->> words-and-phrases
(remove #{"and"})
(map process-phrase)
(remove str/blank?)
(str/join " & ")))
(defn- complete-last-word
"Add wildcards at the end of the final word, so that we match ts completions."
[expression]
(str/replace expression #"(\S+)(?=\s*$)" "$1:*"))
(defn to-tsquery-expr
"Given the user input, construct a query in the Postgres tsvector query language."
[input]
(str
(when input
(let [trimmed (str/trim input)
complete? (not (str/ends-with? trimmed "\""))
;; TODO also only complete if the :context is appropriate
maybe-complete (if complete? complete-last-word identity)]
(->> (str/replace trimmed "\\" "\\\\")
split-preserving-quotes
(remove str/blank?)
(partition-by #{"or"})
(remove #(= (first %) "or"))
(map process-clause)
(str/join " | ")
maybe-complete)))))
(defn weighted-tsvector
"Create a weighted tsvector for Postgres full-text search with the given weight and text."
[weight text]
;; tsvector has a max value size of 1048575 bytes, limit to less than that because the multiple values get concatenated together
[:setweight [:to_tsvector [:inline (tsv-language)] [:cast (u.str/limit-bytes text search.ingestion/max-searchable-value-length) :text]] [:inline weight]])

View File

@@ -24,5 +24,6 @@
"Submit a task to the single thread executor. This will attempt to serialize repeated requests to sync tables. It
obviously cannot work across multiple instances."
^Future [^Callable f]
{:pre [(some? f)]}
(let [task (bound-fn* f)]
(.submit ^ExecutorService @executor ^Callable task)))

View File

@@ -1,82 +1,3 @@
(ns metabase.search.appdb.specializations.postgres-test
(:require
[clojure.test :refer :all]
[metabase.app-db.core :as mdb]
[metabase.search.appdb.specialization.postgres :as search.postgres]
[metabase.test.util :as tu]
[metabase.util.i18n :as i18n]))
(def search-expr #'search.postgres/to-tsquery-expr)
(deftest to-tsquery-expr-test
(is (= "'a' & 'b' & 'c':*"
(search-expr "a b c")))
(is (= "'a' & 'b' & 'c':*"
(search-expr "a AND b AND c")))
(is (= "'a' & 'b' & 'c'"
(search-expr "a b \"c\"")))
(is (= "'a' & 'b' | 'c':*"
(search-expr "a b or c")))
(is (= "'this' & !'that':*"
(search-expr "this -that")))
(testing "hyphens"
(is (= "'[ops' & 'monitoring]' & '-' & 'available':*"
(search-expr "[ops monitoring] - available")))
(is (= "'[ops' & 'monitoring]' & '--' & 'available':*"
(search-expr "[ops monitoring] -- available")))
(is (= "'[ops' & 'monitoring]' & 'not-available':*"
(search-expr "[ops monitoring] not-available"))))
(is (= "'a' & 'b' & 'c' <-> 'd' & 'e' | 'b' & 'e':*"
(search-expr "a b \" c d\" e or b e")))
(is (= "'ab' <-> 'and' <-> 'cde' <-> 'f' | !'abc' & 'def' & 'ghi' | 'jkl' <-> 'mno' <-> 'or' <-> 'pqr'"
(search-expr "\"ab and cde f\" or -abc def AND ghi OR \"jkl mno OR pqr\"")))
(is (= "'big' & 'data' | 'business' <-> 'intelligence' | 'data' & 'wrangling':*"
(search-expr "Big Data oR \"Business Intelligence\" OR data and wrangling")))
(testing "unbalanced quotes"
(is (= "'big' <-> 'data' & 'big' <-> 'mistake':*"
(search-expr "\"Big Data\" \"Big Mistake")))
(is (= "'something'"
(search-expr "something \""))))
(is (= "'partial' <-> 'quoted' <-> 'and' <-> 'or' <-> '-split':*"
(search-expr "\"partial quoted AND OR -split")))
(testing "dangerous characters"
(is (= "'you' & '<-' & 'pointing':*"
(search-expr "you <- pointing"))))
(testing "backslash"
(is (= "'test\\\\':*" (search-expr "test\\"))))
(testing "single quotes"
(is (= "'you''re':*"
(search-expr "you're")))))
(deftest available-tsv-languages
(when (= :postgres (mdb/db-type))
(let [available @#'search.postgres/available-tsv-languages]
(is (= :english (:en available)))
(is (= :german (:de available)))
(is (nil? (:ko available))))))
(deftest tsv-language
(when (= :postgres (mdb/db-type))
(binding [i18n/*site-locale-override* "en"]
(is (= "english" (#'search.postgres/tsv-language))))
(binding [i18n/*site-locale-override* "de"]
(is (= "german" (#'search.postgres/tsv-language))))
(binding [i18n/*site-locale-override* "pt_BR"]
(is (= "portuguese" (#'search.postgres/tsv-language))))
(binding [i18n/*site-locale-override* "ko"]
(is (= "simple" (#'search.postgres/tsv-language))))
(tu/with-temporary-setting-values [search-language "custom-value"]
(is (= "custom-value" (#'search.postgres/tsv-language))))))
[clojure.test :refer :all]))

View File

@@ -1,7 +1,10 @@
(ns metabase.search.util-test
(:require
[clojure.test :refer :all]
[metabase.search.util :as search.util]))
[metabase.app-db.core :as mdb]
[metabase.search.util :as search.util]
[metabase.test.util :as tu]
[metabase.util.i18n :as i18n]))
(def ^:private impossible? search.util/impossible-condition?)
@@ -17,3 +20,82 @@
(is (not (impossible? [:and [:= 1 :this.id] [:!= "card" "dashboard"]])))
(is (not (impossible? [:or [:= 1 :this.id] [:!= "card" "dashboard"]])))
(is (impossible? [:or [:= "oh" "no"] [:= "card" "dashboard"]])))
;;; ============================================================================
;;; Postgres-specific utility tests
;;; ============================================================================
(def search-expr #'search.util/to-tsquery-expr)
(deftest to-tsquery-expr-test
(is (= "'a' & 'b' & 'c':*"
(search-expr "a b c")))
(is (= "'a' & 'b' & 'c':*"
(search-expr "a AND b AND c")))
(is (= "'a' & 'b' & 'c'"
(search-expr "a b \"c\"")))
(is (= "'a' & 'b' | 'c':*"
(search-expr "a b or c")))
(is (= "'this' & !'that':*"
(search-expr "this -that")))
(testing "hyphens"
(is (= "'[ops' & 'monitoring]' & '-' & 'available':*"
(search-expr "[ops monitoring] - available")))
(is (= "'[ops' & 'monitoring]' & '--' & 'available':*"
(search-expr "[ops monitoring] -- available")))
(is (= "'[ops' & 'monitoring]' & 'not-available':*"
(search-expr "[ops monitoring] not-available"))))
(is (= "'a' & 'b' & 'c' <-> 'd' & 'e' | 'b' & 'e':*"
(search-expr "a b \" c d\" e or b e")))
(is (= "'ab' <-> 'and' <-> 'cde' <-> 'f' | !'abc' & 'def' & 'ghi' | 'jkl' <-> 'mno' <-> 'or' <-> 'pqr'"
(search-expr "\"ab and cde f\" or -abc def AND ghi OR \"jkl mno OR pqr\"")))
(is (= "'big' & 'data' | 'business' <-> 'intelligence' | 'data' & 'wrangling':*"
(search-expr "Big Data oR \"Business Intelligence\" OR data and wrangling")))
(testing "unbalanced quotes"
(is (= "'big' <-> 'data' & 'big' <-> 'mistake':*"
(search-expr "\"Big Data\" \"Big Mistake")))
(is (= "'something'"
(search-expr "something \""))))
(is (= "'partial' <-> 'quoted' <-> 'and' <-> 'or' <-> '-split':*"
(search-expr "\"partial quoted AND OR -split")))
(testing "dangerous characters"
(is (= "'you' & '<-' & 'pointing':*"
(search-expr "you <- pointing"))))
(testing "backslash"
(is (= "'test\\\\':*" (search-expr "test\\"))))
(testing "single quotes"
(is (= "'you''re':*"
(search-expr "you're")))))
(deftest available-tsv-languages-test
(when (= :postgres (mdb/db-type))
(let [available @#'search.util/available-tsv-languages]
(is (= :english (:en available)))
(is (= :german (:de available)))
(is (nil? (:ko available))))))
(deftest tsv-language-test
(when (= :postgres (mdb/db-type))
(binding [i18n/*site-locale-override* "en"]
(is (= "english" (#'search.util/tsv-language))))
(binding [i18n/*site-locale-override* "de"]
(is (= "german" (#'search.util/tsv-language))))
(binding [i18n/*site-locale-override* "pt_BR"]
(is (= "portuguese" (#'search.util/tsv-language))))
(binding [i18n/*site-locale-override* "ko"]
(is (= "simple" (#'search.util/tsv-language))))
(tu/with-temporary-setting-values [search-language "custom-value"]
(is (= "custom-value" (#'search.util/tsv-language))))))

View File

@@ -5,15 +5,9 @@
[clojure.string :as str]
[clojure.tools.namespace.find :as ns.find]))
;; This is made `^:const` so it will get calculated when the uberjar is compiled. `find-namespaces` won't work if
;; source is excluded; either way this takes a few seconds, so doing it at compile time speeds up launch as well.
(defonce ^:const ^{:doc "Vector of symbols of all Metabase namespaces, excluding test namespaces. This is intended
(defonce ^{:doc "Vector of symbols of all Metabase namespaces, excluding test namespaces. This is intended
for use by various test routines that need to load related namespaces.
DEPRECATED: Using this is an anti-pattern, it messes up our ability to analyze the code and find dependencies between
namespaces or to topographically sort them correctly during compilation. See
https://metaboat.slack.com/archives/CKZEMT1MJ/p1734635053499399 or ask Cam for more info.
This was moved from metabase.util.jvm for testing purposes only."}
^:deprecated metabase-namespace-symbols
(vec (sort (for [ns-symb (ns.find/find-namespaces (classpath/system-classpath))

View File

@@ -42,6 +42,7 @@ persisted-models-enabled: null
report-timezone: null
report-timezone-long: null
report-timezone-short: null
search-engine: null
search-typeahead-enabled: null
setup-embedding-autoenabled: null
setup-license-active-at-setup: null