@dbx-tools/search

Add agent tools, federated search, index lifecycle helpers, or a Lakebase-backed full-text provider to AppKit AI Search. Use native AppKit aiSearch for Databricks AI Search queries. Add this package when an agent needs search / universal_search, an app needs to create, sync, or seed indexes, or a deployment needs the same AppKit query contract backed by PostgreSQL full-text search instead of Vector Search.

Quick start

Add the extension around an existing AppKit AI Search index:

import { createApp, server } from "@databricks/appkit";
import { aiSearch } from "@databricks/appkit/beta";
import { search } from "@dbx-tools/search";

await createApp({
  plugins: [
    server(),
    aiSearch({ indexes: { docs: { indexName: "main.support.docs" } } }),
    search({ index: "main.support.docs" }),
  ],
});

The extension adds:

• Vector Search reads delegate to AppKit aiSearch, which handles identity selection, caching, reranking, pagination, route validation, and response decoding.
• lakebaseAiSearch() implements the same aiSearch alias, route, client-config, and SearchResponse contract with PostgreSQL tsvector.
• Agent tools for Mastra (searchTool() and related factories) and AppKit agents (through the extension plugin's ToolProvider): search and universal_search reads, plus provider-aware write tools. Lakebase exposes add_documents. Native Vector Search can also expose create_index and sync_index.
• AppKit-compatible query routes under /api/ai-search/:alias. Extension routes under /api/search cover universal search and optional lifecycle operations.
• Configuration can infer columns, page size, mode, and aliases after you name a default index or set DATABRICKS_VECTOR_SEARCH_INDEX.
• Native AppKit aiSearch uses the service principal by default. Set an index's auth to on-behalf-of-user when queries must use the signed-in user's permissions. Each Lakebase full-text index selects the same two modes through the sibling native lakebase routing pool.
• Filters use AppKit's scalar/array shape: { column: valueOrValues }.
• Embedding-model resolution for index creation reuses @dbx-tools/model: a loose name fuzzy-matches the live catalogue, or the best embedding endpoint is chosen automatically.
• Index lifecycle helpers: createIndex / ensureIndex (Delta Sync from a source table, self-managed direct-access with a dimension, or a managed direct-access index that embeds a text column and needs no Delta table or warehouse), provision (ensure + seed in one idempotent call), syncIndex, deleteIndex, listIndexes, and ensureEndpoint. Each infers the endpoint, embedding model, primary key, and columns from package defaults.
• Provision an index on boot with ensureOnSetup: the plugin provisions the endpoint + index and seeds documents in the background using the native AppKit workspace client, so a fresh deployment is searchable with no manual setup.
• An explicit Lakebase full-text provider: register lakebaseAiSearch instead of native aiSearch to serve the same aliases, query routes, filters, and result shape from a Postgres tsvector index.

Choose a search provider

Do not use this package instead of native AppKit for ordinary Vector Search queries. Register aiSearch from @databricks/appkit/beta.

Use this package for capabilities AppKit does not ship: agent tool providers, federated fan-out, index lifecycle and seeding, reusable result components, or the Lakebase full-text implementation of the AppKit AI Search contract.

Register native aiSearch for Databricks Vector Search. Register lakebaseAiSearch for PostgreSQL full-text search with the same AppKit query contract. They both use the aiSearch plugin name, so an application must choose one. Add search() beside the selected provider when agents need tools, federated reads, or lifecycle routes.

Add search to a Mastra agent

import { createApp, server } from "@databricks/appkit";
import { aiSearch } from "@databricks/appkit/beta";
import { search } from "@dbx-tools/search";
import { agents, mastra } from "@dbx-tools/appkit-mastra";

const support = agents.createAgent({
  instructions: "Answer from the docs; use `search` to find them.",
  async tools(plugins) {
    return { ...(await plugins.search.toolkit()) };
  },
});

await createApp({
  plugins: [
    server(),
    aiSearch({
      indexes: {
        docs: {
          indexName: "main.support.docs",
          columns: ["id", "title", "url", "body"],
        },
      },
    }),
    search({
      index: "main.support.docs",
      indexes: [{ name: "main.support.docs", alias: "docs" }],
    }),
    mastra({ agents: support }),
  ],
});

Every search() plugin keeps its resolved config, client, and provider backend separate. App-integrated agents should consume its native toolkit as shown above, which keeps calls attached to that exact plugin. Standalone Mastra factories accept an explicit runtime, or config plus a provider:

import { runtime, tool } from "@dbx-tools/search";

const searchRuntime = runtime.createSearchRuntime({
  config: { index: "main.support.docs" },
  readBackend,
});
const search = tool.searchTool({ runtime: searchRuntime });

Use Lakebase full-text search without changing the AppKit UI hook:

import { lakebase } from "@databricks/appkit";
import { lakebaseAiSearch } from "@dbx-tools/search";

createApp({
  plugins: [
    lakebase(),
    lakebaseAiSearch({
      indexes: {
        docs: {
          indexName: "docs",
          columns: ["id", "title", "body"],
          queryType: "full_text",
        },
      },
    }),
  ],
});

Manage indexes

SearchClient is the lifecycle client. Query execution requires the registered AppKit-compatible provider. A Delta Sync index computes embeddings from a source Delta table and stays synced. The embedding model, endpoint, primary key (id), and text column (text/content/body) are all inferred when omitted.

const client = createSearchClient();

// Ensure the Vector Search endpoint exists (creates a STANDARD one if not).
await client.ensureEndpoint("my-vs-endpoint", { wait: true });

// Delta Sync index from a Delta table (embeddings computed by Databricks).
await client.ensureIndex("main.support.docs", {
  endpoint: "my-vs-endpoint",
  sourceTable: "main.support.docs_source",
  // embeddingModel / primaryKey / embeddingSourceColumn inferred when omitted
});

// Trigger a sync, then later delete.
await client.syncIndex("main.support.docs");
await client.deleteIndex("main.support.docs");

// Managed direct-access index: Databricks embeds a
// text column on write AND query, so no Delta table, no warehouse, no vectors.
await client.createIndex("main.support.docs", {
  endpoint: "my-vs-endpoint",
  // managed by default when no embeddingDimension is given
});
await client
  .index("main.support.docs")
  .addDocuments([{ id: "1", text: "AI Search finds the most relevant documents for a query." }]);

// Self-managed direct-access index you write vectors to yourself.
await client.createIndex("main.support.vectors", { embeddingDimension: 1024 });
await client.index("main.support.vectors").addDocuments([{ id: "1", embedding: [/* … */] }]);

// One call to make an index real AND seeded (idempotent - safe every boot).
await client.provision("main.support.docs", {
  endpoint: "my-vs-endpoint",
  seed: [{ id: "1", text: "Databricks AI Search overview", url: "https://…" }],
});

The SearchIndex handle mirrors these: index.ensure(opts), index.sync(), index.delete(), index.info().

When the plugin's write surface is enabled (search({ allowWrite: true }) or SEARCH_WRITE=true), agents get the same lifecycle as tools: create_index provisions an index with everything inferred (pass a sourceTable for the common Delta Sync case, or an embeddingDimension for a direct-access index), and sync_index refreshes a Delta Sync index from its source table.

allowWrite exposes these tools but doesn't approve them. AppKit marks them with effect: "write" and runs them as the requesting user. The Mastra adapter doesn't suspend routine write effects. If a write requires human approval, register a Mastra tool with requireApproval: true that calls SearchClient after the run resumes instead of exposing the write toolkit directly.

Provision an index on boot

ensureOnSetup provisions an index when the app starts, using the native AppKit workspace client. It creates the endpoint and index when absent and seeds documents only when the index is empty. This work runs in the background, so a slow first-time endpoint build does not block the server. The default is a managed direct-access index, so the seed rows are plain objects and search-by-text works without a Delta table or warehouse.

search({
  index: "main.support.docs",
  endpoint: "my-vs-endpoint",
  ensureOnSetup: {
    embeddingModel: "databricks-gte-large-en", // optional; best embedding endpoint otherwise
    documents: [
      { id: "1", title: "Overview", text: "AI Search finds relevant docs.", url: "https://…" },
      { id: "2", title: "Indexes", text: "Delta Sync vs. direct access." },
    ],
    // schema inferred from the first row; primaryKey/textColumn default to id/text
  },
});

Idempotent: later boots see the endpoint, index, and rows already present and do nothing. Point ensureOnSetup.sourceTable at a Delta table to provision a Delta Sync index instead of a managed direct-access one.

Lakebase full-text provider

lakebaseAiSearch() is an AppKit aiSearch provider backed by a Postgres full-text index. It provisions one table per alias (a generated tsvector column with a GIN index), seeds configured documents, and answers queries with a prefix to_tsquery + ts_rank. The pool comes from the native lakebase plugin's routing pool, so database authentication is not re-implemented. HTTP callers cannot widen configured projection columns. columns overrides remain available only to trusted programmatic query() calls. Set an index's auth to on-behalf-of-user for OBO query/write routes or leave it at service-principal.

Queries are compiled from the search box rather than handed to websearch_to_tsquery, which is too literal for type-ahead in two ways:

• Punctuation: Postgres indexes sample-store-intelligence as the compound lexeme plus its parts, but compiles a hyphenated query to the compound alone. As a result, store-intelligence matches nothing while store intelligence matches everything. The query is split on punctuation, so both spellings (and . / _ in a table reference) behave the same.
• Prefixes: Every term is matched as a prefix, so intel reaches intelligence and store intel finds sample-store-intelligence.

All terms must match. When none do, search relaxes instead of returning an empty box: any single term counts, plus a substring pass that catches a fragment which is not a prefix (telligence). The substring pass cannot use the GIN index, so it only runs after the indexed pass finds nothing.

The provider returns AppKit's SearchResponse and mounts the same /api/ai-search/:alias query surface. AppKit UI's useAiSearchQuery and @dbx-tools/ui/search work without a backend-specific client:

import { createApp, lakebase } from "@databricks/appkit";
import { lakebaseAiSearch } from "@dbx-tools/search";

createApp({
  plugins: [
    lakebase(),
    lakebaseAiSearch({
      indexes: {
        docs: {
          indexName: "docs",
          queryType: "full_text",
          columns: ["id", "title", "text"],
          documents: [{ id: "1", title: "Overview", text: "Search over Postgres full-text." }],
        },
      },
    }),
  ],
});

Register either native aiSearch for Vector Search or lakebaseAiSearch for Postgres full text. Both use the registered plugin name aiSearch, so they are alternatives and must not be registered together.

Configuration

All fields are optional. Precedence is plugin config, then environment, then a default.

Modules

• client - SearchClient, SearchIndex, and the Vector Search lifecycle (createIndex / ensureIndex / syncIndex / deleteIndex / listIndexes / ensureEndpoint), plus provider-backed federated reads.
• plugin - SearchPlugin and the search() factory (ToolProvider, routes, clientConfig, exports).
• tool - the searchTool(), universalSearchTool(), addDocumentsTool(), createIndexTool(), and syncIndexTool() Mastra factories.
• index-tools - toCreateIndexOptions, the shared mapping from the create_index wire request onto SearchClient.createIndex options (used by both the Mastra tool and the plugin route).
• config - resolveSearchConfig, resolveIndexName, SEARCH_CONFIG_SCHEMA, the config env constants, and the SearchPluginConfig / ResolvedSearchConfig types.
• native - adapter from an AppKit-compatible aiSearch provider to the extension tools and universal-search client.
• lakebase-plugin - lakebaseAiSearch, the AppKit-compatible PostgreSQL full-text provider.
• lakebase - LakebaseSearchBackend, the provider's tsvector runtime.
• query - toDocumentArray, shared by write routes and tools.
• runtime - createSearchRuntime / setSearchReadBackend for isolated clients.
• schema - the tool descriptions and re-exported request schemas.

Browser-safe schemas live in @dbx-tools/shared-search. The React search box lives in @dbx-tools/ui/search. Model resolution reuses @dbx-tools/model.