arthur-common 0.0.0.post0__py3-none-any.whl

arthur_common/aggregations/__init__.py

@@ -0,0 +1,2 @@
+from .aggregator import *  # noqa
+from .functions import *  # noqa

arthur_common/aggregations/aggregator.py

@@ -0,0 +1,304 @@
+import os
+import re
+from abc import ABC, abstractmethod
+from base64 import b64encode
+from typing import Any, Type, Union, cast
+
+import pandas as pd
+from datasketches import kll_floats_sketch
+from duckdb import DuckDBPyConnection
+
+from arthur_common.models.metrics import *
+
+
+class AggregationFunction(ABC):
+    FEATURE_FLAG_NAME: str | None = None
+
+    @staticmethod
+    @abstractmethod
+    def id() -> UUID:
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def display_name() -> str:
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def description() -> str:
+        raise NotImplementedError
+
+    @abstractmethod
+    def aggregation_type(self) -> Type[SketchMetric] | Type[NumericMetric]:
+        raise NotImplementedError
+
+    @staticmethod
+    @abstractmethod
+    def reported_aggregations() -> list[BaseReportedAggregation]:
+        """Returns the list of aggregations reported by the aggregate function."""
+        raise NotImplementedError
+
+    @staticmethod
+    def get_innermost_segmentation_columns(segmentation_cols: list[str]) -> list[str]:
+        """
+        Extracts the innermost column name for nested segmentation columns or
+        returns the top-level column name for non-nested segmentation columns.
+        """
+        for i, col in enumerate(segmentation_cols):
+            # extract the innermost column for escaped column names (e.g. '"nested.col"."name"')
+            # otherwise return the name since it's a top-level column
+            if col.startswith('"') and col.endswith('"'):
+                identifier = col[1:-1]
+                identifier_split_in_struct_fields = re.split(r'"\."', identifier)
+
+                # For nested columns, take just the innermost field name
+                # Otherwise for top-level columns, take the whole name
+                if len(identifier_split_in_struct_fields) > 1:
+                    innermost_field = identifier_split_in_struct_fields[-1]
+                    segmentation_cols[i] = innermost_field.replace('""', '"')
+                else:
+                    segmentation_cols[i] = identifier.replace('""', '"')
+            else:
+                segmentation_cols[i] = col
+
+        return segmentation_cols
+
+    @abstractmethod
+    def aggregate(
+        self,
+        ddb_conn: DuckDBPyConnection,
+        *args: Any,
+        **kwargs: Any,
+    ) -> Union[list[SketchMetric], list[NumericMetric]]:
+        raise NotImplementedError
+
+    @staticmethod
+    def string_to_dimension(name: str, value: str | None) -> Dimension:
+        if value is None:
+            value = "null"
+        return Dimension(name=name, value=str(value))
+
+    def is_feature_flag_enabled(self, feature_flag_name: str) -> bool:
+        if feature_flag_name is None:
+            value = os.getenv(self.FEATURE_FLAG_NAME, "false")
+        else:
+            value = os.getenv(feature_flag_name, "false")
+        return value.lower() in ("true", "1", "yes")
+
+
+class NumericAggregationFunction(AggregationFunction, ABC):
+    def aggregation_type(self) -> Type[NumericMetric]:
+        return NumericMetric
+
+    @abstractmethod
+    def aggregate(
+        self,
+        ddb_conn: DuckDBPyConnection,
+        *args: Any,
+        **kwargs: Any,
+    ) -> list[NumericMetric]:
+        raise NotImplementedError
+
+    @staticmethod
+    def group_query_results_to_numeric_metrics(
+        data: pd.DataFrame,
+        value_col: str,
+        dim_columns: list[str],
+        timestamp_col: str,
+    ) -> list[NumericTimeSeries]:
+        """
+        Convert a grouped dataframe with repeated dimensions to internal numeric metric definition.
+
+        At a high level, the query results are already grouped, however,
+        the order isn't guaranteed that groups are sequential (this requires an explicit ORDER BY on the source query.)
+        What this function does is group by the indicated dimensions list, and from each group extract the dimension values once.
+        From there, iterate over the group turning each data point to a *Point. At the end, this single instance of the group metrics
+        and the list of points (values) are merged to one *TimeSeries
+        """
+        if not dim_columns:
+            return [
+                NumericAggregationFunction._dimensionless_query_results_to_numeric_metrics(
+                    data,
+                    value_col,
+                    timestamp_col,
+                ),
+            ]
+
+        # get innermost column name for nested segmentation columns
+        dim_columns = AggregationFunction.get_innermost_segmentation_columns(
+            dim_columns,
+        )
+
+        calculated_metrics: list[NumericTimeSeries] = []
+        # make sure dropna is False or rows with "null" as a dimension value will be dropped
+        groups = data.groupby(dim_columns, dropna=False)
+        for _, group in groups:
+            dimensions: list[Dimension] = []
+            # Get the first row of the group to determine the group level dimensions
+            dims_row = group.iloc[0]
+            for dim in dim_columns:
+                d = AggregationFunction.string_to_dimension(
+                    name=dim,
+                    value=dims_row[dim],
+                )
+                dimensions.append(d)
+
+            values: list[NumericPoint] = []
+            for _, row in group.iterrows():
+                # Skip NaN values
+                if pd.notna(row[value_col]):
+                    values.append(
+                        NumericPoint(
+                            timestamp=row[timestamp_col], value=row[value_col]
+                        ),
+                    )
+            # Only add the series if it has values
+            if values:
+                calculated_metrics.append(
+                    NumericTimeSeries(values=values, dimensions=dimensions),
+                )
+
+        return calculated_metrics
+
+    @staticmethod
+    def _dimensionless_query_results_to_numeric_metrics(
+        data: pd.DataFrame,
+        value_col: str,
+        timestamp_col: str,
+    ) -> NumericTimeSeries:
+        """
+        Convert a dimensionless time / value series to internal numeric metric definition.
+        """
+        values: list[NumericPoint] = []
+        for _, row in data.iterrows():
+            # Skip NaN values
+            if pd.notna(row[value_col]):
+                values.append(
+                    NumericPoint(timestamp=row[timestamp_col], value=row[value_col]),
+                )
+        return NumericTimeSeries(values=values, dimensions=[])
+
+    @staticmethod
+    def series_to_metric(
+        metric_name: str,
+        series: list[NumericTimeSeries],
+    ) -> NumericMetric:
+        return NumericMetric(name=metric_name, numeric_series=series)
+
+
+class SketchAggregationFunction(AggregationFunction, ABC):
+    def aggregation_type(self) -> Type[SketchMetric]:
+        return SketchMetric
+
+    @abstractmethod
+    def aggregate(
+        self,
+        ddb_conn: DuckDBPyConnection,
+        *args: Any,
+        **kwargs: Any,
+    ) -> list[SketchMetric]:
+        raise NotImplementedError
+
+    @staticmethod
+    def group_query_results_to_sketch_metrics(
+        data: pd.DataFrame,
+        value_col: str,
+        dim_columns: list[str],
+        timestamp_col: str,
+    ) -> list[SketchTimeSeries]:
+        """
+        Convert a grouped dataframe with repeated dimensions to internal sketch metric definition.
+
+        For sketch data, what we're doing is grouping the raw row data into the dimensions we care about.
+        Within each group, we extract the dimensions once. Within this single dimension group,
+        we group the data into 5min intervals. Within each interval, the data point we care to sketch is added to the sketch.
+
+        """
+
+        calculated_metrics: list[SketchTimeSeries] = []
+
+        # get innermost column name for nested segmentation columns
+        dim_columns = AggregationFunction.get_innermost_segmentation_columns(
+            dim_columns,
+        )
+
+        if dim_columns:
+            # make sure dropna is False or rows with "null" as a dimension value will be dropped
+            # call _group_to_series for each grouped DF
+            groups = data.groupby(dim_columns, dropna=False)
+            for _, group in groups:
+                calculated_metrics.append(
+                    SketchAggregationFunction._group_to_series(
+                        group,
+                        timestamp_col,
+                        dim_columns,
+                        value_col,
+                    ),
+                )
+        else:
+            calculated_metrics.append(
+                SketchAggregationFunction._group_to_series(
+                    data,
+                    timestamp_col,
+                    dim_columns,
+                    value_col,
+                ),
+            )
+
+        return calculated_metrics
+
+    @staticmethod
+    def _group_to_series(
+        group: pd.DataFrame,
+        timestamp_col: str,
+        dim_columns: list[str],
+        value_col: str,
+    ) -> SketchTimeSeries:
+        def to_sketch(col: pd.Series) -> Optional[kll_floats_sketch]:
+            if not len(col):
+                return None
+            s = kll_floats_sketch()
+            for v in col.values:
+                s.update(v)
+            return s
+
+        dimensions: list[Dimension] = []
+        if dim_columns:
+            # Get the first row of the group to determine the group level dimensions
+            dims_row = group.iloc[0]
+            for dim in dim_columns:
+                d = AggregationFunction.string_to_dimension(
+                    name=dim, value=dims_row[dim]
+                )
+                dimensions.append(d)
+
+        values: list[SketchPoint] = []
+
+        # Group query results into 5min buckets
+        group[timestamp_col] = pd.to_datetime(group[timestamp_col])
+        group.set_index(timestamp_col, inplace=True)
+        # make sure dropna is False or rows with "null" as a dimension value will be dropped
+        time_bucketed_groups = group.groupby(pd.Grouper(freq="5min"), dropna=False)
+
+        for group_timestamp, time_bucket_group in time_bucketed_groups:
+            # Don't generate metrics on empty buckets
+            if time_bucket_group.empty:
+                continue
+            sketch = to_sketch(time_bucket_group[value_col])
+            if sketch is not None:
+                values.append(
+                    SketchPoint(
+                        timestamp=cast(pd.Timestamp, group_timestamp),
+                        value=b64encode(sketch.serialize()).decode(),
+                    ),
+                )
+
+        return SketchTimeSeries(values=values, dimensions=dimensions)
+
+    @staticmethod
+    def series_to_metric(
+        metric_name: str,
+        series: list[SketchTimeSeries],
+    ) -> SketchMetric:
+        return SketchMetric(name=metric_name, sketch_series=series)

