shkit 1.2.0__py3-none-any.whl

healing_kit/__init__.py

@@ -0,0 +1,3 @@
+"""Self-Healing Pipeline Kit — Enterprise AI-driven autonomous remediation for Databricks."""
+
+__version__ = "2.0.0"

healing_kit/auth.py

@@ -0,0 +1,79 @@
+"""Identity & auth (#1) — prefer a service principal (OAuth M2M) over a PAT.
+
+The healer should authenticate as a least-privileged **service principal**, not a
+human PAT. This module mints a short-lived OAuth token via the Databricks OIDC
+`client_credentials` grant. It falls back to a PAT or the notebook's ambient
+token only when no SP credentials are configured, so existing deploys keep working.
+"""
+
+from dataclasses import dataclass
+
+import requests
+
+from healing_kit.clients.databricks_client import DatabricksClient, DatabricksConfig
+
+
+@dataclass
+class AuthResult:
+    """Resolved auth: a bearer token plus how it was obtained (for audit)."""
+
+    config: DatabricksConfig
+    method: str  # "service_principal" | "pat" | "ambient"
+
+
+def get_sp_oauth_token(host: str, client_id: str, client_secret: str, scope: str = "all-apis") -> str:
+    """Exchange SP client credentials for a short-lived OAuth access token."""
+    host = host.rstrip("/")
+    resp = requests.post(
+        f"{host}/oidc/v1/token",
+        auth=(client_id, client_secret),
+        data={"grant_type": "client_credentials", "scope": scope},
+        timeout=30,
+    )
+    resp.raise_for_status()
+    return resp.json()["access_token"]
+
+
+def resolve_auth(
+    host: str,
+    *,
+    sp_client_id: str | None = None,
+    sp_client_secret: str | None = None,
+    pat: str | None = None,
+    ambient_token: str | None = None,
+) -> AuthResult:
+    """Resolve auth in priority order: service principal → PAT → ambient token."""
+    host = host.rstrip("/")
+    if sp_client_id and sp_client_secret:
+        token = get_sp_oauth_token(host, sp_client_id, sp_client_secret)
+        return AuthResult(DatabricksConfig(host=host, token=token), "service_principal")
+    if pat:
+        return AuthResult(DatabricksConfig(host=host, token=pat), "pat")
+    if ambient_token:
+        return AuthResult(DatabricksConfig(host=host, token=ambient_token), "ambient")
+    raise ValueError("No credentials: provide SP client_id/secret, a PAT, or an ambient token")
+
+
+def build_client(auth: AuthResult) -> DatabricksClient:
+    return DatabricksClient(auth.config)
+
+
+def resolve_auth_from_dbutils(dbutils, secret_scope: str, host: str | None = None) -> AuthResult:
+    """Notebook-side helper: read SP creds from a secret scope, fall back to the
+    notebook's ambient token. Never embeds a static PAT in the notebook."""
+    ctx = dbutils.notebook.entry_point.getDbutils().notebook().getContext()
+    host = (host or ctx.apiUrl().get()).rstrip("/")
+    ambient = ctx.apiToken().get()
+
+    def _secret(key):
+        try:
+            return dbutils.secrets.get(scope=secret_scope, key=key)
+        except Exception:
+            return None
+
+    return resolve_auth(
+        host,
+        sp_client_id=_secret("sp_client_id"),
+        sp_client_secret=_secret("sp_client_secret"),
+        ambient_token=ambient,
+    )

healing_kit/clients/__init__.py

@@ -0,0 +1 @@
+"""External service clients."""

healing_kit/clients/databricks_client.py

