openmetadata-ingestion 1.9.12.0__py3-none-any.whl → 1.10.0.0__py3-none-any.whl

metadata/ingestion/source/messaging/kafka/connection.py

@@ -16,8 +16,9 @@ from copy import deepcopy
 from dataclasses import dataclass
 from typing import Optional, Union
 
+from confluent_kafka import DeserializingConsumer
 from confluent_kafka.admin import AdminClient, KafkaException
-from confluent_kafka.avro import AvroConsumer
+from confluent_kafka.schema_registry.avro import AvroDeserializer
 from confluent_kafka.schema_registry.schema_registry_client import SchemaRegistryClient
 
 from metadata.generated.schema.entity.automations.workflow import (
@@ -71,6 +72,7 @@ def get_connection(
     """
     consumer_config = deepcopy(connection.consumerConfig) or {}
     schema_registry_config = deepcopy(connection.schemaRegistryConfig) or {}
+
     if connection.saslUsername or connection.saslPassword or connection.saslMechanism:
         if connection.saslUsername:
             consumer_config["sasl.username"] = connection.saslUsername
@@ -98,19 +100,26 @@ def get_connection(
 
     schema_registry_client = None
     consumer_client = None
+
     if connection.schemaRegistryURL:
         schema_registry_config["url"] = str(connection.schemaRegistryURL)
         schema_registry_client = SchemaRegistryClient(schema_registry_config)
+
         consumer_config["bootstrap.servers"] = connection.bootstrapServers
         if "group.id" not in consumer_config:
             consumer_config["group.id"] = "openmetadata-consumer"
         if "auto.offset.reset" not in consumer_config:
             consumer_config["auto.offset.reset"] = "largest"
         consumer_config["enable.auto.commit"] = False
-        logger.debug(f"Using Kafka consumer config: {consumer_config}")
-        consumer_client = AvroConsumer(
-            consumer_config, schema_registry=schema_registry_client
+
+        avro_deserializer = AvroDeserializer(
+            schema_registry_client=schema_registry_client
         )
+        consumer_config["value.deserializer"] = avro_deserializer
+
+        logger.debug(f"Using Kafka consumer config: {consumer_config}")
+
+        consumer_client = DeserializingConsumer(consumer_config)
 
     return KafkaClient(
         admin_client=admin_client,

metadata/ingestion/source/metadata/amundsen/client.py

@@ -59,16 +59,16 @@ class Neo4jHelper:
         """
         Create a Neo4j connection to Database
         """
-        trust = (
-            neo4j.TRUST_SYSTEM_CA_SIGNED_CERTIFICATES
-            if self.conf.neo4j_validate_ssl
-            else neo4j.TRUST_ALL_CERTIFICATES
-        )
+        if self.conf.neo4j_validate_ssl:
+            trusted_certificates = neo4j.TrustSystemCAs()
+        else:
+            trusted_certificates = neo4j.TrustAll()
+
         return GraphDatabase.driver(
             self.graph_url,
             auth=(self.conf.username, self.conf.password),
             encrypted=self.conf.neo4j_encrypted,
-            trust=trust,
+            trusted_certificates=trusted_certificates,
         )
 
     def _execute_query(self, transaction: Any, query: str) -> Any:

metadata/ingestion/source/mlmodel/mlmodel_service.py

@@ -12,12 +12,13 @@
 Base class for ingesting mlmodel services
 """
 from abc import ABC, abstractmethod
-from typing import Any, Iterable, List, Optional, Set
+from typing import Any, Iterable, List, Optional, Set, Tuple
 
 from pydantic import Field
 from typing_extensions import Annotated
 
 from metadata.generated.schema.api.data.createMlModel import CreateMlModelRequest
+from metadata.generated.schema.api.lineage.addLineage import AddLineageRequest
 from metadata.generated.schema.entity.data.mlmodel import (
     MlFeature,
     MlHyperParameter,
@@ -39,6 +40,7 @@ from metadata.ingestion.api.models import Either
 from metadata.ingestion.api.steps import Source
 from metadata.ingestion.api.topology_runner import TopologyRunnerMixin
 from metadata.ingestion.models.delete_entity import DeleteEntity
+from metadata.ingestion.models.ometa_lineage import OMetaLineageRequest
 from metadata.ingestion.models.topology import (
     NodeStage,
     ServiceTopology,
@@ -92,6 +94,12 @@ class MlModelServiceTopology(ServiceTopology):
                 consumer=["mlmodel_service"],
                 use_cache=True,
             ),
+            NodeStage(
+                type_=AddLineageRequest,
+                processor="yield_mlmodel_lineage",
+                consumer=["mlmodel_service"],
+                nullable=True,
+            ),
         ],
     )
 
@@ -208,3 +216,69 @@ class MlModelServiceSource(TopologyRunnerMixin, Source, ABC):
 
     def prepare(self):
         """By default, nothing to prepare"""
+
+    def get_db_service_prefixes(self) -> List[str]:
+        """
+        Get the list of db service prefixes
+        """
+        return (
+            self.source_config.lineageInformation.dbServicePrefixes or []
+            if hasattr(self.source_config, "lineageInformation")
+            and self.source_config.lineageInformation
+            else []
+        )
+
+    def parse_db_service_prefix(
+        self, db_service_prefix: Optional[str]
+    ) -> Tuple[Optional[str], Optional[str], Optional[str], Optional[str]]:
+        """
+        Parse the db service prefix
+        Returns:
+            Tuple[Optional[str], Optional[str], Optional[str], Optional[str]]: service, database, schema, table
+        """
+        prefix_parts = (db_service_prefix or "").split(".")
+        return tuple(prefix_parts + ([None] * (4 - len(prefix_parts))))
+
+    def yield_mlmodel_lineage_details(
+        self, mlmodel_details: Any, db_service_prefix: Optional[str] = None
+    ) -> Iterable[Either[AddLineageRequest]]:
+        """
+        Get lineage between MLModel and source tables.
+        To be implemented by sources that support lineage.
+        """
+
+    def yield_mlmodel_lineage(
+        self, mlmodel_details: Any
+    ) -> Iterable[Either[OMetaLineageRequest]]:
+        """
+        Yields lineage if config is enabled.
+        We will look for the data in all the services we have informed.
+        """
+        db_service_prefixes = self.get_db_service_prefixes()
+        for db_service_prefix in db_service_prefixes or [None]:
+            for lineage in (
+                self.yield_mlmodel_lineage_details(mlmodel_details, db_service_prefix)
+                or []
+            ):
+                yield from self.yield_lineage_request(lineage)
+
+    def yield_lineage_request(
+        self, lineage: Optional[Either[AddLineageRequest]] = None
+    ) -> Iterable[Either[OMetaLineageRequest]]:
+        """
+        Method to yield lineage request
+        """
+        if lineage:
+            if lineage.right is not None:
+                yield Either(
+                    right=OMetaLineageRequest(
+                        lineage_request=lineage.right,
+                        override_lineage=(
+                            self.source_config.overrideLineage
+                            if hasattr(self.source_config, "overrideLineage")
+                            else False
+                        ),
+                    )
+                )
+            else:
+                yield lineage

metadata/ingestion/source/search/elasticsearch/connection.py

@@ -176,7 +176,7 @@ def get_connection(connection: ElasticsearchConnection) -> Elasticsearch:
 
     return Elasticsearch(
         str(connection.hostPort),
-        http_auth=basic_auth,
+        basic_auth=basic_auth,
         api_key=api_key,
         ssl_context=ssl_context,
         **connection.connectionArguments.root,

metadata/pii/algorithms/classifiers.py

@@ -19,6 +19,7 @@ from typing import (
     Dict,
     Generic,
     Hashable,
+    Iterable,
     Mapping,
     Optional,
     Sequence,
@@ -29,7 +30,7 @@ from typing import (
 
 from presidio_analyzer import AnalyzerEngine
 
-from metadata.generated.schema.entity.data.table import DataType
+from metadata.generated.schema.entity.data.table import Column, DataType
 from metadata.pii.algorithms.column_patterns import get_pii_column_name_patterns
 from metadata.pii.algorithms.feature_extraction import (
     extract_pii_from_column_names,
@@ -48,6 +49,7 @@ from metadata.pii.algorithms.presidio_utils import (
     set_presidio_logger_level,
 )
 from metadata.pii.algorithms.tags import PIISensitivityTag, PIITag
+from metadata.pii.tag_analyzer import TagAnalyzer
 
 T = TypeVar("T", bound=Hashable)
 
@@ -72,6 +74,11 @@ class ColumnClassifier(ABC, Generic[T]):
         higher scores indicate a higher likelihood of the class for the given inputs.
         """
 
+    def classify(self, column: Column, sample_data: Sequence[Any]) -> Mapping[T, float]:
+        return self.predict_scores(
+            sample_data, column_name=column.name.root, column_data_type=column.dataType
+        )
+
 
 @final
 class HeuristicPIIClassifier(ColumnClassifier[PIITag]):
@@ -178,3 +185,58 @@ class PIISensitiveClassifier(ColumnClassifier[PIISensitivityTag]):
                 results[tag] /= counts[tag]
 
         return results
+
+
+@final
+class TagClassifier(ColumnClassifier[str]):
+    """
+    Heuristic PII Column Classifier
+    """
+
+    def __init__(
+        self,
+        *,
+        tag_analyzers: Iterable[TagAnalyzer],
+        column_name_contribution: float = 0.5,
+        score_cutoff: float = 0.1,
+        relative_cardinality_cutoff: float = 0.01,
+    ):
+        set_presidio_logger_level()
+
+        self._analyzers = tag_analyzers
+
+        self._column_name_contribution = column_name_contribution
+        self._score_cutoff = score_cutoff
+        self._relative_cardinality_cutoff = relative_cardinality_cutoff
+
+    def predict_scores(
+        self,
+        sample_data: Sequence[Any],
+        column_name: Optional[str] = None,
+        column_data_type: Optional[DataType] = None,
+    ) -> Mapping[str, float]:
+        str_values = preprocess_values(sample_data)
+
+        if not str_values:
+            return {}
+
+        # Relative cardinality test
+        unique_values = set(str_values)
+        if len(unique_values) / len(str_values) < self._relative_cardinality_cutoff:
+            return {}
+
+        results: dict[str, float] = defaultdict(float)
+        for analyzer in self._analyzers:
+            content_score = analyzer.analyze_content(values=str_values)
+            column_score = 0.0
+            if column_name is not None:
+                column_score = analyzer.analyze_column()
+                column_score *= max(column_score, self._column_name_contribution)
+
+            total_score = content_score + column_score
+            if total_score > self._score_cutoff:
+                results[str(analyzer.tag.fullyQualifiedName)] = (
+                    content_score + column_score
+                )
+
+        return results

metadata/pii/algorithms/presidio_recognizer_factory.py

@@ -0,0 +1,302 @@
+#  Copyright 2025 Collate
+#  Licensed under the Collate Community License, Version 1.0 (the "License");
+#  you may not use this file except in compliance with the License.
+#  You may obtain a copy of the License at
+#  https://github.com/open-metadata/OpenMetadata/blob/main/ingestion/LICENSE
+#  Unless required by applicable law or agreed to in writing, software
+#  distributed under the License is distributed on an "AS IS" BASIS,
+#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+#  See the License for the specific language governing permissions and
+#  limitations under the License.
+"""
+Factory for creating Presidio recognizers from OpenMetadata recognizer configurations.
+"""
+import re
+from typing import Dict, List, Optional, cast
+
+from presidio_analyzer import EntityRecognizer
+from presidio_analyzer import Pattern as PresidioPattern
+from presidio_analyzer import PatternRecognizer as PresidioPatternRecognizer
+from presidio_analyzer import predefined_recognizers
+
+from metadata.generated.schema.entity.classification.tag import Tag
+from metadata.generated.schema.type.contextRecognizer import ContextRecognizer
+from metadata.generated.schema.type.customRecognizer import CustomRecognizer
+from metadata.generated.schema.type.denyListRecognizer import DenyListRecognizer
+from metadata.generated.schema.type.patternRecognizer import PatternRecognizer
+from metadata.generated.schema.type.predefinedRecognizer import PredefinedRecognizer
+from metadata.generated.schema.type.recognizer import Recognizer
+from metadata.generated.schema.type.recognizers.regexFlags import RegexFlags
+from metadata.pii.algorithms.presidio_utils import apply_confidence_threshold
+from metadata.utils.logger import pii_logger
+
+logger = pii_logger()
+
+
+class PresidioRecognizerFactory:
+    """Factory for creating Presidio recognizers from OpenMetadata configurations."""
+
+    @staticmethod
+    def create_recognizer(recognizer_config: Recognizer) -> Optional[EntityRecognizer]:
+        """
+        Create a Presidio recognizer from an OpenMetadata recognizer configuration.
+
+        Args:
+            recognizer_config: The recognizer configuration from OpenMetadata
+            tag_name: The name of the tag this recognizer belongs to
+
+        Returns:
+            A Presidio EntityRecognizer or None if creation fails
+        """
+        if not recognizer_config.enabled:
+            return None
+
+        config = recognizer_config.recognizerConfig.root
+
+        if isinstance(config, PatternRecognizer):
+            recognizer = PresidioRecognizerFactory._create_pattern_recognizer(
+                config, recognizer_config
+            )
+        elif isinstance(config, DenyListRecognizer):
+            recognizer = PresidioRecognizerFactory._create_deny_list_recognizer(
+                config, recognizer_config
+            )
+        elif isinstance(config, ContextRecognizer):
+            recognizer = PresidioRecognizerFactory._create_context_recognizer(
+                config, recognizer_config
+            )
+        elif isinstance(config, CustomRecognizer):
+            recognizer = PresidioRecognizerFactory._create_custom_recognizer(
+                config, recognizer_config
+            )
+        elif isinstance(
+            config, PredefinedRecognizer
+        ):  # pyright: ignore[reportUnnecessaryIsInstance]
+            recognizer = PresidioRecognizerFactory._create_predefined_recognizer(
+                config, recognizer_config
+            )
+        else:
+            logger.warning(f"Unknown recognizer type for {recognizer_config.name}")
+            return None
+
+        if recognizer and (threshold := recognizer_config.confidenceThreshold):
+            patch_analyze = apply_confidence_threshold(threshold)
+            recognizer = patch_analyze(recognizer)
+
+        return recognizer
+
+    @staticmethod
+    def _get_regex_flags(flags: Optional[RegexFlags]) -> Optional[int]:
+        if flags is None:
+            return re.IGNORECASE | re.DOTALL | re.MULTILINE
+
+        re_flags = 0
+        if flags.ignoreCase:
+            re_flags |= re.IGNORECASE
+        if flags.dotAll:
+            re_flags |= re.DOTALL
+        if flags.multiline:
+            re_flags |= re.MULTILINE
+
+        return re_flags
+
+    @staticmethod
+    def _create_pattern_recognizer(
+        config: PatternRecognizer,
+        recognizer_config: Recognizer,
+    ) -> PresidioPatternRecognizer:
+        """Create a pattern-based recognizer."""
+        patterns: List[PresidioPattern] = []
+        for pattern_config in config.patterns:
+            patterns.append(
+                PresidioPattern(
+                    name=pattern_config.name,
+                    regex=pattern_config.regex,
+                    score=pattern_config.score or 0.8,
+                )
+            )
+
+        return PresidioPatternRecognizer(
+            supported_entity=config.supportedEntity.value,
+            patterns=patterns,
+            name=recognizer_config.name.root,
+            supported_language=config.supportedLanguage,
+            global_regex_flags=PresidioRecognizerFactory._get_regex_flags(
+                config.regexFlags
+            ),
+        )
+
+    @staticmethod
+    def _create_deny_list_recognizer(
+        config: DenyListRecognizer, recognizer_config: Recognizer
+    ) -> PresidioPatternRecognizer:
+        """Create a deny list recognizer using patterns."""
+        patterns: List[PresidioPattern] = []
+        for value in config.denyList:
+            # Escape special regex characters in the value
+            escaped_value = re.escape(value)
+
+            patterns.append(
+                PresidioPattern(
+                    name=f"deny_{value}",
+                    regex=escaped_value,
+                    score=0.9,  # High confidence for exact matches
+                )
+            )
+
+        return PresidioPatternRecognizer(
+            supported_entity=config.supportedEntity.value,
+            patterns=patterns,
+            name=recognizer_config.name.root,
+            supported_language=config.supportedLanguage,
+            global_regex_flags=PresidioRecognizerFactory._get_regex_flags(
+                config.regexFlags
+            ),
+        )
+
+    @staticmethod
+    def _create_context_recognizer(
+        config: ContextRecognizer, recognizer_config: Recognizer
+    ) -> PresidioPatternRecognizer:
+        """Create a context-aware recognizer."""
+        # For context recognizers, we can use a pattern recognizer with context words
+        # or implement a custom recognizer that uses NLP
+        context_patterns: List[PresidioPattern] = []
+
+        # Create patterns that look for context words near potential entities
+        for context_word in config.contextWords:
+            # Pattern to match words near context words
+            pattern = f"(?i)(?:{context_word})\\s+\\w+|\\w+\\s+(?:{context_word})"
+            context_patterns.append(
+                PresidioPattern(
+                    name=f"context_{context_word}",
+                    regex=pattern,
+                    score=(config.minScore + config.maxScore) / 2
+                    if config.minScore and config.maxScore
+                    else 0.6,
+                )
+            )
+
+        return PresidioPatternRecognizer(
+            supported_entity=config.supportedEntity.value,
+            patterns=context_patterns,
+            name=recognizer_config.name.root,
+            supported_language=config.supportedLanguage,
+        )
+
+    @staticmethod
+    def _create_custom_recognizer(
+        config: CustomRecognizer,  # pyright: ignore[reportUnusedParameter]
+        recognizer_config: Recognizer,
+    ) -> Optional[EntityRecognizer]:
+        """
+        Create a custom recognizer with user-defined logic.
+
+        For security reasons, we don't execute arbitrary Python code.
+        Instead, we could support a limited set of custom behaviors
+        or require custom recognizers to be implemented as plugins.
+        """
+        logger.warning(
+            f"Custom recognizer {recognizer_config.name} requires implementation. "
+            + "Consider using pattern, deny_list, or context recognizers instead."
+        )
+        return None
+
+    @staticmethod
+    def _create_predefined_recognizer(
+        config: PredefinedRecognizer,
+        recognizer: Recognizer,  # pyright: ignore[reportUnusedParameter]
+    ) -> Optional[EntityRecognizer]:
+        """Create a custom recognizer with user-defined logic."""
+        try:
+            predefined_class = getattr(predefined_recognizers, config.name.value)
+        except AttributeError:
+            logger.error(f"Recognizer {config.name} not found")
+            return None
+
+        args = {}
+        if supported_language := config.supportedLanguage:
+            args["supported_language"] = supported_language
+        if context := config.context:
+            args["context"] = context
+
+        return predefined_class(**args)
+
+    @staticmethod
+    def create_recognizers_for_tag(tag: Tag) -> List[EntityRecognizer]:
+        """
+        Create all enabled recognizers for a given tag.
+
+        Args:
+            tag: The tag containing recognizer configurations
+
+        Returns:
+            List of Presidio EntityRecognizer instances
+        """
+        recognizers: List[EntityRecognizer] = []
+
+        if not tag.autoClassificationEnabled or not tag.recognizers:
+            return recognizers
+
+        for recognizer_config in tag.recognizers:
+            recognizer = PresidioRecognizerFactory.create_recognizer(recognizer_config)
+            if recognizer:
+                recognizers.append(recognizer)
+                logger.info(
+                    f"Created recognizer {recognizer_config.name} for tag {tag.name}"
+                )
+
+        return recognizers
+
+
+class RecognizerRegistry:
+    """Registry for managing custom recognizers from OpenMetadata."""
+
+    def __init__(self):
+        self.recognizers: Dict[str, List[EntityRecognizer]] = {}
+        self.tag_priority: Dict[str, int] = {}
+        self.tag_confidence_threshold: Dict[str, float] = {}
+
+    def register_tag_recognizers(self, tag: Tag) -> None:
+        """
+        Register all recognizers for a tag.
+
+        Args:
+            tag: The tag with recognizer configurations
+        """
+        if not tag.autoClassificationEnabled:
+            return
+
+        tag_fqn = cast(str, tag.fullyQualifiedName)
+        self.recognizers[
+            tag_fqn
+        ] = PresidioRecognizerFactory.create_recognizers_for_tag(tag)
+        self.tag_priority[tag_fqn] = tag.autoClassificationPriority or 50
+
+        # Calculate minimum confidence from all recognizers
+        min_confidence = 1.0
+        for recognizer_config in tag.recognizers or []:
+            if recognizer_config.confidenceThreshold:
+                min_confidence = min(
+                    min_confidence, recognizer_config.confidenceThreshold
+                )
+        self.tag_confidence_threshold[tag_fqn] = min_confidence
+
+    def get_recognizers_for_tag(self, tag_fqn: str) -> List[EntityRecognizer]:
+        """Get all recognizers registered for a tag."""
+        return self.recognizers.get(tag_fqn, [])
+
+    def get_all_recognizers(self) -> List[EntityRecognizer]:
+        """Get all registered recognizers across all tags."""
+        all_recognizers: list[EntityRecognizer] = []
+        for recognizers in self.recognizers.values():
+            all_recognizers.extend(recognizers)
+        return all_recognizers
+
+    def get_tag_priority(self, tag_fqn: str) -> int:
+        """Get the priority for a tag (used in conflict resolution)."""
+        return self.tag_priority.get(tag_fqn, 50)
+
+    def get_tag_confidence_threshold(self, tag_fqn: str) -> float:
+        """Get the minimum confidence threshold for a tag."""
+        return self.tag_confidence_threshold.get(tag_fqn, 0.6)

metadata/pii/algorithms/presidio_utils.py

@@ -13,16 +13,17 @@ Utilities for working with the Presidio Library.
 """
 import inspect
 import logging
-from typing import Iterable, Type, Union
+from typing import Callable, Iterable, List, Type, Union
 
 import spacy
 from presidio_analyzer import (
     AnalyzerEngine,
     EntityRecognizer,
     PatternRecognizer,
+    RecognizerResult,
     predefined_recognizers,
 )
-from presidio_analyzer.nlp_engine import SpacyNlpEngine
+from presidio_analyzer.nlp_engine import NlpArtifacts, SpacyNlpEngine
 from spacy.cli.download import download  # pyright: ignore[reportUnknownVariableType]
 
 from metadata.pii.constants import PRESIDIO_LOGGER, SPACY_EN_MODEL, SUPPORTED_LANG
@@ -31,6 +32,17 @@ from metadata.utils.logger import pii_logger
 logger = pii_logger()
 
 
+def load_nlp_engine(
+    model_name: str = SPACY_EN_MODEL, supported_language: str = SUPPORTED_LANG
+) -> SpacyNlpEngine:
+    _load_spacy_model(model_name)
+    model = {
+        "lang_code": supported_language,
+        "model_name": model_name,
+    }
+    return SpacyNlpEngine(models=[model])
+
+
 def build_analyzer_engine(
     model_name: str = SPACY_EN_MODEL,
 ) -> AnalyzerEngine:
@@ -39,14 +51,9 @@ def build_analyzer_engine(
 
     If the model is not found locally, it will be downloaded.
     """
-    _load_spacy_model(model_name)
-
-    model = {
-        "lang_code": SUPPORTED_LANG,
-        "model_name": model_name,
-    }
-
-    nlp_engine = SpacyNlpEngine(models=[model])
+    nlp_engine = load_nlp_engine(
+        model_name=model_name, supported_language=SUPPORTED_LANG
+    )
     analyzer_engine = AnalyzerEngine(
         nlp_engine=nlp_engine, supported_languages=[SUPPORTED_LANG]
     )
@@ -109,3 +116,24 @@ def _get_all_pattern_recognizers() -> Iterable[EntityRecognizer]:
         elif cls == predefined_recognizers.PhoneRecognizer:
             # Not a pattern recognizer, but pretty much the same
             yield predefined_recognizers.PhoneRecognizer()
+
+
+def apply_confidence_threshold(
+    threshold: float,
+) -> Callable[[EntityRecognizer], EntityRecognizer]:
+    def decorate_entity_recognizer(recognizer: EntityRecognizer) -> EntityRecognizer:
+        original_analyze = recognizer.analyze
+
+        def analyze(
+            instance: EntityRecognizer,  # pyright: ignore[reportUnusedParameter]
+            text: str,
+            entities: List[str],
+            nlp_artifacts: NlpArtifacts,
+        ) -> List[RecognizerResult]:
+            results = original_analyze(text, entities, nlp_artifacts)
+            return [result for result in results if result.score >= threshold]
+
+        recognizer.analyze = analyze.__get__(recognizer, type(recognizer))
+        return recognizer
+
+    return decorate_entity_recognizer

metadata/pii/algorithms/utils.py

@@ -36,3 +36,7 @@ def get_top_classes(scores: Mapping[T, float], n: int, threshold: float) -> Sequ
     sorted_scores = sorted(scores.items(), key=lambda x: x[1], reverse=True)
     top_classes = [key for key, score in sorted_scores if score >= threshold]
     return top_classes[:n]
+
+
+def build_reason(tag_name: str, score: float) -> str:
+    return f"Chose {tag_name} with a classification score of {score:.2f}"

metadata/pii/base_processor.py

@@ -100,7 +100,8 @@ class AutoClassificationProcessor(Processor, ABC):
 
         column_tags = []
 
-        for idx, column in enumerate(record.table.columns):
+        for idx, column_name in enumerate(record.sample_data.data.columns):
+            column = next(c for c in record.table.columns if c.name == column_name)
             try:
                 tags = self.create_column_tag_labels(
                     column=column,