arthur-common 2.4.1__py3-none-any.whl → 2.4.3__py3-none-any.whl

arthur_common/aggregations/aggregator.py

@@ -1,3 +1,4 @@
+import re
 from abc import ABC, abstractmethod
 from base64 import b64encode
 from typing import Any, Type, Union
@@ -35,6 +36,31 @@ class AggregationFunction(ABC):
         """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,
@@ -89,6 +115,11 @@ class NumericAggregationFunction(AggregationFunction, ABC):
                 ),
             ]
 
+        # 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)
@@ -168,12 +199,21 @@ class SketchAggregationFunction(AggregationFunction, ABC):
         """
 
         calculated_metrics: list[SketchTimeSeries] = []
+
+        # get innermost column name for nested segmentation columns
+        dim_columns = AggregationFunction.get_innermost_segmentation_columns(
+            dim_columns,
+        )
+
         # 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:
             calculated_metrics.append(
                 SketchAggregationFunction._group_to_series(
-                    group, timestamp_col, dim_columns, value_col
+                    group,
+                    timestamp_col,
+                    dim_columns,
+                    value_col,
                 ),
             )
 

arthur_common/aggregations/functions/agentic_aggregations.py

@@ -1,5 +1,6 @@
 import json
 import logging
+from datetime import datetime
 from typing import Annotated, Any
 from uuid import UUID
 
@@ -27,6 +28,46 @@ TOOL_SCORE_NO_TOOL_VALUE = 2
 logger = logging.getLogger(__name__)
 
 
+def root_span_in_time_buckets(
+    ddb_conn: DuckDBPyConnection, dataset: DatasetReference
+) -> pd.DataFrame:
+    return ddb_conn.sql(
+        f"""
+            SELECT
+                time_bucket(INTERVAL '5 minutes', start_time) as ts,
+                root_spans
+            FROM {dataset.dataset_table_name}
+            WHERE root_spans IS NOT NULL AND length(root_spans) > 0
+            ORDER BY ts DESC;
+            """,
+    ).df()
+
+
+def span_parser(span_to_parse: str | dict[str, Any]) -> dict[str, Any]:
+    if isinstance(span_to_parse, str):
+        return json.loads(span_to_parse)  # type: ignore[no-any-return]
+
+    return span_to_parse
+
+
+def extract_agent_name_from_span(span: dict[str, Any]) -> str | None:
+    try:
+        raw_data = span.get("raw_data", {})
+        if isinstance(raw_data, str):
+            raw_data = json.loads(raw_data)
+
+        # Try to get agent name from the span's name field
+        agent_name = raw_data.get("name", "unknown")
+        if agent_name != "unknown":
+            return str(agent_name)
+    except (json.JSONDecodeError, KeyError, TypeError):
+        logger.error(
+            f"Error parsing attributes from span (span_id: {span.get('span_id')}) in trace {span.get('trace_id')}",
+        )
+
+    return None
+
+
 # TODO: create TypedDict for span
 def extract_spans_with_metrics_and_agents(
     root_spans: list[str | dict[str, Any]],
@@ -41,39 +82,26 @@ def extract_spans_with_metrics_and_agents(
     # TODO: Improve function so it won't modify variable outside of its scope
     def traverse_spans(
         spans: list[str | dict[str, Any]],
-        current_agent_name: str = "unknown",
+        current_agent: str = "unknown",
     ) -> None:
         for span_to_parse in spans:
-            if isinstance(span_to_parse, str):
-                parsed_span = json.loads(span_to_parse)
-            else:
-                parsed_span = span_to_parse
+            parsed_span = span_parser(span_to_parse)
 
             # Update current agent name if this span is an AGENT
             if parsed_span.get("span_kind") == "AGENT":
-                try:
-                    raw_data = parsed_span.get("raw_data", {})
-                    if isinstance(raw_data, str):
-                        raw_data = json.loads(raw_data)
-
-                    # Try to get agent name from the span's name field
-                    agent_name = raw_data.get("name", "unknown")
-                    if agent_name != "unknown":
-                        current_agent_name = agent_name
-                except (json.JSONDecodeError, KeyError, TypeError):
-                    logger.error(
-                        f"Error parsing attributes from span (span_id: {parsed_span.get('span_id')}) in trace {parsed_span.get('trace_id')}",
-                    )
+                agent_name = extract_agent_name_from_span(parsed_span)
+                if agent_name:
+                    current_agent = agent_name
 
             # Check if this span has metrics
             if parsed_span.get("metric_results", []):
                 spans_with_metrics_and_agents.append(
-                    (parsed_span, current_agent_name),
+                    (parsed_span, current_agent),
                 )
 
             # Recursively traverse children with the current agent name
             if children_span := parsed_span.get("children", []):
-                traverse_spans(children_span, current_agent_name)
+                traverse_spans(children_span, current_agent)
 
     traverse_spans(root_spans)
     return spans_with_metrics_and_agents
@@ -153,16 +181,7 @@ class AgenticMetricsOverTimeAggregation(SketchAggregationFunction):
         ],
     ) -> list[SketchMetric]:
         # Query traces by timestamp
