Skip to content

Commit

Permalink
changes
Browse files Browse the repository at this point in the history
  • Loading branch information
gaurpulkit committed Nov 18, 2024
1 parent c18defa commit 54ba983
Show file tree
Hide file tree
Showing 3 changed files with 40 additions and 11 deletions.
26 changes: 15 additions & 11 deletions src/datapilot/core/platforms/dbt/cli/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,12 @@
import click

from datapilot.clients.altimate.utils import check_token_and_instance
from datapilot.clients.altimate.utils import get_project_governance_llm_checks
from datapilot.clients.altimate.utils import onboard_file
from datapilot.clients.altimate.utils import run_project_governance_llm_checks
from datapilot.clients.altimate.utils import start_dbt_ingestion
from datapilot.clients.altimate.utils import validate_credentials
from datapilot.clients.altimate.utils import validate_permissions
from datapilot.config.config import load_config
from datapilot.core.platforms.dbt.constants import LLM
from datapilot.core.platforms.dbt.constants import MODEL
from datapilot.core.platforms.dbt.constants import PROJECT
from datapilot.core.platforms.dbt.executor import DBTInsightGenerator
Expand All @@ -30,8 +29,8 @@ def dbt():


@dbt.command("project-health")
@click.option("--token", prompt="API Token", help="Your API token for authentication.")
@click.option("--instance-name", prompt="Instance Name", help="Your tenant ID.")
@click.option("--token", required=False, prompt="API Token", help="Your API token for authentication.")
@click.option("--instance-name", required=False, prompt="Instance Name", help="Your tenant ID.")
@click.option(
"--manifest-path",
required=True,
Expand Down Expand Up @@ -70,15 +69,20 @@ def project_health(
manifest = load_manifest(manifest_path)
catalog = load_catalog(catalog_path) if catalog_path else None

llm_checks = get_project_governance_llm_checks(token, instance_name, backend_url)
check_names = [check["name"] for check in llm_checks if check["alias"] not in config.get("disabled_insights", [])]
llm_check_results = run_project_governance_llm_checks(token, instance_name, backend_url, manifest, catalog, check_names)

insight_generator = DBTInsightGenerator(manifest=manifest, catalog=catalog, config=config, selected_models=selected_models)
insight_generator = DBTInsightGenerator(
manifest=manifest,
catalog=catalog,
config=config,
selected_models=selected_models,
token=token,
instance_name=instance_name,
backend_url=backend_url,
)
reports = insight_generator.run()

package_insights = reports[PROJECT]
model_insights = reports[MODEL]
llm_insights = reports[LLM]
model_report = generate_model_insights_table(model_insights)
if len(model_report) > 0:
click.echo("--" * 50)
Expand All @@ -97,11 +101,11 @@ def project_health(
click.echo("--" * 50)
click.echo(tabulate_data(project_report, headers="keys"))

if llm_check_results:
if len(llm_insights):
click.echo("--" * 50)
click.echo("Project Governance LLM Insights")
click.echo("--" * 50)
for check in llm_check_results["results"]:
for check in llm_insights:
click.echo(f"Check: {check['name']}")
for answer in check["answer"]:
click.echo(f"Rule: {answer['Rule']}")
Expand Down
2 changes: 2 additions & 0 deletions src/datapilot/core/platforms/dbt/constants.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,8 @@
MODEL = "model"
SOURCE = "source"

LLM = "llm"


PROJECT = "project"
SQL = "sql"
Expand Down
23 changes: 23 additions & 0 deletions src/datapilot/core/platforms/dbt/executor.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,6 +5,9 @@
from typing import List
from typing import Optional

from datapilot.clients.altimate.utils import get_project_governance_llm_checks
from datapilot.clients.altimate.utils import run_project_governance_llm_checks
from datapilot.core.platforms.dbt.constants import LLM
from datapilot.core.platforms.dbt.constants import MODEL
from datapilot.core.platforms.dbt.constants import PROJECT
from datapilot.core.platforms.dbt.exceptions import AltimateCLIArgumentError
Expand All @@ -29,11 +32,17 @@ def __init__(
target: str = "dev",
selected_models: Optional[str] = None,
selected_model_ids: Optional[List[str]] = None,
token: Optional[str] = None,
instance_name: Optional[str] = None,
backend_url: Optional[str] = None,
):
self.run_results_path = run_results_path
self.target = target
self.env = env
self.config = config or {}
self.token = token
self.instance_name = instance_name
self.backend_url = backend_url

self.manifest_wrapper = DBTFactory.get_manifest_wrapper(manifest)
self.manifest_present = True
Expand Down Expand Up @@ -85,10 +94,19 @@ def _check_if_skipped(self, insight):
return True
return False

def run_llm_checks(self):
llm_checks = get_project_governance_llm_checks(self.token, self.instance_name, self.backend_url)
check_names = [check["name"] for check in llm_checks if check["alias"] not in self.config.get("disabled_insights", [])]
llm_check_results = run_project_governance_llm_checks(
self.token, self.instance_name, self.backend_url, self.manifest, self.catalog, check_names
)
return llm_check_results

def run(self):
reports = {
MODEL: {},
PROJECT: [],
LLM: [],
}
for insight_class in INSIGHTS:
# TODO: Skip insight based on config
Expand Down Expand Up @@ -154,4 +172,9 @@ def run(self):
else:
self.logger.info(color_text(f"Skipping insight {insight_class.NAME} as {message}", YELLOW))

if self.token and self.instance_name and self.backend_url:
llm_check_results = self.run_llm_checks()
if llm_check_results:
reports[LLM].extend(llm_check_results["results"])

return reports

0 comments on commit 54ba983

Please sign in to comment.