AI-Assisted Rules Generation Improvements

demos/dqx_demo_ai_assisted_checks_generation.py

@@ -41,7 +41,7 @@
 
 import os, yaml
 from databricks.labs.dqx.profiler.generator import DQGenerator
-from databricks.labs.dqx.config import LLMModelConfig
+from databricks.labs.dqx.config import LLMModelConfig, InputConfig
 from databricks.labs.dqx.engine import DQEngine
 from databricks.sdk import WorkspaceClient
 
@@ -95,6 +95,6 @@
 
 # COMMAND ----------
 
-checks = generator.generate_dq_rules_ai_assisted(user_input=user_requirement, table_name=table_name)
+checks = generator.generate_dq_rules_ai_assisted(user_input=user_requirement, input_config=InputConfig(location=table_name))
 print("======== Generated checks =========")
 print(checks)

docs/dqx/docs/demos.mdx

@@ -12,6 +12,7 @@ Import the following notebooks in the Databricks workspace to try DQX out:
 * [DQX Demo Notebook](https://github.com/databrickslabs/dqx/blob/v0.10.0/demos/dqx_demo_library.py) - demonstrates how to use DQX as a library.
 * [DQX Demo Notebook for Data Quality Summary Metrics](https://github.com/databrickslabs/dqx/blob/v0.10.0/demos/dqx_demo_summary_metrics.py) - demonstrates how to generate summary-level data quality metrics when validating data with DQX.
 * [DQX Demo Notebook for Profiling and Applying Checks at Scale on Multiple Tables](https://github.com/databrickslabs/dqx/blob/v0.10.0/demos/dqx_multi_table_demo.py) - demonstrates how to use DQX as a library at scale to apply checks on multiple tables.
+* [DQX Demo Notebook for AI-assisted checks generation](https://github.com/databrickslabs/dqx/blob/v0.10.0/demos/dqx_demo_ai_assisted_checks_generation.py) - demonstrates how to generate DQX rules/checks with LLM using natural language.
 * [DQX Demo Notebook for Spark Structured Streaming (Native End-to-End Approach)](https://github.com/databrickslabs/dqx/blob/v0.10.0/demos/dqx_streaming_demo_native.py) - demonstrates how to use DQX as a library with Spark Structured Streaming, using the built-in end-to-end method to handle both reading and writing.
 * [DQX Demo Notebook for Spark Structured Streaming (DIY Approach)](https://github.com/databrickslabs/dqx/blob/v0.10.0/demos/dqx_streaming_demo_diy.py) - demonstrates how to use DQX as a library with Spark Structured Streaming, while handling reading and writing on your own outside DQX using Spark API.
 * [DQX Demo Notebook for Lakeflow Pipelines (formerly DLT)](https://github.com/databrickslabs/dqx/blob/v0.10.0/demos/dqx_dlt_demo.py) - demonstrates how to use DQX as a library with Lakeflow Pipelines.

docs/dqx/docs/guide/ai_assisted_quality_checks_generation.mdx

@@ -14,7 +14,7 @@ DQX provides the capability to generate data quality rule candidates using AI/LL
 
 The AI-assisted quality checks generation is available via the `DQGenerator` class and supports:
 - Analyzing natural language descriptions of data quality requirements.
-- Optionally inspecting table schemas to understand data structure.
+- Optionally inspecting input data schemas to understand data structure.
 - Generating appropriate data quality rules that match the requirements using built-in and custom check functions.
 - Validating the generated rules to ensure they are syntactically correct.
 
@@ -81,21 +81,22 @@ Here's a simple example of generating quality rules from a natural language desc
   </TabItem>
 </Tabs>
 
-### Using with Table Schema
+### Using with Input Data Schema
 
-For better results, you can provide a fully qualified table name of the input data.
-The LLM will analyze the table schema to generate more accurate rules:
+For better results, you can provide a fully qualified table name or path of the input data.
+The LLM will analyze the schema of the input data to generate more accurate rules:
 
 <Tabs>
   <TabItem value="Python" label="Python" default>
     ```python
     from databricks.labs.dqx.profiler.generator import DQGenerator
+    from databricks.labs.dqx.config import InputConfig
     from databricks.sdk import WorkspaceClient
 
     ws = WorkspaceClient()
     generator = DQGenerator(workspace_client=ws)
 
-    # Generate rules with table schema awareness
+    # Generate rules with input data schema awareness
     user_input = """
     All customer records must have complete contact information.
     Email addresses must be valid.
@@ -105,7 +106,7 @@ The LLM will analyze the table schema to generate more accurate rules:
 
     checks = generator.generate_dq_rules_ai_assisted(
       user_input=user_input,
-      table_name="catalog1.schema1.customers"
+      input_config=InputConfig(location="catalog1.schema1.customers")
     )
 
     print(checks)
@@ -209,6 +210,7 @@ Here's a complete workflow showing how to generate, review, and save AI-assisted
     ```python
     from databricks.labs.dqx.profiler.generator import DQGenerator
     from databricks.labs.dqx.engine import DQEngine
+    from databricks.labs.dqx.config import InputConfig
     from databricks.labs.dqx.config import WorkspaceFileChecksStorageConfig
     from databricks.sdk import WorkspaceClient
     import yaml
@@ -233,7 +235,8 @@ Here's a complete workflow showing how to generate, review, and save AI-assisted
     # Generate quality rules using AI
     checks = generator.generate_dq_rules_ai_assisted(
       user_input=business_requirements,
-      table_name="production.users.registrations"
+      # input location could be a table or path
+      input_config=InputConfig(location="production.users.registrations")
     )
 
     # Review the generated checks
@@ -265,6 +268,7 @@ You can combine AI-assisted rules with profiler-generated rules for comprehensiv
     ```python
     from databricks.labs.dqx.profiler.profiler import DQProfiler
     from databricks.labs.dqx.profiler.generator import DQGenerator
+    from databricks.labs.dqx.config import InputConfig
     from databricks.labs.dqx.engine import DQEngine
     from databricks.sdk import WorkspaceClient
 
@@ -291,7 +295,7 @@ You can combine AI-assisted rules with profiler-generated rules for comprehensiv
 
     ai_checks = generator.generate_dq_rules_ai_assisted(
         user_input=business_requirements,
-        table_name="catalog1.schema1.sales_data"
+        input_config=InputConfig(location="catalog1.schema1.sales_data")
     )
 
     # Step 3: Combine both sets of rules
@@ -329,7 +333,7 @@ GDPR Compliance Requirements:
 
 checks = generator.generate_dq_rules_ai_assisted(
   user_input=compliance_requirements,
-  table_name="gdpr.users.personal_data"
+  input_config=InputConfig(location="gdpr.users.personal_data")
 )
 ```
 
@@ -349,7 +353,7 @@ Financial Transaction Rules:
 
 checks = generator.generate_dq_rules_ai_assisted(
   user_input=financial_requirements,
-  table_name="finance.transactions.daily"
+  input_config=InputConfig(location="finance.transactions.daily")
 )
 ```
 
@@ -369,7 +373,7 @@ IoT Sensor Data Quality Rules:
 
 checks = generator.generate_dq_rules_ai_assisted(
   user_input=iot_requirements,
-  table_name="iot.sensors.readings"
+  input_config=InputConfig(location="iot.sensors.readings")
 )
 ```
 
@@ -410,7 +414,7 @@ Use a slash-separated format for specifying the scope and key, such as `api_key:
 ## Best Practices
 
 1. **Be Specific**: Provide clear and specific business requirements in your input. The more detailed your description, the better the generated rules. However, be aware that the model endpoint that you use have token limitations [see Databricks Model APIs limits and quotas](https://docs.databricks.com/aws/en/machine-learning/foundation-model-apis/limits).
-2. **Include Table Name**: When possible, provide the fully qualified table name to help the LLM understand the actual data structure.
+2. **Include Input Config**: When possible, provide input config specifying location of input data as a fully qualified table name or path to help the LLM understand the actual data structure.
 3. **Review Generated Rules**: Always review the generated rules before applying them to production data. The AI may not perfectly understand all nuances of your requirements so treat the generated rules as candidates and review them.
 4. **Combine Approaches**: Use AI-assisted generation for business logic rules and profiler-based generation for statistical (technical) rules.
 5. **Iterate**: If the generated rules don't match your expectations, refine your input description and regenerate, or update manually.
@@ -432,7 +436,7 @@ pip install 'databricks-labs-dqx[llm]'
 
 **Solution**:
 - Make your input description more specific and detailed.
-- Provide the table name so the LLM can analyze the actual schema.
+- Provide the input data location so the LLM can analyze the actual schema.
 - Break complex requirements into simpler, more focused descriptions.
 - Include examples in your description.
 
@@ -452,7 +456,7 @@ pip install 'databricks-labs-dqx[llm]'
 
 **Solution**:
 - Check the validation error messages for specific issues.
-- Verify that the column names mentioned in your requirements exist in the table.
+- Verify that the column names mentioned in your requirements exist in the input data.
 - Ensure the rule types requested are supported by DQX (see [Quality Checks Reference](/docs/reference/quality_checks)).
 - Simplify your requirements and regenerate.