-        results = ddb_conn.sql(
-            f"""
-            SELECT
-                time_bucket(INTERVAL '5 minutes', start_time) as ts,
-                root_spans
-            FROM {dataset.dataset_table_name}
-            WHERE root_spans IS NOT NULL AND length(root_spans) > 0
-            ORDER BY ts DESC;
-            """,
-        ).df()
+        results = root_span_in_time_buckets(ddb_conn, dataset)
 
         # Process traces and extract spans with metrics
         tool_selection_data = []
@@ -409,17 +428,7 @@ class AgenticRelevancePassFailCountAggregation(NumericAggregationFunction):
             ),
         ],
     ) -> list[NumericMetric]:
-        # Query traces by timestamp
-        results = ddb_conn.sql(
-            f"""
-            SELECT
-                time_bucket(INTERVAL '5 minutes', start_time) as ts,
-                root_spans
-            FROM {dataset.dataset_table_name}
-            WHERE root_spans IS NOT NULL AND length(root_spans) > 0
-            ORDER BY ts DESC;
-            """,
-        ).df()
+        results = root_span_in_time_buckets(ddb_conn, dataset)
 
         # Process traces and extract spans with metrics
         processed_data = []
@@ -534,17 +543,7 @@ class AgenticToolPassFailCountAggregation(NumericAggregationFunction):
             ),
         ],
     ) -> list[NumericMetric]:
-        # Query traces by timestamp
-        results = ddb_conn.sql(
-            f"""
-            SELECT
-                time_bucket(INTERVAL '5 minutes', start_time) as ts,
-                root_spans
-            FROM {dataset.dataset_table_name}
-            WHERE root_spans IS NOT NULL AND length(root_spans) > 0
-            ORDER BY ts DESC;
-            """,
-        ).df()
+        results = root_span_in_time_buckets(ddb_conn, dataset)
 
         # Process traces and extract spans with metrics
         processed_data = []
@@ -713,16 +712,7 @@ class AgenticLLMCallCountAggregation(NumericAggregationFunction):
             ),
         ],
     ) -> list[NumericMetric]:
-        results = ddb_conn.sql(
-            f"""
-            SELECT
-                time_bucket(INTERVAL '5 minutes', start_time) as ts,
-                root_spans
-            FROM {dataset.dataset_table_name}
-            WHERE root_spans IS NOT NULL AND length(root_spans) > 0
-            ORDER BY ts DESC;
-            """,
-        ).df()
+        results = root_span_in_time_buckets(ddb_conn, dataset)
 
         # Process traces and count LLM spans
         llm_call_counts = {}
@@ -738,10 +728,7 @@ class AgenticLLMCallCountAggregation(NumericAggregationFunction):
             def count_llm_spans(spans: list[str | dict[str, Any]]) -> int:
                 count = 0
                 for span_to_parse in spans:
-                    if isinstance(span_to_parse, str):
-                        span = json.loads(span_to_parse)
-                    else:
-                        span = span_to_parse
+                    span = span_parser(span_to_parse)
 
                     # Check if this span is an LLM span
                     if span.get("span_kind") == "LLM":
@@ -813,16 +800,7 @@ class AgenticToolSelectionAndUsageByAgentAggregation(NumericAggregationFunction)
         ],
     ) -> list[NumericMetric]:
         # Query traces by timestamp
-        results = ddb_conn.sql(
-            f"""
-            SELECT
-                time_bucket(INTERVAL '5 minutes', start_time) as ts,
-                root_spans
-            FROM {dataset.dataset_table_name}
-            WHERE root_spans IS NOT NULL AND length(root_spans) > 0
-            ORDER BY ts DESC;
-            """,
-        ).df()
+        results = root_span_in_time_buckets(ddb_conn, dataset)
 
         # Process traces and extract spans with metrics
         processed_data = []
@@ -899,3 +877,185 @@ class AgenticToolSelectionAndUsageByAgentAggregation(NumericAggregationFunction)
         )
         metric = self.series_to_metric(self.METRIC_NAME, series)
         return [metric]
