AI Search プラグイン

AppKit アプリケーションから、ハイブリッド検索、リランキング、カーソルページネーションを利用して Databricks Vector Search インデックスにクエリを実行します。

主な機能:

• 複数の Vector Search インデックスを名前付きエイリアスで管理
• ハイブリッド、ANN、全文検索の各クエリモード
• 列単位で制御可能なリランキング (オプション)
• 大量の結果セットに対応するカーソルベースのページネーション
• サービスプリンシパル (既定) および on-behalf-of-user 認証
• カスタム embeddingFn による自己管理型の埋め込みインデックス

基本的な使い方

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

設定オプション

インデックスエイリアス

インデックスエイリアスを使うと、複数の Vector Search インデックスを名前で参照できます。エイリアスは API ルートやプログラムからの呼び出しで使用します。

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

クエリタイプ

• hybrid — ベクトル類似度検索とキーワード検索を組み合わせます。汎用的な検索に最適です。
• ann — 埋め込みのみを使用した近似最近傍探索です。意味的な類似性の検索に最適です。
• full_text — 埋め込みを必要としないキーワードベースの検索です。

再ランキング

再ランキングでは、初期候補に対して第 2 段階のモデルを実行し、結果の関連性を高めます。

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

返される全列を対象に再ランキングするには、reranker: true を指定します。

On-behalf-of-user 認証

デフォルトでは、クエリはアプリのサービスプリンシパルとして実行されます。サインイン中のユーザーとしてクエリを実行する場合は、auth: "on-behalf-of-user" を設定します。

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

ページネーション

カーソルベースのページネーションを有効にすると、大きな結果セットをページ単位で取得できます:

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

pagination が true の場合、endpointName は必須です。2ページ目以降を取得するには /:alias/next-page ルートを使用します。

自己管理型の埋め込みインデックス

埋め込みを自身で管理するインデックスの場合は、クエリ文字列を受け取ってベクトルを返す embeddingFn を指定します。

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 ルート

ルートは /api/ai-search にマウントされます。

インデックスをクエリする

レスポンス:

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

各結果には、関連度を示す score と、data 配下に返却された列が含まれます。nextPageToken は、pagination が有効でさらに結果が存在する場合を除き null になります。

次のページを取得する

インデックス設定の取得

エイリアスに対して解決された IndexConfig を返します (embeddingFn を除く) 。

プログラムからのアクセス

この plugin は、サーバーサイドで使用する query メソッドを公開しています。

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

numResults などの呼び出し単位の設定を変更したい場合は、query の第 2 引数にオプションのオーバーライドを渡します。

キャッシュ

クエリ結果は短いTTL (60秒) でキャッシュされます。そのため、同一のクエリが繰り返された場合 (コンポーネントが再レンダリングされたり2回マウントされたりする場合も含む) 、毎回インデックスにアクセスするのではなく、1回のVector Search呼び出しの結果を再利用します。次ページ用のルートはキャッシュされません。ページトークンは1回限りのカーソルであり、すでに特定のページを一意に指し示しているためです。

キャッシュキーには、結果を左右するすべての要素が含まれます。具体的には、解決済みのインデックス、queryText、queryVector (ハッシュ化) 、queryType、numResults、解決済みのcolumns、filters、そして再ランキングの有効・無効です。これらのいずれかが異なるクエリは、別々にキャッシュされます。

ユーザーごとの分離

auth: "on-behalf-of-user" のインデックスでは、呼び出し元のIDがキャッシュキーに含まれます。そのため、あるユーザーが他のユーザーのキャッシュ結果を目にすることはなく、on-behalf-of-userのクエリがサービスプリンシパルによって登録されたエントリを読み取ることもありません。サービスプリンシパルのインデックスでは、すべての呼び出し元が単一のキャッシュエントリを共有します。

React hook

useAiSearchQuery は plugin のクライアント設定から設定済みのインデックスを読み取り、適切な /:alias/query ルートへ POST するため、UI 側でエイリアスをハードコードする必要はありません。インデックスが 1 つだけ設定されている場合、引数は不要です。特定のインデックスを対象にする場合は { alias } を渡します。

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 は、呼び出しごとに細かく制御できるよう、完全なリクエストオブジェクト ({ queryText, numResults, filters, ... }) も受け付けます。フックの indexes フィールドには設定済みのすべてのインデックスが列挙されるため、インデックス選択UIの構築に利用できます。