datahub-ai-classifier 0.1.0__tar.gz

datahub_ai_classifier-0.1.0/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2026 DataHub Classifier contributors
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

datahub_ai_classifier-0.1.0/PKG-INFO

@@ -0,0 +1,76 @@
+Metadata-Version: 2.4
+Name: datahub-ai-classifier
+Version: 0.1.0
+Summary: Customer-taxonomy LLM classification and tagging for DataHub
+License-Expression: MIT
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: acryl-datahub[datahub-rest]>=1.7.0
+Requires-Dist: openai>=1.40.0
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Dynamic: license-file
+
+# DataHub Classifier
+
+Classify DataHub fields against a customer-provided tag list using an
+OpenAI-compatible LLM, then apply those tags.
+
+Install:
+
+```bash
+pip install datahub-ai-classifier
+```
+
+Example recipe: `datahub_classifier.example.yml`.
+
+Both plugins must be in the same recipe, with the same `taxonomy` map.
+
+```yaml
+source:
+  type: postgres
+  config:
+    classification:
+      enabled: true
+      sample_size: 10
+      max_workers: 1
+      classifiers:
+        - type: datahub_classifier.classifier:AIPIIClassifier
+          config:
+            llm_model: gpt-4.1-mini
+            llm_base_url: https://api.openai.com/v1
+            llm_api_key: "<your-api-key>"
+            llm_timeout_seconds: 60
+            llm_max_retries: 0
+            confidence_threshold: 0.6
+            max_samples: 8
+            taxonomy: &taxonomy
+              C1: >
+                Restricted identifiers. Use for national IDs and tax numbers
+                such as SSN values like 123-45-6789 stored in columns named
+                ssn.
+              C2: >
+                Confidential personal and payment data. Use for a person's
+                full_name, email, phone, home_address, and date_of_birth, and
+                for order_total money amounts.
+              C3: >
+                Internal operational data. Use for customer_id, order_id,
+                loyalty_tier, currency, order status, shipped_at, and
+                created_at timestamps.
+
+transformers:
+  - type: datahub-classifier
+    config:
+      taxonomy: *taxonomy
+```
+
+Missing tags are created. Existing tags are not overwritten.
+Classifier alone leaves glossary terms and does not create tags.
+
+## Development
+
+```bash
+pip install -e '.[dev]'
+pytest
+```

datahub_ai_classifier-0.1.0/README.md

@@ -0,0 +1,62 @@
+# DataHub Classifier
+
+Classify DataHub fields against a customer-provided tag list using an
+OpenAI-compatible LLM, then apply those tags.
+
+Install:
+
+```bash
+pip install datahub-ai-classifier
+```
+
+Example recipe: `datahub_classifier.example.yml`.
+
+Both plugins must be in the same recipe, with the same `taxonomy` map.
+
+```yaml
+source:
+  type: postgres
+  config:
+    classification:
+      enabled: true
+      sample_size: 10
+      max_workers: 1
+      classifiers:
+        - type: datahub_classifier.classifier:AIPIIClassifier
+          config:
+            llm_model: gpt-4.1-mini
+            llm_base_url: https://api.openai.com/v1
+            llm_api_key: "<your-api-key>"
+            llm_timeout_seconds: 60
+            llm_max_retries: 0
+            confidence_threshold: 0.6
+            max_samples: 8
+            taxonomy: &taxonomy
+              C1: >
+                Restricted identifiers. Use for national IDs and tax numbers
+                such as SSN values like 123-45-6789 stored in columns named
+                ssn.
+              C2: >
+                Confidential personal and payment data. Use for a person's
+                full_name, email, phone, home_address, and date_of_birth, and
+                for order_total money amounts.
+              C3: >
+                Internal operational data. Use for customer_id, order_id,
+                loyalty_tier, currency, order status, shipped_at, and
+                created_at timestamps.
+
+transformers:
+  - type: datahub-classifier
+    config:
+      taxonomy: *taxonomy
+```
+
+Missing tags are created. Existing tags are not overwritten.
+Classifier alone leaves glossary terms and does not create tags.
+
+## Development
+
+```bash
+pip install -e '.[dev]'
+pytest
+```