+
+
+class AgenticTraceLatencyAggregation(SketchAggregationFunction):
+    METRIC_NAME = "trace_latency"
+
+    @staticmethod
+    def id() -> UUID:
+        return UUID("00000000-0000-0000-0000-000000000039")
+
+    @staticmethod
+    def display_name() -> str:
+        return "Trace Latency"
+
+    @staticmethod
+    def description() -> str:
+        return "Aggregation that reports the latency of the agentic trace in ms."
+
+    @staticmethod
+    def reported_aggregations() -> list[BaseReportedAggregation]:
+        return [
+            BaseReportedAggregation(
+                metric_name=AgenticTraceLatencyAggregation.METRIC_NAME,
+                description=AgenticTraceLatencyAggregation.description(),
+            ),
+        ]
+
+    def aggregate(
+        self,
+        ddb_conn: DuckDBPyConnection,
+        dataset: Annotated[
+            DatasetReference,
+            MetricDatasetParameterAnnotation(
+                friendly_name="Dataset",
+                description="The agentic trace dataset containing traces with nested spans.",
+                model_problem_type=ModelProblemType.AGENTIC_TRACE,
+            ),
+        ],
+    ) -> list[SketchMetric]:
+        # Query traces by timestamp and calculate latency directly in SQL
+        results = ddb_conn.sql(
+            f"""
+            SELECT
+                time_bucket(INTERVAL '5 minutes', start_time) as ts,
+                CAST(EXTRACT(EPOCH FROM (end_time - start_time)) * 1000 AS INTEGER) as latency_ms
+            FROM {dataset.dataset_table_name}
+            WHERE start_time IS NOT NULL 
+                AND end_time IS NOT NULL
+                AND end_time > start_time
+            ORDER BY ts DESC;
+            """,
+        ).df()
+
+        if results.empty:
+            return []
+
+        df = results
+        # Create a single time series without grouping dimensions
+        # Since we have no dimensions to group by, we create one time series for all data
+        series = [self._group_to_series(df, "ts", [], "latency_ms")]
+        metric = self.series_to_metric(self.METRIC_NAME, series)
+        return [metric]
+
+
+class AgenticSpanLatencyAggregation(SketchAggregationFunction):
+    METRIC_NAME = "span_latency"
+
+    @staticmethod
+    def id() -> UUID:
+        return UUID("00000000-0000-0000-0000-000000000040")
+
+    @staticmethod
+    def display_name() -> str:
+        return "Span Latency"
+
+    @staticmethod
+    def description() -> str:
+        return "Aggregation that reports the latency of the agentic span in ms."
+
+    @staticmethod
+    def reported_aggregations() -> list[BaseReportedAggregation]:
+        return [
+            BaseReportedAggregation(
+                metric_name=AgenticSpanLatencyAggregation.METRIC_NAME,
+                description=AgenticSpanLatencyAggregation.description(),
+            ),
+        ]
+
+    def aggregate(
+        self,
+        ddb_conn: DuckDBPyConnection,
+        dataset: Annotated[
+            DatasetReference,
+            MetricDatasetParameterAnnotation(
+                friendly_name="Dataset",
+                description="The agentic trace dataset containing traces with nested spans.",
+                model_problem_type=ModelProblemType.AGENTIC_TRACE,
+            ),
+        ],
+    ) -> list[SketchMetric]:
+        results = root_span_in_time_buckets(ddb_conn, dataset)
+
+        latency_data = []
+        for _, row in results.iterrows():
+            ts = row["ts"]
+            root_spans = row["root_spans"]
+
+            # Parse root_spans if it's a string
+            if isinstance(root_spans, str):
+                root_spans = json.loads(root_spans)
+
+            # Extract all spans with their timing data
+            spans_with_timing = self._extract_spans_with_timing(root_spans)
+
+            for span_data in spans_with_timing:
+                span, current_agent, latency_ms = span_data
+                span_kind = span.get("span_kind", "unknown")
+
+                if latency_ms is not None and latency_ms > 0:
+                    latency_data.append(
+                        {
+                            "ts": ts,
+                            "latency_ms": latency_ms,
+                            "span_kind": span_kind,
+                            "agent_name": current_agent,
+                        }
+                    )
+
+        if not latency_data:
+            return []
+
+        # Convert to DataFrame and create sketch metrics
+        df = pd.DataFrame(latency_data)
+        series = self.group_query_results_to_sketch_metrics(
+            df,
+            "latency_ms",
+            ["span_kind", "agent_name"],
+            "ts",
+        )
+        metric = self.series_to_metric(self.METRIC_NAME, series)
+        return [metric]
+
+    def _extract_spans_with_timing(
+        self, spans: list[str | dict[str, Any]], current_agent: str = "unknown"
+    ) -> list[tuple[dict[str, Any], str, int | None]]:
+        """Recursively extract spans with calculated latency in milliseconds"""
+        spans_with_timing = []
+
+        for span_to_parse in spans:
+            span = span_parser(span_to_parse)
+
+            # Update current agent name if this span is an AGENT
+            if span.get("span_kind") == "AGENT":
+                agent_name = extract_agent_name_from_span(span)
+                if agent_name:
+                    current_agent = agent_name
+
+            # Calculate latency if both start_time and end_time exist
+            start_time = span.get("start_time")
+            end_time = span.get("end_time")
+            latency_ms = None
+
+            if start_time and end_time:
+                try:
+                    # Parse ISO format timestamps and calculate latency in milliseconds
+                    # Assume same timezone for start and end time, specific TZ not important for latency calculation
+                    start_dt = datetime.fromisoformat(start_time)
+                    end_dt = datetime.fromisoformat(end_time)
+                    latency_ms = int((end_dt - start_dt).total_seconds() * 1000)
+                except (ValueError, TypeError) as e:
+                    logger.warning(
+                        f"Error calculating latency for span {span.get('span_id')}: {e}"
+                    )
+
+            spans_with_timing.append((span, current_agent, latency_ms))
+
+            # Recursively process children
+            if children := span.get("children", []):
+                spans_with_timing.extend(
+                    self._extract_spans_with_timing(children, current_agent)
+                )
+
+        return spans_with_timing

arthur_common/aggregations/functions/categorical_count.py

@@ -18,8 +18,10 @@ from arthur_common.models.schema_definitions import (
     ScalarType,
     ScopeSchemaTag,
 )