@@ -0,0 +1,183 @@
+"""Unified Databricks API client wrapping Jobs, Runs, Workspace, Clusters, UC Lineage, Model Serving."""
+
+import base64
+import time
+from dataclasses import dataclass
+from typing import Optional
+
+import requests
+
+
+@dataclass
+class DatabricksConfig:
+    """Connection configuration."""
+
+    host: str
+    token: str
+
+    @property
+    def headers(self) -> dict:
+        return {"Authorization": f"Bearer {self.token}", "Content-Type": "application/json"}
+
+
+class DatabricksClient:
+    """
+    Unified client for all Databricks REST API interactions.
+    Uses ambient auth via PAT (env var or secret scope).
+    """
+
+    def __init__(self, config: DatabricksConfig):
+        self.config = config
+        self.base_url = config.host.rstrip("/")
+
+    def _get(self, path: str, params: dict = None) -> dict:
+        r = requests.get(f"{self.base_url}{path}", headers=self.config.headers, params=params, timeout=30)
+        r.raise_for_status()
+        return r.json()
+
+    def _post(self, path: str, payload: dict = None) -> dict:
+        r = requests.post(f"{self.base_url}{path}", headers=self.config.headers, json=payload or {}, timeout=60)
+        r.raise_for_status()
+        return r.json() if r.content else {}
+
+    # ─── Jobs API ───
+
+    def list_jobs(self, limit: int = 50) -> list[dict]:
+        data = self._get("/api/2.1/jobs/list", {"limit": limit})
+        return data.get("jobs", [])
+
+    def get_job(self, job_id: int) -> dict:
+        return self._get("/api/2.1/jobs/get", {"job_id": job_id})
+
+    def list_runs(self, job_id: int, limit: int = 5, expand_tasks: bool = True) -> list[dict]:
+        params = {"job_id": job_id, "limit": limit, "expand_tasks": str(expand_tasks).lower()}
+        data = self._get("/api/2.1/jobs/runs/list", params)
+        return data.get("runs", [])
+
+    def get_run_output(self, run_id: int) -> dict:
+        return self._get("/api/2.1/jobs/runs/get-output", {"run_id": run_id})
+
+    def repair_run(self, run_id: int, rerun_tasks: list[str]) -> dict:
+        return self._post("/api/2.1/jobs/runs/repair", {"run_id": run_id, "rerun_tasks": rerun_tasks})
+
+    def create_job(self, settings: dict) -> int:
+        result = self._post("/api/2.1/jobs/create", settings)
+        return result["job_id"]
+
+    def reset_job(self, job_id: int, new_settings: dict) -> None:
+        self._post("/api/2.1/jobs/reset", {"job_id": job_id, "new_settings": new_settings})
+
+    def run_job_now(self, job_id: int, params: dict = None) -> int:
+        payload = {"job_id": job_id}
+        if params:
+            payload["notebook_params"] = params
+        result = self._post("/api/2.1/jobs/run-now", payload)
+        return result.get("run_id")
+
+    # ─── Workspace API ───
+
+    def export_notebook(self, path: str) -> Optional[str]:
+        """Export notebook source code. Returns decoded content or None."""
+        try:
+            data = self._get("/api/2.0/workspace/export", {"path": path, "format": "SOURCE"})
+            return base64.b64decode(data.get("content", "")).decode("utf-8", errors="replace")
+        except Exception:
+            return None
+
+    def import_notebook(self, path: str, content: str, language: str = "PYTHON") -> bool:
+        """Upload a notebook. Returns True on success."""
+        try:
+            # Ensure parent directory exists
+            parent = "/".join(path.split("/")[:-1])
+            self._post("/api/2.0/workspace/mkdirs", {"path": parent})
+            payload = {
+                "path": path,
+                "content": base64.b64encode(content.encode()).decode(),
+                "language": language,
+                "format": "SOURCE",
+                "overwrite": True,
+            }
+            self._post("/api/2.0/workspace/import", payload)
+            return True
+        except Exception:
+            return False
+
+    # ─── Clusters API ───
+
+    def get_cluster_events(self, cluster_id: str, limit: int = 10) -> list[dict]:
+        try:
+            data = self._post("/api/2.0/clusters/events", {"cluster_id": cluster_id, "limit": limit})
+            return data.get("events", [])
+        except Exception:
+            return []
+
+    # ─── Model Serving ───
+
+    def list_serving_endpoints(self) -> list[dict]:
+        data = self._get("/api/2.0/serving-endpoints")
+        return data.get("endpoints", [])
+
+    def invoke_model(self, endpoint_name: str, messages: list[dict], max_tokens: int = 1500, temperature: float = 0.1) -> str:
+        """Call a Model Serving chat endpoint. Returns the response content string."""
+        content, _ = self.invoke_model_full(endpoint_name, messages, max_tokens, temperature)
+        return content
+
+    def invoke_model_full(self, endpoint_name: str, messages: list[dict], max_tokens: int = 1500,
+                          temperature: float = 0.1, retries: int = 4):
+        """Call a Model Serving chat endpoint with backoff on 429/5xx.
+
+        Returns ``(content, usage_dict)`` where usage carries real token counts so
+        the token budget can be enforced. Raises after exhausting retries."""
+        payload = {"messages": messages, "max_tokens": max_tokens, "temperature": temperature}
+        url = f"{self.base_url}/serving-endpoints/{endpoint_name}/invocations"
+        delay, last = 2.0, None
+        for _ in range(retries):
+            r = requests.post(url, headers=self.config.headers, json=payload, timeout=120)
+            last = r
+            if r.status_code == 200:
+                data = r.json()
+                return data["choices"][0]["message"]["content"], (data.get("usage", {}) or {})
+            if r.status_code in (429, 500, 502, 503, 504):
+                time.sleep(delay)
+                delay *= 2
+                continue
+            r.raise_for_status()
+        if last is not None:
+            last.raise_for_status()
+        raise RuntimeError("Model invocation failed with no response")
+
+    # ─── SCIM / Identity ───
+
+    def find_user_by_identity(self, identity: str) -> Optional[dict]:
+        """Look up a workspace user by userName or email via SCIM. Returns the
+        SCIM user object (or None). Used to authorize approvers (#3)."""
+        ident = (identity or "").strip()
+        if not ident:
+            return None
+        for flt in (f'userName eq "{ident}"', f'emails.value eq "{ident}"'):
+            try:
+                data = self._get("/api/2.0/preview/scim/v2/Users", {"filter": flt})
+            except Exception:
+                continue
+            resources = data.get("Resources", [])
+            if resources:
+                return resources[0]
+        return None
+
+    def user_in_group(self, user: dict, group_name: str) -> bool:
+        """True if the SCIM user object lists membership in group_name."""
+        if not group_name:
+            return True
+        for g in user.get("groups", []) or []:
+            if g.get("display") == group_name or g.get("value") == group_name:
+                return True
+        return False
+
+    # ─── Unity Catalog Lineage ───
+
+    def get_table_lineage(self, table_name: str) -> dict:
+        """Get upstream/downstream lineage for a table."""
+        try:
+            return self._get("/api/2.1/unity-catalog/lineage/table-lineage", {"table_name": table_name})
+        except Exception:
+            return {}