datahub_ai_classifier-0.1.0/pyproject.toml

@@ -0,0 +1,24 @@
+[build-system]
+requires = ["setuptools>=68", "wheel"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "datahub-ai-classifier"
+version = "0.1.0"
+description = "Customer-taxonomy LLM classification and tagging for DataHub"
+readme = "README.md"
+requires-python = ">=3.10"
+license = "MIT"
+dependencies = [
+  "acryl-datahub[datahub-rest]>=1.7.0",
+  "openai>=1.40.0",
+]
+
+[project.entry-points."datahub.ingestion.transformer.plugins"]
+datahub-classifier = "datahub_classifier.tagger:PIITagTransformer"
+
+[project.optional-dependencies]
+dev = ["pytest"]
+
+[tool.setuptools.packages.find]
+where = ["src"]

datahub_ai_classifier-0.1.0/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

datahub_ai_classifier-0.1.0/src/datahub_ai_classifier.egg-info/PKG-INFO

@@ -0,0 +1,76 @@
+Metadata-Version: 2.4
+Name: datahub-ai-classifier
+Version: 0.1.0
+Summary: Customer-taxonomy LLM classification and tagging for DataHub
+License-Expression: MIT
+Requires-Python: >=3.10
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: acryl-datahub[datahub-rest]>=1.7.0
+Requires-Dist: openai>=1.40.0
+Provides-Extra: dev
+Requires-Dist: pytest; extra == "dev"
+Dynamic: license-file
+
+# DataHub Classifier
+
+Classify DataHub fields against a customer-provided tag list using an
+OpenAI-compatible LLM, then apply those tags.
+
+Install:
+
+```bash
+pip install datahub-ai-classifier
+```
+
+Example recipe: `datahub_classifier.example.yml`.
+
+Both plugins must be in the same recipe, with the same `taxonomy` map.
+
+```yaml
+source:
+  type: postgres
+  config:
+    classification:
+      enabled: true
+      sample_size: 10
+      max_workers: 1
+      classifiers:
+        - type: datahub_classifier.classifier:AIPIIClassifier
+          config:
+            llm_model: gpt-4.1-mini
+            llm_base_url: https://api.openai.com/v1
+            llm_api_key: "<your-api-key>"
+            llm_timeout_seconds: 60
+            llm_max_retries: 0
+            confidence_threshold: 0.6
+            max_samples: 8
+            taxonomy: &taxonomy
+              C1: >
+                Restricted identifiers. Use for national IDs and tax numbers
+                such as SSN values like 123-45-6789 stored in columns named
+                ssn.
+              C2: >
+                Confidential personal and payment data. Use for a person's
+                full_name, email, phone, home_address, and date_of_birth, and
+                for order_total money amounts.
+              C3: >
+                Internal operational data. Use for customer_id, order_id,
+                loyalty_tier, currency, order status, shipped_at, and
+                created_at timestamps.
+
+transformers:
+  - type: datahub-classifier
+    config:
+      taxonomy: *taxonomy
+```
+
+Missing tags are created. Existing tags are not overwritten.
+Classifier alone leaves glossary terms and does not create tags.
+
+## Development
+
+```bash
+pip install -e '.[dev]'
+pytest
+```

datahub_ai_classifier-0.1.0/src/datahub_ai_classifier.egg-info/SOURCES.txt

@@ -0,0 +1,20 @@
+LICENSE
+README.md
+pyproject.toml
+src/datahub_ai_classifier.egg-info/PKG-INFO
+src/datahub_ai_classifier.egg-info/SOURCES.txt
+src/datahub_ai_classifier.egg-info/dependency_links.txt
+src/datahub_ai_classifier.egg-info/entry_points.txt
+src/datahub_ai_classifier.egg-info/requires.txt
+src/datahub_ai_classifier.egg-info/top_level.txt
+src/datahub_classifier/__init__.py
+src/datahub_classifier/classifier.py
+src/datahub_classifier/llm_client.py
+src/datahub_classifier/tagger.py
+src/datahub_classifier/taxonomy.py
+tests/test_classifier.py
+tests/test_classifier_llm_client.py
+tests/test_datahub_classification.py
+tests/test_plugin_names.py
+tests/test_tagger.py
+tests/test_taxonomy.py

datahub_ai_classifier-0.1.0/src/datahub_ai_classifier.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

datahub_ai_classifier-0.1.0/src/datahub_ai_classifier.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[datahub.ingestion.transformer.plugins]
+datahub-classifier = datahub_classifier.tagger:PIITagTransformer

datahub_ai_classifier-0.1.0/src/datahub_ai_classifier.egg-info/requires.txt

@@ -0,0 +1,5 @@
+acryl-datahub[datahub-rest]>=1.7.0
+openai>=1.40.0
+
+[dev]
+pytest

datahub_ai_classifier-0.1.0/src/datahub_ai_classifier.egg-info/top_level.txt

@@ -0,0 +1 @@
+datahub_classifier

datahub_ai_classifier-0.1.0/src/datahub_classifier/__init__.py

@@ -0,0 +1,5 @@
+from datahub_classifier.classifier import AIPIIClassifier
+from datahub_classifier.tagger import PIITagTransformer
+
+__all__ = ["AIPIIClassifier", "PIITagTransformer"]
+__version__ = "0.1.0"

datahub_ai_classifier-0.1.0/src/datahub_classifier/classifier.py

@@ -0,0 +1,240 @@
+from __future__ import annotations
+
+import logging
+import math
+from typing import Any, Dict, List
+
+from datahub.ingestion.glossary.classification_types import (
+    ColumnInfo,
+    DebugInfo,
+    InfotypeProposal,
+)
+from datahub.ingestion.glossary.classifier import Classifier
+
+from datahub_classifier.llm_client import chat_json, get_llm_client
+from datahub_classifier.taxonomy import parse_taxonomy
+
+logger = logging.getLogger(__name__)
+
+SYSTEM_PROMPT = """You classify database columns against a customer-provided taxonomy.
+Return JSON only:
+{
+  "columns": [
+    {
+      "name": "<column name>",
+      "infotype": "<one of the allowed type ids>",
+      "confidence": 0.0-1.0
+    }
+  ]
+}
+Rules:
+- Use only allowed type ids.
+- Prefer None when evidence is weak.
+- Use the type description to decide.
+- Broad ids such as C1 are valid if they are in the list.
+- Use column name, datatype, and sample values.
+- Do not invent sample values.
+"""
+
+
+def _require_str(config: Dict[str, Any], key: str, *, allow_blank: bool = False) -> str:
+    value = config.get(key)
+    if not isinstance(value, str):
+        raise ValueError(f"{key} is required")
+    if not allow_blank and not value.strip():
+        raise ValueError(f"{key} is required")
+    return value
+
+
+class AIPIIClassifier(Classifier):
+    """DataHub classifier that proposes taxonomy matches via an LLM."""
+
+    def __init__(self, config: Dict[str, Any]) -> None:
+        self.config = config or {}
+        self.llm_model = _require_str(self.config, "llm_model")
+        self.llm_api_key = _require_str(self.config, "llm_api_key", allow_blank=True)
+        self.llm_base_url = _require_str(self.config, "llm_base_url")
+        raw_timeout = self.config.get("llm_timeout_seconds", 60)
+        raw_retries = self.config.get("llm_max_retries", 0)
+        if isinstance(raw_timeout, bool):
+            raise ValueError("llm_timeout_seconds must be a number greater than 0")
+        if isinstance(raw_retries, bool) or not isinstance(raw_retries, int):
+            raise ValueError("llm_max_retries must be a non-negative integer")
+        try:
+            self.llm_timeout_seconds = float(raw_timeout)
+        except (TypeError, ValueError) as exc:
+            raise ValueError(
+                "llm_timeout_seconds must be a number greater than 0"
+            ) from exc
+        if not math.isfinite(self.llm_timeout_seconds) or self.llm_timeout_seconds <= 0:
+            raise ValueError("llm_timeout_seconds must be a number greater than 0")
+        if raw_retries < 0:
+            raise ValueError("llm_max_retries must be a non-negative integer")
+        self.llm_max_retries = raw_retries
+        self.taxonomy = parse_taxonomy(self.config.get("taxonomy"))
+        self.allowed_infotypes = tuple(
+            [entry.id for entry in self.taxonomy] + ["None"]
+        )
+        raw_threshold = self.config.get("confidence_threshold", 0.6)
+        raw_max_samples = self.config.get("max_samples", 8)
+        if isinstance(raw_threshold, bool):
+            raise ValueError("confidence_threshold must be a number between 0 and 1")
+        if isinstance(raw_max_samples, bool) or not isinstance(raw_max_samples, int):
+            raise ValueError("max_samples must be a non-negative integer")
+        try:
+            self.confidence_threshold = float(raw_threshold)
+        except (TypeError, ValueError) as exc:
+            raise ValueError(
+                "confidence_threshold must be a number between 0 and 1"
+            ) from exc
+        self.max_samples = raw_max_samples
+        if not math.isfinite(self.confidence_threshold) or not (
+            0 <= self.confidence_threshold <= 1
+        ):
+            raise ValueError("confidence_threshold must be a number between 0 and 1")
+        if self.max_samples < 0:
+            raise ValueError("max_samples must be a non-negative integer")
+        self._client = None
+
+    @classmethod
+    def create(cls, config_dict: Dict[str, Any]) -> "AIPIIClassifier":
+        return cls(config_dict or {})
+
+    def _client_or_create(self):
+        if self._client is None:
+            self._client = get_llm_client(
+                self.llm_api_key,
+                self.llm_base_url,
+                timeout=self.llm_timeout_seconds,
+                max_retries=self.llm_max_retries,
+            )
+        return self._client
+
+    def classify(self, columns: List[ColumnInfo]) -> List[ColumnInfo]:
+        if not columns:
+            return columns
+
+        payload = []
+        for col in columns:
+            samples = [str(v) for v in (col.values or [])[: self.max_samples]]
+            # DataHub mixin historically populates "DataType"; Metadata reads "Datatype".
+            datatype = col.metadata.datatype or (col.metadata.meta_info or {}).get(
+                "DataType"
+            )
+            payload.append(
+                {
+                    "name": col.metadata.name,
+                    "datatype": datatype,
+                    "description": col.metadata.description,
+                    "dataset": col.metadata.dataset_name,
+                    "samples": samples,
+                }
+            )
+
+        catalog = [
+            {"id": entry.id, "description": entry.description}
+            for entry in self.taxonomy
+        ]
+        user_prompt = (
+            "Allowed type ids: "
+            + ", ".join(self.allowed_infotypes)
+            + "\n\nTaxonomy:\n"
+            + str(catalog)
+            + "\n\nClassify these columns:\n"
+            + str(payload)
+        )
+
+        result = chat_json(
+            self._client_or_create(),
+            system=SYSTEM_PROMPT,
+            user=user_prompt,
+            model=self.llm_model,
+        )
+
+        if not isinstance(result, dict) or not isinstance(result.get("columns"), list):
+            raise ValueError("LLM PII response must contain a columns list")
+
+        known_names = {
+            col.metadata.name for col in columns if col.metadata.name is not None
+        }
+        by_name = {}
+        duplicate_names = set()
+        for item in result["columns"]:
+            if not isinstance(item, dict):
+                logger.warning("Ignoring malformed LLM PII prediction: not an object")
+                continue
+
+            name = item.get("name")
+            infotype = item.get("infotype")
+            confidence = item.get("confidence")
+
+            if not isinstance(name, str) or name not in known_names:
+                logger.warning(
+                    "Ignoring LLM PII prediction for unknown column %r", name
+                )
+                continue
+            if not isinstance(infotype, str) or infotype not in self.allowed_infotypes:
+                logger.warning(
+                    "Ignoring LLM PII prediction for %s: invalid infotype %r",
+                    name,
+                    infotype,
+                )
+                continue
+            if (
+                isinstance(confidence, bool)
+                or not isinstance(confidence, (int, float))
+                or not math.isfinite(float(confidence))
+                or not 0 <= float(confidence) <= 1
+            ):
+                logger.warning(
+                    "Ignoring LLM PII prediction for %s: invalid confidence %r",
+                    name,
+                    confidence,
+                )
+                continue
+            if name in by_name or name in duplicate_names:
+                logger.warning(
+                    "Ignoring duplicate LLM PII predictions for column %s", name
+                )
+                by_name.pop(name, None)
+                duplicate_names.add(name)
+                continue
+
+            by_name[name] = (infotype, float(confidence))
+
+        for col in columns:
+            name = col.metadata.name
+            item = by_name.get(name) if name else None
+            if not item:
+                continue
+
+            infotype, confidence = item
+
+            if infotype == "None":
+                continue
+            if confidence < self.confidence_threshold:
+                logger.debug(
+                    "Skipping %s.%s: confidence %.2f < %.2f",
+                    col.metadata.dataset_name,
+                    name,
+                    confidence,
+                    self.confidence_threshold,
+                )
+                continue
+
+            col.infotype_proposals = [
+                InfotypeProposal(
+                    infotype=infotype,
+                    confidence_level=confidence,
+                    debug_info=DebugInfo(values=confidence, name=confidence),
+                )
+            ]
+            logger.info(
+                "PII proposal %s.%s -> %s (%.2f)",
+                col.metadata.dataset_name,
+                name,
+                infotype,
+                confidence,
+            )
+
+        return columns

datahub_ai_classifier-0.1.0/src/datahub_classifier/llm_client.py

@@ -0,0 +1,93 @@
+from __future__ import annotations
+
+import json
+import logging
+import re
+from typing import Any, Dict, Optional
+
+from openai import OpenAI
+
+logger = logging.getLogger(__name__)
+
+
+def _json_object(content: str) -> Dict[str, Any]:
+    parsed = json.loads(content)
+    if not isinstance(parsed, dict):
+        raise ValueError("LLM response must be a JSON object")
+    return parsed
+
+
+def get_llm_client(
+    api_key: str,
+    base_url: str,
+    *,
+    timeout: float,
+    max_retries: int,
+) -> OpenAI:
+    return OpenAI(
+        api_key=api_key or "not-required",
+        base_url=base_url,
+        timeout=timeout,
+        max_retries=max_retries,
+    )
+
+
+def chat_json(
+    client: OpenAI,
+    *,
+    system: str,
+    user: str,
+    model: str,
+    temperature: Optional[float] = 0.1,
+) -> Dict[str, Any]:
+    messages = [
+        {"role": "system", "content": system},
+        {"role": "user", "content": user},
+    ]
+    kwargs: Dict[str, Any] = {
+        "model": model,
+        "response_format": {"type": "json_object"},
+        "messages": messages,
+    }
+    if temperature is not None:
+        kwargs["temperature"] = temperature
+
+    while True:
+        try:
+            response = client.chat.completions.create(**kwargs)
+            break
+        except Exception as exc:
+            message = str(exc).lower()
+            if "temperature" in kwargs and "temperature" in message:
+                logger.warning(
+                    "Model %s rejected temperature=%s; retrying without it",
+                    model,
+                    kwargs["temperature"],
+                )
+                kwargs.pop("temperature")
+                continue
+            if "response_format" in kwargs and (
+                "response_format" in message
+                or "response format" in message
+                or "json_object" in message
+            ):
+                logger.warning(
+                    "Model %s rejected JSON response format; retrying without it",
+                    model,
+                )
+                kwargs.pop("response_format")
+                continue
+            raise
+
+    content = response.choices[0].message.content or "{}"
+    try:
+        return _json_object(content)
+    except json.JSONDecodeError:
+        match = re.search(r"\{.*\}", content, re.DOTALL)
+        if not match:
+            logger.warning(
+                "Failed to parse LLM JSON response: %s",
+                content[:500],
+            )
+            return {}
+        return _json_object(match.group(0))