LLM assisted rules generation

[souravg-db2]

Changes

• Rules generation based on natural language and optionally schema from input table. Providing support for AI-assisted rules generation using programmatic and no-code (workflows) approach.
• Using Databricks foundational model by default and support providing different LLM models.
• Generated rules are validated for correctness.
• Added demo for AI-assisted rules generation.
• Added new method to the DQGenerator: generate_dq_rules_ai_assisted

Linked issues

Resolves #370 #478 #480

Tests

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

[github-actions]

✅ 409/409 passed, 2 flaky, 39 skipped, 4h26m9s total

Flaky tests:

• 🤪 test_e2e_workflow_for_multiple_run_configs (10m46.953s)
• 🤪 test_e2e_workflow_for_patterns_exclude_patterns (2m19.872s)

_{Running from acceptance #3082}

[Copilot]

Pull Request Overview

This PR introduces LLM-assisted data quality rules generation capabilities to the DQX framework. The implementation leverages DSPy for LLM orchestration and includes comprehensive training examples and validation mechanisms.

• Adds core LLM functionality with DSPy integration for automated data quality rule generation
• Implements training dataset and validation framework for optimizing rule generation accuracy
• Provides utility functions for schema metadata extraction and function documentation

Reviewed Changes

Copilot reviewed 6 out of 6 changed files in this pull request and generated 3 comments.

Show a summary per file

File	Description
src/databricks/labs/dqx/llm/llm_core.py	Core LLM implementation with DSPy module, validation logic, and compiler configuration
src/databricks/labs/dqx/llm/llm_engine.py	High-level interface for generating business rules using LLM
src/databricks/labs/dqx/llm/utils.py	Utility functions for schema extraction, training data loading, and function documentation
src/databricks/labs/dqx/llm/resources/training_examples.yml	Training dataset with business scenarios and expected quality rules
tests/unit/test_llm_utils.py	Unit tests for LLM utility functions and training data validation
pyproject.toml	Adds dspy dependency for LLM functionality

---

_{Tip: Customize your code reviews with copilot-instructions.md. Create the file or learn how to get started.}

[codecov]

Welcome to Codecov 🎉

Once you merge this PR into your default branch, you're all set! Codecov will compare coverage reports and display results in all future pull requests.

Thanks for integrating Codecov - We've got you covered ☂️

[mwojtyczka]

LGTM

[Copilot]

Pull Request Overview

Copilot reviewed 45 out of 45 changed files in this pull request and generated 2 comments.

Comments suppressed due to low confidence (2)

tests/integration/conftest.py:1

• Corrected spelling of 'fo' to 'of'.
  src/databricks/labs/dqx/pii/init.py:1
• Error message split across two lines should be combined into a single string for better readability. Consider using implicit string concatenation or a single multi-line string.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull Request Overview

Copilot reviewed 45 out of 45 changed files in this pull request and generated 4 comments.

Comments suppressed due to low confidence (2)

src/databricks/labs/dqx/pii/init.py:1

• Missing space between sentences. Should be 'PII detection extras not installed. Install additional dependencies...' with proper spacing.
  tests/integration/test_ai_rules_generator.py:1
• Typo in docstring: 'Example fo custom' should be 'Example of custom'.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.

[Copilot]

Pull Request Overview

Copilot reviewed 45 out of 45 changed files in this pull request and generated 5 comments.

Comments suppressed due to low confidence (2)

src/databricks/labs/dqx/profiler/profiler_workflow.py:1

• The secret detection logic using \"/\" in llm_model_config.api_base is fragile. Valid API base URLs typically contain "/" characters (e.g., "https://api.example.com/v1\"), which would incorrectly trigger secret retrieval. Consider using a more explicit pattern, such as checking for the absence of a protocol prefix ("http://" or "https://") or requiring a specific prefix like "secret://" to indicate secret references.
  tests/integration/conftest.py:1
• The custom check function not_ends_with_suffix has a logic error: the condition F.col(column).endswith(suffix) checks if the column does end with the suffix, but the function name suggests it should check if it does not end with the suffix. The condition should be negated: ~F.col(column).endswith(suffix) to match the intended behavior.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.