healing_kit/clients/teams_client.py

@@ -0,0 +1,128 @@
+"""Microsoft Teams webhook client — sends full diagnostic report cards with Approve/Reject."""
+
+from datetime import datetime
+
+import requests
+
+from healing_kit.utils.hmac_tokens import generate_token
+
+
+class TeamsClient:
+    """Sends comprehensive diagnostic report cards to Teams with approval buttons."""
+
+    def __init__(self, webhook_url: str, approval_job_url: str = "", secret_key: bytes = b""):
+        self.webhook_url = webhook_url
+        self.approval_job_url = approval_job_url
+        self.secret_key = secret_key
+
+    def send_diagnosis_report(self, results: list[dict], run_id: str, pipeline_name: str = "",
+                              approver_email: str = "team@company.com") -> bool:
+        """
+        Send ONE consolidated report card with full analysis + Approve/Reject per failure.
+        NEVER auto-fixes. Always waits for human approval.
+
+        Each failure includes:
+        - Root cause analysis
+        - Data evidence (diagnostic query results)
+        - Proposed fix
+        - Fix simulation (impact analysis)
+        - Approve / Reject buttons (signed HMAC URLs)
+        """
+        if not self.webhook_url:
+            return False
+
+        body_blocks = [
+            {"type": "TextBlock", "size": "Large", "weight": "Bolder", "text": "\U0001f916 Pipeline Failure Report — Approval Required"},
+            {"type": "TextBlock", "text": f"**Pipeline:** {pipeline_name} | **Run:** {run_id}", "isSubtle": True, "wrap": True},
+            {"type": "TextBlock", "text": f"**Time:** {datetime.utcnow().strftime('%Y-%m-%d %H:%M UTC')} | **Failures:** {len(results)}", "isSubtle": True, "wrap": True},
+            {"type": "TextBlock", "text": "---", "separator": True},
+        ]
+
+        actions = []
+
+        for i, item in enumerate(results, 1):
+            ctx = item.get("context", {})
+            analysis = item.get("analysis", {})
+            evidence = item.get("evidence")
+            simulation = item.get("simulation", {})
+
+            task_key = ctx.get("task_key", "unknown")
+            confidence = analysis.get("confidence_score", 0)
+            if isinstance(confidence, float) and confidence <= 1.0:
+                confidence = int(confidence * 100)
+
+            # Section header
+            body_blocks.append({"type": "TextBlock", "size": "Medium", "weight": "Bolder",
+                                "text": f"\u26a0\ufe0f Issue {i}: {task_key}", "separator": True, "wrap": True})
+
+            # Root cause
+            body_blocks.append({"type": "TextBlock", "weight": "Bolder", "text": "Root Cause:", "wrap": True})
+            body_blocks.append({"type": "TextBlock", "text": str(analysis.get("root_cause", "Unknown"))[:400], "wrap": True})
+
+            # Data evidence
+            if evidence and "0 rows" not in str(evidence) and "BLOCKED" not in str(evidence):
+                body_blocks.append({"type": "TextBlock", "weight": "Bolder", "text": "Data Evidence:", "wrap": True})
+                body_blocks.append({"type": "TextBlock", "text": f"```\n{str(evidence)[:350]}\n```", "wrap": True, "fontType": "Monospace"})
+
+            # Proposed fix
+            body_blocks.append({"type": "TextBlock", "weight": "Bolder", "text": "Proposed Fix:", "wrap": True})
+            body_blocks.append({"type": "TextBlock", "text": str(analysis.get("reasoning", analysis.get("suggested_fix", "N/A")))[:300], "wrap": True})
+
+            # Fix simulation / impact analysis
+            if simulation:
+                body_blocks.append({"type": "TextBlock", "weight": "Bolder", "text": "Impact Simulation:", "wrap": True})
+                sim_text = f"\u2022 Action: {simulation.get('action_description', 'N/A')}\n"
+                sim_text += f"\u2022 Affected: {simulation.get('affected_tables', 'N/A')}\n"
+                sim_text += f"\u2022 Expected outcome: {simulation.get('expected_outcome', 'N/A')}\n"
+                sim_text += f"\u2022 Risk if fails: {simulation.get('risk_if_fails', 'N/A')}\n"
+                sim_text += f"\u2022 Reversible: {simulation.get('reversible', 'Unknown')}\n"
+                sim_text += f"\u2022 Confidence: {confidence}%"
+                body_blocks.append({"type": "TextBlock", "text": sim_text, "wrap": True})
+            else:
+                body_blocks.append({"type": "FactSet", "facts": [
+                    {"title": "Action", "value": str(analysis.get("action_id", "RETRY"))},
+                    {"title": "Confidence", "value": f"{confidence}%"},
+                    {"title": "Reversible", "value": "Yes (rerun is safe)"},
+                ]})
+
+            # Generate signed approval tokens
+            if self.secret_key and self.approval_job_url:
+                token = generate_token(
+                    run_id=run_id,
+                    approver_email=approver_email,
+                    action_id=str(analysis.get("action_id", "RETRY")),
+                    secret_key=self.secret_key,
+                )
+                approve_url = f"{self.approval_job_url}?token={token}&action=approve&task={task_key}"
+                reject_url = f"{self.approval_job_url}?token={token}&action=reject&task={task_key}"
+
+                actions.append({"type": "Action.OpenUrl", "title": f"\u2705 Approve Fix: {task_key}", "url": approve_url})
+                actions.append({"type": "Action.OpenUrl", "title": f"\u274c Reject: {task_key}", "url": reject_url})
+
+        # Summary footer
+        body_blocks.append({"type": "TextBlock", "text": "---", "separator": True})
+        body_blocks.append({"type": "TextBlock", "text": "\u23f0 **Approval tokens expire in 15 minutes.** No action is taken until you approve.", "wrap": True, "isSubtle": True})
+
+        card = {
+            "type": "message",
+            "attachments": [{
+                "contentType": "application/vnd.microsoft.card.adaptive",
+                "content": {
+                    "$schema": "http://adaptivecards.io/schemas/adaptive-card.json",
+                    "type": "AdaptiveCard",
+                    "version": "1.4",
+                    "body": body_blocks,
+                    "actions": actions if actions else None,
+                },
+            }],
+        }
+
+        # Remove None actions field if empty
+        if not actions:
+            del card["attachments"][0]["content"]["actions"]
+
+        try:
+            r = requests.post(self.webhook_url, json=card, headers={"Content-Type": "application/json"}, timeout=15)
+            return r.status_code in (200, 202)
+        except Exception:
+            return False

