WIP - Feature: Use unity catalog metadata for LLM assisted checks generation

[IvannKurchenko]

Changes

Collects the following additional metadata for profiler for LLM assisted checks generation:

• Unity catalog comments and tags for both tables and columns;
• Upstream lineage with corresponding metadata;
• Upstream external lineage via workspace client;

Additional demo notebook provided;

NOTE: For demo and tests https://docs.databricks.com/aws/en/data-governance/unity-catalog/access-control/privileges-reference#create-external-metadata permissions required.

Demo results - comparison checks generation with and without metadata attached to UC.

Without:

04:30:18     INFO [d.l.d.profiler.generator] Generating DQ rules with LLM for input: 'generate quality checks for an invoice fact table'
04:30:18     INFO [d.l.d.profiler.generator] Using WorkspaceClient to determine the schema info for 'dev.demo.fact_invoice'
 18%|█▊        | 2/11 [00:01<00:05,  1.52it/s]
Bootstrapped 2 full traces after 2 examples for up to 1 rounds, amounting to 2 attempts.
04:30:28     INFO [d.l.d.profiler.generator] Generated 11 rules with LLM: [{'criticality': 'error', 'check': {'function': 'is_not_null', 'arguments': {'column': 'invoice_id'}}}, {'criticality': 'error', 'check': {'function': 'is_not_null', 'arguments': {'column': 'invoice_timestamp'}}}, {'criticality': 'error', 'check': {'function': 'is_not_null', 'arguments': {'column': 'client_id'}}}, {'criticality': 'error', 'check': {'function': 'is_not_null', 'arguments': {'column': 'service_id'}}}, {'criticality': 'error', 'check': {'function': 'is_not_null', 'arguments': {'column': 'amount'}}}, {'criticality': 'error', 'check': {'function': 'is_not_null', 'arguments': {'column': 'currency_code'}}}, {'criticality': 'error', 'check': {'function': 'is_not_null', 'arguments': {'column': 'signed'}}}, {'criticality': 'warn', 'check': {'function': 'is_valid_currency_code', 'arguments': {'column': 'currency_code', 'code_format': 'alphabetic'}}}, {'criticality': 'error', 'check': {'function': 'sql_expression', 'arguments': {'expression': 'amount > 0', 'columns': ['amount']}}}, {'criticality': 'error', 'check': {'function': 'sql_expression', 'arguments': {'expression': 'signed IN (TRUE, FALSE)', 'columns': ['signed']}}}, {'criticality': 'warn', 'check': {'function': 'is_data_fresh', 'arguments': {'column': 'invoice_timestamp', 'max_age_minutes': 10080}}}]
04:30:28     INFO [d.l.d.profiler.generator] LLM reasoning: These quality rules ensure that all required columns are present and not null, and that the data is valid and consistent. The rules also check for valid currency codes and that the amount is greater than 0. The signed column is checked to ensure it only contains TRUE or FALSE values. The data freshness rule checks that the invoice timestamp is within the last week.
04:30:28     INFO [d.l.d.profiler.generator] ✅ AI-Assisted quality rules generation completed. Generated 11 valid rule(s) out of 11 LLM output(s).
- criticality: error
  check:
    function: is_not_null
    arguments:
      column: invoice_id
- criticality: error
  check:
    function: is_not_null
    arguments:
      column: invoice_timestamp
- criticality: error
  check:
    function: is_not_null
    arguments:
      column: client_id
- criticality: error
  check:
    function: is_not_null
    arguments:
      column: service_id
- criticality: error
  check:
    function: is_not_null
    arguments:
      column: amount
- criticality: error
  check:
    function: is_not_null
    arguments:
      column: currency_code
- criticality: error
  check:
    function: is_not_null
    arguments:
      column: signed
- criticality: warn
  check:
    function: is_valid_currency_code
    arguments:
      column: currency_code
      code_format: alphabetic
- criticality: error
  check:
    function: sql_expression
    arguments:
      expression: amount > 0
      columns:
      - amount
- criticality: error
  check:
    function: sql_expression
    arguments:
      expression: signed IN (TRUE, FALSE)
      columns:
      - signed