arthur_common/aggregations/functions/README.md

@@ -0,0 +1,26 @@
+| Class Name                                                                   | UUID                                 | Name                                                                                    |
+|------------------------------------------------------------------------------|--------------------------------------|-----------------------------------------------------------------------------------------|
+| BinaryClassifierCountThresholdClassAggregationFunction                       | 00000000-0000-0000-0000-000000000020 | Binary Classification Count by Class - Probability Threshold                            |
+| BinaryClassifierCountByClassAggregationFunction                              | 00000000-0000-0000-0000-00000000001f | Binary Classification Count by Class - Class Label                                      |
+| BinaryClassifierProbabilityThresholdConfusionMatrixAggregationFunction       | 00000000-0000-0000-0000-00000000001e | Binary Classification Confusion Matrix - Probability Threshold                          |
+| BinaryClassifierStringLabelConfusionMatrixAggregationFunction                | 00000000-0000-0000-0000-00000000001d | Binary Classification Confusion Matrix - String Types                                   |
+| BinaryClassifierIntBoolConfusionMatrixAggregationFunction                    | 00000000-0000-0000-0000-00000000001c | Binary Classification Confusion Matrix - Int/Bool Types                                 |
+| NumericSumAggregationFunction                                                | 00000000-0000-0000-0000-00000000000f | Numeric Sum                                                                             |
+| MeanAbsoluteErrorAggregationFunction                                         | 00000000-0000-0000-0000-00000000000e | Mean Absolute Error                                                                     |
+| MeanSquaredErrorAggregationFunction                                          | 00000000-0000-0000-0000-000000000010 | Mean Squared Error                                                                      |
+| NumericSketchAggregationFunction                                             | 00000000-0000-0000-0000-00000000000d | Numeric Distribution                                                                    |
+| CategoricalCountAggregationFunction                                          | 00000000-0000-0000-0000-00000000000c | Category Count                                                                          |
+| InferenceNullCountAggregationFunction                                        | 00000000-0000-0000-0000-00000000000b | Null Value Count                                                                        |
+| InferenceCountAggregationFunction                                            | 00000000-0000-0000-0000-00000000000a | Inference Count                                                                         |
+| ShieldInferenceRuleLatencyAggregation                                        | 00000000-0000-0000-0000-000000000009 | Rule Latency Distribution                                                               |
+| ShieldInferenceRuleClaimFailCountAggregation                                 | 00000000-0000-0000-0000-000000000008 | Claim Count Distribution - Invalid Claims                                               |
+| ShieldInferenceRuleClaimPassCountAggregation                                 | 00000000-0000-0000-0000-000000000007 | Claim Count Distribution - Valid Claims                                                 |
+| ShieldInferenceRuleClaimCountAggregation                                     | 00000000-0000-0000-0000-000000000006 | Claim Count Distribution                                                                |
+| ShieldInferenceRulePIIDataScoreAggregation                                   | 00000000-0000-0000-0000-000000000005 | PII Score Distribution                                                                  |
+| ShieldInferenceRuleToxicityScoreAggregation                                  | 00000000-0000-0000-0000-000000000004 | Toxicity Distribution                                                                   |
+| ShieldInferenceHallucinationCountAggregation                                 | 00000000-0000-0000-0000-000000000003 | Hallucination Count                                                                     |
+| ShieldInferenceRuleCountAggregation                                          | 00000000-0000-0000-0000-000000000002 | Rule Result Count                                                                       |
+| ShieldInferencePassFailCountAggregation                                      | 00000000-0000-0000-0000-000000000001 | Inference Count                                                                         |
+| ShieldInferenceTokenCountAggregation                                         | 00000000-0000-0000-0000-000000000021 | Token Count                                                                             |
+| MulticlassClassifierCountByClassAggregationFunction                          | 64a338fb-6c99-4c40-ba39-81ab8baa8687 | Multiclass Classification Count by Class - Class Label                                  |
+| MulticlassClassifierStringLabelSingleClassConfusionMatrixAggregationFunction | dc728927-6928-4a3b-b174-8c1ec8b58d62 | Multiclass Classification Confusion Matrix Single Class - String Class Label Prediction |

arthur_common/aggregations/functions/__init__.py

@@ -0,0 +1,25 @@
+import importlib.util
+import inspect
+import os
+
+package_dir = os.path.dirname(__file__)
+
+# Peter 05/08/2024: This is some code I swiped from stackoverflow that iterated through the package directory here looking at .py files
+# It reads each file and imports the classes to add them to the "globals" which we can think of as importing into this namespace
+# By doing that, everything is exported and ready to be read as members of this `functions` package.
+# TLDR: this does what you would think `from . import *` does
+# Benefit here is any file with any class is added to the "exports", so nothing needs to be done after dropping a file in here
+for filename in os.listdir(package_dir):
+    if filename.endswith(".py") and filename != "__init__.py":
+        module_name = filename[:-3]  # Remove the .py extension to get the module name
+        module_path = os.path.join(package_dir, filename)
+
+        spec = importlib.util.spec_from_file_location(module_name, module_path)
+        if not spec:
+            continue
+        module = importlib.util.module_from_spec(spec)
+        if spec.loader:
+            spec.loader.exec_module(module)
+        for name, value in module.__dict__.items():
+            if inspect.isclass(value) and not name.startswith("_"):
+                globals()[name] = value