Skip to main content
Last updated on

ai_extract task for jobs

Beta

This feature is in Beta. To use it, a workspace admin must turn on Lakeflow AI Function Task from the Previews page. See Manage Databricks previews.

To perform information extraction in a job, use an ai_extract task. The task runs ai_extract on each row of a Unity Catalog table. The task extracts the fields that you define from text or parsed documents and writes the result for each row to a Unity Catalog table.

An ai_extract task is a type of AI Function task. In the Jobs API and Declarative Automation Bundles, it's an ai_function_task with the ai_extract function set. For requirements and how AI Function tasks process data, see AI Function tasks for jobs.

Add an ai_extract task​

To add an ai_extract task:

  1. Navigate to the Tasks tab in the Jobs UI.
  2. Click Plus icon. Add task. If the job has no tasks yet, click Plus icon. Add another task type.
  3. Under Document intelligence, select Information Extraction.
  4. Enter a name into the Task name field. It defaults to extraction.
  5. Under Input, select a table and the Column to extract from. To extract from documents, parse them first. See Use with other AI functions.
  6. Under Output, select the Location, and optionally enter a Table name. See Output.
  7. (Optional) Add an ID column. See ID column.
  8. Under Schema, define the fields to extract. See Define the fields to extract.
  9. (Optional) Configure the settings under Optional settings. See Optional settings.
  10. (Optional) To configure retries, run-duration or streaming-backlog thresholds, or notifications, see Advanced task settings.
  11. Click Save task.

The task runs the next time the job runs. To run the job on a schedule, see Run jobs on a schedule. To edit, clone, disable, or delete this task, see Configure and edit tasks in Lakeflow Jobs.

Define the fields to extract​

Under Schema, add each field to extract, with a name, a type, and an optional description. An object field holds nested fields, and an array field holds a list of values. To edit the fields as JSON, click JSON. The JSON uses the format that ai_extract expects, which also lists the supported field types.

Instead of adding the fields yourself, you can do either of the following:

  • Import from Agent: Copy the fields, instructions, and options of an Agent Bricks information extraction agent into the task. The task doesn't stay linked to the agent. To pick up later changes to the agent, import it again.
  • Define with Genie Code: Draft the fields with Genie Code, based on the selected input table and column. This option is available after you select the input and the output Location.

Optional settings​

Under Optional settings, you can configure the following:

  • Precision mode: Improve accuracy for complex fields and long documents.
  • Confidence scores: Include a confidence level for each extracted value.
  • Citations: Include a reference to the source text for each extracted value.
  • Instructions: Guidance applied to every field, such as date formats or content to ignore.
  • Version: The ai_extract version that the task runs, shown for reference. The version is pinned, so scheduled runs don't change when the function is updated.

Define the task in Declarative Automation Bundles​

Declarative Automation Bundles let you define the job in source files, so you can version, review, and deploy changes to it like code. The following example defines a job with an ai_extract task that extracts the counterparty, effective date, and renewal term from each contract in a table, with a citation for each value. The schema field holds the fields to extract as a JSON string. To extract from documents in a volume, see Use with other AI functions. To create the job with the Jobs API instead, send the same job settings to jobs/create.

YAML
# resources/extract_contract_fields.job.yml
resources:
  jobs:
    extract_contract_fields:
      name: extract_contract_fields
      tasks:
        - task_key: extraction
          ai_function_task:
            delta_table:
              table_name: main.legal.contracts
              input_column: contract_text
              merge_key_columns: [contract_id]
            output:
              catalog: main
              schema: legal
            ai_extract:
              version: '2.1'
              schema: |
                {
                  "counterparty": {"type": "string", "description": "Legal name of the other party"},
                  "effective_date": {"type": "string", "description": "Date in YYYY-MM-DD format"},
                  "renewal_term_months": {"type": "integer", "description": "Length of the renewal term in months"}
                }
              instructions: Ignore text in exhibits and appendices.
              enable_citations: true

Output​

The response column of the output table holds the full ai_extract result for each row as a VARIANT, including confidence scores and citations when you turn them on. The task doesn't add a column for each value. To get a value as a column, select it from response in a query. For the format, see ai_extract returns. For the other columns, see Output table.

Additional resources​

On this page