arthur-common 2.1.67__tar.gz → 2.2.0__tar.gz

{arthur_common-2.1.67 → arthur_common-2.2.0}/PKG-INFO

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

{arthur_common-2.1.67 → arthur_common-2.2.0}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "arthur-common"
-version = "2.1.67"
+version = "2.2.0"
 description = "Utility code common to Arthur platform components."
 authors = ["Arthur <engineering@arthur.ai>"]
 license = "MIT"

{arthur_common-2.1.67 → arthur_common-2.2.0}/src/arthur_common/models/connectors.py

@@ -46,9 +46,6 @@ SNOWFLAKE_CONNECTOR_ROLE_FIELD = "role"
 SNOWFLAKE_CONNECTOR_AUTHENTICATOR_FIELD = "authenticator"
 SNOWFLAKE_CONNECTOR_PRIVATE_KEY_FIELD = "private_key"
 SNOWFLAKE_CONNECTOR_PRIVATE_KEY_PASSPHRASE_FIELD = "private_key_passphrase"
-SNOWFLAKE_CONNECTOR_AUTHENTICATOR_KEY_PAIR = "snowflake_key_pair"
-SNOWFLAKE_CONNECTOR_AUTHENTICATOR_PASSWORD = "snowflake_password"
-
 
 # dataset (connector type dependent) constants
 SHIELD_DATASET_TASK_ID_FIELD = "task_id"

{arthur_common-2.1.67 → arthur_common-2.2.0}/src/arthur_common/models/enums.py

@@ -48,6 +48,11 @@ class ModelProblemType(BaseEnum):
     AGENTIC_TRACE = "agentic_trace"
 
 
+class SnowflakeConnectorAuthenticatorMethods(BaseEnum):
+    SNOWFLAKE_KEY_PAIR = "snowflake_key_pair"
+    SNOWFLAKE_PASSWORD = "snowflake_password"
+
+
 # Using version from arthur-engine, which has str and enum type inheritance.
 # Note: These string values are not arbitrary and map to Presidio entity types: https://microsoft.github.io/presidio/supported_entities/
 class PIIEntityTypes(BaseEnum):
@@ -114,9 +119,9 @@ class TokenUsageScope(BaseEnum):
 
 
 class ToolClassEnum(IntEnum):
-    WRONG_TOOL_SELECTED = 0
-    CORRECT_TOOL_SELECTED = 1
-    NO_TOOL_SELECTED = 2
+    INCORRECT = 0
+    CORRECT = 1
+    NA = 2
 
     def __str__(self) -> str:
         return str(self.value)
@@ -142,3 +147,11 @@ class UserPermissionResource(BaseEnum):
     RESPONSES = "responses"
     RULES = "rules"
     TASKS = "tasks"
+
+
+class ComparisonOperatorEnum(BaseEnum):
+    EQUAL = "eq"
+    GREATER_THAN = "gt"
+    GREATER_THAN_OR_EQUAL = "gte"
+    LESS_THAN = "lt"
+    LESS_THAN_OR_EQUAL = "lte"

{arthur_common-2.1.67 → arthur_common-2.2.0}/src/arthur_common/models/request_schemas.py

@@ -1,9 +1,16 @@
 from datetime import datetime
-from typing import Any, Dict, List, Optional, Self, Type, Union
+from typing import Any, Dict, List, Optional, Self, Type
 
 from fastapi import HTTPException
 from openinference.semconv.trace import OpenInferenceSpanKindValues
-from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
+from pydantic import (
+    BaseModel,
+    ConfigDict,
+    Field,
+    ValidationInfo,
+    field_validator,
+    model_validator,
+)
 
 from arthur_common.models.common_schemas import (
     ExamplesConfig,
@@ -25,6 +32,7 @@ from arthur_common.models.enums import (
     PIIEntityTypes,
     RuleScope,
     RuleType,
+    ToolClassEnum,
 )
 from arthur_common.models.metric_schemas import RelevanceMetricConfig
 
@@ -50,12 +58,12 @@ class NewRuleRequest(BaseModel):
         examples=[False],
     )
     config: (
-            KeywordsConfig
-            | RegexConfig
-            | ExamplesConfig
-            | ToxicityConfig
-            | PIIConfig
-            | None
+        KeywordsConfig
+        | RegexConfig
+        | ExamplesConfig
+        | ToxicityConfig
+        | PIIConfig
+        | None
     ) = Field(description="Config of the rule", default=None)
 
     model_config = ConfigDict(
@@ -554,3 +562,250 @@ class SpanQueryRequest(BaseModel):
                 f"Valid values: {', '.join(sorted(valid_span_kinds))}",
             )
         return value
