add llm insights to datapilot

Author: gaurpulkit

src/datapilot/core/platforms/dbt/cli/cli.py

@@ -8,7 +8,6 @@
 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
@@ -82,7 +81,6 @@ def project_health(
 
     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)
@@ -101,18 +99,6 @@ def project_health(
         click.echo("--" * 50)
         click.echo(tabulate_data(project_report, headers="keys"))
 
-    if len(llm_insights):
-        click.echo("--" * 50)
-        click.echo("Project Governance LLM Insights")
-        click.echo("--" * 50)
-        for check in llm_insights:
-            click.echo(f"Check: {check['name']}")
-            for answer in check["answer"]:
-                click.echo(f"Path: {answer['path']}")
-                click.echo(f"Reason to Flag: {answer['reason_to_flag']}")
-                click.echo(f"Recommendation: {answer['recommendation']}")
-                click.echo("\n")
-
 
 @dbt.command("onboard")
 @click.option("--token", prompt="API Token", help="Your API token for authentication.")

src/datapilot/core/platforms/dbt/executor.py

@@ -13,6 +13,8 @@
 from datapilot.core.platforms.dbt.exceptions import AltimateCLIArgumentError
 from datapilot.core.platforms.dbt.factory import DBTFactory
 from datapilot.core.platforms.dbt.insights import INSIGHTS
+from datapilot.core.platforms.dbt.insights.schema import DBTInsightResult
+from datapilot.core.platforms.dbt.insights.schema import DBTModelInsightResponse
 from datapilot.core.platforms.dbt.schemas.manifest import Catalog
 from datapilot.core.platforms.dbt.schemas.manifest import Manifest
 from datapilot.core.platforms.dbt.utils import get_models
@@ -108,7 +110,6 @@ def run(self):
         reports = {
             MODEL: {},
             PROJECT: [],
-            LLM: [],
         }
         for insight_class in INSIGHTS:
             # TODO: Skip insight based on config
@@ -176,7 +177,38 @@ def run(self):
 
         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"])
+            llm_reports = llm_check_results["results"]
+            llm_insights = {}
+            for report in llm_reports:
+                for answer in report["answer"]:
+                    location = answer["unique_id"]
+                    if location not in llm_insights:
+                        llm_insights[location] = []
+                        metadata = answer.get("metadata", {})
+                        metadata["source"] = LLM
+                        metadata["llm_id"] = report["id"]
+                        metadata["catagory"] = report["type"]
+                    llm_insights[location].append(
+                        DBTModelInsightResponse(
+                            insight=DBTInsightResult(
+                                type="Custom",
+                                name=report["name"],
+                                message=answer["message"],
+                                reason_to_flag=answer["reason_to_flag"],
+                                recommendation=answer["recommendation"],
+                                metadata=metadata,
+                            ),
+                            severity=answer["severity"],
+                            path=answer["path"],
+                            original_file_path=answer["original_file_path"],
+                            package_name=answer["package_name"],
+                            unique_id=answer["unique_id"],
+                        )
+                    )
+            for key, value in llm_insights.items():
+                if key in reports[MODEL]:
+                    reports[MODEL][key].extend(value)
+                else:
+                    reports[MODEL][key] = value
 
         return reports