Repository navigation
Conversation
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
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
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 changethe prompt, and results are typed from the task definition.
ai-parse-documentis left for a follow-up. It takes a Unity Catalogvolume path rather than the document itself, so it likely fits best
alongside the files plugin.
Design notes
as const satisfies AiFunctionTasks,validated at startup.
POST /api/ai-functions/:task/invokeacceptsexactly
{ content }or{ state }. The client config exposes only eachtask's name and function.
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 neverwidens an on-behalf-of-user agent run or
runAgent({ caller }); there itruns as the user. A user task's 401 becomes
IDENTITY_EXPIRED.scopes: ["ai-functions"], andai-functionsis added to the shared capability scopes.template/databricks.yml.tmplis unchanged.AppKitErrorwithstatusCodeandisRetryable. Noplugin-specific error class is exported. The upstream 400 detail is
returned to the caller but kept out of spans and logs.
types, plus
extractValues,citedText, andscoreLevel. appkit-ui addsuseAiFunction.Testing
route deprecation check. These include identity-scope tests,
cancel-vs-timeout race tests, and hook-state tests. Full suite: 5,771
passed.
ci.yml(lint, typecheck, knip,licenses, generated files, template sync, unit, Playwright integration,
and template smoke for pnpm and npm).
on-behalf-of-user task routes.
Notes for reviewers
ai-decideis in beta on Databricks and may need enabling from theworkspace Previews page.
custom-plugins.mdlinks toexecution-context.md#telemetry-span-attributes, a heading the recentidentity rewrite removed. The docs build warns about it.
This pull request and its description were written by Isaac.