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 5a34f1797..b80fef165 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..313700e4f 100644 --- a/packages/appkit/src/connectors/ai-search/client.ts +++ b/packages/appkit/src/connectors/ai-search/client.ts @@ -6,11 +6,13 @@ import { SpanStatusCode, TelemetryManager, } from "../../telemetry"; -import type { WorkspaceClient } from "../../workspace-client"; -import { contextFromAbortSignal } from "../context"; +import type { + QueryVectorIndexRequest, + QueryVectorIndexResponse, + WorkspaceClient, +} from "../../workspace-client"; import type { AiSearchConnectorConfig, - UcTableInfo, VsIndexInfo, VsNextPageParams, VsQueryParams, @@ -19,6 +21,32 @@ import type { const logger = createLogger("connectors:ai-search"); +/** + * 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. + */ +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 { private readonly telemetry: TelemetryProvider; @@ -38,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)", @@ -74,26 +102,18 @@ 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 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 = toRawResponse( + await workspaceClient.vectorSearch.queryVectorIndex(req, { + signal, + }), + ); const duration = Date.now() - startTime; span.setAttribute("vs.result_count", response.result.row_count); @@ -153,20 +173,16 @@ export class AiSearchConnector { }, async (span: Span) => { try { - const response = (await workspaceClient.apiClient.request( - { - method: "POST", - path: `/api/2.0/vector-search/indexes/${params.indexName}/query-next-page`, - payload: { - endpoint_name: params.endpointName, - page_token: params.pageToken, + const response = toRawResponse( + await workspaceClient.vectorSearch.queryVectorIndexNextPage( + { + name: params.indexName, + endpointName: params.endpointName, + pageToken: params.pageToken, }, - headers: new Headers({ "Content-Type": "application/json" }), - raw: false, - query: {}, - }, - contextFromAbortSignal(signal), - )) as VsRawResponse; + { signal }, + ), + ); span.setAttribute("vs.result_count", response.result.row_count); span.setStatus({ code: SpanStatusCode.OK }); @@ -193,16 +209,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 +238,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..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,7 +7,7 @@ import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; import { withEnv } from "../../../testing"; import { useTestCache } from "../../../testing/test-cache"; -import { Context } from "../../../workspace-client"; +import { createWorkspaceClient } from "../../../workspace-client"; vi.mock("../../../context", () => ({ getWorkspaceClient: vi.fn(() => mockWorkspaceClient), @@ -119,17 +119,48 @@ const validVsResponse = { debug_info: { response_time: 35 }, }; +// 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 = { - apiClient: { request: mockRequest }, + vectorSearch: { + queryVectorIndex: realVectorSearch.queryVectorIndex.bind(realVectorSearch), + queryVectorIndexNextPage: + realVectorSearch.queryVectorIndexNextPage.bind(realVectorSearch), + getVectorIndex: mockGetVectorIndex, + }, + tables: { getTable: mockGetTable }, }; +/** 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"; describe("AiSearchPlugin", () => { beforeEach(() => { mockRequest.mockClear(); mockRequest.mockResolvedValue(validVsResponse); + mockGetVectorIndex.mockReset(); + mockGetTable.mockReset(); }); describe("setup()", () => { @@ -218,25 +249,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 +273,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 +282,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 +305,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(); }); }); @@ -338,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; @@ -378,16 +426,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 +468,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 +490,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 +514,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 +535,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 +556,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 +663,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 +682,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 +699,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 +711,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 +1062,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 +1144,15 @@ 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), ); + // 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/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..ec1dbe624 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,22 @@ export class AppKitWorkspaceClient implements WorkspaceClient { return this.#jobs; } + // 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 c64136198..4c645ab85 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 @@ -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 @@ -52,6 +58,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"; @@ -591,6 +599,30 @@ export function buildFilesClient(opts: WorkspaceClientOptions): FilesClient { ); } +/** + * 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, +): 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 { @@ -601,6 +633,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 ──────────────────────────────────────────────── @@ -648,3 +682,7 @@ export type { EndpointState, GetWarehouseResponse, } from "@databricks/sdk-warehouses/v1"; +export type { + QueryVectorIndexRequest, + QueryVectorIndexResponse, +} from "@databricks/sdk-vectorsearch/v1"; diff --git a/packages/shared/src/workspace-client/types.ts b/packages/shared/src/workspace-client/types.ts index f32f5b517..4444e5fbe 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/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 3c6b8a5cf..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: @@ -307,6 +308,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(patch_hash=b3d79eea4bba92065adc591ee8e6f1d72c0acdac6c00e0cb522074a5fb14ba8e) '@databricks/sdk-warehouses': specifier: 0.53.0 version: 0.53.0 @@ -649,6 +656,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(patch_hash=b3d79eea4bba92065adc591ee8e6f1d72c0acdac6c00e0cb522074a5fb14ba8e) '@databricks/sdk-warehouses': specifier: 0.53.0 version: 0.53.0 @@ -2057,6 +2070,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 +14428,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(patch_hash=b3d79eea4bba92065adc591ee8e6f1d72c0acdac6c00e0cb522074a5fb14ba8e)': + 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 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