ai_enrich function

Applies to: check marked yes Databricks SQL check marked yes Databricks Runtime

Important

This feature is in Beta. To use it, a workspace admin must turn on Ai Enrich from the Previews page. See Manage Azure Databricks previews.

The ai_enrich() function generates new columns for a row from a schema you define. Given input content and a target schema, the function calls an AI model to fill in each field. It can optionally ground the generated values in one or more knowledge sources like an AI Search index or a search tool on an external Model Context Protocol (MCP) server, so the values reflect your own data or up-to-date information rather than the model's training data alone.

Use ai_enrich to add derived attributes to a table at scale. You can tag and categorize records, fill in missing metadata, or attach researched context to each row from a single SQL function call. By default, each generated field is returned with a short rationale that explains how the value was derived.

Requirements

  • Databricks Runtime 18.2 or above.
  • If you are using Serverless compute, the serverless environment version must be set to 3 or above, as this enables features like VARIANT.
  • To ground enrichment in an AI Search index, you need one or more AI Search indexes to use as knowledge sources.
  • To ground enrichment in an external MCP server, you need an MCP registered in Unity Gateway, EXECUTE on it, and USE CATALOG and USE SCHEMA on its parent catalog and schema. See MCP configuration.
  • The ai_enrich function is available using Azure Databricks notebooks, SQL editor, jobs, or Lakeflow pipelines.

Data security

Your document data is processed within the Databricks security perimeter. Databricks does not store the parameters that are passed into the AI function calls, but does retain metadata run details, such as the Databricks Runtime version used.

If you configure an mcp knowledge source, the search query that ai_enrich builds from each row is sent to the external provider, so treat row content that reaches it as data that leaves Azure Databricks. See MCP configuration.

Syntax

ai_enrich(content, schema [, knowledge_sources] [, options])

Arguments

  • content: A STRING or VARIANT expression. The row to enrich. VARIANT input, such as the output of another AI function like ai_parse_document, is serialized to a JSON string internally.

  • schema: A STRING literal that defines the columns to generate. It uses the same grammar as ai_extract. The schema can be:

    • Simple schema: A JSON array of field names, which are generated as strings.

      ["industry", "headquarters_country", "year_founded"]

    • Advanced schema: A JSON object with type information, descriptions, and nested structures.

      • Supports string, integer, number, boolean, and enum types. Performs type validation. Maximum of 500 enum values.
      • Supports nested objects using "type": "object" with "properties".
      • Supports arrays of primitives or objects using "type": "array" with "items".
      • Optional "description" field for each property to guide the generated value.
      {
  "hq_address": {
    "type": "object",
    "description": "Registered headquarters address",
    "properties": {
      "city": { "type": "string" },
      "country": { "type": "string" }
    }
  },
  "founding_team": { "type": "array", "description": "Full names of the founders", "items": { "type": "string" } },
  "founding_year": { "type": "integer", "description": "Year the company was founded" }
}

  • knowledge_sources: An optional VARIANT or STRING expression containing a JSON array of knowledge source configurations used to ground the generated values. See Knowledge source configuration.

  • options: An optional MAP<STRING, STRING>. Supported keys:

    • 'version': The function version to use.
    • 'instructions': A STRING of up to 20,000 characters. Natural-language guidance that describes the enrichment task. Optional; the schema field names alone can drive the enrichment. For example, 'Infer attributes for each company from its public profile.'
    • 'enableRationale': 'true' (default) or 'false'. When 'true', each generated field is returned as a {rationale, value} object, where rationale explains how the value was derived. Set to 'false' to return {value} only.

Knowledge source configuration

The knowledge_sources argument is a JSON array. Each element is a {type, description, config} envelope. The type field identifies how ai_enrich retrieves grounding context, and the config field contains the source-specific configuration.

Key Required Description
type Yes The knowledge source type. One of vector_search (an AI Search index) or mcp (an external MCP server).
description No A natural-language description of the source.
config Yes An object containing the source-specific configuration. See AI Search index configuration for vector_search and MCP configuration for mcp.

AI Search index configuration

For an AI Search index with type set to vector_search, config accepts the following keys:

Key Required Description
index_name Yes The Unity Catalog three-level name of the AI Search index, for example catalog.schema.my_index.
text_col Yes The column in the index that contains the document text.
doc_uri_col Yes The column in the index that contains the document URI.
filter_columns No A comma-separated string or JSON array of columns available for metadata filtering. When omitted, the list is derived from the index schema, excluding reserved, text, and document URI columns.

You can configure more than one vector_search source in a single call. For an example, see Ground enrichment in an AI Search index.

MCP configuration

An mcp source grounds each row with a search tool on an external MCP server, such as You.com or Exa. Use it to ground rows in a search provider you already have an account with. ai_enrich calls the tool through Unity Gateway as the identity that runs the query. Each mcp source adds up to five result URLs to metadata.sources.

