Skip to main content
Last updated on

Tutorial: Enrich entities extracted from documents

Beta

This feature is in Beta. Workspace admins can control access to this feature 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 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 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