diff --git a/apps/dev-playground/app.yaml b/apps/dev-playground/app.yaml index 36da4ffa..85d20f65 100644 --- a/apps/dev-playground/app.yaml +++ b/apps/dev-playground/app.yaml @@ -5,3 +5,5 @@ env: valueFrom: volume - name: DATABRICKS_VOLUME_OTHER valueFrom: other-volume + - name: DATABRICKS_VS_INDEX_NAME + valueFrom: vs-index diff --git a/apps/dev-playground/client/src/routes/__root.tsx b/apps/dev-playground/client/src/routes/__root.tsx index 35a2282b..bc7f1e34 100644 --- a/apps/dev-playground/client/src/routes/__root.tsx +++ b/apps/dev-playground/client/src/routes/__root.tsx @@ -112,6 +112,14 @@ function RootComponent() { Serving + + + diff --git a/apps/dev-playground/client/src/routes/vector-search.route.tsx b/apps/dev-playground/client/src/routes/vector-search.route.tsx new file mode 100644 index 00000000..bed4885c --- /dev/null +++ b/apps/dev-playground/client/src/routes/vector-search.route.tsx @@ -0,0 +1,154 @@ +import { + Button, + Card, + CardContent, + CardHeader, + CardTitle, + Input, +} from "@databricks/appkit-ui/react"; +import { createFileRoute } from "@tanstack/react-router"; +import { Search } from "lucide-react"; +import { useState } from "react"; +import { Header } from "@/components/layout/header"; + +export const Route = createFileRoute("/vector-search")({ + component: VectorSearchRoute, +}); + +interface SearchResult { + score: number; + data: Record; +} + +interface SearchResponse { + results: SearchResult[]; + totalCount: number; + queryTimeMs: number; + queryType: string; +} + +function VectorSearchRoute() { + const [query, setQuery] = useState(""); + const [loading, setLoading] = useState(false); + const [error, setError] = useState(null); + const [response, setResponse] = useState(null); + + const handleSearch = async () => { + if (!query.trim()) return; + setLoading(true); + setError(null); + setResponse(null); + + try { + const res = await fetch("/api/vector-search/demo/query", { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify({ queryText: query }), + }); + + if (!res.ok) { + const data = await res.json().catch(() => ({})); + throw new Error(data.error ?? `HTTP ${res.status}: ${res.statusText}`); + } + + const data: SearchResponse = await res.json(); + setResponse(data); + } catch (err) { + setError(err instanceof Error ? err.message : String(err)); + } finally { + setLoading(false); + } + }; + + const handleKeyDown = (e: React.KeyboardEvent) => { + if (e.key === "Enter") { + void handleSearch(); + } + }; + + return ( +
+
+
+ +
+
+ setQuery(e.target.value)} + onKeyDown={handleKeyDown} + className="flex-1" + /> + +
+ + {error && ( +
+ {error} +
+ )} + + {response && ( +
+
+ {response.totalCount} result + {response.totalCount !== 1 ? "s" : ""} ·{" "} + {response.queryTimeMs}ms · {response.queryType} +
+ + {response.results.length === 0 ? ( +

+ No results found. +

+ ) : ( + response.results.map((result, index) => ( + + + + Result {index + 1} + + score: {result.score.toFixed(4)} + + + + +
+ {Object.entries(result.data).map(([key, value]) => ( +
+
+ {key} +
+
+ {String(value ?? "")} +
+
+ ))} +
+ + + )) + )} +
+ )} +
+
+
+ ); +} diff --git a/apps/dev-playground/server/index.ts b/apps/dev-playground/server/index.ts index af05b11f..96c1f94f 100644 --- a/apps/dev-playground/server/index.ts +++ b/apps/dev-playground/server/index.ts @@ -8,6 +8,7 @@ import { serving, } from "@databricks/appkit"; import { WorkspaceClient } from "@databricks/sdk-experimental"; +import { vectorSearch } from "../../../packages/appkit/src/plugins/vector-search"; import { lakebaseExamples } from "./lakebase-examples-plugin"; import { reconnect } from "./reconnect-plugin"; import { telemetryExamples } from "./telemetry-example-plugin"; @@ -34,6 +35,16 @@ createApp({ lakebaseExamples(), files(), serving(), + vectorSearch({ + indexes: { + demo: { + indexName: + process.env.DATABRICKS_VS_INDEX_NAME ?? "catalog.schema.index", + columns: ["id", "text", "title"], + queryType: "hybrid", + }, + }, + }), ], ...(process.env.APPKIT_E2E_TEST && { client: createMockClient() }), }).then((appkit) => { diff --git a/docs/docs/plugins/vector-search.md b/docs/docs/plugins/vector-search.md new file mode 100644 index 00000000..5d704641 --- /dev/null +++ b/docs/docs/plugins/vector-search.md @@ -0,0 +1,239 @@ +--- +sidebar_position: 9 +--- + +# Vector 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 + +```ts +import { createApp, vectorSearch, server } from "@databricks/appkit"; + +await createApp({ + plugins: [ + server(), + vectorSearch({ + indexes: { + products: { + indexName: "catalog.schema.products_idx", + columns: ["id", "name", "description"], + queryType: "hybrid", + numResults: 20, + }, + }, + }), + ], +}); +``` + +## Configuration options + +| Option | Type | Default | Description | +|--------|------|---------|-------------| +| `indexes` | `Record` | — | **Required.** Map of alias names to index configurations | +| `timeout` | `number` | `30000` | Query timeout in ms | + +### Index aliases + +Index aliases let you reference multiple Vector Search indexes by name. The alias is used in API routes and programmatic calls: + +```ts +vectorSearch({ + 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 + +| Field | Type | Default | Description | +|-------|------|---------|-------------| +| `indexName` | `string` | — | **Required.** Three-level Unity Catalog name (`catalog.schema.index`) | +| `columns` | `string[]` | — | **Required.** Columns to return in query results | +| `queryType` | `"ann" \| "hybrid" \| "full_text"` | `"hybrid"` | Search mode | +| `numResults` | `number` | `20` | Maximum results per query | +| `reranker` | `boolean \| { columnsToRerank: string[] }` | — | Enable reranking. Pass `true` to rerank all result columns, or specify a subset | +| `auth` | `"service-principal" \| "on-behalf-of-user"` | `"service-principal"` | Authentication mode for query execution | +| `pagination` | `boolean` | — | Enable cursor-based pagination | +| `endpointName` | `string` | — | Vector Search endpoint name. Required when `pagination` is `true` | +| `embeddingFn` | `(text: string) => Promise` | — | Custom embedding function for self-managed embedding indexes | + +### 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: + +```ts +vectorSearch({ + 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: + +```ts +vectorSearch({ + 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: + +```ts +vectorSearch({ + 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: + +```ts +import { embed } from "./my-embedding-client"; + +vectorSearch({ + indexes: { + products: { + indexName: "catalog.schema.products_idx", + columns: ["id", "name", "description"], + queryType: "ann", + embeddingFn: (text) => embed(text), + }, + }, +}); +``` + +## HTTP routes + +Routes are mounted at `/api/vector-search`. + +| Method | Path | Description | +|--------|------|-------------| +| `POST` | `/:alias/query` | Query an index by alias | +| `POST` | `/:alias/next-page` | Fetch the next page of results (requires `pagination: true`) | +| `GET` | `/:alias/config` | Return the resolved config for an index alias | + +### Query an index + +``` +POST /api/vector-search/:alias/query +Content-Type: application/json + +{ + "queryText": "machine learning guide", + "numResults": 10 +} +``` + +Response: + +```json +{ + "results": [ + { "id": "42", "name": "Intro to ML", "description": "..." } + ], + "nextPageToken": "eyJvZmZzZXQiOjEwfQ==" +} +``` + +`nextPageToken` is only present when `pagination` is enabled and more results are available. + +### Fetch the next page + +``` +POST /api/vector-search/:alias/next-page +Content-Type: application/json + +{ + "queryText": "machine learning guide", + "pageToken": "eyJvZmZzZXQiOjEwfQ==" +} +``` + +### Get index config + +``` +GET /api/vector-search/:alias/config +``` + +Returns the resolved `IndexConfig` for the alias (excluding `embeddingFn`). + +## Programmatic access + +The plugin exposes a `query` method for server-side use: + +```ts +const AppKit = await createApp({ + plugins: [ + server(), + vectorSearch({ + indexes: { + products: { + indexName: "catalog.schema.products_idx", + columns: ["id", "name", "description"], + }, + }, + }), + ], +}); + +const result = await AppKit.vectorSearch.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. diff --git a/knip.json b/knip.json index e8eb1eb3..b777d8c2 100644 --- a/knip.json +++ b/knip.json @@ -16,6 +16,7 @@ "**/*.generated.ts", "**/*.example.tsx", "**/*.css", + "packages/appkit/src/plugins/vector-search/**", "template/**", "tools/**", "docs/**" diff --git a/packages/appkit/src/connectors/index.ts b/packages/appkit/src/connectors/index.ts index 41e7748c..54a24fa4 100644 --- a/packages/appkit/src/connectors/index.ts +++ b/packages/appkit/src/connectors/index.ts @@ -3,3 +3,4 @@ export * from "./genie"; export * from "./lakebase"; export * from "./lakebase-v1"; export * from "./sql-warehouse"; +export * from "./vector-search"; diff --git a/packages/appkit/src/connectors/vector-search/client.ts b/packages/appkit/src/connectors/vector-search/client.ts new file mode 100644 index 00000000..ab45b78d --- /dev/null +++ b/packages/appkit/src/connectors/vector-search/client.ts @@ -0,0 +1,176 @@ +import type { WorkspaceClient } from "@databricks/sdk-experimental"; +import { createLogger } from "../../logging/logger"; +import type { TelemetryProvider } from "../../telemetry"; +import { + type Span, + SpanKind, + SpanStatusCode, + TelemetryManager, +} from "../../telemetry"; +import type { + VectorSearchConnectorConfig, + VsNextPageParams, + VsQueryParams, + VsRawResponse, +} from "./types"; + +const logger = createLogger("connectors:vector-search"); + +export class VectorSearchConnector { + private readonly telemetry: TelemetryProvider; + + constructor(config: VectorSearchConnectorConfig = {}) { + this.telemetry = TelemetryManager.getProvider( + "vector-search", + config.telemetry, + ); + } + + async query( + workspaceClient: WorkspaceClient, + params: VsQueryParams, + signal?: AbortSignal, + ): Promise { + if (signal?.aborted) { + throw new Error("Query cancelled before execution"); + } + + const body: Record = { + columns: params.columns, + num_results: params.numResults, + query_type: params.queryType.toUpperCase(), + debug_level: 1, + }; + + if (params.queryText) body.query_text = params.queryText; + if (params.queryVector) body.query_vector = params.queryVector; + if (params.filters && Object.keys(params.filters).length > 0) { + body.filters = params.filters; + } + if (params.reranker) { + body.reranker = { + model: "databricks_reranker", + parameters: { columns_to_rerank: params.reranker.columnsToRerank }, + }; + } + + logger.debug( + "Querying VS index %s (type=%s, num_results=%d)", + params.indexName, + params.queryType, + params.numResults, + ); + + return this.telemetry.startActiveSpan( + "vector-search.query", + { + kind: SpanKind.CLIENT, + attributes: { + "db.system": "databricks", + "vs.index_name": params.indexName, + "vs.query_type": params.queryType, + "vs.num_results": params.numResults, + "vs.has_filters": !!( + params.filters && Object.keys(params.filters).length > 0 + ), + "vs.has_reranker": !!params.reranker, + }, + }, + async (span: Span) => { + const startTime = Date.now(); + try { + const response = (await workspaceClient.apiClient.request({ + method: "POST", + path: `/api/2.0/vector-search/indexes/${params.indexName}/query`, + body, + headers: new Headers({ "Content-Type": "application/json" }), + raw: false, + query: {}, + })) as VsRawResponse; + + const duration = Date.now() - startTime; + span.setAttribute("vs.result_count", response.result.row_count); + span.setAttribute( + "vs.query_time_ms", + response.debug_info?.response_time ?? 0, + ); + span.setAttribute("vs.duration_ms", duration); + span.setStatus({ code: SpanStatusCode.OK }); + + logger.event()?.setContext("vector-search", { + index_name: params.indexName, + query_type: params.queryType, + result_count: response.result.row_count, + query_time_ms: response.debug_info?.response_time ?? 0, + duration_ms: duration, + }); + + return response; + } catch (error) { + span.recordException(error as Error); + span.setStatus({ + code: SpanStatusCode.ERROR, + message: error instanceof Error ? error.message : String(error), + }); + throw error; + } + }, + { name: "vector-search", includePrefix: true }, + ); + } + + async queryNextPage( + workspaceClient: WorkspaceClient, + params: VsNextPageParams, + signal?: AbortSignal, + ): Promise { + if (signal?.aborted) { + throw new Error("Query cancelled before execution"); + } + + logger.debug( + "Fetching next page for index %s (endpoint=%s)", + params.indexName, + params.endpointName, + ); + + return this.telemetry.startActiveSpan( + "vector-search.queryNextPage", + { + kind: SpanKind.CLIENT, + attributes: { + "db.system": "databricks", + "vs.index_name": params.indexName, + "vs.endpoint_name": params.endpointName, + }, + }, + async (span: Span) => { + try { + const response = (await workspaceClient.apiClient.request({ + method: "POST", + path: `/api/2.0/vector-search/indexes/${params.indexName}/query-next-page`, + body: { + endpoint_name: params.endpointName, + page_token: params.pageToken, + }, + headers: new Headers({ "Content-Type": "application/json" }), + raw: false, + query: {}, + })) as VsRawResponse; + + span.setAttribute("vs.result_count", response.result.row_count); + span.setStatus({ code: SpanStatusCode.OK }); + return response; + } catch (error) { + span.recordException(error as Error); + span.setStatus({ + code: SpanStatusCode.ERROR, + message: error instanceof Error ? error.message : String(error), + }); + throw error; + } + }, + { name: "vector-search", includePrefix: true }, + ); + } +} diff --git a/packages/appkit/src/connectors/vector-search/index.ts b/packages/appkit/src/connectors/vector-search/index.ts new file mode 100644 index 00000000..d2ec2302 --- /dev/null +++ b/packages/appkit/src/connectors/vector-search/index.ts @@ -0,0 +1,2 @@ +export * from "./client"; +export * from "./types"; diff --git a/packages/appkit/src/connectors/vector-search/types.ts b/packages/appkit/src/connectors/vector-search/types.ts new file mode 100644 index 00000000..df042e8c --- /dev/null +++ b/packages/appkit/src/connectors/vector-search/types.ts @@ -0,0 +1,42 @@ +import type { TelemetryOptions } from "shared"; + +export interface VectorSearchConnectorConfig { + timeout?: number; + telemetry?: TelemetryOptions; +} + +export interface VsQueryParams { + indexName: string; + queryText?: string; + queryVector?: number[]; + columns: string[]; + numResults: number; + queryType: "ann" | "hybrid" | "full_text"; + filters?: Record; + reranker?: { columnsToRerank: string[] }; +} + +export interface VsNextPageParams { + indexName: string; + endpointName: string; + pageToken: string; +} + +export interface VsRawResponse { + manifest: { + column_count: number; + columns: Array<{ name: string; type?: string }>; + }; + result: { + row_count: number; + data_array: unknown[][]; + }; + next_page_token?: string | null; + debug_info?: { + response_time?: number; + ann_time?: number; + embedding_gen_time?: number; + latency_ms?: number; + [key: string]: unknown; + }; +} diff --git a/packages/appkit/src/plugins/vector-search/defaults.ts b/packages/appkit/src/plugins/vector-search/defaults.ts new file mode 100644 index 00000000..c02b6e80 --- /dev/null +++ b/packages/appkit/src/plugins/vector-search/defaults.ts @@ -0,0 +1,7 @@ +import type { PluginExecuteConfig } from "shared"; + +export const vectorSearchDefaults: PluginExecuteConfig = { + cache: { enabled: false }, + retry: { enabled: true, initialDelay: 1000, attempts: 3 }, + timeout: 30_000, +}; diff --git a/packages/appkit/src/plugins/vector-search/index.ts b/packages/appkit/src/plugins/vector-search/index.ts new file mode 100644 index 00000000..9052cb03 --- /dev/null +++ b/packages/appkit/src/plugins/vector-search/index.ts @@ -0,0 +1,2 @@ +export * from "./vector-search"; +export * from "./types"; diff --git a/packages/appkit/src/plugins/vector-search/manifest.json b/packages/appkit/src/plugins/vector-search/manifest.json new file mode 100644 index 00000000..ed580ecf --- /dev/null +++ b/packages/appkit/src/plugins/vector-search/manifest.json @@ -0,0 +1,69 @@ +{ + "$schema": "https://databricks.github.io/appkit/schemas/plugin-manifest.schema.json", + "name": "vector-search", + "displayName": "Vector Search Plugin", + "description": "Query Databricks Vector Search indexes with built-in hybrid search, reranking, and pagination", + "resources": { + "required": [ + { + "type": "vector_search_index", + "alias": "Vector Search Index", + "resourceKey": "vector-search-index", + "description": "A Databricks Vector Search index to query. Index names configured via plugin config.", + "permission": "SELECT", + "fields": { + "indexName": { + "env": "DATABRICKS_VS_INDEX_NAME", + "description": "Three-level UC name of the default index (catalog.schema.index_name)" + }, + "endpointName": { + "env": "DATABRICKS_VS_ENDPOINT_NAME", + "description": "Vector Search endpoint name (required for pagination)" + } + } + } + ], + "optional": [] + }, + "config": { + "schema": { + "type": "object", + "properties": { + "indexes": { + "type": "object", + "description": "Map of alias names to index configurations", + "additionalProperties": { + "type": "object", + "properties": { + "indexName": { + "type": "string", + "description": "Three-level UC name: catalog.schema.index_name" + }, + "columns": { + "type": "array", + "items": { "type": "string" }, + "description": "Columns to return in results" + }, + "queryType": { + "type": "string", + "enum": ["ann", "hybrid", "full_text"], + "default": "hybrid" + }, + "numResults": { + "type": "number", + "default": 20 + } + }, + "required": ["indexName", "columns"] + } + }, + "timeout": { + "type": "number", + "default": 30000, + "description": "Query execution timeout in milliseconds" + } + }, + "required": ["indexes"] + } + } +} diff --git a/packages/appkit/src/plugins/vector-search/tests/vector-search.test.ts b/packages/appkit/src/plugins/vector-search/tests/vector-search.test.ts new file mode 100644 index 00000000..d3c41f25 --- /dev/null +++ b/packages/appkit/src/plugins/vector-search/tests/vector-search.test.ts @@ -0,0 +1,329 @@ +import { beforeEach, describe, expect, it, vi } from "vitest"; + +vi.mock("../../../context", () => ({ + getWorkspaceClient: vi.fn(() => mockWorkspaceClient), + getCurrentUserId: vi.fn(() => "test-user"), +})); + +vi.mock("../../../logging/logger", () => ({ + createLogger: () => ({ + debug: vi.fn(), + info: vi.fn(), + warn: vi.fn(), + error: vi.fn(), + event: () => ({ + setComponent: vi.fn().mockReturnThis(), + setContext: vi.fn().mockReturnThis(), + }), + }), +})); + +vi.mock("../../../telemetry", () => ({ + TelemetryManager: { + getProvider: () => ({ + getTracer: () => ({}), + getMeter: () => ({ + createCounter: () => ({ add: vi.fn() }), + createHistogram: () => ({ record: vi.fn() }), + }), + startActiveSpan: vi.fn( + ( + _name: string, + _opts: unknown, + fn: (...args: unknown[]) => unknown, + _telemetryOpts?: unknown, + ) => + fn({ + setAttribute: vi.fn(), + setStatus: vi.fn(), + recordException: vi.fn(), + }), + ), + }), + }, + SpanKind: { CLIENT: 3 }, + SpanStatusCode: { OK: 1, ERROR: 2 }, + normalizeTelemetryOptions: () => ({ traces: false, metrics: false }), +})); + +vi.mock("../../../cache", () => ({ + CacheManager: { + getInstanceSync: () => ({ + get: vi.fn(), + set: vi.fn(), + delete: vi.fn(), + generateKey: vi.fn(() => "test-key"), + }), + }, +})); + +vi.mock("../../../app", () => ({ + AppManager: vi.fn().mockImplementation(() => ({})), +})); + +vi.mock("../../../plugin/dev-reader", () => ({ + DevFileReader: { + getInstance: () => ({}), + }, +})); + +vi.mock("../../../stream", () => ({ + StreamManager: vi.fn().mockImplementation(() => ({ + abortAll: vi.fn(), + stream: vi.fn(), + })), +})); + +const validVsResponse = { + manifest: { + column_count: 3, + columns: [{ name: "id" }, { name: "title" }, { name: "score" }], + }, + result: { + row_count: 2, + data_array: [ + [1, "ML Guide", 0.95], + [2, "AI Primer", 0.87], + ], + }, + next_page_token: null, + debug_info: { response_time: 35 }, +}; + +const mockRequest = vi.fn().mockResolvedValue(validVsResponse); +const mockWorkspaceClient = { + apiClient: { request: mockRequest }, +}; + +import { VectorSearchPlugin } from "../vector-search"; + +describe("VectorSearchPlugin", () => { + beforeEach(() => { + mockRequest.mockClear(); + mockRequest.mockResolvedValue(validVsResponse); + }); + + describe("setup()", () => { + it("throws if any index is missing indexName", async () => { + const plugin = new VectorSearchPlugin({ + indexes: { + test: { indexName: "", columns: ["id"] }, + }, + }); + await expect(plugin.setup()).rejects.toThrow("indexName"); + }); + + it("throws if any index is missing columns", async () => { + const plugin = new VectorSearchPlugin({ + indexes: { + test: { indexName: "cat.sch.idx", columns: [] }, + }, + }); + await expect(plugin.setup()).rejects.toThrow("columns"); + }); + + it("throws if pagination enabled but no endpointName", async () => { + const plugin = new VectorSearchPlugin({ + indexes: { + test: { + indexName: "cat.sch.idx", + columns: ["id"], + pagination: true, + }, + }, + }); + await expect(plugin.setup()).rejects.toThrow("endpointName"); + }); + + it("succeeds with valid config", async () => { + const plugin = new VectorSearchPlugin({ + indexes: { + products: { + indexName: "cat.sch.products_idx", + columns: ["id", "name", "description"], + queryType: "hybrid", + numResults: 20, + }, + }, + }); + await expect(plugin.setup()).resolves.not.toThrow(); + }); + }); + + describe("manifest", () => { + it("has correct name", () => { + expect(VectorSearchPlugin.manifest.name).toBe("vector-search"); + }); + }); + + describe("exports()", () => { + it("returns object with query function", () => { + const plugin = new VectorSearchPlugin({ + indexes: { + test: { indexName: "cat.sch.idx", columns: ["id"] }, + }, + }); + const exports = plugin.exports(); + expect(exports).toHaveProperty("query"); + expect(typeof exports.query).toBe("function"); + }); + }); + + describe("query()", () => { + it("calls VS API via connector and parses response", async () => { + const plugin = new VectorSearchPlugin({ + indexes: { + products: { + indexName: "cat.sch.products", + columns: ["id", "title"], + queryType: "hybrid", + }, + }, + }); + await plugin.setup(); + + const result = await plugin.query("products", { + queryText: "machine learning", + }); + + expect(result.results).toHaveLength(2); + expect(result.results[0].score).toBe(0.95); + expect(result.results[0].data).toEqual({ id: 1, title: "ML Guide" }); + expect(result.results[1].score).toBe(0.87); + expect(result.totalCount).toBe(2); + expect(result.queryTimeMs).toBe(35); + }); + + it("constructs correct API request", async () => { + const plugin = new VectorSearchPlugin({ + indexes: { + test: { + indexName: "cat.sch.idx", + columns: ["id", "title"], + queryType: "hybrid", + numResults: 10, + }, + }, + }); + await plugin.setup(); + await plugin.query("test", { queryText: "test query" }); + + expect(mockRequest).toHaveBeenCalledWith( + expect.objectContaining({ + method: "POST", + path: "/api/2.0/vector-search/indexes/cat.sch.idx/query", + }), + ); + + const callBody = mockRequest.mock.calls[0][0].body; + expect(callBody.query_text).toBe("test query"); + expect(callBody.query_type).toBe("HYBRID"); + expect(callBody.num_results).toBe(10); + expect(callBody.columns).toEqual(["id", "title"]); + }); + + it("throws Error for unknown alias", async () => { + const plugin = new VectorSearchPlugin({ + indexes: { + test: { indexName: "cat.sch.idx", columns: ["id"] }, + }, + }); + await plugin.setup(); + + await expect( + plugin.query("unknown", { queryText: "test" }), + ).rejects.toThrow('No index configured with alias "unknown"'); + }); + + it("includes filters when provided", async () => { + const plugin = new VectorSearchPlugin({ + indexes: { + test: { + indexName: "cat.sch.idx", + columns: ["id", "title"], + }, + }, + }); + await plugin.setup(); + await plugin.query("test", { + queryText: "test", + filters: { category: ["books"] }, + }); + + const callBody = mockRequest.mock.calls[0][0].body; + expect(callBody.filters).toEqual({ category: ["books"] }); + }); + + it("includes reranker config when enabled on index", async () => { + const plugin = new VectorSearchPlugin({ + indexes: { + test: { + indexName: "cat.sch.idx", + columns: ["id", "title", "desc"], + reranker: true, + }, + }, + }); + await plugin.setup(); + await plugin.query("test", { queryText: "test" }); + + const callBody = mockRequest.mock.calls[0][0].body; + expect(callBody.reranker.model).toBe("databricks_reranker"); + expect(callBody.reranker.parameters.columns_to_rerank).toEqual([ + "title", + "desc", + ]); + }); + + it("calls embeddingFn for self-managed indexes", async () => { + const mockEmbeddingFn = vi.fn().mockResolvedValue([0.1, 0.2, 0.3]); + const plugin = new VectorSearchPlugin({ + indexes: { + test: { + indexName: "cat.sch.idx", + columns: ["id", "title"], + embeddingFn: mockEmbeddingFn, + }, + }, + }); + await plugin.setup(); + await plugin.query("test", { queryText: "test" }); + + expect(mockEmbeddingFn).toHaveBeenCalledWith("test"); + const callBody = mockRequest.mock.calls[0][0].body; + expect(callBody.query_vector).toEqual([0.1, 0.2, 0.3]); + expect(callBody.query_text).toBeUndefined(); + }); + + it("throws when embeddingFn fails", async () => { + const mockEmbeddingFn = vi + .fn() + .mockRejectedValue(new Error("embedding service unavailable")); + const plugin = new VectorSearchPlugin({ + indexes: { + test: { + indexName: "cat.sch.idx", + columns: ["id", "title"], + embeddingFn: mockEmbeddingFn, + }, + }, + }); + await plugin.setup(); + + await expect(plugin.query("test", { queryText: "test" })).rejects.toThrow( + "Embedding generation failed", + ); + }); + }); + + describe("shutdown()", () => { + it("does not throw", async () => { + const plugin = new VectorSearchPlugin({ + indexes: { + test: { indexName: "cat.sch.idx", columns: ["id"] }, + }, + }); + await expect(plugin.shutdown()).resolves.not.toThrow(); + }); + }); +}); diff --git a/packages/appkit/src/plugins/vector-search/types.ts b/packages/appkit/src/plugins/vector-search/types.ts new file mode 100644 index 00000000..e9380814 --- /dev/null +++ b/packages/appkit/src/plugins/vector-search/types.ts @@ -0,0 +1,67 @@ +import type { BasePluginConfig } from "shared"; + +export interface IVectorSearchConfig extends BasePluginConfig { + timeout?: number; + indexes: Record; +} + +export interface IndexConfig { + /** Three-level UC name: catalog.schema.index_name */ + indexName: string; + /** Columns to return in results */ + columns: string[]; + /** Default search mode */ + queryType?: "ann" | "hybrid" | "full_text"; + /** Max results per query */ + numResults?: number; + /** Enable built-in reranker. Pass true to rerank all non-id columns, or an object for fine control. */ + reranker?: boolean | RerankerConfig; + /** Auth mode — "service-principal" uses the app's SP, "on-behalf-of-user" proxies the logged-in user's token */ + auth?: "service-principal" | "on-behalf-of-user"; + /** Enable cursor pagination */ + pagination?: boolean; + /** VS endpoint name (required when pagination is true) */ + endpointName?: string; + /** + * For self-managed embedding indexes: converts query text to an embedding vector. + * When provided, the plugin calls this function and sends query_vector to VS. + * When omitted, query_text is sent and VS computes embeddings server-side (managed mode). + */ + embeddingFn?: (text: string) => Promise; +} + +export interface RerankerConfig { + columnsToRerank: string[]; +} + +export type SearchFilters = Record< + string, + string | number | boolean | (string | number)[] +>; + +export interface SearchRequest { + queryText?: string; + queryVector?: number[]; + columns?: string[]; + numResults?: number; + queryType?: "ann" | "hybrid" | "full_text"; + filters?: SearchFilters; + reranker?: boolean; +} + +export interface SearchResponse< + T extends Record = Record, +> { + results: SearchResult[]; + totalCount: number; + queryTimeMs: number; + queryType: "ann" | "hybrid" | "full_text"; + nextPageToken: string | null; +} + +export interface SearchResult< + T extends Record = Record, +> { + score: number; + data: T; +} diff --git a/packages/appkit/src/plugins/vector-search/vector-search.ts b/packages/appkit/src/plugins/vector-search/vector-search.ts new file mode 100644 index 00000000..81a66595 --- /dev/null +++ b/packages/appkit/src/plugins/vector-search/vector-search.ts @@ -0,0 +1,367 @@ +import type express from "express"; +import type { IAppRouter, PluginExecutionSettings } from "shared"; +import { VectorSearchConnector } from "../../connectors"; +import type { VsRawResponse } from "../../connectors/vector-search/types"; +import { getWorkspaceClient } from "../../context"; +import { createLogger } from "../../logging/logger"; +import { Plugin, toPlugin } from "../../plugin"; +import type { PluginManifest } from "../../registry"; +import { vectorSearchDefaults } from "./defaults"; +import manifest from "./manifest.json"; +import type { + IndexConfig, + IVectorSearchConfig, + SearchRequest, + SearchResponse, +} from "./types"; + +const logger = createLogger("vector-search"); + +const querySettings: PluginExecutionSettings = { + default: vectorSearchDefaults, +}; + +export class VectorSearchPlugin extends Plugin { + static manifest = manifest as PluginManifest<"vector-search">; + + protected static description = + "Query Databricks Vector Search indexes with hybrid search, reranking, and pagination"; + protected declare config: IVectorSearchConfig; + + private connector: VectorSearchConnector; + + constructor(config: IVectorSearchConfig) { + super(config); + this.config = config; + this.connector = new VectorSearchConnector({ + timeout: config.timeout, + telemetry: config.telemetry, + }); + } + + async setup(): Promise { + for (const [alias, idx] of Object.entries(this.config.indexes)) { + if (!idx.indexName) { + throw new Error( + `Index "${alias}" is missing required field "indexName"`, + ); + } + if (!idx.columns || idx.columns.length === 0) { + throw new Error(`Index "${alias}" is missing required field "columns"`); + } + if (idx.pagination && !idx.endpointName) { + throw new Error( + `Index "${alias}" has pagination enabled but is missing "endpointName"`, + ); + } + } + logger.debug( + "Vector Search plugin configured with %d index(es)", + Object.keys(this.config.indexes).length, + ); + } + + injectRoutes(router: IAppRouter) { + this.route(router, { + name: "query", + method: "post", + path: "/:alias/query", + handler: async (req: express.Request, res: express.Response) => { + const indexConfig = this._resolveIndex(req.params.alias); + if (!indexConfig) { + res.status(404).json({ + error: `No index configured with alias "${req.params.alias}"`, + plugin: this.name, + }); + return; + } + + const body: SearchRequest = req.body; + if (!body.queryText && !body.queryVector) { + res.status(400).json({ + error: "queryText or queryVector is required", + plugin: this.name, + }); + return; + } + + try { + const prepared = await this._prepareQuery(body, indexConfig); + const plugin = + indexConfig.auth === "on-behalf-of-user" ? this.asUser(req) : this; + + const result = await plugin.execute( + async (signal) => + this.connector.query( + getWorkspaceClient(), + { + indexName: indexConfig.indexName, + queryText: prepared.queryText, + queryVector: prepared.queryVector, + columns: prepared.columns, + numResults: prepared.numResults, + queryType: prepared.queryType, + filters: body.filters, + reranker: prepared.rerankerConfig, + }, + signal, + ), + querySettings, + ); + + if (result === undefined) { + res.status(500).json({ error: "Query failed", plugin: this.name }); + return; + } + res.json(this._parseResponse(result, prepared.queryType)); + } catch (error) { + this._handleError(res, error, "Query failed"); + } + }, + }); + + this.route(router, { + name: "queryNextPage", + method: "post", + path: "/:alias/next-page", + handler: async (req: express.Request, res: express.Response) => { + const indexConfig = this._resolveIndex(req.params.alias); + if (!indexConfig) { + res.status(404).json({ + error: `No index configured with alias "${req.params.alias}"`, + plugin: this.name, + }); + return; + } + + if (!indexConfig.pagination) { + res.status(400).json({ + error: `Pagination is not enabled for index "${req.params.alias}"`, + plugin: this.name, + }); + return; + } + + if (!indexConfig.endpointName) { + res.status(400).json({ + error: `Index "${req.params.alias}" is missing endpointName required for pagination`, + plugin: this.name, + }); + return; + } + + const { pageToken } = req.body; + if (!pageToken) { + res.status(400).json({ + error: "pageToken is required", + plugin: this.name, + }); + return; + } + + try { + const plugin = + indexConfig.auth === "on-behalf-of-user" ? this.asUser(req) : this; + + const result = await plugin.execute( + async (signal) => + this.connector.queryNextPage( + getWorkspaceClient(), + { + indexName: indexConfig.indexName, + endpointName: indexConfig.endpointName as string, + pageToken, + }, + signal, + ), + querySettings, + ); + + if (result === undefined) { + res + .status(500) + .json({ error: "Next-page query failed", plugin: this.name }); + return; + } + res.json( + this._parseResponse(result, indexConfig.queryType ?? "hybrid"), + ); + } catch (error) { + this._handleError(res, error, "Next-page query failed"); + } + }, + }); + + this.route(router, { + name: "getConfig", + method: "get", + path: "/:alias/config", + handler: (req: express.Request, res: express.Response) => { + const { alias } = req.params; + const indexConfig = this._resolveIndex(alias); + if (!indexConfig) { + res.status(404).json({ + error: `No index configured with alias "${alias}"`, + plugin: this.name, + }); + return; + } + res.json({ + alias, + columns: indexConfig.columns, + queryType: indexConfig.queryType ?? "hybrid", + numResults: indexConfig.numResults ?? 20, + reranker: !!indexConfig.reranker, + pagination: !!indexConfig.pagination, + }); + }, + }); + } + + /** + * Programmatic query API — available as `appkit.vectorSearch.query()`. + * When called through `asUser(req)`, executes with the user's credentials. + */ + async query(alias: string, request: SearchRequest): Promise { + const indexConfig = this._resolveIndex(alias); + if (!indexConfig) { + throw new Error(`No index configured with alias "${alias}"`); + } + + const prepared = await this._prepareQuery(request, indexConfig); + + const result = await this.execute( + async (signal) => + this.connector.query( + getWorkspaceClient(), + { + indexName: indexConfig.indexName, + queryText: prepared.queryText, + queryVector: prepared.queryVector, + columns: prepared.columns, + numResults: prepared.numResults, + queryType: prepared.queryType, + filters: request.filters, + reranker: prepared.rerankerConfig, + }, + signal, + ), + querySettings, + ); + + if (result === undefined) { + throw new Error(`Vector search query failed for index "${alias}"`); + } + + return this._parseResponse(result, prepared.queryType); + } + + async shutdown(): Promise { + // No streams or persistent connections to clean up + } + + exports() { + return { + query: this.query.bind(this), + }; + } + + private _resolveIndex(alias: string): IndexConfig | undefined { + return this.config.indexes[alias]; + } + + private async _prepareQuery( + request: SearchRequest, + indexConfig: IndexConfig, + ): Promise<{ + queryText: string | undefined; + queryVector: number[] | undefined; + queryType: "ann" | "hybrid" | "full_text"; + columns: string[]; + numResults: number; + rerankerConfig: { columnsToRerank: string[] } | undefined; + }> { + const queryType = request.queryType ?? indexConfig.queryType ?? "hybrid"; + let queryText = request.queryText; + let queryVector = request.queryVector; + + if (indexConfig.embeddingFn && queryText && !queryVector) { + try { + queryVector = await indexConfig.embeddingFn(queryText); + queryText = undefined; + } catch (error) { + throw new Error( + `Embedding generation failed: ${error instanceof Error ? error.message : String(error)}`, + ); + } + } + + const columns = request.columns ?? indexConfig.columns; + return { + queryText, + queryVector, + queryType, + columns, + numResults: request.numResults ?? indexConfig.numResults ?? 20, + rerankerConfig: this._resolveReranker( + request.reranker, + indexConfig, + columns, + ), + }; + } + + private _resolveReranker( + requestReranker: boolean | undefined, + indexConfig: IndexConfig, + columns: string[], + ): { columnsToRerank: string[] } | undefined { + const shouldRerank = requestReranker ?? indexConfig.reranker; + if (!shouldRerank) return undefined; + + if (typeof indexConfig.reranker === "object") { + return indexConfig.reranker; + } + return { columnsToRerank: columns.filter((c) => c !== "id") }; + } + + private _parseResponse( + raw: VsRawResponse, + queryType: "ann" | "hybrid" | "full_text", + ): SearchResponse { + const columnNames = raw.manifest.columns.map((c) => c.name); + const scoreIndex = columnNames.indexOf("score"); + + const results = raw.result.data_array.map((row) => { + const data: Record = {}; + for (let i = 0; i < columnNames.length; i++) { + if (columnNames[i] !== "score") data[columnNames[i]] = row[i]; + } + return { + score: scoreIndex >= 0 ? (row[scoreIndex] as number) : 0, + data, + }; + }); + + return { + results, + totalCount: raw.result.row_count, + queryTimeMs: + raw.debug_info?.response_time ?? raw.debug_info?.latency_ms ?? 0, + queryType, + nextPageToken: raw.next_page_token ?? null, + }; + } + + private _handleError( + res: express.Response, + error: unknown, + fallbackMessage: string, + ): void { + logger.error("%s: %O", fallbackMessage, error); + const message = error instanceof Error ? error.message : fallbackMessage; + res.status(500).json({ error: message, plugin: this.name }); + } +} + +export const vectorSearch = toPlugin(VectorSearchPlugin); diff --git a/template/client/src/App.tsx b/template/client/src/App.tsx index a94bb5bc..e314df82 100644 --- a/template/client/src/App.tsx +++ b/template/client/src/App.tsx @@ -20,6 +20,9 @@ import { FilesPage } from './pages/files/FilesPage'; {{- if .plugins.serving}} import { ServingPage } from './pages/serving/ServingPage'; {{- end}} +{{- if .plugins.vectorSearch}} +import { VectorSearchPage } from './pages/vector-search/VectorSearchPage'; +{{- end}} const navLinkClass = ({ isActive }: { isActive: boolean }) => `px-3 py-1.5 rounded-md text-sm font-medium transition-colors ${ @@ -61,6 +64,11 @@ function Layout() { Serving +{{- end}} +{{- if .plugins.vectorSearch}} + + Vector Search + {{- end}} @@ -91,6 +99,9 @@ const router = createBrowserRouter([ {{- end}} {{- if .plugins.serving}} { path: '/serving', element: }, +{{- end}} +{{- if .plugins.vectorSearch}} + { path: '/vector-search', element: }, {{- end}} ], }, diff --git a/template/client/src/pages/vector-search/VectorSearchPage.tsx b/template/client/src/pages/vector-search/VectorSearchPage.tsx new file mode 100644 index 00000000..f1e5e58f --- /dev/null +++ b/template/client/src/pages/vector-search/VectorSearchPage.tsx @@ -0,0 +1,155 @@ +{{if .plugins.vectorSearch -}} +import { + Button, + Card, + CardContent, + CardHeader, + CardTitle, + Input, + Skeleton, +} from '@databricks/appkit-ui/react'; +import { Search } from 'lucide-react'; +import { useState } from 'react'; + +interface SearchResult { + score: number; + data: Record; +} + +interface SearchResponse { + results: SearchResult[]; + totalCount: number; + queryTimeMs: number; + queryType: string; +} + +export function VectorSearchPage() { + const [query, setQuery] = useState(''); + const [loading, setLoading] = useState(false); + const [error, setError] = useState(null); + const [response, setResponse] = useState(null); + + const handleSearch = async () => { + if (!query.trim()) return; + setLoading(true); + setError(null); + setResponse(null); + + try { + const res = await fetch('/api/vector-search/default/query', { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + body: JSON.stringify({ queryText: query }), + }); + + if (!res.ok) { + const data = await res.json().catch(() => ({})); + throw new Error(data.error ?? `HTTP ${res.status}: ${res.statusText}`); + } + + const data: SearchResponse = await res.json(); + setResponse(data); + } catch (err) { + setError(err instanceof Error ? err.message : String(err)); + } finally { + setLoading(false); + } + }; + + const handleKeyDown = (e: React.KeyboardEvent) => { + if (e.key === 'Enter') { + void handleSearch(); + } + }; + + return ( +
+
+

Vector Search

+

+ Query a Databricks Vector Search index using natural language. +

+
+ +
+ setQuery(e.target.value)} + onKeyDown={handleKeyDown} + className="flex-1" + /> + +
+ + {error && ( +
+ {error} +
+ )} + + {loading && ( +
+ {Array.from({ length: 3 }, (_, i) => ( + + + + + + + ))} +
+ )} + + {response && !loading && ( +
+

+ {response.totalCount} result{response.totalCount !== 1 ? 's' : ''} ·{' '} + {response.queryTimeMs}ms · {response.queryType} +

+ + {response.results.length === 0 ? ( +

No results found.

+ ) : ( + response.results.map((result, index) => ( + + + + Result {index + 1} + + score: {result.score.toFixed(4)} + + + + +
+ {Object.entries(result.data).map(([key, value]) => ( +
+
+ {key} +
+
+ {String(value ?? '')} +
+
+ ))} +
+ + + )) + )} +
+ )} +
+ ); +} +{{- end}}