-
-from arthur_common.tools.duckdb_data_loader import unescape_identifier, escape_str_literal
+from arthur_common.tools.duckdb_data_loader import (
+    escape_str_literal,
+    unescape_identifier,
+)
 
 
 class CategoricalCountAggregationFunction(NumericAggregationFunction):
@@ -103,7 +105,9 @@ class CategoricalCountAggregationFunction(NumericAggregationFunction):
                 group by ts, category
         """
         segmentation_cols = [] if not segmentation_cols else segmentation_cols
-        categorical_col_name_unescaped = escape_str_literal(unescape_identifier(categorical_col))
+        categorical_col_name_unescaped = escape_str_literal(
+            unescape_identifier(categorical_col),
+        )
 
         # build query components with segmentation columns
         all_select_clause_cols = [
@@ -125,11 +129,10 @@ class CategoricalCountAggregationFunction(NumericAggregationFunction):
 
         results = ddb_conn.sql(count_query).df()
 
-        unescaped_segmentation_cols = [unescape_identifier(seg_col) for seg_col in segmentation_cols]
         series = self.group_query_results_to_numeric_metrics(
             results,
             "count",
-            unescaped_segmentation_cols + extra_dims,
+            segmentation_cols + extra_dims,
             timestamp_col="ts",
         )
         metric = self.series_to_metric(self.METRIC_NAME, series)

arthur_common/aggregations/functions/confusion_matrix.py

@@ -20,8 +20,10 @@ from arthur_common.models.schema_definitions import (
     ScalarType,
     ScopeSchemaTag,
 )
-
-from arthur_common.tools.duckdb_data_loader import unescape_identifier, escape_str_literal
+from arthur_common.tools.duckdb_data_loader import (
+    escape_str_literal,
+    unescape_identifier,
+)
 
 
 class ConfusionMatrixAggregationFunction(NumericAggregationFunction):
@@ -97,7 +99,9 @@ class ConfusionMatrixAggregationFunction(NumericAggregationFunction):
                 ORDER BY ts
         """
         segmentation_cols = [] if not segmentation_cols else segmentation_cols