+
+
+class TraceQueryRequest(BaseModel):
+    """Request schema for querying traces with comprehensive filtering."""
+
+    # Required
+    task_ids: list[str] = Field(
+        ...,
+        description="Task IDs to filter on. At least one is required.",
+        min_length=1,
+    )
+
+    # Common optional filters
+    trace_ids: Optional[list[str]] = Field(
+        None,
+        description="Trace IDs to filter on. Optional.",
+    )
+    start_time: Optional[datetime] = Field(
+        None,
+        description="Inclusive start date in ISO8601 string format. Use local time (not UTC).",
+    )
+    end_time: Optional[datetime] = Field(
+        None,
+        description="Exclusive end date in ISO8601 string format. Use local time (not UTC).",
+    )
+
+    # New trace-level filters
+    tool_name: Optional[str] = Field(
+        None,
+        description="Return only results with this tool name.",
+    )
+    span_types: Optional[list[str]] = Field(
+        None,
+        description="Span types to filter on. Optional.",
+    )
+
+    # Query relevance filters
+    query_relevance_eq: Optional[float] = Field(
+        None,
+        ge=0,
+        le=1,
+        description="Equal to this value.",
+    )
+    query_relevance_gt: Optional[float] = Field(
+        None,
+        ge=0,
+        le=1,
+        description="Greater than this value.",
+    )
+    query_relevance_gte: Optional[float] = Field(
+        None,
+        ge=0,
+        le=1,
+        description="Greater than or equal to this value.",
+    )
+    query_relevance_lt: Optional[float] = Field(
+        None,
+        ge=0,
+        le=1,
+        description="Less than this value.",
+    )
+    query_relevance_lte: Optional[float] = Field(
+        None,
+        ge=0,
+        le=1,
+        description="Less than or equal to this value.",
+    )
+
+    # Response relevance filters
+    response_relevance_eq: Optional[float] = Field(
+        None,
+        ge=0,
+        le=1,
+        description="Equal to this value.",
+    )
+    response_relevance_gt: Optional[float] = Field(
+        None,
+        ge=0,
+        le=1,
+        description="Greater than this value.",
+    )
+    response_relevance_gte: Optional[float] = Field(
+        None,
+        ge=0,
+        le=1,
+        description="Greater than or equal to this value.",
+    )
+    response_relevance_lt: Optional[float] = Field(
+        None,
+        ge=0,
+        le=1,
+        description="Less than this value.",
+    )
+    response_relevance_lte: Optional[float] = Field(
+        None,
+        ge=0,
+        le=1,
+        description="Less than or equal to this value.",
+    )
+
+    # Tool classification filters
+    tool_selection: Optional[ToolClassEnum] = Field(
+        None,
+        description="Tool selection evaluation result.",
+    )
+    tool_usage: Optional[ToolClassEnum] = Field(
+        None,
+        description="Tool usage evaluation result.",
+    )
+
+    # Trace duration filters
+    trace_duration_eq: Optional[float] = Field(
+        None,
+        ge=0,
+        description="Duration exactly equal to this value (seconds).",
+    )
+    trace_duration_gt: Optional[float] = Field(
+        None,
+        ge=0,
+        description="Duration greater than this value (seconds).",
+    )
+    trace_duration_gte: Optional[float] = Field(
+        None,
+        ge=0,
+        description="Duration greater than or equal to this value (seconds).",
+    )
+    trace_duration_lt: Optional[float] = Field(
+        None,
+        ge=0,
+        description="Duration less than this value (seconds).",
+    )
+    trace_duration_lte: Optional[float] = Field(
+        None,
+        ge=0,
+        description="Duration less than or equal to this value (seconds).",
+    )
+
+    @field_validator(
+        "query_relevance_eq",
+        "query_relevance_gt",
+        "query_relevance_gte",
+        "query_relevance_lt",
+        "query_relevance_lte",
+        "response_relevance_eq",
+        "response_relevance_gt",
+        "response_relevance_gte",
+        "response_relevance_lt",
+        "response_relevance_lte",
+        mode="before",
+    )
+    @classmethod
+    def validate_relevance_scores(
+        cls,
+        value: Optional[float],
+        info: ValidationInfo,
+    ) -> Optional[float]:
+        """Validate that relevance scores are between 0 and 1 (inclusive)."""
+        if value is not None:
+            if not (0.0 <= value <= 1.0):
+                raise ValueError(
+                    f"{info.field_name} value must be between 0 and 1 (inclusive)",
+                )
+        return value
+
+    @field_validator(
+        "trace_duration_eq",
+        "trace_duration_gt",
+        "trace_duration_gte",
+        "trace_duration_lt",
+        "trace_duration_lte",
+        mode="before",
+    )
+    @classmethod
+    def validate_trace_duration(
+        cls,
+        value: Optional[float],
+        info: ValidationInfo,
+    ) -> Optional[float]:
+        """Validate that trace duration values are non-negative."""
+        if value is not None:
+            if value < 0:
+                raise ValueError(
+                    f"{info.field_name} value must be non-negative (greater than or equal to 0)",
+                )
+        return value
+
+    @field_validator("tool_selection", "tool_usage", mode="before")
+    @classmethod
+    def validate_tool_classification(cls, value: Any) -> Optional[ToolClassEnum]:
+        """Validate tool classification enum values."""
+        if value is not None:
+            # Handle both integer and enum inputs
+            if isinstance(value, int):
+                if value not in [0, 1, 2]:
+                    raise ValueError(
+                        "Tool classification must be 0 (INCORRECT), "
+                        "1 (CORRECT), or 2 (NA)",
+                    )
+                return ToolClassEnum(value)
+            elif isinstance(value, ToolClassEnum):
+                return value
+            else:
+                raise ValueError(
+                    "Tool classification must be an integer (0, 1, 2) or ToolClassEnum instance",
+                )
+        return value
+
+    @field_validator("span_types")
+    @classmethod
+    def validate_span_types(cls, value: Optional[list[str]]) -> Optional[list[str]]:
+        """Validate that all span_types are valid OpenInference span kinds."""
+        if not value:
+            return value
+
+        # Get all valid span kind values
+        valid_span_kinds = [kind.value for kind in OpenInferenceSpanKindValues]
+        invalid_types = [st for st in value if st not in valid_span_kinds]
+
+        if invalid_types:
+            raise ValueError(
+                f"Invalid span_types received: {invalid_types}. "
+                f"Valid values: {', '.join(sorted(valid_span_kinds))}",
+            )
+        return value
+
+    @model_validator(mode="after")
+    def validate_filter_combinations(self) -> Self:
+        """Validate that filter combinations are logically valid."""
+        # Check mutually exclusive filters for each metric type
+        for prefix in ["query_relevance", "response_relevance", "trace_duration"]:
+            eq_field = f"{prefix}_eq"
+            comparison_fields = [f"{prefix}_{op}" for op in ["gt", "gte", "lt", "lte"]]
+
+            if getattr(self, eq_field) and any(
+                getattr(self, field) for field in comparison_fields
+            ):
+                raise ValueError(
+                    f"{eq_field} cannot be combined with other {prefix} comparison operators",
+                )
+
+            # Check for incompatible operator combinations
+            if getattr(self, f"{prefix}_gt") and getattr(self, f"{prefix}_gte"):
+                raise ValueError(f"Cannot combine {prefix}_gt with {prefix}_gte")
+            if getattr(self, f"{prefix}_lt") and getattr(self, f"{prefix}_lte"):
+                raise ValueError(f"Cannot combine {prefix}_lt with {prefix}_lte")
+
+        return self