Skip to content

feat(appkit): add AI Functions beta plugin - #633

Draft
philip wants to merge 1 commit into
databricks:mainfrom
philip:feat/ai-functions-plugin
Draft

philip wants to merge 1 commit into
databricks:mainfrom
philip:feat/ai-functions-plugin

Conversation

@philip

@philip philip commented Oct 6, 2026

Copy link
Copy Markdown
Contributor

Summary

Adds a beta plugin for Databricks AI Functions
(ai-classify, ai-extract, ai-decide), which turn text into labels,
extracted fields, or answers to questions, returned as JSON.

An app defines each job once as a named task, then calls it from React
(useAiFunction), from server code (run), or as an agent tool
(<task>.invoke). The browser sends only the text, so users can't change
the prompt, and results are typed from the task definition.

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.

Design notes

  • Tasks: plain JSON, typed with as const satisfies AiFunctionTasks,
    validated at startup. POST /api/ai-functions/:task/invoke accepts
    exactly { content } or { state }. The client config exposes only each
    task's name and function.
  • Identity: per-task auth ("service-principal" by default, or
    "on-behalf-of-user"), matching AI Search and Files. User-task routes use
    _asUserScoped(req). A service-principal task runs as the app, but never
    widens an on-behalf-of-user agent run or runAgent({ caller }); there it
    runs as the user. A user task's 401 becomes IDENTITY_EXPIRED.
  • Scopes: the manifest declares scopes: ["ai-functions"], and
    ai-functions is added to the shared capability scopes.
    template/databricks.yml.tmpl is unchanged.
  • Errors: an AppKitError with statusCode and isRetryable. No
    plugin-specific error class is exported. The upstream 400 detail is
    returned to the caller but kept out of spans and logs.
  • Public surface: curated beta exports: the task and request/response
    types, plus extractValues, citedText, and scoreLevel. appkit-ui adds
    useAiFunction.

Testing

  • Unit: 229 tests across the plugin, connector, hook, helpers, and the
    route deprecation check. These include identity-scope tests,
    cancel-vs-timeout race tests, and hook-state tests. Full suite: 5,771
    passed.
  • CI, run locally: every job in ci.yml (lint, typecheck, knip,
    licenses, generated files, template sync, unit, Playwright integration,
    and template smoke for pnpm and npm).
  • Live, against a staging workspace:
    • every example in the plugin doc (44 checks);
    • agent tools in mixed and on-behalf-of-user agents;
    • two apps deployed from this branch, exercising service-principal and
      on-behalf-of-user task routes.

Notes for reviewers

  • ai-decide is in beta on Databricks and may need enabling from the
    workspace Previews page.
  • Unrelated, pre-existing: custom-plugins.md links to
    execution-context.md#telemetry-span-attributes, a heading the recent
    identity rewrite removed. The docs build warns about it.

This pull request and its description were written by Isaac.

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 <task>.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<typeof tasks.name>("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 <philip@neon.tech>
Co-authored-by: Isaac <no-reply@databricks.com>

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant