ai_enrich function

Applies to: Databricks SQL Databricks Runtime

Beta

This feature is in Beta. To use it, a workspace admin must turn on Ai Enrich from the Previews page. See Manage 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, a live web search, 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 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 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.

    JSON

    ["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.
    JSON

    {
      "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), web_search, 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, Web search configuration for web_search, and MCP configuration for mcp.

Key	Required	Description
type	Yes	The knowledge source type. One of vector_search (an AI Search index), web_search, 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, Web search configuration for web_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.

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.

Web search configuration

For a web search with type set to web_search, config accepts the following optional keys. Web search runs through web search on Databricks. See Limitations for availability.

Key	Required	Description
allowed_domains	No	A JSON array of domains to restrict the search to. When set, only results from these domains are used.
blocked_domains	No	A JSON array of domains to exclude from the search.

Key	Required	Description
allowed_domains	No	A JSON array of domains to restrict the search to. When set, only results from these domains are used.
blocked_domains	No	A JSON array of domains to exclude from the search.

You can configure at most one web_search source in each call. For an example, see Ground enrichment in web search.

The following example configures an AI Search index and a web search as knowledge sources:

JSON

[
  {
    "type": "vector_search",
    "description": "Internal product catalog",
    "config": {
      "index_name": "prod_catalog.docs.product_catalog",
      "text_col": "description",
      "doc_uri_col": "product_url"
    }
  },
  {
    "type": "web_search",
    "config": {
      "allowed_domains": ["wikipedia.org"]
    }
  }
]

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.

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:

JSON

{
  "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 web search

The following example enriches each company using a simple array schema (each field is generated as a string), grounded in a live web search over the open web:

SQL

SELECT ai_enrich(
  company_name,
  '["industry", "headquarters_country", "year_founded"]',
  PARSE_JSON('[{
    "type": "web_search",
    "config": {}
  }]')
) AS result
FROM sales.accounts.companies;

Restrict web search to specific domains

The following example is the same as above, but restricts the web search to specific sites with allowed_domains, so only results from those domains ground the values:

SQL

SELECT ai_enrich(
  company_name,
  '["industry", "headquarters_country", "year_founded"]',
  PARSE_JSON('[{
    "type": "web_search",
    "config": {"allowed_domains": ["wikipedia.org", "sec.gov"]}
  }]')
) AS result
FROM sales.accounts.companies;

Enforce a typed schema with instructions

The following example uses the same open web search but with a full typed schema, adds instructions to steer the task, and disables rationale so each field returns a plain value:

SQL

SELECT ai_enrich(
  company_name,
  '{
    "industry": {"type": "string", "description": "Primary industry"},
    "headquarters_country": {"type": "string"},
    "year_founded": {"type": "integer"}
  }',
  PARSE_JSON('[{
    "type": "web_search",
    "config": {}
  }]'),
  options => map(
    'instructions', 'Fill each field from current, authoritative public sources. Use null when a value cannot be verified.',
    'enableRationale', 'false'
  )
) AS result
FROM sales.accounts.companies;

Populate a nested schema

The following example populates a nested schema for each company — a structured address object, an array of founder names, a typed year, and a nested funding round — grounded in a live web search so the values reflect current public information rather than the model's training data alone:

SQL

SELECT ai_enrich(
  company_name,
  '{
    "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"},
    "latest_funding_round": {
      "type": "object",
      "properties": {
        "stage": {"type": "string", "description": "Funding stage, for example Seed or Series A"},
        "amount_usd": {"type": "number", "description": "Amount raised in USD"}
      }
    }
  }',
  PARSE_JSON('[{
    "type": "web_search",
    "config": {}
  }]'),
  options => map('instructions', 'Ground each field in current, authoritative public sources. Use null when a value cannot be verified.')
) AS result
FROM sales.accounts.companies;

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:

SQL

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:

SQL

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:

SQL

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

• Grounding with a web_search knowledge source is only available in some regions and workspaces. See web search on Databricks.
• 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.

Related functions

• ai_extract function
• ai_search function
• ai_parse_document function
• ai_prep_search function
• ai_query function