Model Serving plugin

Databricks Model Serving の endpoint への認証済みプロキシを提供し、invoke とストリーミングに対応します。

主な機能:

• 複数のサービング endpoint を名前付きエイリアスで指定可能
• 非ストリーミング (invoke) と SSE ストリーミング (stream) の呼び出し
• リクエスト / レスポンススキーマの OpenAPI 型を自動生成
• endpoint スキーマに基づくリクエストボディのフィルタリング
• ユーザーの代理実行 (On-behalf-of、OBO)

基本的な使い方

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

await createApp({
  plugins: [
    server(),
    serving(),
  ],
});

設定を行わない場合、plugin は環境変数 DATABRICKS_SERVING_ENDPOINT_NAME を読み取り、default というエイリアスで登録します。

設定オプション

Endpoint エイリアス

Endpoint エイリアスを使うと、複数のサービング endpoint を名前で参照できます。

serving({
  endpoints: {
    llm: { env: "DATABRICKS_SERVING_ENDPOINT_NAME" },
    classifier: { env: "DATABRICKS_SERVING_ENDPOINT_CLASSIFIER" },
  },
})

各エイリアスは、実際の endpoint 名を保持する環境変数にマッピングされます。1 つの endpoint が複数のモデルを提供している場合は、servedModel を指定することでトラフィックルーティングを回避し、特定のモデルを直接指定できます。

serving({
  endpoints: {
    llm: { env: "DATABRICKS_SERVING_ENDPOINT_NAME", servedModel: "llama-v2" },
  },
})

型生成

appKitServingTypesPlugin() Vite plugin は、サービング endpoint の OpenAPI スキーマから TypeScript の型を生成します。手動の setup は不要です。AppKit の開発サーバーにはこの plugin が自動的に含まれています。

この plugin は、サーバーファイル (server/index.ts または server/server.ts) から endpoint の設定を自動検出します。

生成される型により、次のことが可能になります。

• バックエンド (AppKit.serving("alias")) とフロントエンドのフック (useServingStream、useServingInvoke) の両方でのエイリアスの自動補完
• OpenAPI スキーマに基づく endpoint ごとの型付きリクエスト／レスポンス／チャンク

endpoint の OpenAPI スキーマが利用できない場合 (deploy されていない、環境変数が未設定など) 、plugin は汎用的なフォールバック型を生成します。この場合でも endpoint は利用可能で、リクエスト／レスポンスが型付けされないだけです。

環境変数

名前付き endpoint を使用する場合は、エイリアスごとにカスタム環境変数を定義します (例: DATABRICKS_SERVING_ENDPOINT_CLASSIFIER) 。

実行コンテキスト

すべてのサービングルートは、デフォルトで認証済みユーザーの代理 (OBO) として実行されます。これは Genie および Files plugin と同じ挙動です。これにより、サービング endpoint 上でユーザーごとの CAN_QUERY permissions が確実に適用されます。

exports() を使ったプログラムからのアクセスでは、.asUser(req) を指定してユーザーコンテキストで実行します:

// サービスプリンシパルのコンテキスト（デフォルト）
const result = await AppKit.serving("llm").invoke({ messages });

// ユーザーコンテキスト（ルートハンドラーでの利用を推奨）
const result = await AppKit.serving("llm").asUser(req).invoke({ messages });

HTTP endpoint

名前付きモード (endpoints 設定を使用)

• POST /api/serving/:alias/invoke — 非ストリーミング呼び出し
• POST /api/serving/:alias/stream — SSE ストリーミング呼び出し

デフォルトモード (endpoints 設定なし)

• POST /api/serving/invoke — 非ストリーミング呼び出し
• POST /api/serving/stream — SSE ストリーミング呼び出し

リクエスト形式

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

このpluginは、サーバーサイドで利用できる invoke メソッドと stream メソッドをエクスポートします。

const AppKit = await createApp({
  plugins: [
    server(),
    serving({
      endpoints: {
        llm: { env: "DATABRICKS_SERVING_ENDPOINT_NAME" },
      },
    }),
  ],
});

// 非ストリーミング
const result = await AppKit.serving("llm").invoke({
  messages: [{ role: "user", content: "Hello" }],
});

// ストリーミング
for await (const chunk of AppKit.serving("llm").stream({
  messages: [{ role: "user", content: "Hello" }],
})) {
  console.log(chunk);
}

フロントエンドフック

@databricks/appkit-ui パッケージは、サービング endpoint 向けの React フックを提供します。

useServingStream

SSE によるストリーミング呼び出し:

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

function ChatStream() {
  const { stream, chunks, streaming, error, reset } = useServingStream(
    { messages: [{ role: "user", content: "Hello" }] },
    {
      alias: "llm",
      onComplete: (finalChunks) => {
        // ストリーム終了時に、蓄積されたすべてのチャンクを引数として呼び出される
        console.log("Stream done, got", finalChunks.length, "chunks");
      },
    },
  );

  return (
    <>
      <button onClick={stream} disabled={streaming}>Send</button>
      <button onClick={reset}>Reset</button>
      {chunks.map((chunk, i) => <pre key={i}>{JSON.stringify(chunk)}</pre>)}
      {error && <p>{error}</p>}
    </>
  );
}

useServingInvoke

非ストリーミングの呼び出しです。invoke() は、レスポンスデータ (エラー時は null) を返す Promise を返します。

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

function Classify() {
  const { invoke, data, loading, error } = useServingInvoke(
    { inputs: ["sample text"] },
    { alias: "classifier" },
  );

  async function handleClick() {
    const result = await invoke();
    if (result) {
      console.log("Classification result:", result);
    }
  }

  return (
    <>
      <button onClick={handleClick} disabled={loading}>Classify</button>
      {data && <pre>{JSON.stringify(data)}</pre>}
      {error && <p>{error}</p>}
    </>
  );
}

どちらのフックも autoStart: true を指定することで、マウント時に自動的に呼び出されます。