Skip to main content
Last updated on

Tutorial: Enrich entities extracted from documents

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.

Open in Databricks

Enterprise document collections often mention companies without including the clean, current attributes needed for analytics and downstream applications. This tutorial builds an end-to-end AI Functions pipeline that turns unstructured contracts into grounded company records. The pipeline uses ai_parse_document and ai_extract to identify the company named in each contract. It then uses ai_enrich to resolve the company and add current information from the web.

The completed pipeline performs these steps:

Text
PDF contracts -> ai_parse_document -> ai_extract -> ai_enrich
                                          |             |
                                   company name    grounded record

Requirements​

  • Databricks Runtime 18.2 or above.
  • If you use Serverless compute, serverless environment version 3 or above.
  • The ai_enrich Beta enabled by a workspace admin from the Previews page.
  • A workspace and region that support web search on Databricks.
  • Access to the samples catalog.

This tutorial uses SEC-filed agreements in /Volumes/samples/sec/contracts/. The samples.sec.contracts volume is available in all workspaces by default. These contracts stand in for the filings, emails, call transcripts, and websites that appear in real document-to-entity pipelines. To process your own PDFs, change SOURCE_PATH to a Unity Catalog volume that contains your files.

Step 1: Ingest sample contracts​

Create a Python notebook and attach it to supported compute. Run the following code to read up to 10 consulting agreements from the sample volume:

Python
from pyspark.sql import functions as F
import json
import uuid

SOURCE_PATH = "/Volumes/samples/sec/contracts/"
TMP_SUFFIX = uuid.uuid4().hex[:8]

raw_contracts_df = (
    spark.read.format("binaryFile")
    .load(SOURCE_PATH)
    .filter(F.lower(F.col("path")).contains("consult"))
    .orderBy("path")
    .limit(10)
)

display(raw_contracts_df.select("path", "length", "modificationTime"))

Step 2: Extract the company name​

First, use ai_parse_document to convert each PDF into a structured VARIANT, and materialize the result in a temporary table so later actions don't invoke the function again.

Python
parsed_contracts_df = raw_contracts_df.select(
    "path",
    F.expr("ai_parse_document(content, MAP('version', '2.0'))").alias("parsed_content"),
)

parsed_table = f"_tmp_ai_enrich_parsed_{TMP_SUFFIX}"
parsed_contracts_df.write.mode("overwrite").saveAsTable(parsed_table)

Then use ai_extract to convert each contract into a sparse entity record. The schema requests the company name and enough contract context to inspect the extraction.

Python
extraction_schema = json.dumps(
    {
        "company_name": {
            "type": "string",
            "description": "Legal name of the company engaging the consultant.",
        },
        "consultant_name": {
            "type": "string",
            "description": "Legal name of the consultant or consulting firm.",
        },
        "effective_date": {
            "type": "string",
            "description": "Contract start date.",
        },
    }
).replace("'", "\\'")

extracted_companies_df = (
    spark.table(parsed_table)
    .filter("TRY_CAST(parsed_content:error_status AS STRING) IS NULL")
    .select(
        "path",
        F.expr(
            f"""
            ai_extract(
              parsed_content,
              '{extraction_schema}',
              MAP('instructions', 'Extract concise values. Return null when a value is absent.')
            )
            """
        ).alias("extracted"),
    )
    .select(
        "path",
        F.expr("extracted:response.company_name::STRING").alias("company_name"),
        F.expr("extracted:response.consultant_name::STRING").alias("consultant_name"),
        F.expr("extracted:response.effective_date::STRING").alias("effective_date"),
    )
    .filter(F.col("company_name").isNotNull())
)

extracted_table = f"_tmp_ai_enrich_entities_{TMP_SUFFIX}"
extracted_companies_df.write.mode("overwrite").saveAsTable(extracted_table)
display(spark.table(extracted_table))

Step 3: Resolve and enrich each company​

Pass the extracted company name to ai_enrich. The typed schema makes the results suitable for downstream processing. The instructions tell the function to resolve the legal entity before generating values and to return null when the available evidence is insufficient.