-        unescaped_prediction_col_name = escape_str_literal(unescape_identifier(prediction_col))
+        unescaped_prediction_col_name = escape_str_literal(
+            unescape_identifier(prediction_col),
+        )
 
         # build query components with segmentation columns
         first_subquery_select_cols = [
@@ -131,29 +135,28 @@ class ConfusionMatrixAggregationFunction(NumericAggregationFunction):
 
         results = ddb_conn.sql(confusion_matrix_query).df()
 
-        unescaped_segmentation_cols = [unescape_identifier(seg_col) for seg_col in segmentation_cols]
         tp = self.group_query_results_to_numeric_metrics(
             results,
             "true_positive_count",
-            dim_columns=unescaped_segmentation_cols + extra_dims,
+            dim_columns=segmentation_cols + extra_dims,
             timestamp_col="ts",
         )
         fp = self.group_query_results_to_numeric_metrics(
             results,
             "false_positive_count",
-            dim_columns=unescaped_segmentation_cols + extra_dims,
+            dim_columns=segmentation_cols + extra_dims,
             timestamp_col="ts",
         )
         fn = self.group_query_results_to_numeric_metrics(
             results,
             "false_negative_count",
-            dim_columns=unescaped_segmentation_cols + extra_dims,
+            dim_columns=segmentation_cols + extra_dims,
             timestamp_col="ts",
         )
         tn = self.group_query_results_to_numeric_metrics(
             results,
             "true_negative_count",
-            dim_columns=unescaped_segmentation_cols + extra_dims,
+            dim_columns=segmentation_cols + extra_dims,
             timestamp_col="ts",
         )
         tp_metric = self.series_to_metric(self.TRUE_POSITIVE_METRIC_NAME, tp)

arthur_common/aggregations/functions/inference_count.py

@@ -18,7 +18,6 @@ from arthur_common.models.schema_definitions import (
     ScalarType,
     ScopeSchemaTag,
 )
-from arthur_common.tools.duckdb_data_loader import unescape_identifier
 
 
 class InferenceCountAggregationFunction(NumericAggregationFunction):
@@ -102,11 +101,11 @@ class InferenceCountAggregationFunction(NumericAggregationFunction):
         """
 
         results = ddb_conn.sql(count_query).df()
-        unescaped_segmentation_cols = [unescape_identifier(seg_col) for seg_col in segmentation_cols]
+
         series = self.group_query_results_to_numeric_metrics(
             results,
             "count",
-            unescaped_segmentation_cols,
+            segmentation_cols,
             "ts",
         )
         metric = self.series_to_metric(self.METRIC_NAME, series)

arthur_common/aggregations/functions/inference_count_by_class.py

@@ -20,7 +20,6 @@ from arthur_common.models.schema_definitions import (
     ScalarType,
     ScopeSchemaTag,
 )
-from arthur_common.tools.duckdb_data_loader import unescape_identifier
 
 
 class BinaryClassifierCountByClassAggregationFunction(NumericAggregationFunction):
@@ -132,11 +131,10 @@ class BinaryClassifierCountByClassAggregationFunction(NumericAggregationFunction
 
         result = ddb_conn.sql(query).df()
 
-        unescaped_segmentation_cols = [unescape_identifier(seg_col) for seg_col in segmentation_cols]
         series = self.group_query_results_to_numeric_metrics(
             result,
             "count",
-            unescaped_segmentation_cols + extra_dims,
+            segmentation_cols + extra_dims,
             "ts",
         )
         metric = self.series_to_metric(self._metric_name(), series)
@@ -278,11 +276,10 @@ class BinaryClassifierCountThresholdClassAggregationFunction(
 
         result = ddb_conn.sql(query).df()
 
-        unescaped_segmentation_cols = [unescape_identifier(seg_col) for seg_col in segmentation_cols]
         series = self.group_query_results_to_numeric_metrics(
             result,
             "count",
-            unescaped_segmentation_cols + extra_dims,
+            segmentation_cols + extra_dims,
             "ts",
         )
         metric = self.series_to_metric(self._metric_name(), series)

arthur_common/aggregations/functions/inference_null_count.py

@@ -114,16 +114,17 @@ class InferenceNullCountAggregationFunction(NumericAggregationFunction):
 
         results = ddb_conn.sql(count_query).df()
 
-        unescaped_segmentation_cols = [unescape_identifier(seg_col) for seg_col in segmentation_cols]
         series = self.group_query_results_to_numeric_metrics(
             results,
             "count",
-            unescaped_segmentation_cols,
+            segmentation_cols,
             "ts",
         )
         # preserve dimension that identifies the name of the nullable column used for the aggregation
         for point in series:
-            point.dimensions.append(Dimension(name="column_name", value=unescape_identifier(nullable_col)))
+            point.dimensions.append(
+                Dimension(name="column_name", value=unescape_identifier(nullable_col)),
+            )
 
         metric = self.series_to_metric(self.METRIC_NAME, series)
         return [metric]

arthur_common/aggregations/functions/mean_absolute_error.py

@@ -19,7 +19,6 @@ from arthur_common.models.schema_definitions import (
     ScalarType,
     ScopeSchemaTag,
 )
-from arthur_common.tools.duckdb_data_loader import unescape_identifier
 
 
 class MeanAbsoluteErrorAggregationFunction(NumericAggregationFunction):
@@ -139,17 +138,17 @@ class MeanAbsoluteErrorAggregationFunction(NumericAggregationFunction):
         """
 
         results = ddb_conn.sql(mae_query).df()
-        unescaped_segmentation_cols = [unescape_identifier(seg_col) for seg_col in segmentation_cols]
+
         count_series = self.group_query_results_to_numeric_metrics(
             results,
             "count",
-            unescaped_segmentation_cols,
+            segmentation_cols,
             "ts",
         )
         absolute_error_series = self.group_query_results_to_numeric_metrics(
             results,
             "ae",
-            unescaped_segmentation_cols,
+            segmentation_cols,
             "ts",
         )
 

arthur_common/aggregations/functions/mean_squared_error.py

@@ -19,7 +19,6 @@ from arthur_common.models.schema_definitions import (
     ScalarType,
     ScopeSchemaTag,
 )
-from arthur_common.tools.duckdb_data_loader import unescape_identifier
 
 
 class MeanSquaredErrorAggregationFunction(NumericAggregationFunction):
@@ -139,17 +138,17 @@ class MeanSquaredErrorAggregationFunction(NumericAggregationFunction):
         """
 
         results = ddb_conn.sql(mse_query).df()
-        unescaped_segmentation_cols = [unescape_identifier(seg_col) for seg_col in segmentation_cols]
+
         count_series = self.group_query_results_to_numeric_metrics(
             results,
             "count",
-            unescaped_segmentation_cols,
+            segmentation_cols,
             "ts",
         )
         squared_error_series = self.group_query_results_to_numeric_metrics(
             results,
             "squared_error",
-            unescaped_segmentation_cols,
+            segmentation_cols,
             "ts",
         )
 

arthur_common/aggregations/functions/multiclass_confusion_matrix.py

@@ -20,8 +20,7 @@ from arthur_common.models.schema_definitions import (
     ScalarType,
     ScopeSchemaTag,
 )
-
-from arthur_common.tools.duckdb_data_loader import escape_str_literal, unescape_identifier
+from arthur_common.tools.duckdb_data_loader import escape_str_literal
 
 
 class MulticlassClassifierStringLabelSingleClassConfusionMatrixAggregationFunction(
@@ -244,30 +243,29 @@ class MulticlassClassifierStringLabelSingleClassConfusionMatrixAggregationFuncti
 """
 
         results = ddb_conn.sql(confusion_matrix_query).df()
-        unescaped_segmentation_cols = [unescape_identifier(seg_col) for seg_col in segmentation_cols]
 
         tp = self.group_query_results_to_numeric_metrics(
             results,
             "true_positive_count",
-            dim_columns=unescaped_segmentation_cols + extra_dims,
+            dim_columns=segmentation_cols + extra_dims,
             timestamp_col="ts",
         )
         fp = self.group_query_results_to_numeric_metrics(
             results,
             "false_positive_count",
-            dim_columns=unescaped_segmentation_cols + extra_dims,
+            dim_columns=segmentation_cols + extra_dims,
             timestamp_col="ts",
         )
         fn = self.group_query_results_to_numeric_metrics(
             results,
             "false_negative_count",
-            dim_columns=unescaped_segmentation_cols + extra_dims,
+            dim_columns=segmentation_cols + extra_dims,
             timestamp_col="ts",
         )
         tn = self.group_query_results_to_numeric_metrics(
             results,
             "true_negative_count",
-            dim_columns=unescaped_segmentation_cols + extra_dims,
+            dim_columns=segmentation_cols + extra_dims,
             timestamp_col="ts",
         )
         tp_metric = self.series_to_metric(

arthur_common/aggregations/functions/numeric_stats.py

@@ -18,8 +18,10 @@ from arthur_common.models.schema_definitions import (
     ScalarType,
     ScopeSchemaTag,
 )
-
-from arthur_common.tools.duckdb_data_loader import unescape_identifier, escape_str_literal
+from arthur_common.tools.duckdb_data_loader import (
+    escape_str_literal,
+    unescape_identifier,
+)
 
 
 class NumericSketchAggregationFunction(SketchAggregationFunction):
@@ -121,12 +123,11 @@ class NumericSketchAggregationFunction(SketchAggregationFunction):
                 """
 
         results = ddb_conn.sql(data_query).df()
-        unescaped_segmentation_cols = [unescape_identifier(seg_col) for seg_col in segmentation_cols]
 
         series = self.group_query_results_to_sketch_metrics(
             results,
             unescape_identifier(numeric_col),
-            unescaped_segmentation_cols + extra_dims,
+            segmentation_cols + extra_dims,
             "ts",
         )
 

arthur_common/aggregations/functions/numeric_sum.py

@@ -118,17 +118,18 @@ class NumericSumAggregationFunction(NumericAggregationFunction):
                 """
 
         results = ddb_conn.sql(query).df()
-        unescaped_segmentation_cols = [unescape_identifier(seg_col) for seg_col in segmentation_cols]
 
         series = self.group_query_results_to_numeric_metrics(
             results,
             "sum",
-            unescaped_segmentation_cols,
+            segmentation_cols,
             "ts",
         )
         # preserve dimension that identifies the name of the numeric column used for the aggregation
         for point in series:
-            point.dimensions.append(Dimension(name="column_name", value=unescape_identifier(numeric_col)))
+            point.dimensions.append(
+                Dimension(name="column_name", value=unescape_identifier(numeric_col)),
+            )
 
         metric = self.series_to_metric(self.METRIC_NAME, series)
         return [metric]

arthur_common/models/connectors.py

@@ -1,4 +1,4 @@
-from pydantic import BaseModel, ConfigDict, Field, computed_field
+from pydantic import BaseModel, Field
 
 
 class ConnectorPaginationOptions(BaseModel):

arthur_common/models/response_schemas.py

@@ -624,6 +624,8 @@ class SpanWithMetricsResponse(BaseModel):
     start_time: datetime
     end_time: datetime
     task_id: Optional[str] = None
+    session_id: Optional[str] = None
+    status_code: str = Field(description="Status code for the span (Unset, Error, Ok)")
     created_at: datetime
     updated_at: datetime
     raw_data: dict[str, Any]
@@ -650,6 +652,8 @@ class NestedSpanWithMetricsResponse(BaseModel):
     start_time: datetime
     end_time: datetime
     task_id: Optional[str] = None
+    session_id: Optional[str] = None
+    status_code: str = Field(description="Status code for the span (Unset, Error, Ok)")
     created_at: datetime
     updated_at: datetime
     raw_data: dict[str, Any]

arthur_common/tools/duckdb_utils.py

@@ -2,7 +2,6 @@ import duckdb
 
 from arthur_common.config.config import Config
 from arthur_common.models.schema_definitions import SEGMENTATION_ALLOWED_DTYPES, DType
-from arthur_common.tools.duckdb_data_loader import escape_identifier
 
 
 def is_column_possible_segmentation(

arthur_common/tools/schema_inferer.py

@@ -40,12 +40,11 @@ class SchemaInferer:
         self.conn.sql(
             f"CREATE OR REPLACE TEMP TABLE {escaped_col} AS SELECT UNNEST({escaped_col}) as {escaped_col} FROM {table}",
         )
-        return self._infer_schema(escaped_col, is_nested_col=True)
+        return self._infer_schema(escaped_col)
 
     def _infer_schema(
         self,
         table: str = "root",
-        is_nested_col: bool = False,
     ) -> DatasetObjectType:
         """is_nested_col indicates whether the function is being called on an unnested/flattened table that represents
         a struct column or list column in the root table."""
@@ -105,9 +104,7 @@ class SchemaInferer:
                         raise NotImplementedError(f"Type {col_type} not mappable.")
 
                 # tag column as a possible segmentation column if it meets criteria
-                # we only support top-level column aggregations right now (ie you can't aggregate on a nested column)
-                # so we don't want to tag nested columns as possible segmentation columns
-                if not is_nested_col and is_column_possible_segmentation(
+                if is_column_possible_segmentation(
                     self.conn,
                     table,
                     escape_identifier(col_name),

{arthur_common-2.4.1.dist-info → arthur_common-2.4.3.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: arthur-common
-Version: 2.4.1
+Version: 2.4.3
 Summary: Utility code common to Arthur platform components.
 License: MIT
 Author: Arthur

{arthur_common-2.4.1.dist-info → arthur_common-2.4.3.dist-info}/RECORD

@@ -1,20 +1,20 @@
 arthur_common/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 arthur_common/aggregations/__init__.py,sha256=vISWyciQAtksa71OKeHNP-QyFGd1NzBKq_LBsG0QSG8,67
-arthur_common/aggregations/aggregator.py,sha256=AhyNqBDEbKtS3ZrnSIT9iZ1SK_TAuiUNg9s9loDvek0,8007
+arthur_common/aggregations/aggregator.py,sha256=1kMyP89biBSR6omD1R9fsAwfjbau0gozcirQOwYVYjg,9680
 arthur_common/aggregations/functions/README.md,sha256=MkZoTAJ94My96R5Z8GAxud7S6vyR0vgVi9gqdt9a4XY,5460
 arthur_common/aggregations/functions/__init__.py,sha256=HqC3UNRURX7ZQHgamTrQvfA8u_FiZGZ4I4eQW7Ooe5o,1299
-arthur_common/aggregations/functions/agentic_aggregations.py,sha256=09th4RPRf-ogtVWpbcqqmITN2UFtfqXhQ7Rr6IBqQHo,33995
-arthur_common/aggregations/functions/categorical_count.py,sha256=_TD0s0JAtqC5RmT6ZNWLEBZm-dU4akm-Aor7EDVazzA,5176
-arthur_common/aggregations/functions/confusion_matrix.py,sha256=n33kyyZuxo8k6jUYnBUsc1fLotTmcw0H8rsX_x_oeJ0,21733
-arthur_common/aggregations/functions/inference_count.py,sha256=D49SpwFywipMqeC93gc3_ZGwBoGL89yKuA9_55dBWBw,3984
-arthur_common/aggregations/functions/inference_count_by_class.py,sha256=mYL6xMTb-_VO6mKGWHOtFAvWzTt-C_4vKf8KgioJGDg,11191
-arthur_common/aggregations/functions/inference_null_count.py,sha256=UlE5EZa3k2nKIv6Yzrnjq1MsZEzrau7Olumny8hsHtg,4672
-arthur_common/aggregations/functions/mean_absolute_error.py,sha256=YzrNHox_4HEGWn33E12d6eiQ8A9Rwct7AW3hOWrTW7I,6544
-arthur_common/aggregations/functions/mean_squared_error.py,sha256=b_is7FKRSninYs1ilAXeLPJFfmyCaiKvCC9Ev_OERio,6565
-arthur_common/aggregations/functions/multiclass_confusion_matrix.py,sha256=e1KEyxIZocWMkDbnW0zfJHd5PUi_kyzwNUVFOD0l5Nk,12359
+arthur_common/aggregations/functions/agentic_aggregations.py,sha256=82OJ174uGcDqf7OLXY7dwnnv1g4kubkjazpc7Yj0xw0,39531
+arthur_common/aggregations/functions/categorical_count.py,sha256=jxV2w2Itmoh02VuazWN5z94PmQ-bRZjZpSoODGeBulQ,5099
+arthur_common/aggregations/functions/confusion_matrix.py,sha256=2fIqo50TcbUlGgPXxGtfFr6ehyZn69R8sphigGuMDgo,21626
+arthur_common/aggregations/functions/inference_count.py,sha256=Pxe5WT_Zgnn_wSDcm48l-flh-M5Zr72SbR4tQyNBk-o,3802
+arthur_common/aggregations/functions/inference_count_by_class.py,sha256=fmzrbRxiWgmutJYrBs7JY1iIRF7F6kozBzcsMypatlE,10896
+arthur_common/aggregations/functions/inference_null_count.py,sha256=X8mfeKb46VxUQFrjukSlVpM9AZCNvStsBHU3LsUbcEM,4591
+arthur_common/aggregations/functions/mean_absolute_error.py,sha256=P9H0rRvpObnWQiu4p7-yW6y6R7_-Ju23y2YlZQgxvHA,6352
+arthur_common/aggregations/functions/mean_squared_error.py,sha256=hZrHzfCscNnGKp_SqOeHEebzjMych1EXtnI1K70EYZE,6373
+arthur_common/aggregations/functions/multiclass_confusion_matrix.py,sha256=eA4y0xJikErkRww5OudUAMG9Y6cYztkO4w561nWVh5w,12195
 arthur_common/aggregations/functions/multiclass_inference_count_by_class.py,sha256=yiMpdz4VuX1ELprXYupFu4B9aDLIhgfEi3ma8jZsT_M,4261
-arthur_common/aggregations/functions/numeric_stats.py,sha256=mMpVH1PvElGaz5mIQWy8sIkKPZ5kyeNOAM2iM2IlBvY,4760
-arthur_common/aggregations/functions/numeric_sum.py,sha256=Vq-dQonKTdLt8pYFwT5tCXyyL_FvVQxb6b3nFNRSqus,4861
+arthur_common/aggregations/functions/numeric_stats.py,sha256=28y0Zdhk3kLFiJYVWq_uev1C1yBZDn1aTUEdvLkqo3k,4660
+arthur_common/aggregations/functions/numeric_sum.py,sha256=TAeVVd5NqF7X9_hnMzbNVOVxdExcra4EZDkubtWHyAs,4780
 arthur_common/aggregations/functions/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 arthur_common/aggregations/functions/shield_aggregations.py,sha256=BzPkpbhZRy16iFOobuusGKHfov5DxnXS2v_WThpw2fk,35659
 arthur_common/aggregations/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -23,7 +23,7 @@ arthur_common/config/config.py,sha256=fcpjOYjPKu4Duk63CuTHrOWKQKAlAhVUR60kF_2_Xo
 arthur_common/config/settings.yaml,sha256=0CrygUwJzC5mGcO5Xnvv2ttp-P7LIsx682jllYA96NQ,161
 arthur_common/models/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 arthur_common/models/common_schemas.py,sha256=31Br7DbIgrwHwzgiyMXrgPYrANhqSqle7kmismcy4TY,6770
-arthur_common/models/connectors.py,sha256=RwjY74cs0KTKw7Opywehg46SZ4vwN3xm6ujHRsRIQ8Y,2292
+arthur_common/models/connectors.py,sha256=gRdX4lNz0ObU64FqMmoffHVBwEgO3JfOf3wjn3tKv0Q,2264
 arthur_common/models/constants.py,sha256=munkU0LrLsDs9BtAfozzw30FCguIowmAUKg_9vqwX24,1049
 arthur_common/models/datasets.py,sha256=7p1tyJEPwXjBs2ZRoai8hTzNl6MK9jU1DluzASApE_4,254
 arthur_common/models/enums.py,sha256=J2beHEMjLfOGgc-vh1aDpE7KmBGKzLoOUGYLtuciJro,3870
@@ -31,7 +31,7 @@ arthur_common/models/metric_schemas.py,sha256=Xf-1RTzg7iYtnBMLkUUUuMPzAujzzNvQx_
 arthur_common/models/metrics.py,sha256=mCa0aN-nuNHYcqGfkyKFeriI0krz0-ScgmXWXHlKoEI,11109
 arthur_common/models/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 arthur_common/models/request_schemas.py,sha256=LEybzPoOzaaTyua48mr9sAVzrBK6dIeYhR158kMp0o8,29749
-arthur_common/models/response_schemas.py,sha256=qc6DDfY4GxtXtUiBllsQglvPnZzhe2Vw8D-S76B_r_0,25393
+arthur_common/models/response_schemas.py,sha256=kY3NZceFaBRkxGDF5-W1CYDjtwFdf3xyLxdzrwHlmJI,25643
 arthur_common/models/schema_definitions.py,sha256=dcUSLjBmvyloStcBFmT_rHdXbKdvA8Yxi_avYUbps3E,16876
 arthur_common/models/task_job_specs.py,sha256=p7jsSb97ylHYNkwoHXNOJvx2zcnh2kxLeh3m0pddo4M,3442
 arthur_common/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -39,11 +39,11 @@ arthur_common/tools/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSu
 arthur_common/tools/aggregation_analyzer.py,sha256=UfMtvFWXV2Dqly8S6nneGgomuvEGN-1tBz81tfkMcAE,11206
 arthur_common/tools/aggregation_loader.py,sha256=3CF46bNi-GdJBNOXkjYfCQ1Aung8lf65L532sdWmR_s,2351
 arthur_common/tools/duckdb_data_loader.py,sha256=A80wpATSc4VJLghoHwxpBEuUsxY93OZS0Qo4cFX7cRw,12462
-arthur_common/tools/duckdb_utils.py,sha256=8l8bUmjqJyj84DXyEOzO_DsD8VsO25DWYK_IYF--Zek,1211
+arthur_common/tools/duckdb_utils.py,sha256=PZ3AKoBUaU6papqNiNQ4Sm2ugg5bGyXfaC_1I-E2q3s,1142
 arthur_common/tools/functions.py,sha256=FWL4eWO5-vLp86WudT-MGUKvf2B8f02IdoXQFKd6d8k,1093
 arthur_common/tools/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-arthur_common/tools/schema_inferer.py,sha256=9teI67umlGn0izp6pZ5UBuWxJthaWEmw3wRj2KPIbf4,5207
+arthur_common/tools/schema_inferer.py,sha256=8ehIqAxuGlgM08RtwPB43a7TfenZyEIf1R0p1RYrkng,4920
 arthur_common/tools/time_utils.py,sha256=4gfiu9NXfvPZltiVNLSIQGylX6h2W0viNi9Kv4bKyfw,1410
-arthur_common-2.4.1.dist-info/METADATA,sha256=LA7R2B8LGE78eJrwpQTHP0nFRPzLTnntZvHQUNtcVm4,2146
-arthur_common-2.4.1.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
-arthur_common-2.4.1.dist-info/RECORD,,
+arthur_common-2.4.3.dist-info/METADATA,sha256=IwWpahpZ5U0mEeZ1YdeR-nBn0tOeYM938q8wIn8Vb-0,2146
+arthur_common-2.4.3.dist-info/WHEEL,sha256=b4K_helf-jlQoXBBETfwnf4B04YC67LOev0jo4fX5m8,88
+arthur_common-2.4.3.dist-info/RECORD,,