diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 212e06a89..fc0fb6272 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -68,7 +68,7 @@ jobs: docs/docs/plugins/lakebase.md \ docs/docs/plugins/model-serving.md \ docs/docs/plugins/server.md \ - docs/docs/plugins/vector-search.md; then + docs/docs/plugins/ai-search.md; then echo "❌ Error: Generated files are out of sync with their source manifests/schemas." echo "" echo "To fix this:" diff --git a/apps/dev-playground/client/src/lib/nav.ts b/apps/dev-playground/client/src/lib/nav.ts index 00f70dfec..86c086805 100644 --- a/apps/dev-playground/client/src/lib/nav.ts +++ b/apps/dev-playground/client/src/lib/nav.ts @@ -114,8 +114,8 @@ export const NAV_GROUPS: ReadonlyArray = [ icon: LineChartIcon, }, { - to: "/vector-search", - label: "Vector Search", + to: "/ai-search", + label: "AI Search", description: "Semantic search backed by Databricks vector indexes, wired into AppKit's retrieval API.", icon: SearchIcon, diff --git a/apps/dev-playground/client/src/routeTree.gen.ts b/apps/dev-playground/client/src/routeTree.gen.ts index 450287592..94034f5f7 100644 --- a/apps/dev-playground/client/src/routeTree.gen.ts +++ b/apps/dev-playground/client/src/routeTree.gen.ts @@ -9,7 +9,6 @@ // Additionally, you should also exclude this file from your linter and/or formatter to prevent it from being checked or modified. import { Route as rootRouteImport } from './routes/__root' -import { Route as VectorSearchRouteRouteImport } from './routes/vector-search.route' import { Route as UiVariantsRouteRouteImport } from './routes/ui-variants.route' import { Route as TypeSafetyRouteRouteImport } from './routes/type-safety.route' import { Route as TelemetryRouteRouteImport } from './routes/telemetry.route' @@ -26,14 +25,10 @@ import { Route as DataVisualizationRouteRouteImport } from './routes/data-visual import { Route as ChartInferenceRouteRouteImport } from './routes/chart-inference.route' import { Route as ArrowAnalyticsRouteRouteImport } from './routes/arrow-analytics.route' import { Route as AnalyticsRouteRouteImport } from './routes/analytics.route' +import { Route as AiSearchRouteRouteImport } from './routes/ai-search.route' import { Route as AgentRouteRouteImport } from './routes/agent.route' import { Route as IndexRouteImport } from './routes/index' -const VectorSearchRouteRoute = VectorSearchRouteRouteImport.update({ - id: '/vector-search', - path: '/vector-search', - getParentRoute: () => rootRouteImport, -} as any) const UiVariantsRouteRoute = UiVariantsRouteRouteImport.update({ id: '/ui-variants', path: '/ui-variants', @@ -114,6 +109,11 @@ const AnalyticsRouteRoute = AnalyticsRouteRouteImport.update({ path: '/analytics', getParentRoute: () => rootRouteImport, } as any) +const AiSearchRouteRoute = AiSearchRouteRouteImport.update({ + id: '/ai-search', + path: '/ai-search', + getParentRoute: () => rootRouteImport, +} as any) const AgentRouteRoute = AgentRouteRouteImport.update({ id: '/agent', path: '/agent', @@ -128,6 +128,7 @@ const IndexRoute = IndexRouteImport.update({ export interface FileRoutesByFullPath { '/': typeof IndexRoute '/agent': typeof AgentRouteRoute + '/ai-search': typeof AiSearchRouteRoute '/analytics': typeof AnalyticsRouteRoute '/arrow-analytics': typeof ArrowAnalyticsRouteRoute '/chart-inference': typeof ChartInferenceRouteRoute @@ -144,11 +145,11 @@ export interface FileRoutesByFullPath { '/telemetry': typeof TelemetryRouteRoute '/type-safety': typeof TypeSafetyRouteRoute '/ui-variants': typeof UiVariantsRouteRoute - '/vector-search': typeof VectorSearchRouteRoute } export interface FileRoutesByTo { '/': typeof IndexRoute '/agent': typeof AgentRouteRoute + '/ai-search': typeof AiSearchRouteRoute '/analytics': typeof AnalyticsRouteRoute '/arrow-analytics': typeof ArrowAnalyticsRouteRoute '/chart-inference': typeof ChartInferenceRouteRoute @@ -165,12 +166,12 @@ export interface FileRoutesByTo { '/telemetry': typeof TelemetryRouteRoute '/type-safety': typeof TypeSafetyRouteRoute '/ui-variants': typeof UiVariantsRouteRoute - '/vector-search': typeof VectorSearchRouteRoute } export interface FileRoutesById { __root__: typeof rootRouteImport '/': typeof IndexRoute '/agent': typeof AgentRouteRoute + '/ai-search': typeof AiSearchRouteRoute '/analytics': typeof AnalyticsRouteRoute '/arrow-analytics': typeof ArrowAnalyticsRouteRoute '/chart-inference': typeof ChartInferenceRouteRoute @@ -187,13 +188,13 @@ export interface FileRoutesById { '/telemetry': typeof TelemetryRouteRoute '/type-safety': typeof TypeSafetyRouteRoute '/ui-variants': typeof UiVariantsRouteRoute - '/vector-search': typeof VectorSearchRouteRoute } export interface FileRouteTypes { fileRoutesByFullPath: FileRoutesByFullPath fullPaths: | '/' | '/agent' + | '/ai-search' | '/analytics' | '/arrow-analytics' | '/chart-inference' @@ -210,11 +211,11 @@ export interface FileRouteTypes { | '/telemetry' | '/type-safety' | '/ui-variants' - | '/vector-search' fileRoutesByTo: FileRoutesByTo to: | '/' | '/agent' + | '/ai-search' | '/analytics' | '/arrow-analytics' | '/chart-inference' @@ -231,11 +232,11 @@ export interface FileRouteTypes { | '/telemetry' | '/type-safety' | '/ui-variants' - | '/vector-search' id: | '__root__' | '/' | '/agent' + | '/ai-search' | '/analytics' | '/arrow-analytics' | '/chart-inference' @@ -252,12 +253,12 @@ export interface FileRouteTypes { | '/telemetry' | '/type-safety' | '/ui-variants' - | '/vector-search' fileRoutesById: FileRoutesById } export interface RootRouteChildren { IndexRoute: typeof IndexRoute AgentRouteRoute: typeof AgentRouteRoute + AiSearchRouteRoute: typeof AiSearchRouteRoute AnalyticsRouteRoute: typeof AnalyticsRouteRoute ArrowAnalyticsRouteRoute: typeof ArrowAnalyticsRouteRoute ChartInferenceRouteRoute: typeof ChartInferenceRouteRoute @@ -274,18 +275,10 @@ export interface RootRouteChildren { TelemetryRouteRoute: typeof TelemetryRouteRoute TypeSafetyRouteRoute: typeof TypeSafetyRouteRoute UiVariantsRouteRoute: typeof UiVariantsRouteRoute - VectorSearchRouteRoute: typeof VectorSearchRouteRoute } declare module '@tanstack/react-router' { interface FileRoutesByPath { - '/vector-search': { - id: '/vector-search' - path: '/vector-search' - fullPath: '/vector-search' - preLoaderRoute: typeof VectorSearchRouteRouteImport - parentRoute: typeof rootRouteImport - } '/ui-variants': { id: '/ui-variants' path: '/ui-variants' @@ -398,6 +391,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof AnalyticsRouteRouteImport parentRoute: typeof rootRouteImport } + '/ai-search': { + id: '/ai-search' + path: '/ai-search' + fullPath: '/ai-search' + preLoaderRoute: typeof AiSearchRouteRouteImport + parentRoute: typeof rootRouteImport + } '/agent': { id: '/agent' path: '/agent' @@ -418,6 +418,7 @@ declare module '@tanstack/react-router' { const rootRouteChildren: RootRouteChildren = { IndexRoute: IndexRoute, AgentRouteRoute: AgentRouteRoute, + AiSearchRouteRoute: AiSearchRouteRoute, AnalyticsRouteRoute: AnalyticsRouteRoute, ArrowAnalyticsRouteRoute: ArrowAnalyticsRouteRoute, ChartInferenceRouteRoute: ChartInferenceRouteRoute, @@ -434,7 +435,6 @@ const rootRouteChildren: RootRouteChildren = { TelemetryRouteRoute: TelemetryRouteRoute, TypeSafetyRouteRoute: TypeSafetyRouteRoute, UiVariantsRouteRoute: UiVariantsRouteRoute, - VectorSearchRouteRoute: VectorSearchRouteRoute, } export const routeTree = rootRouteImport ._addFileChildren(rootRouteChildren) diff --git a/apps/dev-playground/client/src/routes/vector-search.route.tsx b/apps/dev-playground/client/src/routes/ai-search.route.tsx similarity index 60% rename from apps/dev-playground/client/src/routes/vector-search.route.tsx rename to apps/dev-playground/client/src/routes/ai-search.route.tsx index bed4885c7..b64874167 100644 --- a/apps/dev-playground/client/src/routes/vector-search.route.tsx +++ b/apps/dev-playground/client/src/routes/ai-search.route.tsx @@ -6,73 +6,35 @@ import { CardTitle, Input, } from "@databricks/appkit-ui/react"; +import { useAiSearchQuery } from "@databricks/appkit-ui/react/beta"; 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, +export const Route = createFileRoute("/ai-search")({ + component: AiSearchRoute, }); -interface SearchResult { - score: number; - data: Record; -} - -interface SearchResponse { - results: SearchResult[]; - totalCount: number; - queryTimeMs: number; - queryType: string; -} - -function VectorSearchRoute() { +function AiSearchRoute() { 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 { search, data, loading, error } = useAiSearchQuery({ alias: "demo" }); - const data: SearchResponse = await res.json(); - setResponse(data); - } catch (err) { - setError(err instanceof Error ? err.message : String(err)); - } finally { - setLoading(false); - } + const handleSearch = () => { + if (query.trim()) void search(query); }; const handleKeyDown = (e: React.KeyboardEvent) => { - if (e.key === "Enter") { - void handleSearch(); - } + if (e.key === "Enter") handleSearch(); }; return (
@@ -84,10 +46,7 @@ function VectorSearchRoute() { onKeyDown={handleKeyDown} className="flex-1" /> -
)} - {response && ( + {data && (
- {response.totalCount} result - {response.totalCount !== 1 ? "s" : ""} ·{" "} - {response.queryTimeMs}ms · {response.queryType} + {data.totalCount} result + {data.totalCount !== 1 ? "s" : ""} · {data.queryTimeMs}ms + · {data.queryType}
- {response.results.length === 0 ? ( + {data.results.length === 0 ? (

No results found.

) : ( - response.results.map((result, index) => ( + data.results.map((result, index) => ( diff --git a/apps/dev-playground/server/index.ts b/apps/dev-playground/server/index.ts index beb19fead..b30c51684 100644 --- a/apps/dev-playground/server/index.ts +++ b/apps/dev-playground/server/index.ts @@ -14,6 +14,7 @@ import { } from "@databricks/appkit"; import { agents, + aiSearch, createAgent, DatabricksAdapter, supervisorTools, @@ -429,17 +430,14 @@ createApp({ // sense as the user-facing landing agent). defaultAgent: "helper", }), - // 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", - // }, - // }, - // }), + aiSearch({ + indexes: { + demo: { + columns: ["id", "text", "title"], + queryType: "hybrid", + }, + }, + }), ], ...(process.env.APPKIT_E2E_TEST && { client: createMockClient() }), async onPluginsReady(appkit) { diff --git a/docs/docs/api/appkit/Interface.BasePluginConfig.md b/docs/docs/api/appkit/Interface.BasePluginConfig.md index 653df68ce..a109fd560 100644 --- a/docs/docs/api/appkit/Interface.BasePluginConfig.md +++ b/docs/docs/api/appkit/Interface.BasePluginConfig.md @@ -5,6 +5,7 @@ Base configuration interface for AppKit plugins ## Extended by - [`AgentsPluginConfig`](Interface.AgentsPluginConfig.md) +- [`IAiSearchConfig`](Interface.IAiSearchConfig.md) - [`IJobsConfig`](Interface.IJobsConfig.md) ## Indexable diff --git a/docs/docs/api/appkit/Interface.IAiSearchConfig.md b/docs/docs/api/appkit/Interface.IAiSearchConfig.md new file mode 100644 index 000000000..316b4aac3 --- /dev/null +++ b/docs/docs/api/appkit/Interface.IAiSearchConfig.md @@ -0,0 +1,65 @@ +# Interface: IAiSearchConfig + +Base configuration interface for AppKit plugins + +## Extends + +- [`BasePluginConfig`](Interface.BasePluginConfig.md) + +## Indexable + +```ts +[key: string]: unknown +``` + +## Properties + +### host? + +```ts +optional host: string; +``` + +#### Inherited from + +[`BasePluginConfig`](Interface.BasePluginConfig.md).[`host`](Interface.BasePluginConfig.md#host) + +*** + +### indexes? + +```ts +optional indexes: Record; +``` + +*** + +### name? + +```ts +optional name: string; +``` + +#### Inherited from + +[`BasePluginConfig`](Interface.BasePluginConfig.md).[`name`](Interface.BasePluginConfig.md#name) + +*** + +### telemetry? + +```ts +optional telemetry: TelemetryOptions; +``` + +#### Inherited from + +[`BasePluginConfig`](Interface.BasePluginConfig.md).[`telemetry`](Interface.BasePluginConfig.md#telemetry) + +*** + +### timeout? + +```ts +optional timeout: number; +``` diff --git a/docs/docs/api/appkit/Interface.IndexConfig.md b/docs/docs/api/appkit/Interface.IndexConfig.md new file mode 100644 index 000000000..95b6797cc --- /dev/null +++ b/docs/docs/api/appkit/Interface.IndexConfig.md @@ -0,0 +1,110 @@ +# Interface: IndexConfig + +## Properties + +### auth? + +```ts +optional auth: "service-principal" | "on-behalf-of-user"; +``` + +Auth mode for the built-in HTTP routes — "service-principal" (default) +uses the app's SP, "on-behalf-of-user" proxies the logged-in user's token. +Programmatic callers select per call via `appkit.aiSearch.asUser(req)`. + +*** + +### columns? + +```ts +optional columns: string[]; +``` + +Columns to return in results. Optional: in development the plugin +auto-discovers them from the index's source table when omitted (and warns +that they should be set explicitly for production). + +*** + +### embeddingFn()? + +```ts +optional embeddingFn: (text: string) => Promise; +``` + +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). + +#### Parameters + +| Parameter | Type | +| ------ | ------ | +| `text` | `string` | + +#### Returns + +`Promise`\<`number`[]\> + +*** + +### endpointName? + +```ts +optional endpointName: string; +``` + +VS endpoint name (required when pagination is true) + +*** + +### indexName? + +```ts +optional indexName: string; +``` + +Three-level UC name: catalog.schema.index_name. Defaults to the +`DATABRICKS_VS_INDEX_NAME` env var when omitted — so multiple aliases that +omit it all resolve to that same physical index. Set it explicitly per +alias when they should point at distinct indexes. + +*** + +### numResults? + +```ts +optional numResults: number; +``` + +Max results per query + +*** + +### pagination? + +```ts +optional pagination: boolean; +``` + +Enable cursor pagination + +*** + +### queryType? + +```ts +optional queryType: SearchQueryType; +``` + +Default search mode + +*** + +### reranker? + +```ts +optional reranker: boolean | RerankerConfig; +``` + +Enable built-in reranker. Pass true to rerank all non-id columns, or an object for fine control. diff --git a/docs/docs/api/appkit/Interface.RerankerConfig.md b/docs/docs/api/appkit/Interface.RerankerConfig.md new file mode 100644 index 000000000..4012fc326 --- /dev/null +++ b/docs/docs/api/appkit/Interface.RerankerConfig.md @@ -0,0 +1,9 @@ +# Interface: RerankerConfig + +## Properties + +### columnsToRerank + +```ts +columnsToRerank: string[]; +``` diff --git a/docs/docs/api/appkit/Interface.SearchRequest.md b/docs/docs/api/appkit/Interface.SearchRequest.md new file mode 100644 index 000000000..6479ab0da --- /dev/null +++ b/docs/docs/api/appkit/Interface.SearchRequest.md @@ -0,0 +1,57 @@ +# Interface: SearchRequest + +## Properties + +### columns? + +```ts +optional columns: string[]; +``` + +*** + +### filters? + +```ts +optional filters: SearchFilters; +``` + +*** + +### numResults? + +```ts +optional numResults: number; +``` + +*** + +### queryText? + +```ts +optional queryText: string; +``` + +*** + +### queryType? + +```ts +optional queryType: SearchQueryType; +``` + +*** + +### queryVector? + +```ts +optional queryVector: number[]; +``` + +*** + +### reranker? + +```ts +optional reranker: boolean; +``` diff --git a/docs/docs/api/appkit/Interface.SearchResponse.md b/docs/docs/api/appkit/Interface.SearchResponse.md new file mode 100644 index 000000000..40d8941f3 --- /dev/null +++ b/docs/docs/api/appkit/Interface.SearchResponse.md @@ -0,0 +1,47 @@ +# Interface: SearchResponse\ + +## Type Parameters + +| Type Parameter | Default type | +| ------ | ------ | +| `T` *extends* `Record`\<`string`, `unknown`\> | `Record`\<`string`, `unknown`\> | + +## Properties + +### nextPageToken + +```ts +nextPageToken: string | null; +``` + +*** + +### queryTimeMs + +```ts +queryTimeMs: number; +``` + +*** + +### queryType + +```ts +queryType: SearchQueryType; +``` + +*** + +### results + +```ts +results: SearchResult[]; +``` + +*** + +### totalCount + +```ts +totalCount: number; +``` diff --git a/docs/docs/api/appkit/Interface.SearchResult.md b/docs/docs/api/appkit/Interface.SearchResult.md new file mode 100644 index 000000000..0ba3bf264 --- /dev/null +++ b/docs/docs/api/appkit/Interface.SearchResult.md @@ -0,0 +1,23 @@ +# Interface: SearchResult\ + +## Type Parameters + +| Type Parameter | Default type | +| ------ | ------ | +| `T` *extends* `Record`\<`string`, `unknown`\> | `Record`\<`string`, `unknown`\> | + +## Properties + +### data + +```ts +data: T; +``` + +*** + +### score + +```ts +score: number; +``` diff --git a/docs/docs/api/appkit/TypeAlias.SearchFilters.md b/docs/docs/api/appkit/TypeAlias.SearchFilters.md new file mode 100644 index 000000000..4b3149f02 --- /dev/null +++ b/docs/docs/api/appkit/TypeAlias.SearchFilters.md @@ -0,0 +1,5 @@ +# Type Alias: SearchFilters + +```ts +type SearchFilters = Record; +``` diff --git a/docs/docs/api/appkit/Variable.aiSearch.md b/docs/docs/api/appkit/Variable.aiSearch.md new file mode 100644 index 000000000..712d8d94a --- /dev/null +++ b/docs/docs/api/appkit/Variable.aiSearch.md @@ -0,0 +1,5 @@ +# Variable: aiSearch + +```ts +const aiSearch: ToPlugin; +``` diff --git a/docs/docs/api/appkit/index.md b/docs/docs/api/appkit/index.md index 124db7c8f..b7d394176 100644 --- a/docs/docs/api/appkit/index.md +++ b/docs/docs/api/appkit/index.md @@ -51,7 +51,9 @@ surface with `@databricks/appkit/beta`. Not meant for application imports. | [GenerateDatabaseCredentialRequest](Interface.GenerateDatabaseCredentialRequest.md) | Request parameters for generating database OAuth credentials | | [GenerationParams](Interface.GenerationParams.md) | Optional generation parameters forwarded to the OpenAI-compatible serving request body. Names match the serving API wire keys. Only keys that are set are sent — undefined values are omitted so the endpoint applies its own defaults. Ranges are not validated here; the serving endpoint validates. | | [HostedSupervisorTool](Interface.HostedSupervisorTool.md) | Tagged record returned by every [supervisorTools](Variable.supervisorTools.md) factory. The `__kind` discriminator lets the agents plugin (and standalone `runAgent`) classify these tools without a structural match against the wire format — keeps the SA wire shape free to evolve and avoids namespace collisions with MCP hosted tools (which use `type: "genie-space"` hyphenated, vs SA's `type: "genie_space"` underscored). | +| [IAiSearchConfig](Interface.IAiSearchConfig.md) | Base configuration interface for AppKit plugins | | [IJobsConfig](Interface.IJobsConfig.md) | Configuration for the Jobs plugin. | +| [IndexConfig](Interface.IndexConfig.md) | - | | [ITelemetry](Interface.ITelemetry.md) | Plugin-facing interface for OpenTelemetry instrumentation. Provides a thin abstraction over OpenTelemetry APIs for plugins. | | [JobAPI](Interface.JobAPI.md) | User-facing API for a single configured job. | | [JobConfig](Interface.JobConfig.md) | Per-job configuration options. | @@ -67,10 +69,14 @@ surface with `@databricks/appkit/beta`. Not meant for application imports. | [RegisteredAgent](Interface.RegisteredAgent.md) | - | | [RequestedClaims](Interface.RequestedClaims.md) | Optional claims for fine-grained Unity Catalog table permissions When specified, the returned token will be scoped to only the requested tables | | [RequestedResource](Interface.RequestedResource.md) | Resource to request permissions for in Unity Catalog | +| [RerankerConfig](Interface.RerankerConfig.md) | - | | [ResourceEntry](Interface.ResourceEntry.md) | Internal representation of a resource in the registry. Extends ResourceRequirement with resolution state and plugin ownership. | | [ResourceRequirement](Interface.ResourceRequirement.md) | Declares a resource requirement for a plugin. Can be defined statically in a manifest or dynamically via getResourceRequirements(). | | [RunAgentInput](Interface.RunAgentInput.md) | - | | [RunAgentResult](Interface.RunAgentResult.md) | - | +| [SearchRequest](Interface.SearchRequest.md) | - | +| [SearchResponse](Interface.SearchResponse.md) | - | +| [SearchResult](Interface.SearchResult.md) | - | | [ServingEndpointEntry](Interface.ServingEndpointEntry.md) | Shape of a single registry entry. | | [ServingEndpointRegistry](Interface.ServingEndpointRegistry.md) | Registry interface for serving endpoint type generation. Empty by default — augmented by the Vite type generator's `.d.ts` output via module augmentation. When populated, provides autocomplete for alias names and typed request/response/chunk per endpoint. | | [StreamExecutionSettings](Interface.StreamExecutionSettings.md) | Execution settings for streaming endpoints. Extends PluginExecutionSettings with SSE stream configuration. | @@ -112,6 +118,7 @@ surface with `@databricks/appkit/beta`. Not meant for application imports. | [ResolvedToolEntry](TypeAlias.ResolvedToolEntry.md) | Internal tool-index entry after a tool record has been resolved to a dispatchable form. | | [ResourceFieldEntry](TypeAlias.ResourceFieldEntry.md) | - | | [ResourcePermission](TypeAlias.ResourcePermission.md) | Union of all possible permission levels across all resource types. | +| [SearchFilters](TypeAlias.SearchFilters.md) | - | | [ServingFactory](TypeAlias.ServingFactory.md) | Factory function returned by `AppKit.serving`. | | [SupervisorTool](TypeAlias.SupervisorTool.md) | Tools supported by the Databricks AI Gateway Responses API. The shapes match the wire format the endpoint expects, so the adapter passes the array straight into the request body. | | [ToolRegistry](TypeAlias.ToolRegistry.md) | - | @@ -122,6 +129,7 @@ surface with `@databricks/appkit/beta`. Not meant for application imports. | Variable | Description | | ------ | ------ | | [agents](Variable.agents.md) | Plugin factory for the agents plugin. Reads `config/agents/*.md` by default, resolves toolkits/tools from registered plugins, exposes `appkit.agents.*` runtime API and mounts `POST /invocations` and `POST /responses` (aliased non-streaming invoke endpoints) plus `POST /chat` (streaming, HITL-capable). | +| [aiSearch](Variable.aiSearch.md) | - | | [READ\_ACTIONS](Variable.READ_ACTIONS.md) | Actions that only read data. | | [sql](Variable.sql.md) | SQL helper namespace | | [SUPERVISOR\_EXTENSION\_KEY](Variable.SUPERVISOR_EXTENSION_KEY.md) | Namespace key under which the adapter reads its hosted-tool payload from [AgentInput.extensions](Interface.AgentInput.md#extensions). Exported so the agents plugin and standalone `runAgent` (the producers) can write under the same key the adapter reads. | diff --git a/docs/docs/api/appkit/typedoc-sidebar.ts b/docs/docs/api/appkit/typedoc-sidebar.ts index 1c510b7a8..489424a2e 100644 --- a/docs/docs/api/appkit/typedoc-sidebar.ts +++ b/docs/docs/api/appkit/typedoc-sidebar.ts @@ -187,11 +187,21 @@ const typedocSidebar: SidebarsConfig = { id: "api/appkit/Interface.HostedSupervisorTool", label: "HostedSupervisorTool" }, + { + type: "doc", + id: "api/appkit/Interface.IAiSearchConfig", + label: "IAiSearchConfig" + }, { type: "doc", id: "api/appkit/Interface.IJobsConfig", label: "IJobsConfig" }, + { + type: "doc", + id: "api/appkit/Interface.IndexConfig", + label: "IndexConfig" + }, { type: "doc", id: "api/appkit/Interface.ITelemetry", @@ -267,6 +277,11 @@ const typedocSidebar: SidebarsConfig = { id: "api/appkit/Interface.RequestedResource", label: "RequestedResource" }, + { + type: "doc", + id: "api/appkit/Interface.RerankerConfig", + label: "RerankerConfig" + }, { type: "doc", id: "api/appkit/Interface.ResourceEntry", @@ -287,6 +302,21 @@ const typedocSidebar: SidebarsConfig = { id: "api/appkit/Interface.RunAgentResult", label: "RunAgentResult" }, + { + type: "doc", + id: "api/appkit/Interface.SearchRequest", + label: "SearchRequest" + }, + { + type: "doc", + id: "api/appkit/Interface.SearchResponse", + label: "SearchResponse" + }, + { + type: "doc", + id: "api/appkit/Interface.SearchResult", + label: "SearchResult" + }, { type: "doc", id: "api/appkit/Interface.ServingEndpointEntry", @@ -473,6 +503,11 @@ const typedocSidebar: SidebarsConfig = { id: "api/appkit/TypeAlias.ResourcePermission", label: "ResourcePermission" }, + { + type: "doc", + id: "api/appkit/TypeAlias.SearchFilters", + label: "SearchFilters" + }, { type: "doc", id: "api/appkit/TypeAlias.ServingFactory", @@ -504,6 +539,11 @@ const typedocSidebar: SidebarsConfig = { id: "api/appkit/Variable.agents", label: "agents" }, + { + type: "doc", + id: "api/appkit/Variable.aiSearch", + label: "aiSearch" + }, { type: "doc", id: "api/appkit/Variable.READ_ACTIONS", diff --git a/docs/docs/plugins/vector-search.md b/docs/docs/plugins/ai-search.md similarity index 66% rename from docs/docs/plugins/vector-search.md rename to docs/docs/plugins/ai-search.md index 5d704641d..397e00f3a 100644 --- a/docs/docs/plugins/vector-search.md +++ b/docs/docs/plugins/ai-search.md @@ -2,7 +2,13 @@ sidebar_position: 9 --- -# Vector Search plugin +# AI Search plugin + + +:::warning Beta plugin +This plugin is currently **beta**. APIs may change between minor releases. Import from `@databricks/appkit/beta`. See [Plugin Stability Tiers](./stability.md). +::: + Query Databricks Vector Search indexes with hybrid search, reranking, and cursor pagination from your AppKit application. @@ -17,12 +23,13 @@ Query Databricks Vector Search indexes with hybrid search, reranking, and cursor ## Basic usage ```ts -import { createApp, vectorSearch, server } from "@databricks/appkit"; +import { createApp, server } from "@databricks/appkit"; +import { aiSearch } from "@databricks/appkit/beta"; await createApp({ plugins: [ server(), - vectorSearch({ + aiSearch({ indexes: { products: { indexName: "catalog.schema.products_idx", @@ -48,7 +55,7 @@ await createApp({ Index aliases let you reference multiple Vector Search indexes by name. The alias is used in API routes and programmatic calls: ```ts -vectorSearch({ +aiSearch({ indexes: { products: { indexName: "catalog.schema.products_idx", @@ -63,12 +70,19 @@ vectorSearch({ }); ``` +:::note +An alias without its own `indexName` falls back to the `DATABRICKS_VS_INDEX_NAME` +env var. If several aliases omit `indexName`, they all resolve to that one +physical index (with their own per-alias `columns`, `queryType`, etc.). Give +each alias an explicit `indexName` when you mean distinct indexes. +::: + ## 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 | +| `indexName` | `string` | `DATABRICKS_VS_INDEX_NAME` | Three-level Unity Catalog name (`catalog.schema.index`). Defaults to the `DATABRICKS_VS_INDEX_NAME` env var when omitted. | +| `columns` | `string[]` | auto-discovered in dev | Columns to return in query results. Optional in development — when omitted, the plugin reads them from the index's source table and warns. **Set explicitly for production**, where a missing value is not auto-filled. | | `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 | @@ -88,7 +102,7 @@ vectorSearch({ Reranking improves result relevance by running a second-stage model over the initial candidates: ```ts -vectorSearch({ +aiSearch({ indexes: { products: { indexName: "catalog.schema.products_idx", @@ -106,7 +120,7 @@ Pass `reranker: true` to rerank across all returned columns. 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({ +aiSearch({ indexes: { documents: { indexName: "catalog.schema.documents_idx", @@ -122,7 +136,7 @@ vectorSearch({ Enable cursor pagination to page through large result sets: ```ts -vectorSearch({ +aiSearch({ indexes: { products: { indexName: "catalog.schema.products_idx", @@ -143,7 +157,7 @@ For indexes that manage their own embeddings, provide an `embeddingFn` that take ```ts import { embed } from "./my-embedding-client"; -vectorSearch({ +aiSearch({ indexes: { products: { indexName: "catalog.schema.products_idx", @@ -157,7 +171,7 @@ vectorSearch({ ## HTTP routes -Routes are mounted at `/api/vector-search`. +Routes are mounted at `/api/ai-search`. | Method | Path | Description | |--------|------|-------------| @@ -168,7 +182,7 @@ Routes are mounted at `/api/vector-search`. ### Query an index ``` -POST /api/vector-search/:alias/query +POST /api/ai-search/:alias/query Content-Type: application/json { @@ -182,18 +196,24 @@ Response: ```json { "results": [ - { "id": "42", "name": "Intro to ML", "description": "..." } + { + "score": 0.87, + "data": { "id": "42", "name": "Intro to ML", "description": "..." } + } ], + "totalCount": 1, + "queryTimeMs": 35, + "queryType": "hybrid", "nextPageToken": "eyJvZmZzZXQiOjEwfQ==" } ``` -`nextPageToken` is only present when `pagination` is enabled and more results are available. +Each result carries its relevance `score` and the returned columns under `data`. `nextPageToken` is `null` unless `pagination` is enabled and more results are available. ### Fetch the next page ``` -POST /api/vector-search/:alias/next-page +POST /api/ai-search/:alias/next-page Content-Type: application/json { @@ -205,7 +225,7 @@ Content-Type: application/json ### Get index config ``` -GET /api/vector-search/:alias/config +GET /api/ai-search/:alias/config ``` Returns the resolved `IndexConfig` for the alias (excluding `embeddingFn`). @@ -215,10 +235,13 @@ Returns the resolved `IndexConfig` for the alias (excluding `embeddingFn`). The plugin exposes a `query` method for server-side use: ```ts +import { createApp, server } from "@databricks/appkit"; +import { aiSearch } from "@databricks/appkit/beta"; + const AppKit = await createApp({ plugins: [ server(), - vectorSearch({ + aiSearch({ indexes: { products: { indexName: "catalog.schema.products_idx", @@ -229,7 +252,7 @@ const AppKit = await createApp({ ], }); -const result = await AppKit.vectorSearch.query("products", { +const result = await AppKit.aiSearch.query("products", { queryText: "machine learning guide", }); @@ -237,3 +260,27 @@ console.log(result.results); ``` Pass optional overrides as a second argument to `query` to adjust `numResults` or other per-call settings. + +## React hook + +`useAiSearchQuery` reads the configured indexes from the plugin's client config and posts to the right `/:alias/query` route, so the UI never hardcodes an alias. With one index configured it needs no arguments; pass `{ alias }` to target a specific one. + +```tsx +import { useAiSearchQuery } from "@databricks/appkit-ui/react/beta"; + +function Search() { + const { search, data, loading, error } = useAiSearchQuery(); + + return ( + <> + e.key === "Enter" && search(e.currentTarget.value)} /> + {error &&

{error}

} + {data?.results.map((r, i) => ( +
{JSON.stringify(r.data)}
+ ))} + + ); +} +``` + +`search` also accepts a full request object (`{ queryText, numResults, filters, ... }`) for per-call control. The hook's `indexes` field lists every configured index, which you can use to build an index picker. diff --git a/docs/static/appkit-ui/styles.gen.css b/docs/static/appkit-ui/styles.gen.css index 58ecf8142..9e1d5c0c0 100644 --- a/docs/static/appkit-ui/styles.gen.css +++ b/docs/static/appkit-ui/styles.gen.css @@ -2773,6 +2773,12 @@ border-color: var(--ring); } } + .focus-visible\:shadow-none { + &:focus-visible { + --tw-shadow: 0 0 #0000; + box-shadow: var(--tw-inset-shadow), var(--tw-inset-ring-shadow), var(--tw-ring-offset-shadow), var(--tw-ring-shadow), var(--tw-shadow); + } + } .focus-visible\:ring-0 { &:focus-visible { --tw-ring-shadow: var(--tw-ring-inset,) 0 0 0 calc(0px + var(--tw-ring-offset-width)) var(--tw-ring-color, currentcolor); diff --git a/docs/static/schemas/plugin-manifest.schema.json b/docs/static/schemas/plugin-manifest.schema.json index 3e2279e73..c964aafed 100644 --- a/docs/static/schemas/plugin-manifest.schema.json +++ b/docs/static/schemas/plugin-manifest.schema.json @@ -10,8 +10,8 @@ }, "name": { "type": "string", - "pattern": "^[a-z][a-z0-9-]*$", - "description": "Plugin identifier. Must be lowercase, start with a letter, and contain only letters, numbers, and hyphens." + "pattern": "^[a-z][a-zA-Z0-9]*$", + "description": "Plugin identifier and JS binding. Must start with a lowercase letter; camelCase for multi-word names (e.g. aiSearch)." }, "displayName": { "type": "string", diff --git a/docs/static/schemas/template-plugins.schema.json b/docs/static/schemas/template-plugins.schema.json index fb7c32a3f..61ed50f88 100644 --- a/docs/static/schemas/template-plugins.schema.json +++ b/docs/static/schemas/template-plugins.schema.json @@ -23,8 +23,8 @@ "properties": { "name": { "type": "string", - "pattern": "^[a-z][a-z0-9-]*$", - "description": "Plugin identifier. Must be lowercase, start with a letter, and contain only letters, numbers, and hyphens." + "pattern": "^[a-z][a-zA-Z0-9]*$", + "description": "Plugin identifier and JS binding. Must start with a lowercase letter; camelCase for multi-word names (e.g. aiSearch)." }, "displayName": { "type": "string", diff --git a/knip.json b/knip.json index 251cc61eb..0e96b7df5 100644 --- a/knip.json +++ b/knip.json @@ -15,7 +15,7 @@ "**/*.generated.ts", "**/*.example.tsx", "**/*.css", - "packages/appkit/src/plugins/vector-search/**", + "packages/appkit/src/plugins/ai-search/**", "packages/appkit/src/plugin/index.ts", "packages/appkit/src/plugin/to-plugin.ts", "packages/appkit/src/plugins/agents/**", diff --git a/packages/appkit-ui/src/react/beta.ts b/packages/appkit-ui/src/react/beta.ts index 0405e50de..992a79635 100644 --- a/packages/appkit-ui/src/react/beta.ts +++ b/packages/appkit-ui/src/react/beta.ts @@ -1,2 +1,18 @@ // Beta React components -- APIs may change between minor releases. // Import from '@databricks/appkit-ui/react' once graduated to stable. + +// AI Search hook + types. Tracks the `aiSearch` plugin, which ships at beta +// from '@databricks/appkit/beta'. +export type { + AiSearchClientConfig, + AiSearchIndexSummary, + AiSearchQueryType, + AiSearchRequest, + AiSearchResponse, + AiSearchResult, +} from "./hooks/types"; +export { + type UseAiSearchQueryOptions, + type UseAiSearchQueryResult, + useAiSearchQuery, +} from "./hooks/use-ai-search-query"; diff --git a/packages/appkit-ui/src/react/hooks/__tests__/use-ai-search-query.test.ts b/packages/appkit-ui/src/react/hooks/__tests__/use-ai-search-query.test.ts new file mode 100644 index 000000000..e208f47ac --- /dev/null +++ b/packages/appkit-ui/src/react/hooks/__tests__/use-ai-search-query.test.ts @@ -0,0 +1,199 @@ +import { act, renderHook, waitFor } from "@testing-library/react"; +import { afterEach, beforeEach, describe, expect, test, vi } from "vitest"; + +const mockUsePluginClientConfig = vi.fn(); + +vi.mock("../use-plugin-config", () => ({ + usePluginClientConfig: (...args: unknown[]) => + mockUsePluginClientConfig(...args), +})); + +import { useAiSearchQuery } from "../use-ai-search-query"; + +const RESPONSE = { + results: [{ score: 0.9, data: { id: "1", text: "hi" } }], + totalCount: 1, + queryTimeMs: 12, + queryType: "hybrid", + nextPageToken: null, +}; + +describe("useAiSearchQuery", () => { + beforeEach(() => { + mockUsePluginClientConfig.mockReturnValue({ + indexes: [{ alias: "demo", queryType: "hybrid", pagination: false }], + }); + vi.spyOn(globalThis, "fetch").mockResolvedValue( + new Response(JSON.stringify(RESPONSE), { status: 200 }), + ); + }); + + afterEach(() => { + vi.restoreAllMocks(); + }); + + test("defaults to the first configured index", async () => { + const fetchSpy = vi.spyOn(globalThis, "fetch"); + const { result } = renderHook(() => useAiSearchQuery()); + + expect(result.current.alias).toBe("demo"); + expect(result.current.error).toBeNull(); + + act(() => { + void result.current.search("hello"); + }); + + await waitFor(() => { + expect(fetchSpy).toHaveBeenCalledWith( + "/api/ai-search/demo/query", + expect.objectContaining({ + method: "POST", + body: JSON.stringify({ queryText: "hello" }), + }), + ); + }); + }); + + test("uses the provided alias", async () => { + mockUsePluginClientConfig.mockReturnValue({ + indexes: [ + { alias: "demo", queryType: "hybrid", pagination: false }, + { alias: "docs", queryType: "ann", pagination: false }, + ], + }); + const fetchSpy = vi.spyOn(globalThis, "fetch"); + + const { result } = renderHook(() => useAiSearchQuery({ alias: "docs" })); + + act(() => { + void result.current.search("hello"); + }); + + await waitFor(() => { + expect(fetchSpy).toHaveBeenCalledWith( + "/api/ai-search/docs/query", + expect.any(Object), + ); + }); + }); + + test("forwards a full request object", async () => { + const fetchSpy = vi.spyOn(globalThis, "fetch"); + const { result } = renderHook(() => useAiSearchQuery()); + + act(() => { + void result.current.search({ queryText: "hi", numResults: 5 }); + }); + + await waitFor(() => { + expect(fetchSpy).toHaveBeenCalledWith( + "/api/ai-search/demo/query", + expect.objectContaining({ + body: JSON.stringify({ queryText: "hi", numResults: 5 }), + }), + ); + }); + }); + + test("errors when no indexes are configured", () => { + mockUsePluginClientConfig.mockReturnValue({ indexes: [] }); + + const { result } = renderHook(() => useAiSearchQuery()); + + expect(result.current.alias).toBeNull(); + expect(result.current.error).toBe("No AI Search indexes are configured."); + }); + + test("errors for an unknown alias without calling fetch", async () => { + const fetchSpy = vi.spyOn(globalThis, "fetch"); + + const { result } = renderHook(() => useAiSearchQuery({ alias: "nope" })); + + expect(result.current.error).toBe( + 'Unknown AI Search index "nope". Available: demo', + ); + + let returnValue: unknown; + act(() => { + returnValue = result.current.search("hello"); + }); + + expect(await returnValue).toBeNull(); + expect(fetchSpy).not.toHaveBeenCalled(); + }); + + test("sets data on a successful search", async () => { + const { result } = renderHook(() => useAiSearchQuery()); + + act(() => { + void result.current.search("hello"); + }); + + await waitFor(() => { + expect(result.current.data).toEqual(RESPONSE); + expect(result.current.loading).toBe(false); + }); + }); + + test("surfaces the server error message", async () => { + vi.spyOn(globalThis, "fetch").mockResolvedValue( + new Response(JSON.stringify({ error: "index not ready" }), { + status: 500, + }), + ); + + const { result } = renderHook(() => useAiSearchQuery()); + + await act(async () => { + void result.current.search("hello"); + await new Promise((r) => setTimeout(r, 10)); + }); + + await waitFor(() => { + expect(result.current.error).toBe("index not ready"); + expect(result.current.loading).toBe(false); + }); + }); + + test("clears stale data when the alias changes", async () => { + mockUsePluginClientConfig.mockReturnValue({ + indexes: [ + { alias: "demo", queryType: "hybrid", pagination: false }, + { alias: "docs", queryType: "ann", pagination: false }, + ], + }); + + const { result, rerender } = renderHook( + ({ alias }) => useAiSearchQuery({ alias }), + { initialProps: { alias: "demo" } }, + ); + + act(() => { + void result.current.search("hello"); + }); + await waitFor(() => expect(result.current.data).toEqual(RESPONSE)); + + // Switch index: prior results must not linger under the new alias. + rerender({ alias: "docs" }); + expect(result.current.alias).toBe("docs"); + expect(result.current.data).toBeNull(); + }); + + test("re-syncs the error when the alias becomes unknown", async () => { + mockUsePluginClientConfig.mockReturnValue({ + indexes: [{ alias: "demo", queryType: "hybrid", pagination: false }], + }); + + const { result, rerender } = renderHook( + ({ alias }) => useAiSearchQuery({ alias }), + { initialProps: { alias: "demo" } }, + ); + + expect(result.current.error).toBeNull(); + + rerender({ alias: "nope" }); + expect(result.current.error).toBe( + 'Unknown AI Search index "nope". Available: demo', + ); + }); +}); diff --git a/packages/appkit-ui/src/react/hooks/types.ts b/packages/appkit-ui/src/react/hooks/types.ts index aa0df8905..7f4dba04e 100644 --- a/packages/appkit-ui/src/react/hooks/types.ts +++ b/packages/appkit-ui/src/react/hooks/types.ts @@ -197,6 +197,51 @@ export interface ServingClientConfig { aliases: string[]; } +// ============================================================================ +// AI Search +// ============================================================================ + +export type AiSearchQueryType = "ann" | "hybrid" | "full_text"; + +/** One configured index, as exposed by the ai-search plugin's `clientConfig()`. */ +export interface AiSearchIndexSummary { + alias: string; + queryType: AiSearchQueryType; + pagination: boolean; +} + +/** Shape of the ai-search plugin's client config on `window.__appkit__`. */ +export interface AiSearchClientConfig { + indexes: AiSearchIndexSummary[]; +} + +export interface AiSearchRequest { + queryText?: string; + queryVector?: number[]; + columns?: string[]; + numResults?: number; + queryType?: AiSearchQueryType; + filters?: Record; + reranker?: boolean; +} + +export interface AiSearchResult< + T extends Record = Record, +> { + score: number; + data: T; +} + +export interface AiSearchResponse< + T extends Record = Record, +> { + results: AiSearchResult[]; + totalCount: number; + queryTimeMs: number; + queryType: AiSearchQueryType; + nextPageToken: string | null; +} + // ============================================================================ // Serving Endpoint Registry // ============================================================================ diff --git a/packages/appkit-ui/src/react/hooks/use-ai-search-query.ts b/packages/appkit-ui/src/react/hooks/use-ai-search-query.ts new file mode 100644 index 000000000..e15559588 --- /dev/null +++ b/packages/appkit-ui/src/react/hooks/use-ai-search-query.ts @@ -0,0 +1,125 @@ +import { useCallback, useEffect, useRef, useState } from "react"; +import type { + AiSearchClientConfig, + AiSearchIndexSummary, + AiSearchRequest, + AiSearchResponse, +} from "./types"; +import { usePluginClientConfig } from "./use-plugin-config"; + +export interface UseAiSearchQueryOptions { + /** + * Index alias to query. Defaults to the first index exposed by the plugin's + * `clientConfig()`, so a single-index app needs no alias. + */ + alias?: string; +} + +export interface UseAiSearchQueryResult< + T extends Record = Record, +> { + /** Run a search. Pass query text, or a full request for filters/paging control. */ + search: ( + query: string | AiSearchRequest, + ) => Promise | null>; + /** Latest response, null until the first successful search. */ + data: AiSearchResponse | null; + /** Whether a search is in progress. */ + loading: boolean; + /** Error message, if any. */ + error: string | null; + /** The resolved alias this hook queries. */ + alias: string | null; + /** All configured indexes, for building a selector. */ + indexes: AiSearchIndexSummary[]; +} + +/** + * Hook for querying a Databricks AI Search index. Reads the available indexes + * from the ai-search plugin's `clientConfig()` and POSTs to + * `/api/ai-search/{alias}/query`, so the UI never hardcodes an endpoint alias. + */ +export function useAiSearchQuery< + T extends Record = Record, +>(options: UseAiSearchQueryOptions = {}): UseAiSearchQueryResult { + const config = usePluginClientConfig("aiSearch"); + const indexes = config.indexes ?? []; + + const alias = options.alias ?? indexes[0]?.alias ?? null; + + // Config validation error; null for a valid alias. + let aliasError: string | null = null; + if (!alias) { + aliasError = "No AI Search indexes are configured."; + } else if (options.alias && !indexes.some((i) => i.alias === options.alias)) { + const available = indexes.map((i) => i.alias).join(", ") || "none"; + aliasError = `Unknown AI Search index "${options.alias}". Available: ${available}`; + } + + const [data, setData] = useState | null>(null); + const [loading, setLoading] = useState(false); + const [fetchError, setFetchError] = useState(null); + const error = fetchError ?? aliasError; + const abortControllerRef = useRef(null); + + const search = useCallback( + (query: string | AiSearchRequest): Promise | null> => { + if (aliasError || !alias) { + return Promise.resolve(null); + } + + abortControllerRef.current?.abort(); + const abortController = new AbortController(); + abortControllerRef.current = abortController; + + setLoading(true); + setFetchError(null); + setData(null); + + const body: AiSearchRequest = + typeof query === "string" ? { queryText: query } : query; + + return fetch(`/api/ai-search/${encodeURIComponent(alias)}/query`, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(body), + signal: abortController.signal, + }) + .then(async (res) => { + if (!res.ok) { + const errorBody = await res.json().catch(() => null); + throw new Error(errorBody?.error || `HTTP ${res.status}`); + } + return res.json(); + }) + .then((result: AiSearchResponse) => { + if (abortController.signal.aborted) return null; + setData(result); + setLoading(false); + return result; + }) + .catch((err: Error) => { + if (abortController.signal.aborted) return null; + setFetchError(err.message || "Search failed"); + setLoading(false); + return null; + }); + }, + [alias, aliasError], + ); + + // Reset when the target alias changes: abort any in-flight request (its + // result would otherwise land under the new alias) and clear stale data + + // fetch error. `error` re-derives from the new alias's validation state. + // biome-ignore lint/correctness/useExhaustiveDependencies: `alias` is the intended trigger — the effect runs to reset on every alias change, not because it reads alias. + useEffect(() => { + abortControllerRef.current?.abort(); + setData(null); + setLoading(false); + setFetchError(null); + }, [alias]); + + useEffect(() => () => abortControllerRef.current?.abort(), []); + + return { search, data, loading, error, alias, indexes }; +} diff --git a/packages/appkit/src/beta.ts b/packages/appkit/src/beta.ts index b20833fb6..d94b4e0d4 100644 --- a/packages/appkit/src/beta.ts +++ b/packages/appkit/src/beta.ts @@ -89,5 +89,15 @@ export { loadAgentFromFile, loadAgentsFromDir, } from "./plugins/agents"; - +// AI Search plugin config and query types (the `aiSearch` binding +// itself is exported via the generated barrel above). +export type { + IAiSearchConfig, + IndexConfig, + RerankerConfig, + SearchFilters, + SearchRequest, + SearchResponse, + SearchResult, +} from "./plugins/ai-search/types"; export * from "./plugins/beta-exports.generated"; diff --git a/packages/appkit/src/connectors/vector-search/client.ts b/packages/appkit/src/connectors/ai-search/client.ts similarity index 60% rename from packages/appkit/src/connectors/vector-search/client.ts rename to packages/appkit/src/connectors/ai-search/client.ts index f424b17e4..f4ed44b15 100644 --- a/packages/appkit/src/connectors/vector-search/client.ts +++ b/packages/appkit/src/connectors/ai-search/client.ts @@ -7,21 +7,24 @@ import { TelemetryManager, } from "../../telemetry"; import type { WorkspaceClient } from "../../workspace-client"; +import { contextFromAbortSignal } from "../context"; import type { - VectorSearchConnectorConfig, + AiSearchConnectorConfig, + UcTableInfo, + VsIndexInfo, VsNextPageParams, VsQueryParams, VsRawResponse, } from "./types"; -const logger = createLogger("connectors:vector-search"); +const logger = createLogger("connectors:ai-search"); -export class VectorSearchConnector { +export class AiSearchConnector { private readonly telemetry: TelemetryProvider; - constructor(config: VectorSearchConnectorConfig = {}) { + constructor(config: AiSearchConnectorConfig = {}) { this.telemetry = TelemetryManager.getProvider( - "vector-search", + "ai-search", config.telemetry, ); } @@ -45,7 +48,8 @@ export class VectorSearchConnector { 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; + // VS silently ignores an object under `filters`; it wants a JSON string. + body.filters_json = JSON.stringify(params.filters); } if (params.reranker) { body.reranker = { @@ -62,7 +66,7 @@ export class VectorSearchConnector { ); return this.telemetry.startActiveSpan( - "vector-search.query", + "ai-search.query", { kind: SpanKind.CLIENT, attributes: { @@ -79,14 +83,17 @@ export class VectorSearchConnector { 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: {}, - })) as VsRawResponse; + 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 duration = Date.now() - startTime; span.setAttribute("vs.result_count", response.result.row_count); @@ -97,7 +104,7 @@ export class VectorSearchConnector { span.setAttribute("vs.duration_ms", duration); span.setStatus({ code: SpanStatusCode.OK }); - logger.event()?.setContext("vector-search", { + logger.event()?.setContext("ai-search", { index_name: params.indexName, query_type: params.queryType, result_count: response.result.row_count, @@ -115,7 +122,7 @@ export class VectorSearchConnector { throw error; } }, - { name: "vector-search", includePrefix: true }, + { name: "ai-search", includePrefix: true }, ); } @@ -135,7 +142,7 @@ export class VectorSearchConnector { ); return this.telemetry.startActiveSpan( - "vector-search.queryNextPage", + "ai-search.queryNextPage", { kind: SpanKind.CLIENT, attributes: { @@ -146,17 +153,20 @@ export class VectorSearchConnector { }, 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 = (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, + }, + headers: new Headers({ "Content-Type": "application/json" }), + raw: false, + query: {}, }, - headers: new Headers({ "Content-Type": "application/json" }), - raw: false, - query: {}, - })) as VsRawResponse; + contextFromAbortSignal(signal), + )) as VsRawResponse; span.setAttribute("vs.result_count", response.result.row_count); span.setStatus({ code: SpanStatusCode.OK }); @@ -170,7 +180,50 @@ export class VectorSearchConnector { throw error; } }, - { name: "vector-search", includePrefix: true }, + { name: "ai-search", includePrefix: true }, ); } + + /** + * Fetches index metadata (index type, source table). Used to auto-discover + * returnable columns when they aren't configured. No warehouse required. + */ + async getIndex( + workspaceClient: WorkspaceClient, + 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: {}, + }, + contextFromAbortSignal(signal), + )) as VsIndexInfo; + } + + /** + * Lists a Unity Catalog table's column names via the tables REST API + * (no warehouse required). + */ + async getSourceColumns( + workspaceClient: WorkspaceClient, + 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); + } } diff --git a/packages/appkit/src/connectors/vector-search/index.ts b/packages/appkit/src/connectors/ai-search/index.ts similarity index 100% rename from packages/appkit/src/connectors/vector-search/index.ts rename to packages/appkit/src/connectors/ai-search/index.ts diff --git a/packages/appkit/src/connectors/vector-search/types.ts b/packages/appkit/src/connectors/ai-search/types.ts similarity index 66% rename from packages/appkit/src/connectors/vector-search/types.ts rename to packages/appkit/src/connectors/ai-search/types.ts index df042e8c6..0d3e175b2 100644 --- a/packages/appkit/src/connectors/vector-search/types.ts +++ b/packages/appkit/src/connectors/ai-search/types.ts @@ -1,6 +1,6 @@ import type { TelemetryOptions } from "shared"; -export interface VectorSearchConnectorConfig { +export interface AiSearchConnectorConfig { timeout?: number; telemetry?: TelemetryOptions; } @@ -22,6 +22,20 @@ export interface VsNextPageParams { pageToken: string; } +/** Subset of the get-index response used for column auto-discovery. */ +export interface VsIndexInfo { + index_type?: "DELTA_SYNC" | "DIRECT_ACCESS"; + delta_sync_index_spec?: { + source_table?: string; + embedding_vector_columns?: Array<{ name: string }>; + }; +} + +/** 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/connectors/context.ts b/packages/appkit/src/connectors/context.ts new file mode 100644 index 000000000..d1cdd82dd --- /dev/null +++ b/packages/appkit/src/connectors/context.ts @@ -0,0 +1,52 @@ +import { type CancellationToken, Context } from "../workspace-client"; + +/** + * Bridges {@link AbortSignal} to the SDK's {@link CancellationToken} so + * `apiClient.request` can abort the outbound HTTP request (and stop pulling a + * response body) when the caller aborts. + */ +function cancellationTokenFromAbortSignal( + signal: AbortSignal, +): CancellationToken { + const listeners = new Set<() => void>(); + signal.addEventListener( + "abort", + () => { + for (const cb of listeners) { + try { + cb(); + } catch { + // ignore listener failures — abort must stay best-effort + } + } + }, + { passive: true }, + ); + + return { + get isCancellationRequested() { + return signal.aborted; + }, + onCancellationRequested(callback: (e?: unknown) => unknown) { + listeners.add(callback as () => void); + if (signal.aborted) { + void callback(); + } + }, + }; +} + +/** + * Wraps an optional {@link AbortSignal} in an SDK {@link Context} for the + * second argument of `apiClient.request`. Returns `undefined` when no signal + * is given, so callers can pass the result straight through. + */ +export function contextFromAbortSignal( + signal?: AbortSignal, +): InstanceType | undefined { + return signal + ? new Context({ + cancellationToken: cancellationTokenFromAbortSignal(signal), + }) + : undefined; +} diff --git a/packages/appkit/src/connectors/index.ts b/packages/appkit/src/connectors/index.ts index 5fad31d99..438d334af 100644 --- a/packages/appkit/src/connectors/index.ts +++ b/packages/appkit/src/connectors/index.ts @@ -1,7 +1,7 @@ +export * from "./ai-search"; export * from "./files"; export * from "./genie"; export * from "./jobs"; export * from "./lakebase"; export * from "./mcp"; export * from "./sql-warehouse"; -export * from "./vector-search"; diff --git a/packages/appkit/src/connectors/serving/client.ts b/packages/appkit/src/connectors/serving/client.ts index ba54b6687..de9d0465c 100644 --- a/packages/appkit/src/connectors/serving/client.ts +++ b/packages/appkit/src/connectors/serving/client.ts @@ -1,46 +1,9 @@ import { createLogger } from "../../logging/logger"; -import { - type CancellationToken, - Context, - type serving, - type WorkspaceClient, -} from "../../workspace-client"; +import type { serving, WorkspaceClient } from "../../workspace-client"; +import { contextFromAbortSignal } from "../context"; const logger = createLogger("connectors:serving"); -/** - * Bridges {@link AbortSignal} to the SDK's {@link CancellationToken} so - * `apiClient.request` can abort the outbound HTTP request (and stop pulling - * the SSE body) when the agent run is cancelled. - */ -function cancellationTokenFromAbortSignal( - signal: AbortSignal, -): CancellationToken { - const listeners = new Set<() => void>(); - const fire = () => { - for (const cb of listeners) { - try { - cb(); - } catch { - // ignore listener failures — abort must stay best-effort - } - } - }; - signal.addEventListener("abort", fire, { passive: true }); - - return { - get isCancellationRequested() { - return signal.aborted; - }, - onCancellationRequested(callback: (e?: unknown) => unknown) { - listeners.add(callback as () => void); - if (signal.aborted) { - void callback(); - } - }, - }; -} - /** * Structural shape of a Databricks SDK client we need for the low-level * `apiClient.request` call. Lets `streamPath` be reused by adapters that @@ -115,11 +78,7 @@ export async function streamPath( ): Promise> { logger.debug("Streaming from path %s", path); - const context = signal - ? new Context({ - cancellationToken: cancellationTokenFromAbortSignal(signal), - }) - : undefined; + const context = contextFromAbortSignal(signal); const response = (await client.apiClient.request( { diff --git a/packages/appkit/src/connectors/tests/context.test.ts b/packages/appkit/src/connectors/tests/context.test.ts new file mode 100644 index 000000000..58eb034fc --- /dev/null +++ b/packages/appkit/src/connectors/tests/context.test.ts @@ -0,0 +1,61 @@ +import { describe, expect, it, vi } from "vitest"; +import { Context } from "../../workspace-client"; +import { contextFromAbortSignal } from "../context"; + +describe("contextFromAbortSignal", () => { + it("returns undefined when no signal is given", () => { + expect(contextFromAbortSignal()).toBeUndefined(); + }); + + it("wraps a signal in an SDK Context", () => { + const ctx = contextFromAbortSignal(new AbortController().signal); + expect(ctx).toBeInstanceOf(Context); + expect(ctx?.cancellationToken).toBeDefined(); + }); + + it("reflects the signal's aborted state via isCancellationRequested", () => { + const controller = new AbortController(); + const token = contextFromAbortSignal(controller.signal)?.cancellationToken; + + expect(token?.isCancellationRequested).toBe(false); + controller.abort(); + expect(token?.isCancellationRequested).toBe(true); + }); + + it("fires registered callbacks when the signal aborts", () => { + const controller = new AbortController(); + const token = contextFromAbortSignal(controller.signal)?.cancellationToken; + + const cb = vi.fn(); + token?.onCancellationRequested(cb); + expect(cb).not.toHaveBeenCalled(); + + controller.abort(); + expect(cb).toHaveBeenCalledTimes(1); + }); + + it("fires immediately when registering on an already-aborted signal", () => { + const controller = new AbortController(); + controller.abort(); + const token = contextFromAbortSignal(controller.signal)?.cancellationToken; + + const cb = vi.fn(); + token?.onCancellationRequested(cb); + expect(cb).toHaveBeenCalledTimes(1); + }); + + it("isolates callback failures so abort stays best-effort", () => { + const controller = new AbortController(); + const token = contextFromAbortSignal(controller.signal)?.cancellationToken; + + const bad = vi.fn(() => { + throw new Error("listener boom"); + }); + const good = vi.fn(); + token?.onCancellationRequested(bad); + token?.onCancellationRequested(good); + + expect(() => controller.abort()).not.toThrow(); + expect(good).toHaveBeenCalledTimes(1); + }); +}); diff --git a/packages/appkit/src/core/appkit.ts b/packages/appkit/src/core/appkit.ts index 54f107568..201acf190 100644 --- a/packages/appkit/src/core/appkit.ts +++ b/packages/appkit/src/core/appkit.ts @@ -107,6 +107,8 @@ export class AppKit { const self = this; + // The manifest `name` is camelCase, so it doubles as the public handle key + // (`appkit.aiSearch`). The kebab HTTP route is derived separately. Object.defineProperty(this, name, { get() { const plugin = self.#pluginInstances[name]; diff --git a/packages/appkit/src/core/tests/databricks.test.ts b/packages/appkit/src/core/tests/databricks.test.ts index 7db561cdb..b3abc5bea 100644 --- a/packages/appkit/src/core/tests/databricks.test.ts +++ b/packages/appkit/src/core/tests/databricks.test.ts @@ -472,7 +472,8 @@ describe("AppKit", () => { plugins: [{ plugin: NormalTestPlugin, config: {}, name: "normalTest" }], })) as any; - expect(instance["ui-variants"]).toBeDefined(); + // Exposed under the camelCase handle key, not the kebab manifest name. + expect(instance.uiVariants).toBeDefined(); }); test("drops the default ui-variants in production", async () => { @@ -484,7 +485,7 @@ describe("AppKit", () => { // The default is devOnly, so the same guard that skips user devOnly // plugins strips it from a deployed app. - expect(Object.keys(instance)).not.toContain("ui-variants"); + expect(Object.keys(instance)).not.toContain("uiVariants"); expect(instance.normalTest).toBeDefined(); }); @@ -496,14 +497,32 @@ describe("AppKit", () => { plugins: [explicit], })) as any; - expect(instance["ui-variants"]).toBeDefined(); + expect(instance.uiVariants).toBeDefined(); // Only one instance was constructed for the single name. expect( - Object.keys(instance).filter((k) => k === "ui-variants"), + Object.keys(instance).filter((k) => k === "uiVariants"), ).toHaveLength(1); }); }); + describe("plugin accessor", () => { + class MultiWordPlugin extends NormalTestPlugin { + static manifest = createTestManifest("aiSearch"); + name = "aiSearch"; + } + + test("exposes a plugin under its camelCase manifest name", async () => { + const instance = (await createApp({ + plugins: [{ plugin: MultiWordPlugin, config: {}, name: "aiSearch" }], + })) as any; + + // The camelCase name is the accessor key verbatim (no transform). + expect(instance.aiSearch).toBeDefined(); + expect(instance.aiSearch.setupCalled).toBe(true); + expect(Object.keys(instance)).toContain("aiSearch"); + }); + }); + describe("preparePlugins", () => { test("should transform plugin data array to plugin map", () => { const pluginData = [ diff --git a/packages/appkit/src/plugin/plugin.ts b/packages/appkit/src/plugin/plugin.ts index 84c2e247c..ee403a924 100644 --- a/packages/appkit/src/plugin/plugin.ts +++ b/packages/appkit/src/plugin/plugin.ts @@ -12,6 +12,7 @@ import type { StreamExecuteHandler, StreamExecutionSettings, } from "shared"; +import { camelToKebab } from "shared"; import { AppManager } from "../app"; import { CacheManager } from "../cache"; import { getCurrentUserId, runInUserContext, ServiceContext } from "../context"; @@ -669,7 +670,7 @@ export abstract class Plugin< router[method](path, forwardAsyncErrors(handler)); - const fullPath = `/api/${this.name}${path}`; + const fullPath = `/api/${camelToKebab(this.name)}${path}`; this.registerEndpoint(name, fullPath); if (config.skipBodyParsing) { diff --git a/packages/appkit/src/plugins/ai-search/ai-search.ts b/packages/appkit/src/plugins/ai-search/ai-search.ts new file mode 100644 index 000000000..345961672 --- /dev/null +++ b/packages/appkit/src/plugins/ai-search/ai-search.ts @@ -0,0 +1,471 @@ +import type express from "express"; +import type { IAppRouter, PluginExecutionSettings } from "shared"; +import { AiSearchConnector } from "../../connectors/ai-search/client"; +import type { + VsQueryParams, + VsRawResponse, +} from "../../connectors/ai-search/types"; +import { getWorkspaceClient } from "../../context"; +import { createLogger } from "../../logging/logger"; +import { Plugin, toPlugin } from "../../plugin"; +import type { PluginManifest } from "../../registry"; +import { formatWarningBanner } from "../../utils/banner"; +import { aiSearchDefaults } from "./defaults"; +import manifest from "./manifest.json"; +import type { + IAiSearchConfig, + IndexConfig, + IndexSummary, + SearchQueryType, + SearchRequest, + SearchResponse, + SearchResult, +} from "./types"; + +const logger = createLogger("ai-search"); + +const querySettings: PluginExecutionSettings = { + default: aiSearchDefaults, +}; + +export class AiSearchPlugin extends Plugin { + static manifest = manifest as PluginManifest<"aiSearch">; + + protected static description = + "Query Databricks Vector Search indexes with hybrid search, reranking, and pagination"; + protected declare config: IAiSearchConfig; + + private connector: AiSearchConnector; + + constructor(config: IAiSearchConfig) { + super(config); + this.config = { + ...config, + indexes: config.indexes ?? this._defaultIndexes(), + }; + this.connector = new AiSearchConnector({ + timeout: config.timeout, + telemetry: config.telemetry, + }); + } + + /** + * Seeds a `default` index from `DATABRICKS_VS_INDEX_NAME` when no `indexes` + * are configured, so `aiSearch()` works with just the env var. + */ + private _defaultIndexes(): Record { + const indexName = process.env.DATABRICKS_VS_INDEX_NAME; + return indexName ? { default: { indexName } } : {}; + } + + async setup(): Promise { + // pagination needs an endpointName the framework's resource validation + // can't see, so check it here. + for (const [alias, idx] of Object.entries(this.config.indexes ?? {})) { + if (idx.pagination && !idx.endpointName) { + throw new Error( + `Index "${alias}" has pagination enabled but is missing "endpointName"`, + ); + } + } + + // Dev fills in missing `columns` from the source table; prod can't query + // without them (VS requires `columns`), so fail fast at boot. + if (process.env.NODE_ENV === "development") { + await this._autoDiscoverColumns(); + } else { + for (const [alias, idx] of Object.entries(this.config.indexes ?? {})) { + if (!idx.columns || idx.columns.length === 0) { + throw new Error( + `Index "${alias}" has no columns configured. Vector Search queries require "columns"; set them explicitly (auto-discovered only in development).`, + ); + } + } + } + } + + /** + * For each configured index missing `columns`, fill them from its Delta-Sync + * source table (all source columns minus embedding vectors). Best-effort: + * failures are logged and skipped, never thrown. A partial `columns_to_sync` + * isn't honored, so the discovered list is a starting point to trim. + */ + private async _autoDiscoverColumns(): Promise { + const discovered: Record = {}; + for (const [alias, idx] of Object.entries(this.config.indexes ?? {})) { + if (idx.columns && idx.columns.length > 0) continue; + const indexName = idx.indexName ?? process.env.DATABRICKS_VS_INDEX_NAME; + if (!indexName) continue; + try { + const client = getWorkspaceClient(); + const info = await this.connector.getIndex(client, indexName); + const sourceTable = info.delta_sync_index_spec?.source_table; + if (!sourceTable) continue; + const excluded = new Set( + (info.delta_sync_index_spec?.embedding_vector_columns ?? []).map( + (c) => c.name, + ), + ); + const columns = ( + await this.connector.getSourceColumns(client, sourceTable) + ).filter((c) => !excluded.has(c)); + if (columns.length > 0) { + idx.columns = columns; + discovered[alias] = columns; + } + } catch (error) { + logger.warn( + 'Could not auto-discover columns for index "%s": %s', + alias, + error instanceof Error ? error.message : String(error), + ); + } + } + if (Object.keys(discovered).length > 0) { + logger.warn("\n%s", this._formatColumnDiscoveryBanner(discovered)); + } + } + + private _formatColumnDiscoveryBanner( + discovered: Record, + ): string { + const lines = [ + "AI SEARCH: columns auto-discovered (dev mode — would fail in production)", + "", + ]; + for (const [alias, columns] of Object.entries(discovered)) { + lines.push(` ${alias}: ${columns.join(", ")}`); + } + lines.push(""); + lines.push( + "Set `columns` explicitly in the plugin config before deploying.", + ); + + return formatWarningBanner(lines); + } + + injectRoutes(router: IAppRouter) { + this.route(router, { + name: "query", + method: "post", + path: "/:alias/query", + handler: async (req: express.Request, res: express.Response) => { + const indexConfig = this._resolveOr404(req, res); + if (!indexConfig) 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; + } + + // Drop client-supplied `columns` so an HTTP caller can't widen the + // projection past what the app configured. (query() callers are + // trusted and keep the override.) + const { columns: _clientColumns, ...safeBody } = body; + const plugin = + indexConfig.auth === "on-behalf-of-user" ? this.asUser(req) : this; + const queryType = + safeBody.queryType ?? indexConfig.queryType ?? "hybrid"; + + try { + // Prepare inside execute so a self-managed embeddingFn runs in the + // same OBO context as the query, not as the service principal. + const result = await plugin.execute(async (signal) => { + const prepared = await this._prepareQuery(safeBody, indexConfig); + return this.connector.query( + getWorkspaceClient(), + { indexName: indexConfig.indexName, ...prepared }, + signal, + ); + }, querySettings); + + this._sendResult(res, result, 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._resolveOr404(req, res); + if (!indexConfig) 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, queryType } = req.body; + if (!pageToken) { + res.status(400).json({ + error: "pageToken is required", + plugin: this.name, + }); + return; + } + // Echo the original query's queryType so paged responses stay + // consistent with page 1; fall back to the index default. + const pageQueryType = queryType ?? indexConfig.queryType ?? "hybrid"; + + 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, + ); + + this._sendResult(res, result, pageQueryType); + } catch (error) { + this._handleError(res, error, "Next-page query failed"); + } + }, + }); + + this.route(router, { + name: "getConfig", + method: "get", + path: "/:alias/config", + handler: async (req: express.Request, res: express.Response) => { + const indexConfig = this._resolveOr404(req, res); + if (!indexConfig) return; + res.json({ + alias: req.params.alias, + columns: indexConfig.columns, + queryType: indexConfig.queryType ?? "hybrid", + numResults: indexConfig.numResults ?? 20, + reranker: !!indexConfig.reranker, + pagination: !!indexConfig.pagination, + }); + }, + }); + } + + /** + * Index aliases + non-sensitive query metadata, serialized to the client so + * the UI can discover available indexes instead of hardcoding an alias. + */ + clientConfig(): { indexes: IndexSummary[] } { + const indexes = Object.entries(this.config.indexes ?? {}).map( + ([alias, idx]) => ({ + alias, + queryType: idx.queryType ?? "hybrid", + pagination: !!idx.pagination, + }), + ); + return { indexes }; + } + + /** + * Programmatic query API — available as `appkit.aiSearch.query()`. + * When called through `asUser(req)`, executes with the user's credentials. + * + * @remarks `T` types each result's `data` but is an unchecked cast — the row + * shape isn't validated at runtime. + */ + async query = Record>( + 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, ...prepared }, + signal, + ), + querySettings, + ); + + if (!result.ok) { + throw new Error( + `Vector search query failed for index "${alias}": ${result.message}`, + ); + } + + return this._parseResponse(result.data, 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 & { indexName: string }) | undefined { + const idx = this.config.indexes?.[alias]; + if (!idx) return undefined; + const indexName = idx.indexName ?? process.env.DATABRICKS_VS_INDEX_NAME; + if (!indexName) return undefined; + return { ...idx, indexName }; + } + + /** Resolve an index by route alias, or send a 404 and return null. */ + private _resolveOr404( + req: express.Request, + res: express.Response, + ): (IndexConfig & { indexName: string }) | null { + 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 null; + } + return indexConfig; + } + + /** Send an execution result as JSON, or its error status/message. */ + private _sendResult( + res: express.Response, + result: Awaited>>, + queryType: SearchQueryType, + ): void { + if (!result.ok) { + res + .status(result.status) + .json({ error: result.message, plugin: this.name }); + return; + } + res.json(this._parseResponse(result.data, queryType)); + } + + private async _prepareQuery( + request: SearchRequest, + indexConfig: IndexConfig, + ): Promise> { + const queryType = request.queryType ?? indexConfig.queryType ?? "hybrid"; + let queryText = request.queryText; + let queryVector = request.queryVector; + + // full_text uses no vector; hybrid keeps the text for its keyword half. + if ( + indexConfig.embeddingFn && + queryText && + !queryVector && + queryType !== "full_text" + ) { + try { + queryVector = await indexConfig.embeddingFn(queryText); + if (queryType === "ann") 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, + filters: request.filters, + reranker: 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; + } + // Auto-derive from returnable columns (excluding the id). With no columns + // resolved there's nothing to rerank on, so skip it. + const columnsToRerank = columns.filter((c) => c !== "id"); + return columnsToRerank.length > 0 ? { columnsToRerank } : undefined; + } + + private _parseResponse< + T extends Record = Record, + >(raw: VsRawResponse, queryType: SearchQueryType): SearchResponse { + const columnNames = raw.manifest.columns.map((c) => c.name); + const scoreIndex = columnNames.indexOf("score"); + + const results: SearchResult[] = raw.result.data_array.map((row) => { + const data: Record = {}; + for (let i = 0; i < columnNames.length; i++) { + if (i !== scoreIndex) data[columnNames[i]] = row[i]; + } + return { + score: scoreIndex >= 0 ? (row[scoreIndex] as number) : 0, + data: data as T, + }; + }); + + 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); + // Match Plugin.execute(): the raw message is only exposed outside production. + const isDev = process.env.NODE_ENV !== "production"; + const message = + isDev && error instanceof Error ? error.message : fallbackMessage; + res.status(500).json({ error: message, plugin: this.name }); + } +} + +export const aiSearch = toPlugin(AiSearchPlugin); diff --git a/packages/appkit/src/plugins/vector-search/defaults.ts b/packages/appkit/src/plugins/ai-search/defaults.ts similarity index 73% rename from packages/appkit/src/plugins/vector-search/defaults.ts rename to packages/appkit/src/plugins/ai-search/defaults.ts index c02b6e804..d927c4490 100644 --- a/packages/appkit/src/plugins/vector-search/defaults.ts +++ b/packages/appkit/src/plugins/ai-search/defaults.ts @@ -1,6 +1,6 @@ import type { PluginExecuteConfig } from "shared"; -export const vectorSearchDefaults: PluginExecuteConfig = { +export const aiSearchDefaults: PluginExecuteConfig = { cache: { enabled: false }, retry: { enabled: true, initialDelay: 1000, attempts: 3 }, timeout: 30_000, diff --git a/packages/appkit/src/plugins/ai-search/index.ts b/packages/appkit/src/plugins/ai-search/index.ts new file mode 100644 index 000000000..8577b6f08 --- /dev/null +++ b/packages/appkit/src/plugins/ai-search/index.ts @@ -0,0 +1,2 @@ +export * from "./ai-search"; +export * from "./types"; diff --git a/packages/appkit/src/plugins/vector-search/manifest.json b/packages/appkit/src/plugins/ai-search/manifest.json similarity index 78% rename from packages/appkit/src/plugins/vector-search/manifest.json rename to packages/appkit/src/plugins/ai-search/manifest.json index a4451b1af..849c3315a 100644 --- a/packages/appkit/src/plugins/vector-search/manifest.json +++ b/packages/appkit/src/plugins/ai-search/manifest.json @@ -1,8 +1,8 @@ { "$schema": "https://databricks.github.io/appkit/schemas/plugin-manifest.schema.json", - "name": "vector-search", - "displayName": "Vector Search Plugin", - "hidden": true, + "name": "aiSearch", + "displayName": "AI Search Plugin", + "stability": "beta", "description": "Query Databricks Vector Search indexes with built-in hybrid search, reranking, and pagination", "resources": { "required": [ @@ -13,19 +13,24 @@ "description": "A Databricks Vector Search index to query. Index names configured via plugin config.", "permission": "SELECT", "fields": { - "indexName": { + "id": { "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": [] }, + "scaffolding": { + "rules": { + "must": [ + "List the vector search index's source-table columns (unity-catalog tables API) so the user can pick which to return", + "Ask the user which of those columns the search should return", + "Write aiSearch({ indexes: { default: { columns: [...] } } }) in the server file; the client queries the 'default' alias" + ] + } + }, "config": { "schema": { "type": "object", @@ -54,8 +59,7 @@ "type": "number", "default": 20 } - }, - "required": ["indexName", "columns"] + } } }, "timeout": { 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 new file mode 100644 index 000000000..3a8efc0ab --- /dev/null +++ b/packages/appkit/src/plugins/ai-search/tests/ai-search.test.ts @@ -0,0 +1,898 @@ +import { + createMockRequest, + createMockResponse, + createMockRouter, +} from "@tools/test-helpers"; +import { afterEach, beforeEach, describe, expect, it, vi } from "vitest"; +import { Context } from "../../../workspace-client"; + +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(), + setExecution: 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 { AiSearchPlugin } from "../ai-search"; + +describe("AiSearchPlugin", () => { + beforeEach(() => { + mockRequest.mockClear(); + mockRequest.mockResolvedValue(validVsResponse); + }); + + describe("setup()", () => { + const originalIndexEnv = process.env.DATABRICKS_VS_INDEX_NAME; + afterEach(() => { + if (originalIndexEnv === undefined) { + delete process.env.DATABRICKS_VS_INDEX_NAME; + } else { + process.env.DATABRICKS_VS_INDEX_NAME = originalIndexEnv; + } + }); + + it("defaults indexName from DATABRICKS_VS_INDEX_NAME when omitted", async () => { + process.env.DATABRICKS_VS_INDEX_NAME = "cat.sch.from_env"; + const plugin = new AiSearchPlugin({ + indexes: { + test: { columns: ["id"] }, + }, + }); + await expect(plugin.setup()).resolves.not.toThrow(); + + await plugin.query("test", { queryText: "q" }); + expect(mockRequest.mock.calls[0][0].path).toBe( + "/api/2.0/vector-search/indexes/cat.sch.from_env/query", + ); + }); + + it("seeds a 'default' index from the env var when no indexes are configured", async () => { + process.env.DATABRICKS_VS_INDEX_NAME = "cat.sch.from_env"; + // Bare aiSearch() — no indexes config. + const plugin = new AiSearchPlugin({}); + + await plugin.query("default", { queryText: "q", columns: ["id"] }); + expect(mockRequest.mock.calls[0][0].path).toBe( + "/api/2.0/vector-search/indexes/cat.sch.from_env/query", + ); + }); + + it("throws if pagination enabled but no endpointName", async () => { + const plugin = new AiSearchPlugin({ + 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 AiSearchPlugin({ + indexes: { + products: { + indexName: "cat.sch.products_idx", + columns: ["id", "name", "description"], + queryType: "hybrid", + numResults: 20, + }, + }, + }); + await expect(plugin.setup()).resolves.not.toThrow(); + }); + + it("throws outside development when an index has no columns", async () => { + const originalNodeEnv = process.env.NODE_ENV; + process.env.NODE_ENV = "production"; + try { + const plugin = new AiSearchPlugin({ + indexes: { docs: { indexName: "cat.sch.idx" } }, + }); + await expect(plugin.setup()).rejects.toThrow( + 'Index "docs" has no columns configured', + ); + } finally { + process.env.NODE_ENV = originalNodeEnv; + } + }); + + it("does not throw outside development when columns are configured", async () => { + const originalNodeEnv = process.env.NODE_ENV; + process.env.NODE_ENV = "production"; + try { + const plugin = new AiSearchPlugin({ + indexes: { docs: { indexName: "cat.sch.idx", columns: ["id"] } }, + }); + await expect(plugin.setup()).resolves.not.toThrow(); + } finally { + process.env.NODE_ENV = originalNodeEnv; + } + }); + }); + + 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" }], + }, + }); + } + if (opts.path.startsWith("/api/2.1/unity-catalog/tables/")) { + return Promise.resolve({ + columns: [{ name: "id" }, { name: "body" }, { name: "__vec" }], + }); + } + return Promise.resolve(validVsResponse); + }; + + afterEach(() => { + process.env.NODE_ENV = originalNodeEnv; + }); + + it("fills columns from the source table in development and warns", async () => { + process.env.NODE_ENV = "development"; + mockRequest.mockImplementation(routeByPath); + const plugin = new AiSearchPlugin({ + indexes: { docs: { indexName: "cat.sch.idx" } }, + }); + + await plugin.setup(); + + // 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(queryCall?.[0].payload.columns).toEqual(["id", "body"]); + }); + + it("does not discover columns outside development", async () => { + process.env.NODE_ENV = "production"; + mockRequest.mockImplementation(routeByPath); + // 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({ + indexes: { docs: { indexName: "cat.sch.idx", columns: ["id"] } }, + }); + + 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); + }); + + it("skips (does not throw) when an index already has columns", async () => { + process.env.NODE_ENV = "development"; + mockRequest.mockImplementation(routeByPath); + 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); + }); + }); + + describe("manifest", () => { + it("has correct name", () => { + expect(AiSearchPlugin.manifest.name).toBe("aiSearch"); + }); + }); + + describe("exports()", () => { + it("returns object with query function", () => { + const plugin = new AiSearchPlugin({ + 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 AiSearchPlugin({ + 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("types result.data via the generic parameter", async () => { + interface Doc extends Record { + id: number; + title: string; + } + const plugin = new AiSearchPlugin({ + indexes: { + products: { indexName: "cat.sch.products", columns: ["id", "title"] }, + }, + }); + await plugin.setup(); + + const result = await plugin.query("products", { + queryText: "machine learning", + }); + + // `data` is typed as Doc — these fields resolve without a cast. + const first: Doc = result.results[0].data; + expect(first.id).toBe(1); + expect(first.title).toBe("ML Guide"); + }); + + it("constructs correct API request", async () => { + const plugin = new AiSearchPlugin({ + 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", + }), + // 2nd arg is the SDK Context bridging the execution's abort signal. + expect.any(Context), + ); + + 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); + expect(callBody.columns).toEqual(["id", "title"]); + }); + + it("throws Error for unknown alias", async () => { + const plugin = new AiSearchPlugin({ + 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 AiSearchPlugin({ + 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].payload; + // VS expects a JSON-encoded string under `filters_json`; a raw object + // under `filters` is silently ignored by the API. + expect(callBody.filters).toBeUndefined(); + expect(callBody.filters_json).toBe( + JSON.stringify({ category: ["books"] }), + ); + }); + + it("includes reranker config when enabled on index", async () => { + const plugin = new AiSearchPlugin({ + 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].payload; + expect(callBody.reranker.model).toBe("databricks_reranker"); + expect(callBody.reranker.parameters.columns_to_rerank).toEqual([ + "title", + "desc", + ]); + }); + + it("calls embeddingFn and drops query_text for ann (vector-only)", async () => { + const mockEmbeddingFn = vi.fn().mockResolvedValue([0.1, 0.2, 0.3]); + const plugin = new AiSearchPlugin({ + indexes: { + test: { + indexName: "cat.sch.idx", + columns: ["id", "title"], + queryType: "ann", + embeddingFn: mockEmbeddingFn, + }, + }, + }); + await plugin.setup(); + await plugin.query("test", { queryText: "test" }); + + expect(mockEmbeddingFn).toHaveBeenCalledWith("test"); + const callBody = mockRequest.mock.calls[0][0].payload; + expect(callBody.query_vector).toEqual([0.1, 0.2, 0.3]); + expect(callBody.query_text).toBeUndefined(); + }); + + it("keeps query_text alongside the embedded vector for hybrid", async () => { + const mockEmbeddingFn = vi.fn().mockResolvedValue([0.1, 0.2, 0.3]); + const plugin = new AiSearchPlugin({ + indexes: { + test: { + indexName: "cat.sch.idx", + columns: ["id", "title"], + queryType: "hybrid", + embeddingFn: mockEmbeddingFn, + }, + }, + }); + await plugin.setup(); + await plugin.query("test", { queryText: "test" }); + + expect(mockEmbeddingFn).toHaveBeenCalledWith("test"); + const callBody = mockRequest.mock.calls[0][0].payload; + expect(callBody.query_vector).toEqual([0.1, 0.2, 0.3]); + expect(callBody.query_text).toBe("test"); + }); + + it("skips embeddingFn for full_text and sends query_text only", async () => { + const mockEmbeddingFn = vi.fn().mockResolvedValue([0.1, 0.2, 0.3]); + const plugin = new AiSearchPlugin({ + indexes: { + test: { + indexName: "cat.sch.idx", + columns: ["id", "title"], + queryType: "full_text", + embeddingFn: mockEmbeddingFn, + }, + }, + }); + await plugin.setup(); + await plugin.query("test", { queryText: "test" }); + + expect(mockEmbeddingFn).not.toHaveBeenCalled(); + const callBody = mockRequest.mock.calls[0][0].payload; + expect(callBody.query_text).toBe("test"); + expect(callBody.query_vector).toBeUndefined(); + }); + + it("throws when embeddingFn fails", async () => { + const mockEmbeddingFn = vi + .fn() + .mockRejectedValue(new Error("embedding service unavailable")); + const plugin = new AiSearchPlugin({ + 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 AiSearchPlugin({ + indexes: { + test: { indexName: "cat.sch.idx", columns: ["id"] }, + }, + }); + await expect(plugin.shutdown()).resolves.not.toThrow(); + }); + }); + + describe("_parseResponse edge cases", () => { + it("defaults score to 0 when the index returns no score column", async () => { + mockRequest.mockResolvedValueOnce({ + manifest: { column_count: 2, columns: [{ name: "id" }, { name: "t" }] }, + result: { row_count: 1, data_array: [[1, "hi"]] }, + next_page_token: null, + }); + const plugin = new AiSearchPlugin({ + indexes: { test: { indexName: "cat.sch.idx", columns: ["id", "t"] } }, + }); + await plugin.setup(); + const res = await plugin.query("test", { queryText: "q" }); + + expect(res.results[0].score).toBe(0); + expect(res.results[0].data).toEqual({ id: 1, t: "hi" }); + }); + + it("propagates a non-null next_page_token", async () => { + mockRequest.mockResolvedValueOnce({ + ...validVsResponse, + next_page_token: "tok-123", + }); + const plugin = new AiSearchPlugin({ + indexes: { test: { indexName: "cat.sch.idx", columns: ["id"] } }, + }); + await plugin.setup(); + const res = await plugin.query("test", { queryText: "q" }); + + expect(res.nextPageToken).toBe("tok-123"); + }); + + it("falls back to latency_ms for queryTimeMs when response_time is absent", async () => { + mockRequest.mockResolvedValueOnce({ + manifest: { column_count: 1, columns: [{ name: "id" }] }, + result: { row_count: 0, data_array: [] }, + next_page_token: null, + debug_info: { latency_ms: 42 }, + }); + const plugin = new AiSearchPlugin({ + indexes: { test: { indexName: "cat.sch.idx", columns: ["id"] } }, + }); + await plugin.setup(); + const res = await plugin.query("test", { queryText: "q" }); + + expect(res.queryTimeMs).toBe(42); + expect(res.results).toEqual([]); + expect(res.totalCount).toBe(0); + }); + }); + + describe("query() overrides and reranker", () => { + it("lets the request override index queryType, numResults, and columns", async () => { + const plugin = new AiSearchPlugin({ + indexes: { + test: { + indexName: "cat.sch.idx", + columns: ["id", "title"], + queryType: "hybrid", + numResults: 10, + }, + }, + }); + await plugin.setup(); + await plugin.query("test", { + queryText: "q", + queryType: "ann", + numResults: 5, + columns: ["id"], + }); + + const callBody = mockRequest.mock.calls[0][0].payload; + expect(callBody.query_type).toBe("ANN"); + expect(callBody.num_results).toBe(5); + expect(callBody.columns).toEqual(["id"]); + }); + + it("passes an object reranker through untouched", async () => { + const plugin = new AiSearchPlugin({ + indexes: { + test: { + indexName: "cat.sch.idx", + columns: ["id", "title", "body"], + reranker: { columnsToRerank: ["title"] }, + }, + }, + }); + await plugin.setup(); + await plugin.query("test", { queryText: "q" }); + + const callBody = mockRequest.mock.calls[0][0].payload; + expect(callBody.reranker.parameters.columns_to_rerank).toEqual(["title"]); + }); + + it("lets request.reranker=false suppress an index-enabled reranker", async () => { + const plugin = new AiSearchPlugin({ + indexes: { + test: { + indexName: "cat.sch.idx", + columns: ["id", "title"], + reranker: true, + }, + }, + }); + await plugin.setup(); + await plugin.query("test", { queryText: "q", reranker: false }); + + const callBody = mockRequest.mock.calls[0][0].payload; + expect(callBody.reranker).toBeUndefined(); + }); + + it("skips the reranker when enabled but no columns are resolved", async () => { + // Query-time behavior only; skip setup() (its prod guard rejects the + // deliberately column-less config used to exercise this path). + const plugin = new AiSearchPlugin({ + indexes: { test: { indexName: "cat.sch.idx", reranker: true } }, + }); + await plugin.query("test", { queryText: "q" }); + + const callBody = mockRequest.mock.calls[0][0].payload; + expect(callBody.reranker).toBeUndefined(); + expect(callBody.columns).toEqual([]); + }); + + it("throws a wrapped error when the connector query fails", async () => { + // Persistent reject so the retry interceptor exhausts its attempts and + // execute() surfaces a failed result, driving the !result.ok branch. + mockRequest.mockRejectedValue(new Error("VS 503")); + const plugin = new AiSearchPlugin({ + indexes: { products: { indexName: "cat.sch.p", columns: ["id"] } }, + }); + await plugin.setup(); + + await expect( + plugin.query("products", { queryText: "q" }), + ).rejects.toThrow(/Vector search query failed for index "products"/); + }); + }); + + describe("injectRoutes", () => { + const makePlugin = () => + new AiSearchPlugin({ + indexes: { + demo: { + indexName: "cat.sch.idx", + columns: ["id", "title"], + queryType: "hybrid", + }, + paged: { + indexName: "cat.sch.paged", + columns: ["id"], + pagination: true, + endpointName: "ep", + }, + }, + }); + + it("registers the three routes", () => { + const plugin = makePlugin(); + const { router } = createMockRouter(); + plugin.injectRoutes(router); + + expect(router.post).toHaveBeenCalledWith( + "/:alias/query", + expect.any(Function), + ); + expect(router.post).toHaveBeenCalledWith( + "/:alias/next-page", + expect.any(Function), + ); + expect(router.get).toHaveBeenCalledWith( + "/:alias/config", + expect.any(Function), + ); + }); + + describe("/:alias/query", () => { + it("404s an unknown alias", async () => { + const plugin = makePlugin(); + const { router, getHandler } = createMockRouter(); + plugin.injectRoutes(router); + + const res = createMockResponse(); + await getHandler("POST", "/:alias/query")( + createMockRequest({ params: { alias: "nope" }, body: {} }), + res, + ); + + expect(res.status).toHaveBeenCalledWith(404); + }); + + it("400s when neither queryText nor queryVector is provided", async () => { + const plugin = makePlugin(); + const { router, getHandler } = createMockRouter(); + plugin.injectRoutes(router); + + const res = createMockResponse(); + await getHandler("POST", "/:alias/query")( + createMockRequest({ params: { alias: "demo" }, body: {} }), + res, + ); + + expect(res.status).toHaveBeenCalledWith(400); + }); + + it("returns the parsed response on success", async () => { + const plugin = makePlugin(); + const { router, getHandler } = createMockRouter(); + plugin.injectRoutes(router); + + const res = createMockResponse(); + await getHandler("POST", "/:alias/query")( + createMockRequest({ + params: { alias: "demo" }, + body: { queryText: "hi" }, + }), + res, + ); + + expect(res.json).toHaveBeenCalledWith( + expect.objectContaining({ totalCount: 2, queryType: "hybrid" }), + ); + }); + + it("ignores a client-supplied columns override and uses the configured projection", async () => { + const plugin = makePlugin(); + const { router, getHandler } = createMockRouter(); + plugin.injectRoutes(router); + + const res = createMockResponse(); + await getHandler("POST", "/:alias/query")( + createMockRequest({ + params: { alias: "demo" }, + body: { queryText: "hi", columns: ["ssn", "internal_notes"] }, + }), + res, + ); + + // demo is configured with columns ["id", "title"]; the request's + // columns must not widen the projection. + const callBody = mockRequest.mock.calls[0][0].payload; + expect(callBody.columns).toEqual(["id", "title"]); + }); + + it("500s when query preparation throws", async () => { + // Query prep (embeddingFn) runs inside execute() so it shares the OBO + // context; a failure surfaces as a non-ok result → 500. + const plugin = new AiSearchPlugin({ + indexes: { + demo: { + indexName: "cat.sch.idx", + columns: ["id", "title"], + queryType: "ann", + embeddingFn: vi.fn().mockRejectedValue(new Error("embed down")), + }, + }, + }); + const { router, getHandler } = createMockRouter(); + plugin.injectRoutes(router); + + const res = createMockResponse(); + await getHandler("POST", "/:alias/query")( + createMockRequest({ + params: { alias: "demo" }, + body: { queryText: "hi" }, + }), + res, + ); + + expect(res.status).toHaveBeenCalledWith(500); + }); + }); + + describe("/:alias/next-page", () => { + it("400s when pagination is not enabled", async () => { + const plugin = makePlugin(); + const { router, getHandler } = createMockRouter(); + plugin.injectRoutes(router); + + const res = createMockResponse(); + await getHandler("POST", "/:alias/next-page")( + createMockRequest({ + params: { alias: "demo" }, + body: { pageToken: "t" }, + }), + res, + ); + + expect(res.status).toHaveBeenCalledWith(400); + }); + + it("400s when pageToken is missing", async () => { + const plugin = makePlugin(); + const { router, getHandler } = createMockRouter(); + plugin.injectRoutes(router); + + const res = createMockResponse(); + await getHandler("POST", "/:alias/next-page")( + createMockRequest({ params: { alias: "paged" }, body: {} }), + res, + ); + + expect(res.status).toHaveBeenCalledWith(400); + }); + + it("fetches the next page on success", async () => { + const plugin = makePlugin(); + const { router, getHandler } = createMockRouter(); + plugin.injectRoutes(router); + + const res = createMockResponse(); + await getHandler("POST", "/:alias/next-page")( + createMockRequest({ + params: { alias: "paged" }, + body: { pageToken: "t" }, + }), + res, + ); + + 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" }, + }), + expect.any(Context), + ); + expect(res.json).toHaveBeenCalled(); + }); + }); + + describe("/:alias/config", () => { + it("404s an unknown alias", async () => { + const plugin = makePlugin(); + const { router, getHandler } = createMockRouter(); + plugin.injectRoutes(router); + + const res = createMockResponse(); + await getHandler("GET", "/:alias/config")( + createMockRequest({ params: { alias: "nope" } }), + res, + ); + + expect(res.status).toHaveBeenCalledWith(404); + }); + + it("returns resolved config with defaults", async () => { + const plugin = makePlugin(); + const { router, getHandler } = createMockRouter(); + plugin.injectRoutes(router); + + const res = createMockResponse(); + await getHandler("GET", "/:alias/config")( + createMockRequest({ params: { alias: "demo" } }), + res, + ); + + expect(res.json).toHaveBeenCalledWith({ + alias: "demo", + columns: ["id", "title"], + queryType: "hybrid", + numResults: 20, + reranker: false, + pagination: false, + }); + }); + }); + }); +}); diff --git a/packages/appkit/src/plugins/vector-search/types.ts b/packages/appkit/src/plugins/ai-search/types.ts similarity index 54% rename from packages/appkit/src/plugins/vector-search/types.ts rename to packages/appkit/src/plugins/ai-search/types.ts index a2760fced..93dba6063 100644 --- a/packages/appkit/src/plugins/vector-search/types.ts +++ b/packages/appkit/src/plugins/ai-search/types.ts @@ -1,22 +1,38 @@ import type { BasePluginConfig } from "shared"; -export interface IVectorSearchConfig extends BasePluginConfig { +/** Vector Search query mode: semantic (`ann`), keyword+semantic (`hybrid`), or keyword-only (`full_text`). */ +export type SearchQueryType = "ann" | "hybrid" | "full_text"; + +export interface IAiSearchConfig 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[]; + /** + * Three-level UC name: catalog.schema.index_name. Defaults to the + * `DATABRICKS_VS_INDEX_NAME` env var when omitted — so multiple aliases that + * omit it all resolve to that same physical index. Set it explicitly per + * alias when they should point at distinct indexes. + */ + indexName?: string; + /** + * Columns to return in results. Optional: in development the plugin + * auto-discovers them from the index's source table when omitted (and warns + * that they should be set explicitly for production). + */ + columns?: string[]; /** Default search mode */ - queryType?: "ann" | "hybrid" | "full_text"; + queryType?: SearchQueryType; /** 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 mode for the built-in HTTP routes — "service-principal" (default) + * uses the app's SP, "on-behalf-of-user" proxies the logged-in user's token. + * Programmatic callers select per call via `appkit.aiSearch.asUser(req)`. + */ auth?: "service-principal" | "on-behalf-of-user"; /** Enable cursor pagination */ pagination?: boolean; @@ -34,6 +50,13 @@ export interface RerankerConfig { columnsToRerank: string[]; } +/** Public summary of a configured index, exposed to the client via `clientConfig()`. */ +export interface IndexSummary { + alias: string; + queryType: SearchQueryType; + pagination: boolean; +} + export type SearchFilters = Record< string, string | number | boolean | (string | number)[] @@ -44,7 +67,7 @@ export interface SearchRequest { queryVector?: number[]; columns?: string[]; numResults?: number; - queryType?: "ann" | "hybrid" | "full_text"; + queryType?: SearchQueryType; filters?: SearchFilters; reranker?: boolean; } @@ -55,7 +78,7 @@ export interface SearchResponse< results: SearchResult[]; totalCount: number; queryTimeMs: number; - queryType: "ann" | "hybrid" | "full_text"; + queryType: SearchQueryType; nextPageToken: string | null; } diff --git a/packages/appkit/src/plugins/beta-exports.generated.ts b/packages/appkit/src/plugins/beta-exports.generated.ts index 82f6c4a78..7e556ebd9 100644 --- a/packages/appkit/src/plugins/beta-exports.generated.ts +++ b/packages/appkit/src/plugins/beta-exports.generated.ts @@ -6,3 +6,4 @@ // manifests and the synced appkit.plugins.json. export { agents } from "./agents"; +export { aiSearch } from "./ai-search"; diff --git a/packages/appkit/src/plugins/server/index.ts b/packages/appkit/src/plugins/server/index.ts index 968f7bc78..bf1e58091 100644 --- a/packages/appkit/src/plugins/server/index.ts +++ b/packages/appkit/src/plugins/server/index.ts @@ -5,6 +5,7 @@ import dotenv from "dotenv"; import express from "express"; import getPort, { portNumbers } from "get-port"; import type { PluginClientConfigs, PluginPhase } from "shared"; +import { camelToKebab } from "shared"; import { AppKitError, ServerError } from "../../errors"; import { TelemetryReporter } from "../../internal-telemetry"; import { createLogger } from "../../logging/logger"; @@ -272,7 +273,7 @@ export class ServerPlugin extends Plugin { plugin.injectRoutes(router); - const basePath = `/api/${plugin.name}`; + const basePath = `/api/${camelToKebab(plugin.name)}`; this.serverApplication.use(basePath, router); endpoints[plugin.name] = plugin.getEndpoints(); diff --git a/packages/appkit/src/plugins/ui-variants/index.ts b/packages/appkit/src/plugins/ui-variants/index.ts index e53a7ff6d..c33f2ae3e 100644 --- a/packages/appkit/src/plugins/ui-variants/index.ts +++ b/packages/appkit/src/plugins/ui-variants/index.ts @@ -25,7 +25,7 @@ interface ConfirmRequestBody { * per `` id — and the plugin only records; it never edits source. */ class UiVariantsPlugin extends Plugin { - static manifest = manifest as PluginManifest<"ui-variants">; + static manifest = manifest as PluginManifest<"uiVariants">; protected static description = "Dev-only recorder for the UI picker"; diff --git a/packages/appkit/src/plugins/ui-variants/manifest.json b/packages/appkit/src/plugins/ui-variants/manifest.json index 2c9fbc27c..94cde02c8 100644 --- a/packages/appkit/src/plugins/ui-variants/manifest.json +++ b/packages/appkit/src/plugins/ui-variants/manifest.json @@ -1,6 +1,6 @@ { "$schema": "https://databricks.github.io/appkit/schemas/plugin-manifest.schema.json", - "name": "ui-variants", + "name": "uiVariants", "displayName": "UI Variants Plugin", "description": "Dev-only recorder for the UI picker: records the developer's chosen variant so a coding agent can finalize the component source", "hidden": true, diff --git a/packages/appkit/src/plugins/vector-search/index.ts b/packages/appkit/src/plugins/vector-search/index.ts deleted file mode 100644 index d733a0f27..000000000 --- a/packages/appkit/src/plugins/vector-search/index.ts +++ /dev/null @@ -1,2 +0,0 @@ -export * from "./types"; -export * from "./vector-search"; 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 deleted file mode 100644 index 104eb1cb6..000000000 --- a/packages/appkit/src/plugins/vector-search/tests/vector-search.test.ts +++ /dev/null @@ -1,330 +0,0 @@ -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(), - setExecution: 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].payload; - 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].payload; - 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].payload; - 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].payload; - 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/vector-search.ts b/packages/appkit/src/plugins/vector-search/vector-search.ts deleted file mode 100644 index fefc3f409..000000000 --- a/packages/appkit/src/plugins/vector-search/vector-search.ts +++ /dev/null @@ -1,376 +0,0 @@ -import type express from "express"; -import type { IAppRouter, PluginExecutionSettings } from "shared"; -import { VectorSearchConnector } from "../../connectors/vector-search/client"; -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 { - 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( - `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.ok) { - res - .status(result.status) - .json({ error: result.message, plugin: this.name }); - return; - } - res.json(this._parseResponse(result.data, 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.ok) { - res - .status(result.status) - .json({ error: result.message, plugin: this.name }); - return; - } - res.json( - this._parseResponse(result.data, indexConfig.queryType ?? "hybrid"), - ); - } catch (error) { - this._handleError(res, error, "Next-page query failed"); - } - }, - }); - - this.route(router, { - name: "getConfig", - method: "get", - path: "/:alias/config", - handler: async (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.ok) { - throw new Error( - `Vector search query failed for index "${alias}": ${result.message}`, - ); - } - - return this._parseResponse(result.data, 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/packages/appkit/src/registry/resource-registry.ts b/packages/appkit/src/registry/resource-registry.ts index 7b2025b21..489faa536 100644 --- a/packages/appkit/src/registry/resource-registry.ts +++ b/packages/appkit/src/registry/resource-registry.ts @@ -13,6 +13,7 @@ import type { BasePluginConfig, PluginConstructor, PluginData } from "shared"; import { ConfigurationError } from "../errors"; import { createLogger } from "../logging/logger"; +import { formatWarningBanner } from "../utils/banner"; import { getPluginManifest } from "./manifest-loader"; import type { ResourceEntry, @@ -477,11 +478,6 @@ export class ResourceRegistry { "Add these to your .env file or environment to suppress this warning.", ); - const maxLen = Math.max(...contentLines.map((l) => l.length)); - const border = "=".repeat(maxLen + 4); - - const boxed = contentLines.map((line) => `| ${line.padEnd(maxLen)} |`); - - return [border, ...boxed, border].join("\n"); + return formatWarningBanner(contentLines); } } diff --git a/packages/appkit/src/utils/banner.ts b/packages/appkit/src/utils/banner.ts new file mode 100644 index 000000000..251d2d865 --- /dev/null +++ b/packages/appkit/src/utils/banner.ts @@ -0,0 +1,11 @@ +/** + * Frames content lines in an ASCII box (`===` border, `| … |` sides) padded to + * the widest line. Used for prominent dev-mode warnings that must not be missed + * in a noisy console. + */ +export function formatWarningBanner(lines: string[]): string { + const maxLen = Math.max(...lines.map((l) => l.length)); + const border = "=".repeat(maxLen + 4); + const boxed = lines.map((line) => `| ${line.padEnd(maxLen)} |`); + return [border, ...boxed, border].join("\n"); +} diff --git a/packages/appkit/src/utils/index.ts b/packages/appkit/src/utils/index.ts index c0b1b55bd..d0ebedc44 100644 --- a/packages/appkit/src/utils/index.ts +++ b/packages/appkit/src/utils/index.ts @@ -1,3 +1,4 @@ +export * from "./banner"; export * from "./merge"; export * from "./path-exclusions"; export * from "./vite-config-merge"; diff --git a/packages/shared/src/cli/commands/plugin/create/create.ts b/packages/shared/src/cli/commands/plugin/create/create.ts index 5917cfe11..db9e659b4 100644 --- a/packages/shared/src/cli/commands/plugin/create/create.ts +++ b/packages/shared/src/cli/commands/plugin/create/create.ts @@ -13,6 +13,7 @@ import { text, } from "@clack/prompts"; import { Command, Option } from "commander"; +import { PLUGIN_NAME_PATTERN } from "../../../../naming"; import { promptOneResource } from "./prompt-resource"; import { DEFAULT_PERMISSION_BY_TYPE, @@ -25,7 +26,6 @@ import { import { resolveTargetDir, scaffoldPlugin } from "./scaffold"; import type { CreateAnswers, Placement, SelectedResource } from "./types"; -const NAME_PATTERN = /^[a-z][a-z0-9-]*$/; const DEFAULT_VERSION = "0.1.0"; const VALID_PLACEMENTS: Placement[] = ["in-repo", "isolated"]; const REQUIRED_FLAGS = ["placement", "path", "name", "description"] as const; @@ -42,17 +42,9 @@ interface CreateOptions { } function deriveDisplayName(name: string): string { - return name - .split("-") - .map((s) => s.charAt(0).toUpperCase() + s.slice(1)) - .join(" "); -} - -function deriveExportName(name: string): string { - return name - .split("-") - .map((s, i) => (i === 0 ? s : s.charAt(0).toUpperCase() + s.slice(1))) - .join(""); + // camelCase -> Title Case words (e.g. "myPlugin" -> "My Plugin"). + const spaced = name.replace(/([a-z0-9])([A-Z])/g, "$1 $2"); + return spaced.charAt(0).toUpperCase() + spaced.slice(1); } function buildResourceFromType(type: string): SelectedResource { @@ -138,7 +130,8 @@ function printNextSteps(answers: CreateAnswers, targetDir: string): void { const importPath = relativePath.startsWith(".") ? relativePath : `./${relativePath}`; - const exportName = deriveExportName(answers.name); + // The camelCase plugin name is already the JS binding. + const exportName = answers.name; console.log("\nNext steps:\n"); if (answers.placement === "in-repo") { @@ -169,7 +162,7 @@ function runNonInteractive(opts: CreateOptions): void { ); console.error(`Missing: ${missing.map((f) => `--${f}`).join(", ")}`); console.error( - ' appkit plugin create --placement in-repo --path plugins/my-plugin --name my-plugin --description "Does X"', + ' appkit plugin create --placement in-repo --path plugins/my-plugin --name myPlugin --description "Does X"', ); process.exit(1); } @@ -194,9 +187,9 @@ function runNonInteractive(opts: CreateOptions): void { } const name = opts.name as string; - if (!NAME_PATTERN.test(name)) { + if (!PLUGIN_NAME_PATTERN.test(name)) { console.error( - "Error: --name must be lowercase, start with a letter, and use only letters, numbers, and hyphens.", + "Error: --name must start with a lowercase letter and be camelCase (letters and numbers only, e.g. myPlugin).", ); process.exit(1); } @@ -289,11 +282,11 @@ async function runInteractive(): Promise { const name = await text({ message: "Plugin name (id)", - placeholder: "my-plugin", + placeholder: "myPlugin", validate(value) { if (!value?.trim()) return "Name is required."; - if (!NAME_PATTERN.test(value as string)) { - return "Must be lowercase, start with a letter, and use only letters, numbers, and hyphens."; + if (!PLUGIN_NAME_PATTERN.test(value as string)) { + return "Must start with a lowercase letter and be camelCase (letters and numbers only, e.g. myPlugin)."; } return undefined; }, @@ -445,7 +438,7 @@ async function runPluginCreate(opts: CreateOptions): Promise { `Error: Non-interactive mode requires: ${REQUIRED_FLAGS.map((f) => `--${f}`).join(", ")}`, ); console.error( - ' appkit plugin create --placement in-repo --path plugins/my-plugin --name my-plugin --description "Does X"', + ' appkit plugin create --placement in-repo --path plugins/my-plugin --name myPlugin --description "Does X"', ); process.exit(1); } @@ -478,8 +471,8 @@ export const pluginCreateCommand = new Command("create") ` Examples: $ appkit plugin create - $ appkit plugin create --placement in-repo --path plugins/my-plugin --name my-plugin --description "Does X" - $ appkit plugin create --placement in-repo --path plugins/my-plugin --name my-plugin --description "Does X" --resources sql_warehouse,volume --force + $ appkit plugin create --placement in-repo --path plugins/my-plugin --name myPlugin --description "Does X" + $ appkit plugin create --placement in-repo --path plugins/my-plugin --name myPlugin --description "Does X" --resources sql_warehouse,volume --force $ appkit plugin create --placement isolated --path appkit-plugin-ml --name ml --description "ML" --resources-json '[{"type":"serving_endpoint"}]'`, ) .action((opts) => diff --git a/packages/shared/src/cli/commands/plugin/list/list.test.ts b/packages/shared/src/cli/commands/plugin/list/list.test.ts index e7fe8887f..1ce35aa92 100644 --- a/packages/shared/src/cli/commands/plugin/list/list.test.ts +++ b/packages/shared/src/cli/commands/plugin/list/list.test.ts @@ -46,7 +46,7 @@ const TEMPLATE_MANIFEST_JSON = { const PLUGIN_MANIFEST_JSON = { $schema: "https://databricks.github.io/appkit/schemas/plugin-manifest.schema.json", - name: "my-feature", + name: "myFeature", displayName: "My Feature", description: "A test plugin", resources: { required: [], optional: [] }, @@ -144,7 +144,7 @@ describe("list", () => { plugins: { ...TEMPLATE_MANIFEST_JSON.plugins, beta: { - name: "beta-plugin", + name: "betaPlugin", displayName: "Beta Plugin", package: "@databricks/appkit", stability: "beta", @@ -156,7 +156,7 @@ describe("list", () => { fs.writeFileSync(manifestPath, JSON.stringify(manifest, null, 2)); const rows = listFromManifestFile(manifestPath); - const betaRow = rows.find((r) => r.name === "beta-plugin"); + const betaRow = rows.find((r) => r.name === "betaPlugin"); const gaRow = rows.find((r) => r.name === "server"); expect(betaRow?.stability).toBe("beta"); @@ -178,7 +178,7 @@ describe("list", () => { const rows = await listFromDirectory(tmp, path.dirname(tmp)); expect(rows).toHaveLength(1); - expect(rows[0].name).toBe("my-feature"); + expect(rows[0].name).toBe("myFeature"); expect(rows[0].displayName).toBe("My Feature"); expect(rows[0].package).toContain("my-feature"); expect(rows[0].required).toBe(0); @@ -212,7 +212,7 @@ describe("list", () => { const rows = await listFromDirectory(tmp, path.dirname(tmp)); expect(rows).toHaveLength(1); - expect(rows[0].name).toBe("my-feature"); + expect(rows[0].name).toBe("myFeature"); }); it("reads stability from manifest in directory scan", async () => { @@ -224,7 +224,7 @@ describe("list", () => { path.join(pluginDir, "manifest.json"), JSON.stringify({ ...PLUGIN_MANIFEST_JSON, - name: "beta-feature", + name: "betaFeature", stability: "beta", }), ); @@ -277,7 +277,7 @@ describe("list", () => { const rows = await listFromDirectory(tmp, path.dirname(tmp), true); expect(rows).toHaveLength(1); - expect(rows[0].name).toBe("my-feature"); + expect(rows[0].name).toBe("myFeature"); }); it("loads JS manifests from trusted node_modules packages by default", async () => { @@ -298,7 +298,7 @@ describe("list", () => { const rows = await listFromDirectory(tmp, tmp); expect(rows).toHaveLength(1); - expect(rows[0].name).toBe("my-feature"); + expect(rows[0].name).toBe("myFeature"); }); it("does not load JS manifests from untrusted node_modules packages by default", async () => { diff --git a/packages/shared/src/cli/commands/plugin/promote/promote.ts b/packages/shared/src/cli/commands/plugin/promote/promote.ts index 39d62d602..6f6506fc7 100644 --- a/packages/shared/src/cli/commands/plugin/promote/promote.ts +++ b/packages/shared/src/cli/commands/plugin/promote/promote.ts @@ -2,6 +2,7 @@ import fs from "node:fs"; import path from "node:path"; import process from "node:process"; import { Command } from "commander"; +import { kebabToCamel } from "../../../../naming"; import { resolveManifestInDir } from "../manifest-resolve"; import { isWithinDirectory } from "../sync/sync"; import { shouldAllowJsManifestForDir } from "../trusted-js-manifest"; @@ -172,18 +173,6 @@ function escapeRegex(s: string): string { return s.replace(/[.*+?^${}()|[\]\\]/g, "\\$&"); } -/** - * Convert a kebab-case manifest name to its camelCase JS identifier form - * (e.g. `vector-search` -> `vectorSearch`). Mirrors the convention used by - * first-party plugin index files: a manifest's `name` field may be - * kebab-case (the schema permits `^[a-z][a-z0-9-]*$`), but the actual - * exported binding is always a JS identifier. We try both forms when - * matching specifiers in user code. - */ -function manifestNameToBinding(pluginName: string): string { - return pluginName.replace(/-+([a-z0-9])/g, (_, c: string) => c.toUpperCase()); -} - /** * Returns true when the named import specifier `spec` resolves to either * `pluginName` itself or its kebab-to-camelCase JS-identifier form. @@ -196,7 +185,7 @@ function specifierMatchesPlugin(spec: string, pluginName: string): boolean { const stripped = spec.replace(/^type\s+/, "").trim(); const head = stripped.split(/\s+as\s+/)[0]?.trim(); if (!head) return false; - return head === pluginName || head === manifestNameToBinding(pluginName); + return head === pluginName || head === kebabToCamel(pluginName); } /** diff --git a/packages/shared/src/cli/commands/plugin/validate/validate-manifest.test.ts b/packages/shared/src/cli/commands/plugin/validate/validate-manifest.test.ts index 70490747e..ee448beea 100644 --- a/packages/shared/src/cli/commands/plugin/validate/validate-manifest.test.ts +++ b/packages/shared/src/cli/commands/plugin/validate/validate-manifest.test.ts @@ -15,7 +15,7 @@ import { const VALID_MANIFEST = { $schema: "https://databricks.github.io/appkit/schemas/plugin-manifest.schema.json", - name: "test-plugin", + name: "testPlugin", displayName: "Test Plugin", description: "A test plugin", resources: { @@ -90,7 +90,7 @@ describe("validate-manifest", () => { const result = validateManifest(VALID_MANIFEST); expect(result.valid).toBe(true); expect(result.manifest).toBeDefined(); - expect(result.manifest?.name).toBe("test-plugin"); + expect(result.manifest?.name).toBe("testPlugin"); }); it("validates a manifest with resources", () => { diff --git a/packages/shared/src/naming.ts b/packages/shared/src/naming.ts new file mode 100644 index 000000000..9f0f037b2 --- /dev/null +++ b/packages/shared/src/naming.ts @@ -0,0 +1,27 @@ +/** + * Plugin-name casing helpers. A dependency-free leaf module so both the + * runtime (appkit) and the bundled CLI can import it without pulling in the + * heavier `plugin.ts` graph. + */ + +/** + * Canonical plugin-name charset: a lowercase-initial camelCase JS identifier + * (e.g. `aiSearch`). The name doubles as the JS binding and the accessor key, + * so kebab is not allowed. Single source of truth for the manifest schema and + * the plugin-tree generators. + */ +export const PLUGIN_NAME_PATTERN = /^[a-z][a-zA-Z0-9]*$/; + +/** kebab-case to camelCase (e.g. `"ai-search"` -> `"aiSearch"`). */ +export function kebabToCamel(name: string): string { + return name.replace(/-+([a-z0-9])/g, (_, c: string) => c.toUpperCase()); +} + +/** + * camelCase to kebab-case (e.g. `"aiSearch"` -> `"ai-search"`). Used to derive + * HTTP route prefixes and folder paths from the canonical camelCase plugin + * name. A no-op for single-word names. + */ +export function camelToKebab(name: string): string { + return name.replace(/[A-Z]/g, (m) => `-${m.toLowerCase()}`); +} diff --git a/packages/shared/src/plugin.ts b/packages/shared/src/plugin.ts index 9800c261f..46f1186b9 100644 --- a/packages/shared/src/plugin.ts +++ b/packages/shared/src/plugin.ts @@ -11,6 +11,8 @@ import type { // Sourced from `./schemas/manifest` (the Zod canonical) so `DiscoveryDescriptor` // stays the discriminated union shape rather than the free-form predecessor. export type { ResourceFieldEntry, DiscoveryDescriptor, PluginScaffoldingRules }; +// Re-export the naming helpers so `shared` consumers import them from here. +export { camelToKebab, kebabToCamel } from "./naming"; /** Base plugin interface. */ export interface BasePlugin { @@ -250,11 +252,13 @@ export type WithAsUser = SDK extends (...args: any[]) => any /** * Maps plugin names to their exported types (with asUser automatically added). - * Each plugin exposes its public API via the exports() method, - * and AppKit wraps it with asUser() for user-scoped execution. + * Each plugin exposes its public API via the exports() method, and AppKit + * wraps it with asUser() for user-scoped execution. Callable exports + * (functions) are passed through without wrapping, as they manage their own + * `asUser` pattern (e.g. files plugin). * - * Callable exports (functions) are passed through without wrapping, - * as they manage their own `asUser` pattern (e.g. files plugin). + * The key is the plugin's manifest `name`, which is camelCase by convention + * (`appkit.aiSearch`), so it doubles as a valid JS accessor. */ export type PluginMap< U extends readonly PluginData[], diff --git a/packages/shared/src/schemas/manifest.ts b/packages/shared/src/schemas/manifest.ts index ceb0412ab..bf41293f7 100644 --- a/packages/shared/src/schemas/manifest.ts +++ b/packages/shared/src/schemas/manifest.ts @@ -29,6 +29,7 @@ */ import { z } from "zod"; +import { PLUGIN_NAME_PATTERN } from "../naming"; // ── Resource type + per-type permission enums ──────────────────────────── @@ -555,7 +556,7 @@ export const configSchemaPropertySchema: z.ZodType = z.lazy(() => maxLength: z.number().int().min(0).optional(), required: z.array(z.string()).optional(), // `additionalProperties` is a standard JSON Schema keyword used by core - // plugin manifests (e.g., serving, vector-search, genie) to constrain + // plugin manifests (e.g., serving, ai-search, genie) to constrain // dictionary-shaped properties. Allowed on nested property entries as // either a boolean or a sub-schema, mirroring JSON Schema semantics. additionalProperties: z @@ -672,9 +673,9 @@ export const pluginManifestSchema = z .describe("Reference to the JSON Schema for validation"), name: z .string() - .regex(/^[a-z][a-z0-9-]*$/) + .regex(PLUGIN_NAME_PATTERN) .describe( - "Plugin identifier. Must be lowercase, start with a letter, and contain only letters, numbers, and hyphens.", + "Plugin identifier and JS binding. Must start with a lowercase letter; camelCase for multi-word names (e.g. aiSearch).", ), displayName: z .string() @@ -908,9 +909,9 @@ export const templatePluginSchema = z .object({ name: z .string() - .regex(/^[a-z][a-z0-9-]*$/) + .regex(PLUGIN_NAME_PATTERN) .describe( - "Plugin identifier. Must be lowercase, start with a letter, and contain only letters, numbers, and hyphens.", + "Plugin identifier and JS binding. Must start with a lowercase letter; camelCase for multi-word names (e.g. aiSearch).", ), displayName: z .string() diff --git a/template/appkit.plugins.json b/template/appkit.plugins.json index a4aa1c26e..078ab524a 100644 --- a/template/appkit.plugins.json +++ b/template/appkit.plugins.json @@ -28,6 +28,41 @@ }, "stability": "beta" }, + "aiSearch": { + "name": "aiSearch", + "displayName": "AI Search Plugin", + "description": "Query Databricks Vector Search indexes with built-in hybrid search, reranking, and pagination", + "package": "@databricks/appkit", + "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": { + "id": { + "env": "DATABRICKS_VS_INDEX_NAME", + "description": "Three-level UC name of the default index (catalog.schema.index_name)", + "origin": "user" + } + } + } + ], + "optional": [] + }, + "stability": "beta", + "scaffolding": { + "rules": { + "must": [ + "List the vector search index's source-table columns (unity-catalog tables API) so the user can pick which to return", + "Ask the user which of those columns the search should return", + "Write aiSearch({ indexes: { default: { columns: [...] } } }) in the server file; the client queries the 'default' alias" + ] + } + } + }, "analytics": { "name": "analytics", "displayName": "Analytics Plugin", diff --git a/template/client/src/App.tsx b/template/client/src/App.tsx index 2e9f9a9c6..3b4f0b664 100644 --- a/template/client/src/App.tsx +++ b/template/client/src/App.tsx @@ -31,8 +31,8 @@ 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'; +{{- if .plugins.aiSearch}} +import { AiSearchPage } from './pages/ai-search/AiSearchPage'; {{- end}} {{- if .plugins.jobs}} import { JobsPage } from './pages/jobs/JobsPage'; @@ -90,9 +90,9 @@ function NavLinks({ className, linkClass, onClick }: { className?: string; linkC Serving {{- end}} -{{- if .plugins.vectorSearch}} - - Vector Search +{{- if .plugins.aiSearch}} + + AI Search {{- end}} {{- if .plugins.jobs}} @@ -166,8 +166,8 @@ const router = createBrowserRouter([ {{- if .plugins.serving}} { path: '/serving', element: }, {{- end}} -{{- if .plugins.vectorSearch}} - { path: '/vector-search', element: }, +{{- if .plugins.aiSearch}} + { path: '/ai-search', element: }, {{- end}} {{- if .plugins.jobs}} { path: '/jobs', element: }, diff --git a/template/client/src/pages/vector-search/VectorSearchPage.tsx b/template/client/src/pages/ai-search/AiSearchPage.tsx similarity index 62% rename from template/client/src/pages/vector-search/VectorSearchPage.tsx rename to template/client/src/pages/ai-search/AiSearchPage.tsx index f1e5e58fe..5d07b335a 100644 --- a/template/client/src/pages/vector-search/VectorSearchPage.tsx +++ b/template/client/src/pages/ai-search/AiSearchPage.tsx @@ -1,4 +1,4 @@ -{{if .plugins.vectorSearch -}} +{{if .plugins.aiSearch -}} import { Button, Card, @@ -8,66 +8,29 @@ import { Input, Skeleton, } from '@databricks/appkit-ui/react'; +import { useAiSearchQuery } from '@databricks/appkit-ui/react/beta'; 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() { +export function AiSearchPage() { 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}`); - } + // Queries the first configured index. Pass `{ alias }` to target another. + const { search, data, loading, error } = useAiSearchQuery(); - const data: SearchResponse = await res.json(); - setResponse(data); - } catch (err) { - setError(err instanceof Error ? err.message : String(err)); - } finally { - setLoading(false); - } + const handleSearch = () => { + if (query.trim()) void search(query); }; const handleKeyDown = (e: React.KeyboardEvent) => { - if (e.key === 'Enter') { - void handleSearch(); - } + if (e.key === 'Enter') handleSearch(); }; return (
-

Vector Search

+

AI Search

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

@@ -79,7 +42,7 @@ export function VectorSearchPage() { onKeyDown={handleKeyDown} className="flex-1" /> -
)} - {response && !loading && ( + {data && !loading && (

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

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

No results found.

) : ( - response.results.map((result, index) => ( + data.results.map((result, index) => ( diff --git a/tools/dist-appkit.ts b/tools/dist-appkit.ts index abedf4754..9005e8905 100644 --- a/tools/dist-appkit.ts +++ b/tools/dist-appkit.ts @@ -90,6 +90,17 @@ if (fs.existsSync(sharedBin)) { fs.cpSync(sharedCliDist, tmpCliDist, { recursive: true }); } + // The CLI imports leaf modules that live outside dist/cli (e.g. + // `naming.ts`, referenced by the `plugin promote` command). Copy them to + // tmp/dist so the CLI's relative imports resolve in the published tarball. + const sharedNaming = path.join( + __dirname, + "../packages/shared/dist/naming.js", + ); + if (fs.existsSync(sharedNaming)) { + fs.copyFileSync(sharedNaming, "tmp/dist/naming.js"); + } + // Copy JSON schemas so CLI (e.g. plugin validate/sync) can load them at runtime. // Place in both dist/schemas and dist/cli/schemas so resolution works whether // the running module's __dirname is under dist/ or dist/cli/ (e.g. after bundling). diff --git a/tools/generate-plugin-doc-banners.ts b/tools/generate-plugin-doc-banners.ts index 470729f6c..59812a0bd 100644 --- a/tools/generate-plugin-doc-banners.ts +++ b/tools/generate-plugin-doc-banners.ts @@ -16,18 +16,16 @@ import fs from "node:fs"; import path from "node:path"; import { fileURLToPath } from "node:url"; +import { + camelToKebab, + PLUGIN_NAME_PATTERN, +} from "../packages/shared/src/naming"; const __dirname = path.dirname(fileURLToPath(import.meta.url)); const REPO_ROOT = path.join(__dirname, ".."); const PLUGINS_DIR = path.join(REPO_ROOT, "packages/appkit/src/plugins"); const DOCS_DIR = path.join(REPO_ROOT, "docs/docs/plugins"); -/** - * Same as `plugin-manifest.schema.json` `name` pattern; keeps `path.join` targets - * under `docs/docs/plugins` (defense in depth vs path traversal in `name`). - */ -const SCHEMA_NAME_PATTERN = /^[a-z][a-z0-9-]*$/; - /** * Checks whether a resolved file path is within a given directory boundary. */ @@ -114,9 +112,9 @@ function readPluginInfos(): PluginInfo[] { continue; // not a valid plugin manifest, skip silently } - if (!SCHEMA_NAME_PATTERN.test(manifest.name)) { + if (!PLUGIN_NAME_PATTERN.test(manifest.name)) { throw new Error( - `Manifest name "${manifest.name}" in ${manifestPath} doesn't match the plugin manifest schema pattern ^[a-z][a-z0-9-]*$.`, + `Manifest name "${manifest.name}" in ${manifestPath} doesn't match the plugin manifest schema pattern ${PLUGIN_NAME_PATTERN.source}.`, ); } @@ -128,7 +126,7 @@ function readPluginInfos(): PluginInfo[] { } const docBasename = - DOC_FILE_OVERRIDES[manifest.name] ?? `${manifest.name}.md`; + DOC_FILE_OVERRIDES[manifest.name] ?? `${camelToKebab(manifest.name)}.md`; if ( docBasename.includes("..") || docBasename !== path.basename(docBasename) || @@ -191,7 +189,7 @@ function main(): void { for (const s of summary) { const rel = path.relative(REPO_ROOT, DOCS_DIR); - const docName = DOC_FILE_OVERRIDES[s.name] ?? `${s.name}.md`; + const docName = DOC_FILE_OVERRIDES[s.name] ?? `${camelToKebab(s.name)}.md`; if (s.action === "missing") { console.warn( ` warn: ${s.name} — no doc page at ${rel}/${docName} (skipping)`, diff --git a/tools/generate-plugin-entries.ts b/tools/generate-plugin-entries.ts index c2aa6ec7a..f0ce6980e 100644 --- a/tools/generate-plugin-entries.ts +++ b/tools/generate-plugin-entries.ts @@ -14,6 +14,10 @@ import fs from "node:fs"; import path from "node:path"; import { fileURLToPath } from "node:url"; +import { + kebabToCamel, + PLUGIN_NAME_PATTERN, +} from "../packages/shared/src/naming"; import { formatWithBiome } from "./format-with-biome.ts"; const __dirname = path.dirname(fileURLToPath(import.meta.url)); @@ -31,45 +35,40 @@ const HEADER = `// AUTO-GENERATED from packages/appkit/src/plugins//manife `; interface PluginInfo { - name: string; + /** camelCase JS-identifier binding emitted into the barrel. */ + binding: string; folder: string; stability: "beta" | "ga"; } /** - * Mirrors `^[a-z][a-z0-9-]*$` from `plugin-manifest.schema.json`. Catches - * malformed manifests that bypassed `appkit plugin validate`. + * Charsets the manifest `name` (camelCase, from `plugin-manifest.schema.json`) + * and the folder name (kebab) must match. Both flow into generated TS source, + * so these double as a code-injection gate (CWE-94): the charsets forbid + * quotes, semicolons, braces, backslashes, and newlines, so neither can break + * out of the string/identifier context it lands in. */ -const SCHEMA_NAME_PATTERN = /^[a-z][a-z0-9-]*$/; +const FOLDER_NAME_PATTERN = /^[a-z][a-z0-9-]*$/; /** - * Generator-only: the `name` field is interpolated unescaped into a TS - * `export { } from "./";` template, so it MUST be a valid - * JavaScript identifier. The schema accepts hyphens (e.g. "my-plugin"), - * which would produce `export { my-plugin }` — a TypeScript syntax error. - * - * This is also a defense-in-depth gate against code-injection (CWE-94) - * via a malicious `name` containing `}`, `;`, quotes, newlines, etc. - * - * Restricted to camelCase / underscore identifiers starting with a lowercase - * letter to match the existing built-in plugins (`analytics`, `lakebase`, - * `vectorSearch`, …) and the schema's lowercase-first rule. + * The barrel exports each plugin under a JS-identifier binding + * (`export { } from "./";`). The camelCase `name` is the + * binding; `kebabToCamel` is a no-op for it but normalizes the kebab folder + * name. This pattern is the final assertion that the derived binding is safe + * to interpolate unescaped. */ const JS_IDENTIFIER_PATTERN = /^[a-z][a-zA-Z0-9_]*$/; -function validateIdentifier( +function validateSchemaName( value: string, kind: "manifest name" | "folder name", manifestPath: string, ): void { - if (!SCHEMA_NAME_PATTERN.test(value)) { - throw new Error( - `${kind} "${value}" in ${manifestPath} doesn't match the plugin manifest schema pattern ^[a-z][a-z0-9-]*$. Run \`appkit plugin validate\` to catch this earlier.`, - ); - } - if (!JS_IDENTIFIER_PATTERN.test(value)) { + const pattern = + kind === "manifest name" ? PLUGIN_NAME_PATTERN : FOLDER_NAME_PATTERN; + if (!pattern.test(value)) { throw new Error( - `${kind} "${value}" in ${manifestPath} is not a valid JavaScript identifier (must match ^[a-z][a-zA-Z0-9_]*$). The generator interpolates this name into \`export { ${value} } from "./";\` and would emit invalid TypeScript. Rename the plugin folder + manifest \`name\` to camelCase, or set \`hidden: true\` to exclude it from the auto-generated barrels.`, + `${kind} "${value}" in ${manifestPath} doesn't match ${pattern.source}. Run \`appkit plugin validate\` to catch this earlier.`, ); } } @@ -102,12 +101,21 @@ function readPluginInfos(): PluginInfo[] { throw new Error(`Manifest missing "name": ${manifestPath}`); } - // Both the manifest `name` (used as the exported binding) and the + // Both the manifest `name` (source of the exported binding) and the // folder name (used as the `from` path) flow into a TS source file - // unescaped. Validate both against the schema and the JS-identifier - // rule before we emit anything. - validateIdentifier(manifest.name, "manifest name", manifestPath); - validateIdentifier(entry.name, "folder name", manifestPath); + // unescaped, so both must match the schema charset before we emit + // anything. + validateSchemaName(manifest.name, "manifest name", manifestPath); + validateSchemaName(entry.name, "folder name", manifestPath); + + // The schema permits kebab-case names, but the barrel binding must be a + // valid JS identifier, so derive it via kebab->camelCase and assert. + const binding = kebabToCamel(manifest.name); + if (!JS_IDENTIFIER_PATTERN.test(binding)) { + throw new Error( + `Manifest name "${manifest.name}" in ${manifestPath} does not convert to a valid JavaScript identifier (got "${binding}"). The generator emits \`export { ${binding} } from "./";\`, which would be invalid TypeScript. Rename the plugin so its name is kebab-case or camelCase, or set \`hidden: true\` to exclude it from the auto-generated barrels.`, + ); + } const tier = manifest.stability ?? "ga"; if (tier !== "ga" && tier !== "beta") { @@ -117,14 +125,14 @@ function readPluginInfos(): PluginInfo[] { } infos.push({ - name: manifest.name, + binding, folder: entry.name, stability: tier, }); } // Deterministic order so re-runs produce reproducible diffs. - infos.sort((a, b) => a.name.localeCompare(b.name)); + infos.sort((a, b) => a.binding.localeCompare(b.binding)); return infos; } @@ -132,7 +140,9 @@ function renderBarrel(infos: PluginInfo[]): string { if (infos.length === 0) { return `${HEADER}\nexport {};\n`; } - const lines = infos.map((p) => `export { ${p.name} } from "./${p.folder}";`); + const lines = infos.map( + (p) => `export { ${p.binding} } from "./${p.folder}";`, + ); return `${HEADER}\n${lines.join("\n")}\n`; }