Before you use an mcp source:

  1. Register the server as an MCP. See External MCP servers. For a server listed in Databricks Marketplace, such as You.com, installing the listing creates the MCP for you. See Get access to external MCP servers.
  2. Use a connection with a shared credential, such as the provider's API key as a bearer token. ai_enrich can't prompt for sign-in, so if the connection requires each user to sign in, the mcp source fails for users who haven't signed in.
  3. Grant EXECUTE on the MCP, and USE CATALOG and USE SCHEMA on its parent catalog and schema, to the identity that runs the query: each user, or the run-as identity of a job or pipeline. See Share an MCP.

For an mcp source, config accepts the following keys:

Key Required Description
mcp_service Yes The full name of the MCP, <catalog>.<schema>.<service>.
tool Yes The search tool to call, for example you-search for You.com or web_search_exa for Exa. The tool must be one that the MCP exposes.
query_argument No The tool argument that receives the search query. Defaults to query.
arguments No A JSON object of extra tool arguments to send on every call. For example, Exa's web_search_exa tool lists objective as a required argument.
max_query_chars No The maximum length of the search query, in characters. Defaults to 400. Longer queries are truncated.

For each row with non-empty content, ai_enrich builds a short search query from the row content followed by the requested fields. Each top-level field contributes its description, or its name if it has no description. The query is cut to max_query_chars, so for a long row the field text can be cut off. ai_enrich calls the tool of every mcp source one time with that query. Each call times out after 60 seconds. A call that times out or fails with a 429, 500, 502, 503, or 504 error is retried up to three times. Because retries can repeat a call, use a read-only search tool. Calls to the provider count against your provider account's pricing and rate limits.

If a call fails and no other knowledge source returns results, error_message is set and response is null for that row. For example, a call fails if the provider rejects the API key or the identity that runs the query lacks EXECUTE on the MCP. If another knowledge source returns results, the row is grounded in those results only and no error is reported. If the tool returns no results and no other knowledge source does either, the row is enriched from its content alone, without grounding.

If mcp knowledge sources aren't enabled for the workspace, every row with non-empty content returns null for response and the error_message mcp knowledge source is not available: it is not enabled for this workspace., even when the call also configures other knowledge sources.

You can configure more than one mcp source in a single call. For an example, see Ground enrichment in an external MCP search provider.

Returns

A VARIANT with the following schema:

{
  "response": { ... },     // Generated columns matching the provided schema. Each leaf is returned as an object (see below).
  "error_message": null,   // null on success, or an error message on failure
  "metadata": { ... }      // Metadata about the response, including grounding sources.
}

The response field contains the generated columns:

  • Field names and types match the schema definition. Nested objects and arrays keep their original shape.
  • By default (enableRationale is 'true'), each leaf is a {rationale, value} object, where rationale is a short explanation of how the value was derived and value is the generated value, typed according to the schema. When enableRationale is 'false', each leaf is a {value} object.
  • A field's value is null when it cannot be generated.

The metadata field contains metadata about the response. When the row is grounded by a knowledge source, metadata.sources is an array of the source document identifiers that grounded the row. Grounding is row-level, so sources applies to the whole row rather than to individual fields.

If content is NULL, the result is NULL.

Examples

Ground enrichment in an AI Search index

The following example enriches each support ticket with fields grounded in an AI Search index of product documentation, so the generated values are drawn from your own content:

SELECT
  ticket_id,
  ai_enrich(
    customer_description,
    '{
      "affected_product": {"type": "string"},
      "suggested_resolution": {"type": "string"},
      "documentation_url": {"type": "string"}
    }',
    PARSE_JSON('[{
      "type": "vector_search",
      "description": "Product documentation and troubleshooting guides",
      "config": {
        "index_name": "support.docs.product_documentation",
        "text_col": "content",
        "doc_uri_col": "doc_url"
      }
    }]')
  ) AS result
FROM support.tickets.open_tickets;

Ground enrichment in an external MCP search provider

The following example enriches each company with fields grounded in You.com search results. It calls the you-search tool on an MCP named main.default.youcom:

SELECT ai_enrich(
  company_name,
  '["industry", "headquarters_country", "year_founded"]',
  PARSE_JSON('[{
    "type": "mcp",
    "description": "You.com web search",
    "config": {
      "mcp_service": "main.default.youcom",
      "tool": "you-search"
    }
  }]')
) AS result
FROM sales.accounts.companies;

The following example uses Exa instead. Exa's web_search_exa tool lists objective as a required argument, so the example passes it in arguments:

SELECT ai_enrich(
  company_name,
  '["industry", "headquarters_country", "year_founded"]',
  PARSE_JSON('[{
    "type": "mcp",
    "description": "Exa web search",
    "config": {
      "mcp_service": "main.default.exa",
      "tool": "web_search_exa",
      "arguments": {"objective": "Find public company profile facts"}
    }
  }]')
) AS result
FROM sales.accounts.companies;

Limitations

  • ai_enrich does not support the web_search knowledge source. To ground rows in a search provider, use an mcp source. See MCP configuration.
  • An mcp knowledge source can't prompt users to sign in to the provider. If the connection requires each user to sign in, calls fail for users who haven't signed in.
  • An mcp knowledge source doesn't support service policies that return ASK. ai_enrich can't show the approval request, so calls that require approval fail.
  • The instructions option is limited to 20,000 characters.
  • Last updated on