AI Search plugin

Query Databricks Vector Search indexes with hybrid search, reranking, and cursor pagination from your AppKit application.

Key features:

• Named index aliases for multiple Vector Search indexes
• Hybrid, ANN, and full-text query modes
• Optional reranking with column-level control
• Cursor-based pagination for large result sets
• Service principal (default) and on-behalf-of-user auth
• Self-managed embedding indexes via custom embeddingFn

Basic usage

import { createApp, server } from "@databricks/appkit";
import { aiSearch } from "@databricks/appkit/beta";

await createApp({
  plugins: [
    server(),
    aiSearch({
      indexes: {
        products: {
          indexName: "catalog.schema.products_idx",
          columns: ["id", "name", "description"],
          queryType: "hybrid",
          numResults: 20,
        },
      },
    }),
  ],
});

Configuration options

Index aliases

Index aliases let you reference multiple Vector Search indexes by name. The alias is used in API routes and programmatic calls:

aiSearch({
  indexes: {
    products: {
      indexName: "catalog.schema.products_idx",
      columns: ["id", "name", "description"],
    },
    docs: {
      indexName: "catalog.schema.docs_idx",
      columns: ["id", "title", "content", "url"],
      queryType: "full_text",
    },
  },
});

IndexConfig

Query types

• hybrid — Combines vector similarity and keyword search. Best for general-purpose retrieval.
• ann — Approximate nearest neighbor search using embeddings only. Best for semantic similarity.
• full_text — Keyword-based search with no embedding required.

Reranking

Reranking improves result relevance by running a second-stage model over the initial candidates:

aiSearch({
  indexes: {
    products: {
      indexName: "catalog.schema.products_idx",
      columns: ["id", "name", "description", "category"],
      reranker: { columnsToRerank: ["name", "description"] },
    },
  },
});

Pass reranker: true to rerank across all returned columns.

On-behalf-of-user auth

By default, queries run as the app's service principal. Set auth: "on-behalf-of-user" to execute queries as the signed-in user instead:

aiSearch({
  indexes: {
    documents: {
      indexName: "catalog.schema.documents_idx",
      columns: ["id", "title", "body"],
      auth: "on-behalf-of-user",
    },
  },
});

Pagination

Enable cursor pagination to page through large result sets:

aiSearch({
  indexes: {
    products: {
      indexName: "catalog.schema.products_idx",
      columns: ["id", "name", "description"],
      pagination: true,
      endpointName: "my-vector-search-endpoint",
    },
  },
});

endpointName is required when pagination is true. Use the /:alias/next-page route to fetch subsequent pages.

Self-managed embedding indexes

For indexes that manage their own embeddings, provide an embeddingFn that takes a query string and returns a vector:

import { embed } from "./my-embedding-client";

aiSearch({
  indexes: {
    products: {
      indexName: "catalog.schema.products_idx",
      columns: ["id", "name", "description"],
      queryType: "ann",
      embeddingFn: (text) => embed(text),
    },
  },
});

HTTP routes

Routes are mounted at /api/ai-search.

Query an index

Response:

{
  "results": [
    {
      "score": 0.87,
      "data": { "id": "42", "name": "Intro to ML", "description": "..." }
    }
  ],
  "totalCount": 1,
  "queryTimeMs": 35,
  "queryType": "hybrid",
  "nextPageToken": "eyJvZmZzZXQiOjEwfQ=="
}

Each result carries its relevance score and the returned columns under data. nextPageToken is null unless pagination is enabled and more results are available.

Fetch the next page

Get index config

Returns the resolved IndexConfig for the alias (excluding embeddingFn).

Programmatic access

The plugin exposes a query method for server-side use:

import { createApp, server } from "@databricks/appkit";
import { aiSearch } from "@databricks/appkit/beta";

const AppKit = await createApp({
  plugins: [
    server(),
    aiSearch({
      indexes: {
        products: {
          indexName: "catalog.schema.products_idx",
          columns: ["id", "name", "description"],
        },
      },
    }),
  ],
});

const result = await AppKit.aiSearch.query("products", {
  queryText: "machine learning guide",
});

console.log(result.results);

Pass optional overrides as a second argument to query to adjust numResults or other per-call settings.

Caching

Query results are cached with a short TTL (60s) so repeated identical queries — including a component that re-renders or mounts twice — reuse a single Vector Search call instead of hitting the index each time. The next-page route is not cached: a page token is a single-use cursor and already identifies the exact page.

The cache key covers everything that changes results: the resolved index, queryText, queryVector (hashed), queryType, numResults, the resolved columns, filters, and whether reranking is on. Two queries that differ in any of these are cached separately.

Per-user isolation

For auth: "on-behalf-of-user" indexes the caller's identity is part of the cache key, so one user never sees another user's cached results — and an on-behalf-of-user query never reads a service-principal-populated entry. Service-principal indexes share a single cache entry across callers.

React hook

useAiSearchQuery reads the configured indexes from the plugin's client config and posts to the right /:alias/query route, so the UI never hardcodes an alias. With one index configured it needs no arguments; pass { alias } to target a specific one.

import { useAiSearchQuery } from "@databricks/appkit-ui/react/beta";

function Search() {
  const { search, data, loading, error } = useAiSearchQuery();

  return (
    <>
      <input onKeyDown={(e) => e.key === "Enter" && search(e.currentTarget.value)} />
      {error && <p>{error}</p>}
      {data?.results.map((r, i) => (
        <div key={i}>{JSON.stringify(r.data)}</div>
      ))}
    </>
  );
}

search also accepts a full request object ({ queryText, numResults, filters, ... }) for per-call control. The hook's indexes field lists every configured index, which you can use to build an index picker.