healing_kit/models/__init__.py

@@ -0,0 +1 @@
+"""Data models for the healing pipeline."""

healing_kit/models/diagnosis.py

@@ -0,0 +1,45 @@
+"""Diagnosis response model from the AI engine."""
+
+from dataclasses import dataclass, field
+from enum import Enum
+
+
+class ActionId(str, Enum):
+    """Deterministic action identifiers for resolution routing."""
+
+    CLUSTER_RESIZE = "CLUSTER_RESIZE"
+    RETRY = "RETRY"
+    SCHEMA_FIX = "SCHEMA_FIX"
+    CODE_PATCH = "CODE_PATCH"
+    ESCALATE = "ESCALATE"
+    UNKNOWN = "UNKNOWN"
+    CACHE_HIT = "CACHE_HIT"
+
+
+@dataclass
+class DiagnosisResponse:
+    """Structured response from the AI Diagnosis Engine."""
+
+    root_cause: str
+    confidence_score: float  # 0.0 to 1.0
+    action_id: ActionId
+    action_params: dict = field(default_factory=dict)
+    reasoning: str = ""
+    evidence_used: list[str] = field(default_factory=list)
+
+    def __post_init__(self):
+        """Validate constraints."""
+        self.confidence_score = max(0.0, min(1.0, float(self.confidence_score)))
+        if isinstance(self.action_id, str):
+            self.action_id = ActionId(self.action_id)
+
+
+@dataclass
+class ExecutionResult:
+    """Result of executing a remediation action."""
+
+    action_id: str
+    success: bool
+    run_id: str
+    details: dict = field(default_factory=dict)
+    error: str = ""