Python
enrichment_schema = json.dumps(
    {
        "industry": {"type": "string", "description": "Primary industry."},
        "headquarters_country": {
            "type": "string",
            "description": "Country of the current headquarters.",
        },
        "official_website": {
            "type": "string",
            "description": "Canonical URL of the official company website.",
        },
        "is_currently_active": {
            "type": "boolean",
            "description": "Whether the legal entity or its clear successor is currently operating.",
        },
    }
).replace("'", "\\'")

enriched_companies_df = spark.table(extracted_table).select(
    "path",
    "company_name",
    "consultant_name",
    "effective_date",
    F.expr(
        f"""
        ai_enrich(
          company_name,
          '{enrichment_schema}',
          PARSE_JSON('[{{"type":"web_search","config":{{}}}}]'),
          MAP(
            'instructions',
            'Resolve the exact legal entity before enriching it. Prefer official and authoritative sources. If identity is ambiguous or evidence is insufficient, return null rather than guessing.'
          )
        )
        """
    ).alias("enrichment"),
)

enriched_table = f"_tmp_ai_enrich_results_{TMP_SUFFIX}"
enriched_companies_df.write.mode("overwrite").saveAsTable(enriched_table)

Step 4: Inspect values and grounding sources​

By default, each field in response contains a typed value and a rationale. For grounded rows, metadata.sources contains the source document identifiers used for the row. Grounding provenance applies to the entire row, not to individual fields.

Python
final_df = spark.table(enriched_table).select(
    "path",
    "company_name",
    F.expr("enrichment:response.industry.value::STRING").alias("industry"),
    F.expr("enrichment:response.headquarters_country.value::STRING").alias("headquarters_country"),
    F.expr("enrichment:response.official_website.value::STRING").alias("official_website"),
    F.expr("enrichment:response.is_currently_active.value::BOOLEAN").alias("is_currently_active"),
    F.expr("enrichment:response.official_website.rationale::STRING").alias("website_rationale"),
    F.expr("enrichment:metadata.sources").alias("grounding_sources"),
    F.expr("enrichment:error_message::STRING").alias("error_message"),
)

display(final_df)

When the function cannot support a field value, value is null. Preserve that distinction instead of replacing it with a guessed value.

Add an external search provider as a knowledge source​

To ground the enrichment in a search provider that you already have an account with, such as You.com or Exa, add an mcp entry to the knowledge sources array, alongside or in place of the web_search entry. An mcp source calls a search tool on an external Model Context Protocol (MCP) server through Unity Gateway. Configure the full name of the MCP and the search tool to call:

JSON
[
  {
    "type": "mcp",
    "description": "You.com web search",
    "config": {
      "mcp_service": "main.default.youcom",
      "tool": "you-search"
    }
  }
]

Before you run the enrichment with an mcp source:

  1. Register the provider's MCP server in Unity Gateway. See External MCP servers. For You.com, installing the Databricks Marketplace listing creates the MCP for you.
  2. Use a connection with a shared credential, such as the provider's API key. ai_enrich can't prompt users to sign in to the provider, 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 enrichment: you, or the run-as identity of a job or pipeline. See Share an MCP.

ai_enrich calls the tool one time for each row with non-empty content. The search query is built from the row content and is sent to the provider, so treat that content as data that leaves Databricks. For all configuration keys, including the objective argument for Exa, and for failure behavior, see MCP configuration.

Add governed data as a knowledge source​

To ground the same enrichment in proprietary data, add a vector_search entry to the knowledge sources array. Configure the three-level name of an AI Search index and its text and document URI columns:

JSON
[
  { "type": "web_search", "config": {} },
  {
    "type": "vector_search",
    "description": "Governed company profiles",
    "config": {
      "index_name": "prod_catalog.crm.company_kb",
      "text_col": "profile_text",
      "doc_uri_col": "source_url"
    }
  }
]

You can configure multiple AI Search indexes and MCP sources, but at most one web search source, in a call. For production workloads, review the returned rationales and sources, and monitor null rates, errors, source quality, and entity-resolution accuracy.

On this page