- criticality: warn
  check:
    function: is_data_fresh
    arguments:
      column: invoice_timestamp
      max_age_minutes: 10080

With:

04:31:12     INFO [d.l.d.profiler.generator] Generated 13 rules with LLM: [{'criticality': 'error', 'check': {'function': 'is_not_null', 'arguments': {'column': 'invoice_id'}}}, {'criticality': 'error', 'check': {'function': 'is_not_null', 'arguments': {'column': 'invoice_timestamp'}}}, {'criticality': 'error', 'check': {'function': 'is_not_null', 'arguments': {'column': 'client_id'}}}, {'criticality': 'error', 'check': {'function': 'is_not_null', 'arguments': {'column': 'service_id'}}}, {'criticality': 'error', 'check': {'function': 'is_not_null', 'arguments': {'column': 'amount'}}}, {'criticality': 'error', 'check': {'function': 'is_not_null', 'arguments': {'column': 'currency_code'}}}, {'criticality': 'error', 'check': {'function': 'is_not_null', 'arguments': {'column': 'signed'}}}, {'criticality': 'warn', 'check': {'function': 'is_valid_date', 'arguments': {'column': 'invoice_timestamp'}}}, {'criticality': 'warn', 'check': {'function': 'is_valid_currency_code', 'arguments': {'column': 'currency_code', 'code_format': 'alphabetic'}}}, {'criticality': 'error', 'check': {'function': 'is_not_less_than', 'arguments': {'column': 'amount', 'limit': 0}}}, {'criticality': 'error', 'check': {'function': 'is_in_list', 'arguments': {'column': 'currency_code', 'allowed': ['USD', 'EUR', 'GBP']}}}, {'criticality': 'error', 'check': {'function': 'sql_expression', 'arguments': {'expression': 'signed = TRUE', 'columns': ['signed']}}}, {'criticality': 'warn', 'check': {'function': 'is_data_fresh', 'arguments': {'column': 'invoice_timestamp', 'max_age_minutes': 10080}}}]
04:31:12     INFO [d.l.d.profiler.generator] LLM reasoning: These quality rules are designed to ensure the integrity and consistency of the invoice data. The rules check for null values, valid dates, currency codes, and amounts. The rules also check if the signed column is True, indicating that the invoice has been signed. The data freshness rule checks if the invoice timestamp is within a certain time range.
04:31:12     INFO [d.l.d.profiler.generator] ✅ AI-Assisted quality rules generation completed. Generated 13 valid rule(s) out of 13 LLM output(s).
- criticality: error
  check:
    function: is_not_null
    arguments:
      column: invoice_id
- criticality: error
  check:
    function: is_not_null
    arguments:
      column: invoice_timestamp
- criticality: error
  check:
    function: is_not_null
    arguments:
      column: client_id
- criticality: error
  check:
    function: is_not_null
    arguments:
      column: service_id
- criticality: error
  check:
    function: is_not_null
    arguments:
      column: amount
- criticality: error
  check:
    function: is_not_null
    arguments:
      column: currency_code
- criticality: error
  check:
    function: is_not_null
    arguments:
      column: signed
- criticality: warn
  check:
    function: is_valid_date
    arguments:
      column: invoice_timestamp
- criticality: warn
  check:
    function: is_valid_currency_code
    arguments:
      column: currency_code
      code_format: alphabetic
- criticality: error
  check:
    function: is_not_less_than
    arguments:
      column: amount
      limit: 0
- criticality: error
  check:
    function: is_in_list
    arguments:
      column: currency_code
      allowed:
      - USD
      - EUR
      - GBP
- criticality: error
  check:
    function: sql_expression
    arguments:
      expression: signed = TRUE
      columns:
      - signed
- criticality: warn
  check:
    function: is_data_fresh
    arguments:
      column: invoice_timestamp
      max_age_minutes: 10080

Model used: databricks/databricks-meta-llama-3-3-70b-instruct

Linked issues

Resolves #1329

Tests

• manually tested
• added unit tests
• added integration tests
• added end-to-end tests
• added performance tests

Documentation and Demos

• added/updated demos
• added/updated docs
• added/updated agent skills

[IvannKurchenko]

@ghanse Could you review this one when you will have a chance, please? Thanks.