healing_kit/models/events.py

@@ -0,0 +1,30 @@
+"""Event data models for failure detection and batching."""
+
+from dataclasses import dataclass, field
+from datetime import datetime
+
+
+@dataclass
+class FailureEvent:
+    """Raw failure event from a webhook or API poll."""
+
+    run_id: str
+    job_id: str
+    task_key: str
+    workspace_id: str
+    failure_time: datetime
+    error_message: str = ""
+    error_trace: str = ""
+
+
+@dataclass
+class RootCauseEvent:
+    """Deduplicated root-cause event after dependency graph analysis."""
+
+    root_run_id: str
+    root_job_id: str
+    root_task_key: str
+    error_hash: str
+    derived_failures: list[FailureEvent] = field(default_factory=list)
+    total_downstream_impact: int = 0
+    detected_at: datetime = field(default_factory=datetime.utcnow)

healing_kit/models/evidence.py

@@ -0,0 +1,83 @@
+"""Structured evidence package assembled by the Context Agent."""
+
+from dataclasses import dataclass, field
+from datetime import datetime
+from typing import Optional
+
+
+@dataclass
+class SparkLogEvidence:
+    """Evidence from Spark event logs."""
+
+    first_errors: list[str] = field(default_factory=list)
+    oom_events: list[str] = field(default_factory=list)
+    task_failure_traces: list[str] = field(default_factory=list)
+
+
+@dataclass
+class DriverLogEvidence:
+    """Evidence from driver stdout."""
+
+    schema_errors: list[str] = field(default_factory=list)
+    data_source_exceptions: list[str] = field(default_factory=list)
+    missing_table_errors: list[str] = field(default_factory=list)
+
+
+@dataclass
+class TaskMetricsEvidence:
+    """Evidence from task runtime metrics."""
+
+    spill_to_disk_bytes: int = 0
+    gc_overhead_pct: float = 0.0
+    shuffle_read_bytes: int = 0
+    shuffle_write_bytes: int = 0
+
+
+@dataclass
+class LineageEvidence:
+    """Evidence from Unity Catalog lineage."""
+
+    upstream_tables: list[dict] = field(default_factory=list)  # [{name, last_modified}]
+
+
+@dataclass
+class SchemaHistoryEvidence:
+    """Evidence from DESCRIBE HISTORY."""
+
+    recent_changes: list[dict] = field(default_factory=list)  # Last 5 schema changes
+
+
+@dataclass
+class GitContextEvidence:
+    """Evidence from job git metadata."""
+
+    last_commit_message: str = ""
+    last_commit_author: str = ""
+    last_commit_timestamp: Optional[datetime] = None
+
+
+@dataclass
+class ClusterEventsEvidence:
+    """Evidence from cluster events API."""
+
+    termination_reason: str = ""
+    autoscaling_events: list[str] = field(default_factory=list)
+    spot_preemptions: int = 0
+
+
+@dataclass
+class EvidencePackage:
+    """Complete structured evidence package for AI diagnosis."""
+
+    run_id: str
+    job_id: str
+    task_key: str
+    spark_event_logs: Optional[SparkLogEvidence] = None
+    driver_stdout: Optional[DriverLogEvidence] = None
+    task_metrics: Optional[TaskMetricsEvidence] = None
+    lineage: Optional[LineageEvidence] = None
+    schema_history: Optional[SchemaHistoryEvidence] = None
+    git_context: Optional[GitContextEvidence] = None
+    cluster_events: Optional[ClusterEventsEvidence] = None
+    missing_sources: list[str] = field(default_factory=list)
+    collected_at: datetime = field(default_factory=datetime.utcnow)

healing_kit/runtime/__init__.py

@@ -0,0 +1,6 @@
+"""Runtime orchestrators imported by the thin Databricks notebook drivers (#5).
+
+All healing logic lives here in the packaged wheel so it is the single, tested
+source of truth — the notebooks are now just entrypoints that pass `spark`/
+`dbutils` in. This replaces the previously inlined copies in the notebooks.
+"""