diff --git a/apps/dev-playground/server/index.ts b/apps/dev-playground/server/index.ts index 96c1f94f..913f547c 100644 --- a/apps/dev-playground/server/index.ts +++ b/apps/dev-playground/server/index.ts @@ -8,7 +8,8 @@ import { serving, } from "@databricks/appkit"; import { WorkspaceClient } from "@databricks/sdk-experimental"; -import { vectorSearch } from "../../../packages/appkit/src/plugins/vector-search"; +// TODO: re-enable once vector-search is exported from @databricks/appkit +// import { vectorSearch } from "@databricks/appkit"; import { lakebaseExamples } from "./lakebase-examples-plugin"; import { reconnect } from "./reconnect-plugin"; import { telemetryExamples } from "./telemetry-example-plugin"; @@ -35,16 +36,17 @@ createApp({ lakebaseExamples(), files(), serving(), - vectorSearch({ - indexes: { - demo: { - indexName: - process.env.DATABRICKS_VS_INDEX_NAME ?? "catalog.schema.index", - columns: ["id", "text", "title"], - queryType: "hybrid", - }, - }, - }), + // TODO: re-enable once vector-search is exported from @databricks/appkit + // 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/packages/appkit/src/cache/index.ts b/packages/appkit/src/cache/index.ts index aec14c8a..37bd1659 100644 --- a/packages/appkit/src/cache/index.ts +++ b/packages/appkit/src/cache/index.ts @@ -1,7 +1,7 @@ import { createHash } from "node:crypto"; import { ApiError, WorkspaceClient } from "@databricks/sdk-experimental"; import type { CacheConfig, CacheStorage } from "shared"; -import { createLakebasePool } from "@/connectors/lakebase"; +import { createLakebasePool } from "../connectors/lakebase"; import { AppKitError, ExecutionError, InitializationError } from "../errors"; import { createLogger } from "../logging/logger"; import type { Counter, TelemetryProvider } from "../telemetry"; diff --git a/packages/appkit/src/connectors/lakebase-v1/client.ts b/packages/appkit/src/connectors/lakebase-v1/client.ts index ed3e2fea..9f734f94 100644 --- a/packages/appkit/src/connectors/lakebase-v1/client.ts +++ b/packages/appkit/src/connectors/lakebase-v1/client.ts @@ -2,13 +2,6 @@ import { randomUUID } from "node:crypto"; import type { WorkspaceClient } from "@databricks/sdk-experimental"; import { ApiClient, Config } from "@databricks/sdk-experimental"; import pg from "pg"; -import { - type Counter, - type Histogram, - SpanStatusCode, - TelemetryManager, - type TelemetryProvider, -} from "@/telemetry"; import { AppKitError, AuthenticationError, @@ -17,6 +10,13 @@ import { ValidationError, } from "../../errors"; import { createLogger } from "../../logging/logger"; +import { + type Counter, + type Histogram, + SpanStatusCode, + TelemetryManager, + type TelemetryProvider, +} from "../../telemetry"; import { deepMerge } from "../../utils"; import { lakebaseV1Defaults } from "./defaults"; import type { diff --git a/packages/appkit/src/connectors/lakebase/index.ts b/packages/appkit/src/connectors/lakebase/index.ts index adb09713..c58b7a8c 100644 --- a/packages/appkit/src/connectors/lakebase/index.ts +++ b/packages/appkit/src/connectors/lakebase/index.ts @@ -3,7 +3,7 @@ import { type LakebasePoolConfig, } from "@databricks/lakebase"; import type { Pool } from "pg"; -import { createLogger } from "@/logging/logger"; +import { createLogger } from "../../logging/logger"; /** * Create a Lakebase pool with appkit's logger integration. diff --git a/packages/appkit/src/connectors/vector-search/client.ts b/packages/appkit/src/connectors/vector-search/client.ts index ab45b78d..a20a59ed 100644 --- a/packages/appkit/src/connectors/vector-search/client.ts +++ b/packages/appkit/src/connectors/vector-search/client.ts @@ -82,7 +82,7 @@ export class VectorSearchConnector { const response = (await workspaceClient.apiClient.request({ method: "POST", path: `/api/2.0/vector-search/indexes/${params.indexName}/query`, - body, + payload: body, headers: new Headers({ "Content-Type": "application/json" }), raw: false, query: {}, @@ -149,7 +149,7 @@ export class VectorSearchConnector { const response = (await workspaceClient.apiClient.request({ method: "POST", path: `/api/2.0/vector-search/indexes/${params.indexName}/query-next-page`, - body: { + payload: { endpoint_name: params.endpointName, page_token: params.pageToken, }, diff --git a/packages/appkit/src/plugin/dev-reader.ts b/packages/appkit/src/plugin/dev-reader.ts index e2333629..8aa4ba6a 100644 --- a/packages/appkit/src/plugin/dev-reader.ts +++ b/packages/appkit/src/plugin/dev-reader.ts @@ -1,8 +1,8 @@ import { randomUUID } from "node:crypto"; import type { TunnelConnection } from "shared"; -import { isRemoteTunnelAllowedByEnv } from "@/plugins/server/remote-tunnel/gate"; import { TunnelError } from "../errors"; import { createLogger } from "../logging/logger"; +import { isRemoteTunnelAllowedByEnv } from "../plugins/server/remote-tunnel/gate"; const logger = createLogger("plugin:dev-reader"); diff --git a/packages/appkit/src/plugins/server/vite-dev-server.ts b/packages/appkit/src/plugins/server/vite-dev-server.ts index 1edaa070..0340bbde 100644 --- a/packages/appkit/src/plugins/server/vite-dev-server.ts +++ b/packages/appkit/src/plugins/server/vite-dev-server.ts @@ -2,11 +2,11 @@ import fs from "node:fs"; import path from "node:path"; import type express from "express"; import type { ViteDevServer as ViteDevServerType } from "vite"; -import { mergeConfigDedup } from "@/utils"; import { ServerError } from "../../errors"; import { createLogger } from "../../logging/logger"; import { appKitServingTypesPlugin } from "../../type-generator/serving/vite-plugin"; import { appKitTypesPlugin } from "../../type-generator/vite-plugin"; +import { mergeConfigDedup } from "../../utils"; import { BaseServer } from "./base-server"; import type { PluginClientConfigs, PluginEndpoints } from "./utils"; diff --git a/packages/appkit/src/plugins/vector-search/index.ts b/packages/appkit/src/plugins/vector-search/index.ts index 9052cb03..d733a0f2 100644 --- a/packages/appkit/src/plugins/vector-search/index.ts +++ b/packages/appkit/src/plugins/vector-search/index.ts @@ -1,2 +1,2 @@ -export * from "./vector-search"; export * from "./types"; +export * from "./vector-search"; diff --git a/packages/appkit/src/plugins/vector-search/manifest.json b/packages/appkit/src/plugins/vector-search/manifest.json index ed580ecf..a4451b1a 100644 --- a/packages/appkit/src/plugins/vector-search/manifest.json +++ b/packages/appkit/src/plugins/vector-search/manifest.json @@ -2,6 +2,7 @@ "$schema": "https://databricks.github.io/appkit/schemas/plugin-manifest.schema.json", "name": "vector-search", "displayName": "Vector Search Plugin", + "hidden": true, "description": "Query Databricks Vector Search indexes with built-in hybrid search, reranking, and pagination", "resources": { "required": [ 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 index d3c41f25..104eb1cb 100644 --- a/packages/appkit/src/plugins/vector-search/tests/vector-search.test.ts +++ b/packages/appkit/src/plugins/vector-search/tests/vector-search.test.ts @@ -14,6 +14,7 @@ vi.mock("../../../logging/logger", () => ({ event: () => ({ setComponent: vi.fn().mockReturnThis(), setContext: vi.fn().mockReturnThis(), + setExecution: vi.fn().mockReturnThis(), }), }), })); @@ -215,7 +216,7 @@ describe("VectorSearchPlugin", () => { }), ); - const callBody = mockRequest.mock.calls[0][0].body; + const callBody = mockRequest.mock.calls[0][0].payload; expect(callBody.query_text).toBe("test query"); expect(callBody.query_type).toBe("HYBRID"); expect(callBody.num_results).toBe(10); @@ -250,7 +251,7 @@ describe("VectorSearchPlugin", () => { filters: { category: ["books"] }, }); - const callBody = mockRequest.mock.calls[0][0].body; + const callBody = mockRequest.mock.calls[0][0].payload; expect(callBody.filters).toEqual({ category: ["books"] }); }); @@ -267,7 +268,7 @@ describe("VectorSearchPlugin", () => { await plugin.setup(); await plugin.query("test", { queryText: "test" }); - const callBody = mockRequest.mock.calls[0][0].body; + const callBody = mockRequest.mock.calls[0][0].payload; expect(callBody.reranker.model).toBe("databricks_reranker"); expect(callBody.reranker.parameters.columns_to_rerank).toEqual([ "title", @@ -290,7 +291,7 @@ describe("VectorSearchPlugin", () => { await plugin.query("test", { queryText: "test" }); expect(mockEmbeddingFn).toHaveBeenCalledWith("test"); - const callBody = mockRequest.mock.calls[0][0].body; + const callBody = mockRequest.mock.calls[0][0].payload; expect(callBody.query_vector).toEqual([0.1, 0.2, 0.3]); expect(callBody.query_text).toBeUndefined(); }); diff --git a/packages/appkit/src/plugins/vector-search/types.ts b/packages/appkit/src/plugins/vector-search/types.ts index e9380814..a2760fce 100644 --- a/packages/appkit/src/plugins/vector-search/types.ts +++ b/packages/appkit/src/plugins/vector-search/types.ts @@ -2,7 +2,7 @@ import type { BasePluginConfig } from "shared"; export interface IVectorSearchConfig extends BasePluginConfig { timeout?: number; - indexes: Record; + indexes?: Record; } export interface IndexConfig { diff --git a/packages/appkit/src/plugins/vector-search/vector-search.ts b/packages/appkit/src/plugins/vector-search/vector-search.ts index 81a66595..fefc3f40 100644 --- a/packages/appkit/src/plugins/vector-search/vector-search.ts +++ b/packages/appkit/src/plugins/vector-search/vector-search.ts @@ -1,6 +1,6 @@ import type express from "express"; import type { IAppRouter, PluginExecutionSettings } from "shared"; -import { VectorSearchConnector } from "../../connectors"; +import { VectorSearchConnector } from "../../connectors/vector-search/client"; import type { VsRawResponse } from "../../connectors/vector-search/types"; import { getWorkspaceClient } from "../../context"; import { createLogger } from "../../logging/logger"; @@ -40,6 +40,11 @@ export class VectorSearchPlugin extends Plugin { } async setup(): Promise { + if (!this.config.indexes || Object.keys(this.config.indexes).length === 0) { + throw new Error( + 'VectorSearchPlugin requires at least one index in "indexes" config', + ); + } for (const [alias, idx] of Object.entries(this.config.indexes)) { if (!idx.indexName) { throw new Error( @@ -109,11 +114,13 @@ export class VectorSearchPlugin extends Plugin { querySettings, ); - if (result === undefined) { - res.status(500).json({ error: "Query failed", plugin: this.name }); + if (!result.ok) { + res + .status(result.status) + .json({ error: result.message, plugin: this.name }); return; } - res.json(this._parseResponse(result, prepared.queryType)); + res.json(this._parseResponse(result.data, prepared.queryType)); } catch (error) { this._handleError(res, error, "Query failed"); } @@ -177,14 +184,14 @@ export class VectorSearchPlugin extends Plugin { querySettings, ); - if (result === undefined) { + if (!result.ok) { res - .status(500) - .json({ error: "Next-page query failed", plugin: this.name }); + .status(result.status) + .json({ error: result.message, plugin: this.name }); return; } res.json( - this._parseResponse(result, indexConfig.queryType ?? "hybrid"), + this._parseResponse(result.data, indexConfig.queryType ?? "hybrid"), ); } catch (error) { this._handleError(res, error, "Next-page query failed"); @@ -196,7 +203,7 @@ export class VectorSearchPlugin extends Plugin { name: "getConfig", method: "get", path: "/:alias/config", - handler: (req: express.Request, res: express.Response) => { + handler: async (req: express.Request, res: express.Response) => { const { alias } = req.params; const indexConfig = this._resolveIndex(alias); if (!indexConfig) { @@ -249,11 +256,13 @@ export class VectorSearchPlugin extends Plugin { querySettings, ); - if (result === undefined) { - throw new Error(`Vector search query failed for index "${alias}"`); + if (!result.ok) { + throw new Error( + `Vector search query failed for index "${alias}": ${result.message}`, + ); } - return this._parseResponse(result, prepared.queryType); + return this._parseResponse(result.data, prepared.queryType); } async shutdown(): Promise { @@ -267,7 +276,7 @@ export class VectorSearchPlugin extends Plugin { } private _resolveIndex(alias: string): IndexConfig | undefined { - return this.config.indexes[alias]; + return this.config.indexes?.[alias]; } private async _prepareQuery(