From 75607a0854dfe021a2d10d3245de71af4ccc2089 Mon Sep 17 00:00:00 2001 From: Philip Olson Date: Mon, 5 Oct 2026 12:16:07 -0700 Subject: [PATCH] feat(appkit): add AI Functions beta plugin Add a beta plugin for Databricks AI Functions, which turn text into labels, extracted fields, or answers to questions, returned as JSON. REST reference: https://docs.databricks.com/api/ai-functions/v1 It covers ai-classify, ai-extract, and ai-decide. The fourth function, ai-parse-document, is left for a follow-up. It takes a Unity Catalog volume path rather than the document itself, so it likely fits best alongside the files plugin, for example uploading to a volume, parsing, then extracting from the result. An app defines each job once as a named task, such as routing a support ticket or pulling the total from an invoice, then calls it from React, from server code, or as an agent tool. The browser sends only the text, so users can't change the prompt, and results are typed from the task. Details: - aiFunctions({ tasks }): each task is an AI Functions request minus its input, with per-task auth (service principal by default, or on-behalf-of-user). Tasks are plain JSON, typed with `as const satisfies AiFunctionTasks`, and validated at startup. - Each task is reachable at POST /api/ai-functions/:task/invoke (body is exactly { content } or { state }), from the server with run(task, input), and as the agent tool .invoke. The client config exposes only each task's name and function. - Server code can also call classify, extract, and decide directly, without defining a task. - appkit-ui: useAiFunction hook, typed by passing the task's type (useAiFunction("name")). It aborts the call in flight on re-invoke, unmount, and task change, and resets its state when the task changes. - Result helpers, exported for server and React: extractValues (schema-guided), citedText, and scoreLevel. - Errors are AppKitErrors with statusCode and isRetryable. 503 and 429 are retried; the Databricks SDK retries 429 itself, so a sustained rate limit surfaces as a 504 timeout. A caller's own cancel reports 499, and a user task's expired token reports IDENTITY_EXPIRED. The upstream 400 detail is appended outside execute(), so it stays out of spans and logs. - Identity follows AppKit's execution scopes: user task routes use the plugin's internal user scope, and service-principal tasks run as the app, except inside an on-behalf-of-user agent run or runAgent with a caller, where they run as the user and never widen it. - The manifest declares the ai-functions user API scope, and the app template lists the plugin. Docs cover the plugin page and execution identity, plus a dev-playground page and the CI generated-docs check. Signed-off-by: Philip Olson Co-authored-by: Isaac --- .claude/commands/create-core-plugin.md | 1 + .github/workflows/ci.yml | 3 +- apps/dev-playground/client/src/lib/nav.ts | 8 + .../client/src/routeTree.gen.ts | 21 + .../client/src/routes/ai-functions.route.tsx | 773 +++++++++++ apps/dev-playground/client/tsconfig.app.json | 6 +- apps/dev-playground/server/index.ts | 73 +- .../shared/ai-functions-examples.ts | 288 ++++ apps/dev-playground/shared/ai-tasks.ts | 42 + docs/docs/api/appkit/Function.citedText.md | 24 + .../docs/api/appkit/Function.extractValues.md | 26 + docs/docs/api/appkit/Function.scoreLevel.md | 38 + .../api/appkit/Interface.BasePluginConfig.md | 1 + .../api/appkit/Interface.ClassifyRequest.md | 31 + .../api/appkit/Interface.ClassifyResponse.md | 23 + .../api/appkit/Interface.DecideRequest.md | 31 + .../api/appkit/Interface.DecideResponse.md | 31 + .../api/appkit/Interface.ExtractRequest.md | 31 + .../api/appkit/Interface.ExtractResponse.md | 23 + .../appkit/Interface.IAiFunctionsConfig.md | 111 ++ docs/docs/api/appkit/Interface.IndexConfig.md | 2 +- .../api/appkit/Interface.PluginManifest.md | 1 + .../api/appkit/TypeAlias.AiFunctionTask.md | 7 + .../appkit/TypeAlias.AiFunctionTaskInput.md | 19 + .../appkit/TypeAlias.AiFunctionTaskResult.md | 22 + .../api/appkit/TypeAlias.AiFunctionTasks.md | 7 + .../docs/api/appkit/TypeAlias.ExtractField.md | 39 + docs/docs/api/appkit/Variable.aiFunctions.md | 5 + docs/docs/api/appkit/index.md | 16 + docs/docs/api/appkit/typedoc-sidebar.ts | 80 ++ docs/docs/plugins/agents.md | 4 +- docs/docs/plugins/ai-functions.md | 274 ++++ docs/docs/plugins/execution-context.md | 8 +- .../schemas/plugin-manifest.schema.json | 1 + .../schemas/template-plugins.schema.json | 1 + packages/appkit-ui/src/react/beta.ts | 15 + .../hooks/__tests__/use-ai-function.test.ts | 244 ++++ .../src/react/hooks/use-ai-function.ts | 111 ++ packages/appkit/src/beta.ts | 17 + .../src/connectors/ai-functions/client.ts | 259 ++++ .../src/connectors/ai-functions/index.ts | 1 + .../ai-functions/tests/client.test.ts | 408 ++++++ packages/appkit/src/connectors/index.ts | 1 + .../src/plugins/ai-functions/ai-functions.ts | 621 +++++++++ .../src/plugins/ai-functions/defaults.ts | 9 + .../appkit/src/plugins/ai-functions/errors.ts | 29 + .../appkit/src/plugins/ai-functions/index.ts | 3 + .../src/plugins/ai-functions/manifest.json | 85 ++ .../src/plugins/ai-functions/schemas.ts | 314 +++++ .../appkit/src/plugins/ai-functions/tasks.ts | 144 ++ .../ai-functions/tests/ai-functions.test.ts | 1184 +++++++++++++++++ .../ai-functions/tests/schemas.test.ts | 663 +++++++++ .../plugins/ai-functions/tests/tasks.test.ts | 193 +++ .../plugins/ai-functions/tests/tools.test.ts | 329 +++++ .../plugins/ai-functions/tests/types.test.ts | 386 ++++++ .../appkit/src/plugins/ai-functions/types.ts | 36 + .../src/plugins/beta-exports.generated.ts | 1 + .../tests/core-obo-no-deprecation.test.ts | 33 + packages/shared/src/ai-functions-helpers.ts | 113 ++ packages/shared/src/ai-functions.ts | 322 +++++ packages/shared/src/index.ts | 2 + packages/shared/src/schemas/manifest.ts | 1 + .../src/tests/ai-functions-helpers.test.ts | 131 ++ template/appkit.plugins.json | 25 + 64 files changed, 7742 insertions(+), 9 deletions(-) create mode 100644 apps/dev-playground/client/src/routes/ai-functions.route.tsx create mode 100644 apps/dev-playground/shared/ai-functions-examples.ts create mode 100644 apps/dev-playground/shared/ai-tasks.ts create mode 100644 docs/docs/api/appkit/Function.citedText.md create mode 100644 docs/docs/api/appkit/Function.extractValues.md create mode 100644 docs/docs/api/appkit/Function.scoreLevel.md create mode 100644 docs/docs/api/appkit/Interface.ClassifyRequest.md create mode 100644 docs/docs/api/appkit/Interface.ClassifyResponse.md create mode 100644 docs/docs/api/appkit/Interface.DecideRequest.md create mode 100644 docs/docs/api/appkit/Interface.DecideResponse.md create mode 100644 docs/docs/api/appkit/Interface.ExtractRequest.md create mode 100644 docs/docs/api/appkit/Interface.ExtractResponse.md create mode 100644 docs/docs/api/appkit/Interface.IAiFunctionsConfig.md create mode 100644 docs/docs/api/appkit/TypeAlias.AiFunctionTask.md create mode 100644 docs/docs/api/appkit/TypeAlias.AiFunctionTaskInput.md create mode 100644 docs/docs/api/appkit/TypeAlias.AiFunctionTaskResult.md create mode 100644 docs/docs/api/appkit/TypeAlias.AiFunctionTasks.md create mode 100644 docs/docs/api/appkit/TypeAlias.ExtractField.md create mode 100644 docs/docs/api/appkit/Variable.aiFunctions.md create mode 100644 docs/docs/plugins/ai-functions.md create mode 100644 packages/appkit-ui/src/react/hooks/__tests__/use-ai-function.test.ts create mode 100644 packages/appkit-ui/src/react/hooks/use-ai-function.ts create mode 100644 packages/appkit/src/connectors/ai-functions/client.ts create mode 100644 packages/appkit/src/connectors/ai-functions/index.ts create mode 100644 packages/appkit/src/connectors/ai-functions/tests/client.test.ts create mode 100644 packages/appkit/src/plugins/ai-functions/ai-functions.ts create mode 100644 packages/appkit/src/plugins/ai-functions/defaults.ts create mode 100644 packages/appkit/src/plugins/ai-functions/errors.ts create mode 100644 packages/appkit/src/plugins/ai-functions/index.ts create mode 100644 packages/appkit/src/plugins/ai-functions/manifest.json create mode 100644 packages/appkit/src/plugins/ai-functions/schemas.ts create mode 100644 packages/appkit/src/plugins/ai-functions/tasks.ts create mode 100644 packages/appkit/src/plugins/ai-functions/tests/ai-functions.test.ts create mode 100644 packages/appkit/src/plugins/ai-functions/tests/schemas.test.ts create mode 100644 packages/appkit/src/plugins/ai-functions/tests/tasks.test.ts create mode 100644 packages/appkit/src/plugins/ai-functions/tests/tools.test.ts create mode 100644 packages/appkit/src/plugins/ai-functions/tests/types.test.ts create mode 100644 packages/appkit/src/plugins/ai-functions/types.ts create mode 100644 packages/shared/src/ai-functions-helpers.ts create mode 100644 packages/shared/src/ai-functions.ts create mode 100644 packages/shared/src/tests/ai-functions-helpers.test.ts diff --git a/.claude/commands/create-core-plugin.md b/.claude/commands/create-core-plugin.md index 90a7a5624..6c7c2877e 100644 --- a/.claude/commands/create-core-plugin.md +++ b/.claude/commands/create-core-plugin.md @@ -385,6 +385,7 @@ If the plugin performs operations on behalf of the logged-in user via `this.asUs | Scope | When to use | |-------|-------------| | `sql` | Plugin executes SQL queries on behalf of the user (e.g. analytics `.obo.sql` files) | +| `ai-functions` | Plugin calls AI Functions on behalf of the user | | `dashboards.genie` | Plugin interacts with Genie AI/BI spaces on behalf of the user | | `files.files` | Plugin reads/writes files on behalf of the user | diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml index 2eb25ba49..e2a1a69eb 100644 --- a/.github/workflows/ci.yml +++ b/.github/workflows/ci.yml @@ -68,7 +68,8 @@ jobs: docs/docs/plugins/lakebase.md \ docs/docs/plugins/model-serving.md \ docs/docs/plugins/server.md \ - docs/docs/plugins/ai-search.md; then + docs/docs/plugins/ai-search.md \ + docs/docs/plugins/ai-functions.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 5e6eda2a9..e4951193b 100644 --- a/apps/dev-playground/client/src/lib/nav.ts +++ b/apps/dev-playground/client/src/lib/nav.ts @@ -16,6 +16,7 @@ import { ServerIcon, ShieldIcon, SigmaIcon, + TagsIcon, Wand2Icon, ZapIcon, } from "lucide-react"; @@ -144,6 +145,13 @@ export const NAV_GROUPS: ReadonlyArray = [ "Semantic search backed by Databricks vector indexes, wired into AppKit's retrieval API.", icon: SearchIcon, }, + { + to: "/ai-functions", + label: "AI Functions", + description: + "Classify, extract, and decide over text, as the signed-in user or the service principal.", + icon: TagsIcon, + }, { to: "/serving", label: "Serving", diff --git a/apps/dev-playground/client/src/routeTree.gen.ts b/apps/dev-playground/client/src/routeTree.gen.ts index 8658c1952..02380e8e8 100644 --- a/apps/dev-playground/client/src/routeTree.gen.ts +++ b/apps/dev-playground/client/src/routeTree.gen.ts @@ -29,6 +29,7 @@ import { Route as ChartInferenceRouteRouteImport } from './routes/chart-inferenc 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 AiFunctionsRouteRouteImport } from './routes/ai-functions.route' import { Route as AgentRouteRouteImport } from './routes/agent.route' import { Route as IndexRouteImport } from './routes/index' @@ -132,6 +133,11 @@ const AiSearchRouteRoute = AiSearchRouteRouteImport.update({ path: '/ai-search', getParentRoute: () => rootRouteImport, } as any) +const AiFunctionsRouteRoute = AiFunctionsRouteRouteImport.update({ + id: '/ai-functions', + path: '/ai-functions', + getParentRoute: () => rootRouteImport, +} as any) const AgentRouteRoute = AgentRouteRouteImport.update({ id: '/agent', path: '/agent', @@ -146,6 +152,7 @@ const IndexRoute = IndexRouteImport.update({ export interface FileRoutesByFullPath { '/': typeof IndexRoute '/agent': typeof AgentRouteRoute + '/ai-functions': typeof AiFunctionsRouteRoute '/ai-search': typeof AiSearchRouteRoute '/analytics': typeof AnalyticsRouteRoute '/arrow-analytics': typeof ArrowAnalyticsRouteRoute @@ -170,6 +177,7 @@ export interface FileRoutesByFullPath { export interface FileRoutesByTo { '/': typeof IndexRoute '/agent': typeof AgentRouteRoute + '/ai-functions': typeof AiFunctionsRouteRoute '/ai-search': typeof AiSearchRouteRoute '/analytics': typeof AnalyticsRouteRoute '/arrow-analytics': typeof ArrowAnalyticsRouteRoute @@ -195,6 +203,7 @@ export interface FileRoutesById { __root__: typeof rootRouteImport '/': typeof IndexRoute '/agent': typeof AgentRouteRoute + '/ai-functions': typeof AiFunctionsRouteRoute '/ai-search': typeof AiSearchRouteRoute '/analytics': typeof AnalyticsRouteRoute '/arrow-analytics': typeof ArrowAnalyticsRouteRoute @@ -221,6 +230,7 @@ export interface FileRouteTypes { fullPaths: | '/' | '/agent' + | '/ai-functions' | '/ai-search' | '/analytics' | '/arrow-analytics' @@ -245,6 +255,7 @@ export interface FileRouteTypes { to: | '/' | '/agent' + | '/ai-functions' | '/ai-search' | '/analytics' | '/arrow-analytics' @@ -269,6 +280,7 @@ export interface FileRouteTypes { | '__root__' | '/' | '/agent' + | '/ai-functions' | '/ai-search' | '/analytics' | '/arrow-analytics' @@ -294,6 +306,7 @@ export interface FileRouteTypes { export interface RootRouteChildren { IndexRoute: typeof IndexRoute AgentRouteRoute: typeof AgentRouteRoute + AiFunctionsRouteRoute: typeof AiFunctionsRouteRoute AiSearchRouteRoute: typeof AiSearchRouteRoute AnalyticsRouteRoute: typeof AnalyticsRouteRoute ArrowAnalyticsRouteRoute: typeof ArrowAnalyticsRouteRoute @@ -458,6 +471,13 @@ declare module '@tanstack/react-router' { preLoaderRoute: typeof AiSearchRouteRouteImport parentRoute: typeof rootRouteImport } + '/ai-functions': { + id: '/ai-functions' + path: '/ai-functions' + fullPath: '/ai-functions' + preLoaderRoute: typeof AiFunctionsRouteRouteImport + parentRoute: typeof rootRouteImport + } '/agent': { id: '/agent' path: '/agent' @@ -478,6 +498,7 @@ declare module '@tanstack/react-router' { const rootRouteChildren: RootRouteChildren = { IndexRoute: IndexRoute, AgentRouteRoute: AgentRouteRoute, + AiFunctionsRouteRoute: AiFunctionsRouteRoute, AiSearchRouteRoute: AiSearchRouteRoute, AnalyticsRouteRoute: AnalyticsRouteRoute, ArrowAnalyticsRouteRoute: ArrowAnalyticsRouteRoute, diff --git a/apps/dev-playground/client/src/routes/ai-functions.route.tsx b/apps/dev-playground/client/src/routes/ai-functions.route.tsx new file mode 100644 index 000000000..bc86c4cdd --- /dev/null +++ b/apps/dev-playground/client/src/routes/ai-functions.route.tsx @@ -0,0 +1,773 @@ +import { + Badge, + Button, + Card, + CardContent, + Checkbox, + Progress, + Select, + SelectContent, + SelectItem, + SelectTrigger, + SelectValue, + Tabs, + TabsList, + TabsTrigger, + Textarea, + ToggleGroup, + ToggleGroupItem, +} from "@databricks/appkit-ui/react"; +import { useAiFunction } from "@databricks/appkit-ui/react/beta"; +import { createFileRoute } from "@tanstack/react-router"; +import { Play } from "lucide-react"; +import { useEffect, useRef, useState } from "react"; + +import { Header } from "@/components/layout/header"; + +import { + EXAMPLES, + type Example, + FUNCTION_INFO, + type FunctionName, +} from "../../../shared/ai-functions-examples"; + +export const Route = createFileRoute("/ai-functions")({ + component: AiFunctionsRoute, +}); + +type RunAs = "user" | "sp"; +type RunMode = RunAs | "both"; + +const RUN_MODE_HINT: Record = { + user: "Runs the example's named task as you (POST /api/ai-functions/:task/invoke). Edited requests use a demo-only route.", + sp: "Sends the request to a demo-only route that calls appkit.aiFunctions as the app service principal.", + both: "Runs both ways so you can compare the results.", +}; + +interface Identity { + user: string | null; + userTokenPresent: boolean; +} + +interface RunResult { + fn: FunctionName; + runAs: RunAs; + url: string; + request: Record; + status: number; + ms: number; + body: Record | null; +} + +type Target = + | { kind: "task"; runAs: RunAs; task: string; body: Record } + | { + kind: "demo"; + runAs: RunAs; + mode: "as-user" | "as-sp"; + body: Record; + }; + +function targetUrl(fn: FunctionName, target: Target): string { + return target.kind === "task" + ? `/api/ai-functions/${encodeURIComponent(target.task)}/invoke` + : `/api/ai-functions-demo/${fn}/${target.mode}`; +} + +/** Edited requests and the service-principal mode use the demo routes. */ +function targetFor( + fn: FunctionName, + runAs: RunAs, + example: Example, + edited: boolean, + text: string, + request: Record, +): Target { + if (runAs === "sp") + return { kind: "demo", runAs, mode: "as-sp", body: request }; + if (edited) return { kind: "demo", runAs, mode: "as-user", body: request }; + return { + kind: "task", + runAs, + task: example.id, + body: { [FUNCTION_INFO[fn].textField]: text }, + }; +} + +async function callAiFunction( + fn: FunctionName, + target: Target, + signal: AbortSignal, +): Promise { + const url = targetUrl(fn, target); + const started = performance.now(); + const result = { fn, runAs: target.runAs, url, request: target.body }; + try { + const res = await fetch(url, { + method: "POST", + headers: { "Content-Type": "application/json" }, + body: JSON.stringify(target.body), + signal, + }); + const body = await res.json().catch(() => null); + return { + ...result, + status: res.status, + ms: Math.round(performance.now() - started), + body, + }; + } catch (error) { + if (signal.aborted) throw error; + // Network failure: report it like an HTTP error so the result pane shows it. + return { + ...result, + status: 0, + ms: Math.round(performance.now() - started), + body: { + error: error instanceof Error ? error.message : "Network error", + }, + }; + } +} + +function toJson(value: unknown): string { + return JSON.stringify(value, null, 2); +} + +function AiFunctionsRoute() { + const [fn, setFn] = useState("classify"); + const [example, setExample] = useState(EXAMPLES.classify[0]); + const [text, setText] = useState(example.text); + const [optionsJson, setOptionsJson] = useState(toJson(example.options)); + const [identity, setIdentity] = useState(null); + const [runAs, setRunAs] = useState("user"); + const [results, setResults] = useState([]); + // The prompt-injection check is a plain named task, so it uses the hook. + const injectionCheck = useAiFunction("promptInjectionCheck"); + const [showInjection, setShowInjection] = useState(false); + const runController = useRef(null); + + // Cancel an in-flight run when the page unmounts. + useEffect(() => () => runController.current?.abort(), []); + const [checkInjection, setCheckInjection] = useState(false); + const [loading, setLoading] = useState(false); + const [parseError, setParseError] = useState(null); + + useEffect(() => { + fetch("/api/ai-functions-demo/identity") + .then((res) => res.json()) + .then(setIdentity) + .catch(() => setIdentity(null)); + }, []); + + function selectExample(next: Example) { + setExample(next); + setText(next.text); + setOptionsJson(toJson(next.options)); + setResults([]); + setShowInjection(false); + setParseError(null); + } + + function selectFunction(next: FunctionName) { + setFn(next); + selectExample(EXAMPLES[next][0]); + } + + async function run() { + const targets: RunAs[] = runAs === "both" ? ["user", "sp"] : [runAs]; + let options: Record; + try { + options = JSON.parse(optionsJson); + } catch { + setParseError("Request options are not valid JSON."); + return; + } + setParseError(null); + setLoading(true); + // A new run replaces any earlier one, so stale responses never land. + runController.current?.abort(); + const controller = new AbortController(); + runController.current = controller; + const { signal } = controller; + const request = { [FUNCTION_INFO[fn].textField]: text, ...options }; + try { + const edited = optionsJson !== toJson(example.options); + const [runs] = await Promise.all([ + Promise.all( + targets.map((t) => + callAiFunction( + fn, + targetFor(fn, t, example, edited, text, request), + signal, + ), + ), + ), + checkInjection ? injectionCheck.invoke({ content: text }) : null, + ]); + setShowInjection(checkInjection); + setResults(runs); + } catch (error) { + if (!signal.aborted) throw error; + } finally { + if (runController.current === controller) setLoading(false); + } + } + + const info = FUNCTION_INFO[fn]; + + return ( +
+
+
+ +
+ + +
+ + selectFunction(value as FunctionName) + } + > + + Classify + Extract + Decide + + +

{info.summary}

+
+ +
+ Example + + + {example.shows} + +
+ +
+ +