From a51e148cd3db4e54e69cdf79e4692cdcaecc90fa Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Thu, 8 Oct 2026 16:12:50 +0200 Subject: [PATCH 1/3] feat(appkit): migrate ai-search connector to modular sdk-vectorsearch/uc-tables Add VectorSearchClient (sdk-vectorsearch 0.51.0) and TablesClient (sdk-uc-tables 0.52.0) to the workspace-client facade. - getIndex uses vectorSearch.getVectorIndex and maps back to VsIndexInfo. - getSourceColumns uses tables.getTable. - query and query-next-page use the raw request() seam: the generated model has no debug_level / debug_info, so typed unmarshal would drop the timings behind queryTimeMs. Wire shape and errors are unchanged. Co-authored-by: Isaac Signed-off-by: MarioCadenas # Conflicts: # packages/shared/src/workspace-client/client.ts # packages/shared/src/workspace-client/modular.ts # packages/shared/src/workspace-client/types.ts --- knip.json | 2 + packages/appkit/package.json | 2 + .../appkit/src/connectors/ai-search/client.ts | 102 +++++++++-------- .../appkit/src/connectors/ai-search/types.ts | 5 - .../plugins/ai-search/tests/ai-search.test.ts | 108 ++++++++++-------- .../src/testing/mock-workspace-client.ts | 2 + packages/shared/package.json | 2 + .../shared/src/workspace-client/client.ts | 23 ++++ .../shared/src/workspace-client/modular.ts | 33 +++++- packages/shared/src/workspace-client/types.ts | 8 ++ pnpm-lock.yaml | 38 ++++++ 11 files changed, 224 insertions(+), 101 deletions(-) diff --git a/knip.json b/knip.json index e42dd21c6..3cec36a06 100644 --- a/knip.json +++ b/knip.json @@ -16,6 +16,8 @@ "@databricks/sdk-files", "@databricks/sdk-genie", "@databricks/sdk-jobs", + "@databricks/sdk-uc-tables", + "@databricks/sdk-vectorsearch", "@databricks/sdk-options", "@databricks/sdk-scim", "@databricks/sdk-statementexecution", diff --git a/packages/appkit/package.json b/packages/appkit/package.json index 29799f1eb..40a5d6b8f 100644 --- a/packages/appkit/package.json +++ b/packages/appkit/package.json @@ -77,6 +77,8 @@ "@databricks/sdk-files": "0.52.0", "@databricks/sdk-genie": "0.54.0", "@databricks/sdk-jobs": "0.57.0", + "@databricks/sdk-uc-tables": "0.52.0", + "@databricks/sdk-vectorsearch": "0.51.0", "@databricks/sdk-options": "0.51.0", "@databricks/sdk-scim": "0.51.0", "@databricks/sdk-statementexecution": "0.52.0", diff --git a/packages/appkit/src/connectors/ai-search/client.ts b/packages/appkit/src/connectors/ai-search/client.ts index f4ed44b15..4f49f6064 100644 --- a/packages/appkit/src/connectors/ai-search/client.ts +++ b/packages/appkit/src/connectors/ai-search/client.ts @@ -7,10 +7,8 @@ import { TelemetryManager, } from "../../telemetry"; import type { WorkspaceClient } from "../../workspace-client"; -import { contextFromAbortSignal } from "../context"; import type { AiSearchConnectorConfig, - UcTableInfo, VsIndexInfo, VsNextPageParams, VsQueryParams, @@ -19,6 +17,28 @@ import type { const logger = createLogger("connectors:ai-search"); +/** + * POST through the raw `request()` seam, not the typed `VectorSearchClient`: + * its generated model has no `debug_level` (request) or `debug_info` (response), + * so unmarshal would strip the timings the plugin reports, and its reranker + * shape differs. Raw JSON keeps the exact wire shape `VsRawResponse` describes. + */ +async function postJson( + workspaceClient: WorkspaceClient, + path: string, + body: unknown, + signal?: AbortSignal, +): Promise { + const res = await workspaceClient.request({ + method: "POST", + path, + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(body), + signal, + }); + return (await res.json()) as T; +} + export class AiSearchConnector { private readonly telemetry: TelemetryProvider; @@ -83,17 +103,12 @@ export class AiSearchConnector { 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`, - payload: body, - headers: new Headers({ "Content-Type": "application/json" }), - raw: false, - query: {}, - }, - contextFromAbortSignal(signal), - )) as VsRawResponse; + const response = await postJson( + workspaceClient, + `/api/2.0/vector-search/indexes/${params.indexName}/query`, + body, + signal, + ); const duration = Date.now() - startTime; span.setAttribute("vs.result_count", response.result.row_count); @@ -153,20 +168,15 @@ export class AiSearchConnector { }, async (span: Span) => { try { - const response = (await workspaceClient.apiClient.request( + const response = await postJson( + workspaceClient, + `/api/2.0/vector-search/indexes/${params.indexName}/query-next-page`, { - method: "POST", - path: `/api/2.0/vector-search/indexes/${params.indexName}/query-next-page`, - payload: { - endpoint_name: params.endpointName, - page_token: params.pageToken, - }, - headers: new Headers({ "Content-Type": "application/json" }), - raw: false, - query: {}, + endpoint_name: params.endpointName, + page_token: params.pageToken, }, - contextFromAbortSignal(signal), - )) as VsRawResponse; + signal, + ); span.setAttribute("vs.result_count", response.result.row_count); span.setStatus({ code: SpanStatusCode.OK }); @@ -193,16 +203,24 @@ export class AiSearchConnector { indexName: string, signal?: AbortSignal, ): Promise { - return (await workspaceClient.apiClient.request( - { - method: "GET", - path: `/api/2.0/vector-search/indexes/${indexName}`, - headers: new Headers({ "Content-Type": "application/json" }), - raw: false, - query: {}, + const index = await workspaceClient.vectorSearch.getVectorIndex( + { name: indexName }, + { signal }, + ); + // Map the camelCase model back to the snake_case `VsIndexInfo` callers read. + const spec = + index.indexSpec?.$case === "deltaSyncIndexSpec" + ? index.indexSpec.deltaSyncIndexSpec + : undefined; + return { + index_type: index.indexType as VsIndexInfo["index_type"], + delta_sync_index_spec: spec && { + source_table: spec.sourceTable, + embedding_vector_columns: spec.embeddingVectorColumns?.map((c) => ({ + name: c.name ?? "", + })), }, - contextFromAbortSignal(signal), - )) as VsIndexInfo; + }; } /** @@ -214,16 +232,10 @@ export class AiSearchConnector { sourceTable: string, signal?: AbortSignal, ): Promise { - const table = (await workspaceClient.apiClient.request( - { - method: "GET", - path: `/api/2.1/unity-catalog/tables/${sourceTable}`, - headers: new Headers({ "Content-Type": "application/json" }), - raw: false, - query: {}, - }, - contextFromAbortSignal(signal), - )) as UcTableInfo; - return (table.columns ?? []).map((c) => c.name); + const table = await workspaceClient.tables.getTable( + { fullNameArg: sourceTable }, + { signal }, + ); + return (table.columns ?? []).flatMap((c) => (c.name ? [c.name] : [])); } } diff --git a/packages/appkit/src/connectors/ai-search/types.ts b/packages/appkit/src/connectors/ai-search/types.ts index 0d3e175b2..f87fefafb 100644 --- a/packages/appkit/src/connectors/ai-search/types.ts +++ b/packages/appkit/src/connectors/ai-search/types.ts @@ -31,11 +31,6 @@ export interface VsIndexInfo { }; } -/** Subset of the Unity Catalog get-table response used for column discovery. */ -export interface UcTableInfo { - columns?: Array<{ name: string }>; -} - export interface VsRawResponse { manifest: { column_count: number; diff --git a/packages/appkit/src/plugins/ai-search/tests/ai-search.test.ts b/packages/appkit/src/plugins/ai-search/tests/ai-search.test.ts index f9277932f..b6a2307a0 100644 --- a/packages/appkit/src/plugins/ai-search/tests/ai-search.test.ts +++ b/packages/appkit/src/plugins/ai-search/tests/ai-search.test.ts @@ -7,7 +7,6 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { withEnv } from "../../../testing"; import { useTestCache } from "../../../testing/test-cache"; -import { Context } from "../../../workspace-client"; vi.mock("../../../context", () => ({ getWorkspaceClient: vi.fn(() => mockWorkspaceClient), @@ -119,17 +118,28 @@ const validVsResponse = { debug_info: { response_time: 35 }, }; +// Records each raw `WorkspaceRequest`; resolves to the JSON body the facade's +// `request()` would return as a `Response`. const mockRequest = vi.fn().mockResolvedValue(validVsResponse); +const mockGetVectorIndex = vi.fn(); +const mockGetTable = vi.fn(); const mockWorkspaceClient = { - apiClient: { request: mockRequest }, + request: async (req: unknown) => Response.json(await mockRequest(req)), + vectorSearch: { getVectorIndex: mockGetVectorIndex }, + tables: { getTable: mockGetTable }, }; +/** JSON body of the n-th raw request sent through `request()`. */ +const bodyOf = (n = 0) => JSON.parse(mockRequest.mock.calls[n][0].body); + import { AiSearchPlugin } from "../ai-search"; describe("AiSearchPlugin", () => { beforeEach(() => { mockRequest.mockClear(); mockRequest.mockResolvedValue(validVsResponse); + mockGetVectorIndex.mockReset(); + mockGetTable.mockReset(); }); describe("setup()", () => { @@ -218,25 +228,22 @@ describe("AiSearchPlugin", () => { describe("setup() column auto-discovery", () => { const originalNodeEnv = process.env.NODE_ENV; - // Route GET metadata calls to discovery fixtures; POST queries stay on the - // default validVsResponse. - const routeByPath = (opts: { method: string; path: string }) => { - if (opts.path.endsWith("/query")) return Promise.resolve(validVsResponse); - if (opts.path.startsWith("/api/2.0/vector-search/indexes/")) { - return Promise.resolve({ - index_type: "DELTA_SYNC", - delta_sync_index_spec: { - source_table: "cat.sch.src", - embedding_vector_columns: [{ name: "__vec" }], + // Discovery goes through the typed modular clients (camelCase models); + // queries stay on the raw `request()` mock and its default validVsResponse. + const mockDiscovery = () => { + mockGetVectorIndex.mockResolvedValue({ + indexType: "DELTA_SYNC", + indexSpec: { + $case: "deltaSyncIndexSpec", + deltaSyncIndexSpec: { + sourceTable: "cat.sch.src", + embeddingVectorColumns: [{ name: "__vec" }], }, - }); - } - if (opts.path.startsWith("/api/2.1/unity-catalog/tables/")) { - return Promise.resolve({ - columns: [{ name: "id" }, { name: "body" }, { name: "__vec" }], - }); - } - return Promise.resolve(validVsResponse); + }, + }); + mockGetTable.mockResolvedValue({ + columns: [{ name: "id" }, { name: "body" }, { name: "__vec" }], + }); }; afterEach(() => { @@ -245,7 +252,7 @@ describe("AiSearchPlugin", () => { it("fills columns from the source table in development and warns", async () => { process.env.NODE_ENV = "development"; - mockRequest.mockImplementation(routeByPath); + mockDiscovery(); const plugin = new AiSearchPlugin({ indexes: { docs: { indexName: "cat.sch.idx" } }, }); @@ -254,15 +261,20 @@ describe("AiSearchPlugin", () => { // Discovered columns, minus the embedding vector column. await plugin.query("docs", { queryText: "q" }); - const queryCall = mockRequest.mock.calls.find((c) => - c[0].path.endsWith("/query"), + expect(mockGetVectorIndex).toHaveBeenCalledWith( + { name: "cat.sch.idx" }, + expect.anything(), + ); + expect(mockGetTable).toHaveBeenCalledWith( + { fullNameArg: "cat.sch.src" }, + expect.anything(), ); - expect(queryCall?.[0].payload.columns).toEqual(["id", "body"]); + expect(bodyOf().columns).toEqual(["id", "body"]); }); it("does not discover columns outside development", async () => { process.env.NODE_ENV = "production"; - mockRequest.mockImplementation(routeByPath); + mockDiscovery(); // Columns set so the prod no-columns guard doesn't fire; this test only // asserts discovery doesn't run outside development. const plugin = new AiSearchPlugin({ @@ -272,25 +284,21 @@ describe("AiSearchPlugin", () => { await plugin.setup(); // No get-index / get-table calls were made. - const metadataCalls = mockRequest.mock.calls.filter( - (c) => !c[0].path.endsWith("/query"), - ); - expect(metadataCalls).toHaveLength(0); + expect(mockGetVectorIndex).not.toHaveBeenCalled(); + expect(mockGetTable).not.toHaveBeenCalled(); }); it("skips (does not throw) when an index already has columns", async () => { process.env.NODE_ENV = "development"; - mockRequest.mockImplementation(routeByPath); + mockDiscovery(); const plugin = new AiSearchPlugin({ indexes: { docs: { indexName: "cat.sch.idx", columns: ["id"] } }, }); await plugin.setup(); - const metadataCalls = mockRequest.mock.calls.filter( - (c) => !c[0].path.endsWith("/query"), - ); - expect(metadataCalls).toHaveLength(0); + expect(mockGetVectorIndex).not.toHaveBeenCalled(); + expect(mockGetTable).not.toHaveBeenCalled(); }); }); @@ -378,16 +386,18 @@ describe("AiSearchPlugin", () => { expect.objectContaining({ method: "POST", path: "/api/2.0/vector-search/indexes/cat.sch.idx/query", + signal: expect.any(AbortSignal), }), - // 2nd arg is the SDK Context bridging the execution's abort signal. - expect.any(Context), ); - const callBody = mockRequest.mock.calls[0][0].payload; + const callBody = bodyOf(); 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"]); + // Not in the modular request model; the raw `request()` path must keep it + // so the response carries `debug_info` (queryTimeMs). + expect(callBody.debug_level).toBe(1); }); it("throws Error for unknown alias", async () => { @@ -418,7 +428,7 @@ describe("AiSearchPlugin", () => { filters: { category: ["books"] }, }); - const callBody = mockRequest.mock.calls[0][0].payload; + const callBody = bodyOf(); // VS expects a JSON-encoded string under `filters_json`; a raw object // under `filters` is silently ignored by the API. expect(callBody.filters).toBeUndefined(); @@ -440,7 +450,7 @@ describe("AiSearchPlugin", () => { await plugin.setup(); await plugin.query("test", { queryText: "test" }); - const callBody = mockRequest.mock.calls[0][0].payload; + const callBody = bodyOf(); expect(callBody.reranker.model).toBe("databricks_reranker"); expect(callBody.reranker.parameters.columns_to_rerank).toEqual([ "title", @@ -464,7 +474,7 @@ describe("AiSearchPlugin", () => { await plugin.query("test", { queryText: "test" }); expect(mockEmbeddingFn).toHaveBeenCalledWith("test"); - const callBody = mockRequest.mock.calls[0][0].payload; + const callBody = bodyOf(); expect(callBody.query_vector).toEqual([0.1, 0.2, 0.3]); expect(callBody.query_text).toBeUndefined(); }); @@ -485,7 +495,7 @@ describe("AiSearchPlugin", () => { await plugin.query("test", { queryText: "test" }); expect(mockEmbeddingFn).toHaveBeenCalledWith("test"); - const callBody = mockRequest.mock.calls[0][0].payload; + const callBody = bodyOf(); expect(callBody.query_vector).toEqual([0.1, 0.2, 0.3]); expect(callBody.query_text).toBe("test"); }); @@ -506,7 +516,7 @@ describe("AiSearchPlugin", () => { await plugin.query("test", { queryText: "test" }); expect(mockEmbeddingFn).not.toHaveBeenCalled(); - const callBody = mockRequest.mock.calls[0][0].payload; + const callBody = bodyOf(); expect(callBody.query_text).toBe("test"); expect(callBody.query_vector).toBeUndefined(); }); @@ -613,7 +623,7 @@ describe("AiSearchPlugin", () => { columns: ["id"], }); - const callBody = mockRequest.mock.calls[0][0].payload; + const callBody = bodyOf(); expect(callBody.query_type).toBe("ANN"); expect(callBody.num_results).toBe(5); expect(callBody.columns).toEqual(["id"]); @@ -632,7 +642,7 @@ describe("AiSearchPlugin", () => { await plugin.setup(); await plugin.query("test", { queryText: "q" }); - const callBody = mockRequest.mock.calls[0][0].payload; + const callBody = bodyOf(); expect(callBody.reranker.parameters.columns_to_rerank).toEqual(["title"]); }); @@ -649,7 +659,7 @@ describe("AiSearchPlugin", () => { await plugin.setup(); await plugin.query("test", { queryText: "q", reranker: false }); - const callBody = mockRequest.mock.calls[0][0].payload; + const callBody = bodyOf(); expect(callBody.reranker).toBeUndefined(); }); @@ -661,7 +671,7 @@ describe("AiSearchPlugin", () => { }); await plugin.query("test", { queryText: "q" }); - const callBody = mockRequest.mock.calls[0][0].payload; + const callBody = bodyOf(); expect(callBody.reranker).toBeUndefined(); expect(callBody.columns).toEqual([]); }); @@ -1012,7 +1022,7 @@ describe("AiSearchPlugin", () => { // demo is configured with columns ["id", "title"]; the request's // columns must not widen the projection. - const callBody = mockRequest.mock.calls[0][0].payload; + const callBody = bodyOf(); expect(callBody.columns).toEqual(["id", "title"]); }); @@ -1094,10 +1104,10 @@ describe("AiSearchPlugin", () => { expect(mockRequest).toHaveBeenCalledWith( expect.objectContaining({ path: "/api/2.0/vector-search/indexes/cat.sch.paged/query-next-page", - payload: { endpoint_name: "ep", page_token: "t" }, + signal: expect.any(AbortSignal), }), - expect.any(Context), ); + expect(bodyOf()).toEqual({ endpoint_name: "ep", page_token: "t" }); expect(res.json).toHaveBeenCalled(); }); }); diff --git a/packages/appkit/src/testing/mock-workspace-client.ts b/packages/appkit/src/testing/mock-workspace-client.ts index bd65eff3d..3b21896a7 100644 --- a/packages/appkit/src/testing/mock-workspace-client.ts +++ b/packages/appkit/src/testing/mock-workspace-client.ts @@ -69,6 +69,8 @@ const FACADE_SERVICES = [ "warehouses", "genie", "jobs", + "vectorSearch", + "tables", "statementExecution", "servingEndpoints", "currentUser", diff --git a/packages/shared/package.json b/packages/shared/package.json index 451fbcbc0..33df9896a 100644 --- a/packages/shared/package.json +++ b/packages/shared/package.json @@ -54,6 +54,8 @@ "@databricks/sdk-files": "0.52.0", "@databricks/sdk-genie": "0.54.0", "@databricks/sdk-jobs": "0.57.0", + "@databricks/sdk-uc-tables": "0.52.0", + "@databricks/sdk-vectorsearch": "0.51.0", "@databricks/sdk-options": "0.51.0", "@databricks/sdk-scim": "0.51.0", "@databricks/sdk-statementexecution": "0.52.0", diff --git a/packages/shared/src/workspace-client/client.ts b/packages/shared/src/workspace-client/client.ts index ff0074ca5..e51635507 100644 --- a/packages/shared/src/workspace-client/client.ts +++ b/packages/shared/src/workspace-client/client.ts @@ -24,6 +24,10 @@ import { type JobsClient, buildFilesClient, type FilesClient, + buildTablesClient, + type TablesClient, + buildVectorSearchClient, + type VectorSearchClient, type StatementExecutionClient, type WarehousesClient, type WorkspaceAuth, @@ -41,6 +45,8 @@ export class AppKitWorkspaceClient implements WorkspaceClient { #genie?: GenieClient; #jobs?: JobsClient; #files?: FilesClient; + #vectorSearch?: VectorSearchClient; + #tables?: TablesClient; constructor(opts: WorkspaceClientOptions) { this.#opts = opts; @@ -78,6 +84,23 @@ export class AppKitWorkspaceClient implements WorkspaceClient { return this.#jobs; } + // Migrated to the modular SDK — built lazily, independent of the legacy client. + // Migrated to the modular SDK — built lazily, independent of the legacy client. + get vectorSearch(): VectorSearchClient { + if (!this.#vectorSearch) { + this.#vectorSearch = buildVectorSearchClient(this.#opts); + } + return this.#vectorSearch; + } + + // Migrated to the modular SDK — built lazily, independent of the legacy client. + get tables(): TablesClient { + if (!this.#tables) { + this.#tables = buildTablesClient(this.#opts); + } + return this.#tables; + } + // Migrated to the modular SDK — built lazily, independent of the legacy client. get statementExecution(): StatementExecutionClient { if (!this.#statementExecution) { diff --git a/packages/shared/src/workspace-client/modular.ts b/packages/shared/src/workspace-client/modular.ts index d07a3699d..deb17e81f 100644 --- a/packages/shared/src/workspace-client/modular.ts +++ b/packages/shared/src/workspace-client/modular.ts @@ -8,8 +8,8 @@ * * Migrated services are built here as per-service clients; the facade delegates * their accessors to these instead of the legacy monolithic client. Currently - * `warehouses`, `statementExecution`, `currentUser` (SCIM), `genie` and `jobs` - * are migrated, plus the auth + raw-request seam ({@link buildWorkspaceAuth}); + * `warehouses`, `statementExecution`, `currentUser` (SCIM), `genie`, `jobs`, `vectorSearch` + * and `tables` (UC) are migrated, plus the auth + raw-request seam ({@link buildWorkspaceAuth}); * every other service still routes through `legacy.ts`. * * NOTE: statementExecution relies on a pinned pnpm patch @@ -50,6 +50,8 @@ import type { JobsClient } from "@databricks/sdk-jobs/v2"; import type { ClientOptions } from "@databricks/sdk-options/client"; import { ScimClient } from "@databricks/sdk-scim/v1"; import { StatementExecutionClient } from "@databricks/sdk-statementexecution/v1"; +import type { TablesClient } from "@databricks/sdk-uc-tables/v1"; +import type { VectorSearchClient } from "@databricks/sdk-vectorsearch/v1"; import { WarehousesClient } from "@databricks/sdk-warehouses/v1"; import { ApiError } from "./errors"; @@ -484,6 +486,31 @@ export function buildFilesClient(opts: WorkspaceClientOptions): FilesClient { ); } +/** + * Build a modular Vector Search client from wrapper options. Only index metadata + * (`getVectorIndex`) goes typed: the query endpoints use the raw `request()` + * seam because the generated model drops `debug_level` / `debug_info`. + */ +export function buildVectorSearchClient( + opts: WorkspaceClientOptions, +): VectorSearchClient { + return lazyClient( + () => + import("@databricks/sdk-vectorsearch/v1").then( + (m) => m.VectorSearchClient, + ), + mapToClientOptions(opts), + ); +} + +/** Build a modular Unity Catalog Tables client from wrapper options. */ +export function buildTablesClient(opts: WorkspaceClientOptions): TablesClient { + return lazyClient( + () => import("@databricks/sdk-uc-tables/v1").then((m) => m.TablesClient), + mapToClientOptions(opts), + ); +} + // ── Client type re-exports (for the facade accessor types) ─────────────── export type { FilesClient } from "@databricks/sdk-files/v2"; export type { @@ -494,6 +521,8 @@ export type { GenieClient } from "@databricks/sdk-genie/v1"; export type { JobsClient } from "@databricks/sdk-jobs/v2"; export type { ScimClient } from "@databricks/sdk-scim/v1"; export type { StatementExecutionClient } from "@databricks/sdk-statementexecution/v1"; +export type { TablesClient } from "@databricks/sdk-uc-tables/v1"; +export type { VectorSearchClient } from "@databricks/sdk-vectorsearch/v1"; export type { WarehousesClient } from "@databricks/sdk-warehouses/v1"; // ── Model type re-exports ──────────────────────────────────────────────── diff --git a/packages/shared/src/workspace-client/types.ts b/packages/shared/src/workspace-client/types.ts index cebab1935..dc381eef3 100644 --- a/packages/shared/src/workspace-client/types.ts +++ b/packages/shared/src/workspace-client/types.ts @@ -22,6 +22,8 @@ import type { GenieClient, JobsClient, FilesClient, + TablesClient, + VectorSearchClient, } from "./modular"; // Legacy SDK type namespaces for un-migrated services, re-exported so AppKit @@ -57,6 +59,12 @@ export interface WorkspaceClient extends WorkspaceAuth { /** Jobs (modular SDK, Jobs API 2.2). */ readonly jobs: JobsClient; + /** Vector Search (modular SDK; index metadata only, queries use `request()`). */ + readonly vectorSearch: VectorSearchClient; + + /** Unity Catalog Tables (modular SDK). */ + readonly tables: TablesClient; + /** Statement Execution (modular SDK). */ readonly statementExecution: StatementExecutionClient; diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index 3c6b8a5cf..c72820060 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -307,6 +307,12 @@ importers: '@databricks/sdk-statementexecution': specifier: 0.52.0 version: 0.52.0(patch_hash=9377a46b883b548d47116662bf73b159a4da51ec28b9ff1896a7e4a7841c240c) + '@databricks/sdk-uc-tables': + specifier: 0.52.0 + version: 0.52.0 + '@databricks/sdk-vectorsearch': + specifier: 0.51.0 + version: 0.51.0 '@databricks/sdk-warehouses': specifier: 0.53.0 version: 0.53.0 @@ -649,6 +655,12 @@ importers: '@databricks/sdk-statementexecution': specifier: 0.52.0 version: 0.52.0(patch_hash=9377a46b883b548d47116662bf73b159a4da51ec28b9ff1896a7e4a7841c240c) + '@databricks/sdk-uc-tables': + specifier: 0.52.0 + version: 0.52.0 + '@databricks/sdk-vectorsearch': + specifier: 0.51.0 + version: 0.51.0 '@databricks/sdk-warehouses': specifier: 0.53.0 version: 0.53.0 @@ -2057,6 +2069,14 @@ packages: resolution: {integrity: sha512-ymDGDW67PJ/0VujM89mUrE6mbbrabsKJP5LvPv8ZbgzPR0bm0DwXMcQzvS0XUt0UGITjV2gqbx5jQwP3ijtOmg==} engines: {node: '>=22.0.0'} + '@databricks/sdk-uc-tables@0.52.0': + resolution: {integrity: sha512-Dlqu3+O1krk6ogBX/uJGUDn1xMdt0gVMvrpmf16XkPelmgzwrbdWWtU2LU/zlcCBrb3XvhPPUpxVVHmyNS8Pjg==} + engines: {node: '>=22.0.0'} + + '@databricks/sdk-vectorsearch@0.51.0': + resolution: {integrity: sha512-Hr5z5lbT4pt0FV7L+gvKIFdp0PlWtyH14bs6+rJsCxcgwV6BQhuuVecJv8M4FB8QdsDZd2Us/JVIa5LwHHVORA==} + engines: {node: '>=22.0.0'} + '@databricks/sdk-warehouses@0.53.0': resolution: {integrity: sha512-+NlLn/HP3+5YUxTKlp0+zwsbAGaU2GHl8C+8JlWk04VLmBZhJg/PJnr4WlNrQVTAzqH01dy4DmK7MfqyOetHRg==} engines: {node: '>=22.0.0'} @@ -14407,6 +14427,24 @@ snapshots: json-bigint: 1.0.0 zod: 4.3.6 + '@databricks/sdk-uc-tables@0.52.0': + dependencies: + '@databricks/sdk-auth': 0.51.0 + '@databricks/sdk-core': 0.51.0 + '@databricks/sdk-options': 0.51.0 + '@js-temporal/polyfill': 0.5.1 + json-bigint: 1.0.0 + zod: 4.3.6 + + '@databricks/sdk-vectorsearch@0.51.0': + dependencies: + '@databricks/sdk-auth': 0.51.0 + '@databricks/sdk-core': 0.51.0 + '@databricks/sdk-options': 0.51.0 + '@js-temporal/polyfill': 0.5.1 + json-bigint: 1.0.0 + zod: 4.3.6 + '@databricks/sdk-warehouses@0.53.0': dependencies: '@databricks/sdk-auth': 0.51.0 From c3ce2dee5b819837262833a8d64e3caf0fc4ea9b Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Tue, 13 Oct 2026 15:37:01 +0200 Subject: [PATCH 2/3] chore(shared): drop a duplicated comment on the vectorSearch accessor Co-authored-by: Isaac Signed-off-by: MarioCadenas --- packages/shared/src/workspace-client/client.ts | 1 - 1 file changed, 1 deletion(-) diff --git a/packages/shared/src/workspace-client/client.ts b/packages/shared/src/workspace-client/client.ts index e51635507..ec1dbe624 100644 --- a/packages/shared/src/workspace-client/client.ts +++ b/packages/shared/src/workspace-client/client.ts @@ -84,7 +84,6 @@ export class AppKitWorkspaceClient implements WorkspaceClient { return this.#jobs; } - // Migrated to the modular SDK — built lazily, independent of the legacy client. // Migrated to the modular SDK — built lazily, independent of the legacy client. get vectorSearch(): VectorSearchClient { if (!this.#vectorSearch) { From c357e17a9c73dd76076080f9c66fc9887f62b47f Mon Sep 17 00:00:00 2001 From: MarioCadenas Date: Tue, 13 Oct 2026 18:44:07 +0200 Subject: [PATCH 3/3] feat(appkit): run ai-search queries through the typed vector search client Replace the raw request() path for query and query-next-page with vectorSearch.queryVectorIndex / queryVectorIndexNextPage, backed by a pnpm patch on @databricks/sdk-vectorsearch@0.51.0 that: - adds the undocumented debug_level request / debug_info response fields (source of the plugin's queryTimeMs) - accepts bigint row values (the SDK parses with json-bigint, and the unpatched schema rejected ints over 2^53) - accepts a null next_page_token The connector maps the camelCase model back to VsRawResponse and coerces bigint rows to number, so the public response shape is unchanged. Tests now drive a real VectorSearchClient over a stubbed fetch. Co-authored-by: Isaac Signed-off-by: MarioCadenas --- .../appkit/src/connectors/ai-search/client.ts | 108 +++++++++--------- .../plugins/ai-search/tests/ai-search.test.ts | 57 ++++++++- .../shared/src/workspace-client/modular.ts | 15 ++- ...@databricks__sdk-vectorsearch@0.51.0.patch | 67 +++++++++++ pnpm-lock.yaml | 7 +- pnpm-workspace.yaml | 1 + 6 files changed, 192 insertions(+), 63 deletions(-) create mode 100644 patches/@databricks__sdk-vectorsearch@0.51.0.patch diff --git a/packages/appkit/src/connectors/ai-search/client.ts b/packages/appkit/src/connectors/ai-search/client.ts index 4f49f6064..313700e4f 100644 --- a/packages/appkit/src/connectors/ai-search/client.ts +++ b/packages/appkit/src/connectors/ai-search/client.ts @@ -6,7 +6,11 @@ import { SpanStatusCode, TelemetryManager, } from "../../telemetry"; -import type { WorkspaceClient } from "../../workspace-client"; +import type { + QueryVectorIndexRequest, + QueryVectorIndexResponse, + WorkspaceClient, +} from "../../workspace-client"; import type { AiSearchConnectorConfig, VsIndexInfo, @@ -18,25 +22,29 @@ import type { const logger = createLogger("connectors:ai-search"); /** - * POST through the raw `request()` seam, not the typed `VectorSearchClient`: - * its generated model has no `debug_level` (request) or `debug_info` (response), - * so unmarshal would strip the timings the plugin reports, and its reranker - * shape differs. Raw JSON keeps the exact wire shape `VsRawResponse` describes. + * Map the typed (camelCase) response back to the snake_case wire shape the + * plugin reads. `debugInfo` relies on the sdk-vectorsearch pnpm patch. Rows + * may hold `bigint` (the SDK parses with json-bigint); coerce to `number` so + * results stay JSON-serializable, matching legacy `JSON.parse` precision. */ -async function postJson( - workspaceClient: WorkspaceClient, - path: string, - body: unknown, - signal?: AbortSignal, -): Promise { - const res = await workspaceClient.request({ - method: "POST", - path, - headers: { "Content-Type": "application/json" }, - body: JSON.stringify(body), - signal, - }); - return (await res.json()) as T; +function toRawResponse(res: QueryVectorIndexResponse): VsRawResponse { + return { + manifest: { + column_count: res.manifest?.columnCount ?? 0, + columns: (res.manifest?.columns ?? []).map((c) => ({ + name: c.name ?? "", + type: c.typeText, + })), + }, + result: { + row_count: res.result?.rowCount ?? 0, + data_array: (res.result?.dataArray ?? []).map((row) => + row.map((v) => (typeof v === "bigint" ? Number(v) : v)), + ), + }, + next_page_token: res.nextPageToken, + debug_info: res.debugInfo as VsRawResponse["debug_info"], + }; } export class AiSearchConnector { @@ -58,25 +66,25 @@ export class AiSearchConnector { throw new Error("Query cancelled before execution"); } - const body: Record = { + const hasFilters = !!( + params.filters && Object.keys(params.filters).length > 0 + ); + const req: QueryVectorIndexRequest = { + name: params.indexName, 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) { + numResults: params.numResults, + queryType: params.queryType.toUpperCase(), + // Makes the response carry `debug_info` timings (queryTimeMs). + debugLevel: 1, + queryText: params.queryText || undefined, + queryVector: params.queryVector, // VS silently ignores an object under `filters`; it wants a JSON string. - body.filters_json = JSON.stringify(params.filters); - } - if (params.reranker) { - body.reranker = { + filtersJson: hasFilters ? JSON.stringify(params.filters) : undefined, + reranker: params.reranker && { model: "databricks_reranker", - parameters: { columns_to_rerank: params.reranker.columnsToRerank }, - }; - } + parameters: { columnsToRerank: params.reranker.columnsToRerank }, + }, + }; logger.debug( "Querying VS index %s (type=%s, num_results=%d)", @@ -94,20 +102,17 @@ export class AiSearchConnector { "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_filters": hasFilters, "vs.has_reranker": !!params.reranker, }, }, async (span: Span) => { const startTime = Date.now(); try { - const response = await postJson( - workspaceClient, - `/api/2.0/vector-search/indexes/${params.indexName}/query`, - body, - signal, + const response = toRawResponse( + await workspaceClient.vectorSearch.queryVectorIndex(req, { + signal, + }), ); const duration = Date.now() - startTime; @@ -168,14 +173,15 @@ export class AiSearchConnector { }, async (span: Span) => { try { - const response = await postJson( - workspaceClient, - `/api/2.0/vector-search/indexes/${params.indexName}/query-next-page`, - { - endpoint_name: params.endpointName, - page_token: params.pageToken, - }, - signal, + const response = toRawResponse( + await workspaceClient.vectorSearch.queryVectorIndexNextPage( + { + name: params.indexName, + endpointName: params.endpointName, + pageToken: params.pageToken, + }, + { signal }, + ), ); span.setAttribute("vs.result_count", response.result.row_count); diff --git a/packages/appkit/src/plugins/ai-search/tests/ai-search.test.ts b/packages/appkit/src/plugins/ai-search/tests/ai-search.test.ts index b6a2307a0..8b3eff3c6 100644 --- a/packages/appkit/src/plugins/ai-search/tests/ai-search.test.ts +++ b/packages/appkit/src/plugins/ai-search/tests/ai-search.test.ts @@ -7,6 +7,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { withEnv } from "../../../testing"; import { useTestCache } from "../../../testing/test-cache"; +import { createWorkspaceClient } from "../../../workspace-client"; vi.mock("../../../context", () => ({ getWorkspaceClient: vi.fn(() => mockWorkspaceClient), @@ -118,18 +119,38 @@ const validVsResponse = { debug_info: { response_time: 35 }, }; -// Records each raw `WorkspaceRequest`; resolves to the JSON body the facade's -// `request()` would return as a `Response`. +// Queries run through a REAL modular VectorSearchClient over a stubbed `fetch`, +// so the patched marshal/unmarshal (debug_level/debug_info, bigint rows) is +// exercised. `mockRequest` records each wire request and resolves to the JSON +// body to respond with. const mockRequest = vi.fn().mockResolvedValue(validVsResponse); +vi.stubGlobal("fetch", async (url: string, init: RequestInit) => { + const body = await mockRequest({ + method: init.method, + path: new URL(url).pathname, + body: init.body, + signal: init.signal, + }); + // Raw text so tests can send ints beyond 2^53 that JSON.stringify can't. + return new Response(typeof body === "string" ? body : JSON.stringify(body)); +}); const mockGetVectorIndex = vi.fn(); const mockGetTable = vi.fn(); +const realVectorSearch = createWorkspaceClient({ + host: "https://test.cloud.databricks.com", + token: "test-token", +}).vectorSearch; const mockWorkspaceClient = { - request: async (req: unknown) => Response.json(await mockRequest(req)), - vectorSearch: { getVectorIndex: mockGetVectorIndex }, + vectorSearch: { + queryVectorIndex: realVectorSearch.queryVectorIndex.bind(realVectorSearch), + queryVectorIndexNextPage: + realVectorSearch.queryVectorIndexNextPage.bind(realVectorSearch), + getVectorIndex: mockGetVectorIndex, + }, tables: { getTable: mockGetTable }, }; -/** JSON body of the n-th raw request sent through `request()`. */ +/** JSON body of the n-th wire request. */ const bodyOf = (n = 0) => JSON.parse(mockRequest.mock.calls[n][0].body); import { AiSearchPlugin } from "../ai-search"; @@ -346,6 +367,25 @@ describe("AiSearchPlugin", () => { expect(result.queryTimeMs).toBe(35); }); + it("coerces int64 row values the SDK parses as bigint to number", async () => { + // The SDK parses with json-bigint, so ints past 2^53 arrive as bigint; + // the patched schema accepts them and the connector coerces, keeping the + // result JSON-serializable (legacy JSON.parse precision). + mockRequest.mockResolvedValueOnce( + `{"manifest":{"column_count":2,"columns":[{"name":"id"},{"name":"score"}]},` + + `"result":{"row_count":1,"data_array":[[9007199254740993,0.5]]}}`, + ); + const plugin = new AiSearchPlugin({ + indexes: { big: { indexName: "cat.sch.big", columns: ["id"] } }, + }); + await plugin.setup(); + + const result = await plugin.query("big", { queryText: "q" }); + + expect(result.results[0].data.id).toBe(9007199254740992); + expect(() => JSON.stringify(result)).not.toThrow(); + }); + it("types result.data via the generic parameter", async () => { interface Doc extends Record { id: number; @@ -1107,7 +1147,12 @@ describe("AiSearchPlugin", () => { signal: expect.any(AbortSignal), }), ); - expect(bodyOf()).toEqual({ endpoint_name: "ep", page_token: "t" }); + // The SDK also sends `name` in the body; the path carries it either way. + expect(bodyOf()).toEqual({ + name: "cat.sch.paged", + endpoint_name: "ep", + page_token: "t", + }); expect(res.json).toHaveBeenCalled(); }); }); diff --git a/packages/shared/src/workspace-client/modular.ts b/packages/shared/src/workspace-client/modular.ts index deb17e81f..fa4221025 100644 --- a/packages/shared/src/workspace-client/modular.ts +++ b/packages/shared/src/workspace-client/modular.ts @@ -17,6 +17,12 @@ * undocumented Reyden `attachment` response field, which the SDK's generated * unmarshal transform would otherwise strip. * + * NOTE: vectorSearch relies on a pinned pnpm patch + * (`patches/@databricks__sdk-vectorsearch@0.51.0.patch`): it adds the + * undocumented query `debug_level` request / `debug_info` response fields (the + * source of ai-search `queryTimeMs`), and lets row values be `bigint`, since the + * SDK parses with json-bigint and the unpatched schema rejects ints over 2^53. + * * NOTE: genie relies on a pinned pnpm patch * (`patches/@databricks__sdk-genie@0.54.0.patch`): the generated model types the * query result's `data_array` as protobuf `ListValue[]` (`{ values: [...] }`), but @@ -487,9 +493,8 @@ export function buildFilesClient(opts: WorkspaceClientOptions): FilesClient { } /** - * Build a modular Vector Search client from wrapper options. Only index metadata - * (`getVectorIndex`) goes typed: the query endpoints use the raw `request()` - * seam because the generated model drops `debug_level` / `debug_info`. + * Build a modular Vector Search client from wrapper options. Query calls rely + * on the sdk-vectorsearch pnpm patch (`debugLevel` / `debugInfo`, `bigint` rows). */ export function buildVectorSearchClient( opts: WorkspaceClientOptions, @@ -570,3 +575,7 @@ export type { EndpointState, GetWarehouseResponse, } from "@databricks/sdk-warehouses/v1"; +export type { + QueryVectorIndexRequest, + QueryVectorIndexResponse, +} from "@databricks/sdk-vectorsearch/v1"; diff --git a/patches/@databricks__sdk-vectorsearch@0.51.0.patch b/patches/@databricks__sdk-vectorsearch@0.51.0.patch new file mode 100644 index 000000000..58991c28b --- /dev/null +++ b/patches/@databricks__sdk-vectorsearch@0.51.0.patch @@ -0,0 +1,67 @@ +diff --git a/dist/v1/model.d.ts b/dist/v1/model.d.ts +index aeb9a49c9dd947b7275b391d09b9b5594282cb1c..00a7dad9c3086458d2a262e2e0502ee17c31aa89 100644 +--- a/dist/v1/model.d.ts ++++ b/dist/v1/model.d.ts +@@ -465,6 +465,8 @@ interface QueryVectorIndexRequest { + * `TOP` and `BUCKETS` are case-insensitive. A column may appear at most once. + */ + facets?: string[] | undefined; ++ /** Undocumented: `1` makes the response include `debug_info` timings. */ ++ debugLevel?: number | undefined; + } + interface QueryVectorIndexResponse { + /** Metadata about the result set. */ +@@ -479,6 +481,8 @@ interface QueryVectorIndexResponse { + nextPageToken?: string | undefined; + /** Facet aggregation rows returned by a query. */ + facetResult?: FacetResultData | undefined; ++ /** Undocumented: query timings (`response_time`, `latency_ms`, ...) when `debugLevel` is set. */ ++ debugInfo?: Record | undefined; + } + interface RerankerConfig { + /** +diff --git a/dist/v1/model.js b/dist/v1/model.js +index 3cafbc0454e697f5b4f72d6adbc5df23ca9b430d..94bf38e9bf96b326e2f3e359ee3d9a197b24154c 100644 +--- a/dist/v1/model.js ++++ b/dist/v1/model.js +@@ -4,6 +4,7 @@ import { z } from "zod"; + const jsonValueSchema = z.lazy(() => z.union([ + z.null(), + z.number(), ++ z.bigint(), + z.string(), + z.boolean(), + z.record(z.string(), jsonValueSchema), +@@ -288,13 +289,15 @@ const unmarshalPatchEndpointBudgetPolicyResponseSchema = z.object({ + const unmarshalQueryVectorIndexResponseSchema = z.object({ + manifest: z.lazy(() => unmarshalResultManifestSchema).optional(), + result: z.lazy(() => unmarshalResultDataSchema).optional(), +- next_page_token: z.string().optional(), +- facet_result: z.lazy(() => unmarshalFacetResultDataSchema).optional() ++ next_page_token: z.string().nullish(), ++ facet_result: z.lazy(() => unmarshalFacetResultDataSchema).optional(), ++ debug_info: z.record(z.string(), z.unknown()).optional() + }).transform((d) => ({ + manifest: d.manifest, + result: d.result, +- nextPageToken: d.next_page_token, +- facetResult: d.facet_result ++ nextPageToken: d.next_page_token ?? void 0, ++ facetResult: d.facet_result, ++ debugInfo: d.debug_info + })); + const unmarshalResultDataSchema = z.object({ + row_count: z.number().optional(), +@@ -561,9 +564,11 @@ const marshalQueryVectorIndexRequestSchema = z.object({ + reranker: z.lazy(() => marshalRerankerConfigSchema).optional(), + queryColumns: z.array(z.string()).optional(), + sortColumns: z.array(z.string()).optional(), +- facets: z.array(z.string()).optional() ++ facets: z.array(z.string()).optional(), ++ debugLevel: z.number().optional() + }).transform((d) => ({ + name: d.name, ++ debug_level: d.debugLevel, + num_results: d.numResults, + columns: d.columns, + filters_json: d.filtersJson, diff --git a/pnpm-lock.yaml b/pnpm-lock.yaml index c72820060..87927c2ad 100644 --- a/pnpm-lock.yaml +++ b/pnpm-lock.yaml @@ -14,6 +14,7 @@ overrides: patchedDependencies: '@databricks/sdk-genie@0.54.0': 664f3e13b155eb1b5499486a4f6fc1ab9a6a502e2c6b200d06749bd25dd9ae54 '@databricks/sdk-statementexecution@0.52.0': 9377a46b883b548d47116662bf73b159a4da51ec28b9ff1896a7e4a7841c240c + '@databricks/sdk-vectorsearch@0.51.0': b3d79eea4bba92065adc591ee8e6f1d72c0acdac6c00e0cb522074a5fb14ba8e importers: @@ -312,7 +313,7 @@ importers: version: 0.52.0 '@databricks/sdk-vectorsearch': specifier: 0.51.0 - version: 0.51.0 + version: 0.51.0(patch_hash=b3d79eea4bba92065adc591ee8e6f1d72c0acdac6c00e0cb522074a5fb14ba8e) '@databricks/sdk-warehouses': specifier: 0.53.0 version: 0.53.0 @@ -660,7 +661,7 @@ importers: version: 0.52.0 '@databricks/sdk-vectorsearch': specifier: 0.51.0 - version: 0.51.0 + version: 0.51.0(patch_hash=b3d79eea4bba92065adc591ee8e6f1d72c0acdac6c00e0cb522074a5fb14ba8e) '@databricks/sdk-warehouses': specifier: 0.53.0 version: 0.53.0 @@ -14436,7 +14437,7 @@ snapshots: json-bigint: 1.0.0 zod: 4.3.6 - '@databricks/sdk-vectorsearch@0.51.0': + '@databricks/sdk-vectorsearch@0.51.0(patch_hash=b3d79eea4bba92065adc591ee8e6f1d72c0acdac6c00e0cb522074a5fb14ba8e)': dependencies: '@databricks/sdk-auth': 0.51.0 '@databricks/sdk-core': 0.51.0 diff --git a/pnpm-workspace.yaml b/pnpm-workspace.yaml index 3f48fc9af..3d74b5aaf 100644 --- a/pnpm-workspace.yaml +++ b/pnpm-workspace.yaml @@ -21,3 +21,4 @@ allowBuilds: patchedDependencies: '@databricks/sdk-genie@0.54.0': patches/@databricks__sdk-genie@0.54.0.patch '@databricks/sdk-statementexecution@0.52.0': patches/@databricks__sdk-statementexecution@0.52.0.patch + '@databricks/sdk-vectorsearch@0.51.0': patches/@databricks__sdk-vectorsearch@0.51.0.patch