arthur-common 2.1.58__py3-none-any.whl → 2.4.13__py3-none-any.whl

arthur_common/models/common_schemas.py

@@ -0,0 +1,214 @@
+from __future__ import annotations
+
+import math
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, ConfigDict, Field, field_validator
+
+from arthur_common.models.constants import (
+    DEFAULT_PII_RULE_CONFIDENCE_SCORE_THRESHOLD,
+    DEFAULT_TOXICITY_RULE_THRESHOLD,
+    NEGATIVE_BLOOD_EXAMPLE,
+)
+from arthur_common.models.enums import (
+    PaginationSortMethod,
+    PIIEntityTypes,
+    UserPermissionAction,
+    UserPermissionResource,
+)
+
+
+class AuthUserRole(BaseModel):
+    id: str | None = None
+    name: str
+    description: str
+    composite: bool
+
+
+class ExampleConfig(BaseModel):
+    example: str = Field(description="Custom example for the sensitive data")
+    result: bool = Field(
+        description="Boolean value representing if the example passes or fails the the sensitive "
+        "data rule ",
+    )
+
+    model_config = ConfigDict(
+        json_schema_extra={
+            "example": {"example": NEGATIVE_BLOOD_EXAMPLE, "result": True},
+        },
+    )
+
+
+class ExamplesConfig(BaseModel):
+    examples: List[ExampleConfig] = Field(
+        description="List of all the examples for Sensitive Data Rule",
+    )
+
+    model_config = ConfigDict(
+        json_schema_extra={
+            "example": {
+                "examples": [
+                    {"example": NEGATIVE_BLOOD_EXAMPLE, "result": True},
+                    {
+                        "example": "Most of the people have A positive blood group",
+                        "result": False,
+                    },
+                ],
+                "hint": "specific individual's blood type",
+            },
+        },
+    )
+    hint: Optional[str] = Field(
+        description="Optional. Hint added to describe what Sensitive Data Rule should be checking for",
+        default=None,
+    )
+
+    def to_dict(self) -> Dict[str, Any]:
+        d = self.__dict__
+        d["examples"] = [ex.__dict__ for ex in self.examples]
+        d["hint"] = self.hint
+        return d
+
+
+class KeywordsConfig(BaseModel):
+    keywords: List[str] = Field(description="List of Keywords")
+
+    model_config = ConfigDict(
+        json_schema_extra={
+            "example": {"keywords": ["Blocked_Keyword_1", "Blocked_Keyword_2"]},
+        },
+    )
+
+
+class LLMTokenConsumption(BaseModel):
+    prompt_tokens: int
+    completion_tokens: int
+
+    def total_tokens(self) -> int:
+        return self.prompt_tokens + self.completion_tokens
+
+    def add(self, token_consumption: LLMTokenConsumption) -> "LLMTokenConsumption":
+        self.prompt_tokens += token_consumption.prompt_tokens
+        self.completion_tokens += token_consumption.completion_tokens
+        return self
+
+
+class PaginationParameters(BaseModel):
+    sort: Optional[PaginationSortMethod] = PaginationSortMethod.DESCENDING
+    page_size: int = 10
+    page: int = 0
+
+    def calculate_total_pages(self, total_items_count: int) -> int:
+        return math.ceil(total_items_count / self.page_size)
+
+
+class PIIConfig(BaseModel):
+    disabled_pii_entities: Optional[list[str]] = Field(
+        description=f"Optional. List of PII entities to disable. Valid values are: {PIIEntityTypes.to_string()}",
+        default=None,
+    )
+
+    confidence_threshold: Optional[float] = Field(
+        description=f"Optional. Float (0, 1) indicating the level of tolerable PII to consider the rule passed or failed. Min: 0 (less confident) Max: 1 (very confident). Default: {DEFAULT_PII_RULE_CONFIDENCE_SCORE_THRESHOLD}",
+        default=DEFAULT_PII_RULE_CONFIDENCE_SCORE_THRESHOLD,
+        json_schema_extra={"deprecated": True},
+    )
+
+    allow_list: Optional[list[str]] = Field(
+        description="Optional. List of strings to pass PII validation.",
+        default=None,
+    )
+
+    @field_validator("disabled_pii_entities")
+    def validate_pii_entities(cls, v: list[str] | None) -> list[str] | None:
+        if v:
+            entities_passed = set(v)
+            entities_supported = set(PIIEntityTypes.values())
+            invalid_entities = entities_passed - entities_supported
+            if invalid_entities:
+                raise ValueError(
+                    f"The following values are not valid PII entities: {invalid_entities}",
+                )
+
+            # Fail the case where they are trying to disable all PII entity types
+            if (not invalid_entities) & (
+                len(entities_passed) == len(entities_supported)
+            ):
+                raise ValueError(
+                    f"Cannot disable all supported PII entities on PIIDataRule",
+                )
+            return v
+        else:
+            return v
+
+    @field_validator("confidence_threshold")
+    def validate_confidence_threshold(cls, v: float | None) -> float | None:
+        if v:
+            if (v < 0) | (v > 1):
+                raise ValueError(f'"confidence_threshold" must be between 0 and 1')
+            return v
+        else:
+            return v
+
+    model_config = ConfigDict(
+        json_schema_extra={
+            "example": {
+                "disabled_pii_entities": ["PERSON", "URL"],
+                "confidence_threshold": "0.5",
+                "allow_list": ["arthur.ai", "Arthur"],
+            },
+        },
+        extra="forbid",
+    )
+
+
+class RegexConfig(BaseModel):
+    regex_patterns: List[str] = Field(
+        description="List of Regex patterns to be used for validation. Be sure to encode requests in JSON and account for escape characters.",
+    )
+
+    model_config = ConfigDict(
+        json_schema_extra={
+            "example": {
+                "regex_patterns": ["\\d{3}-\\d{2}-\\d{4}", "\\d{5}-\\d{6}-\\d{7}"],
+            },
+        },
+        extra="forbid",
+    )
+
+
+class ToxicityConfig(BaseModel):
+    threshold: float = Field(
+        default=DEFAULT_TOXICITY_RULE_THRESHOLD,
+        description=f"Optional. Float (0, 1) indicating the level of tolerable toxicity to consider the rule passed or failed. Min: 0 (no toxic language) Max: 1 (very toxic language). Default: {DEFAULT_TOXICITY_RULE_THRESHOLD}",
+    )
+
+    model_config = ConfigDict(
+        extra="forbid",
+        json_schema_extra={"example": {"threshold": DEFAULT_TOXICITY_RULE_THRESHOLD}},
+    )
+
+    @field_validator("threshold", mode="before")
+    @classmethod
+    def validate_toxicity_threshold(cls, v: float | None) -> float:
+        if v is None:
+            return float(DEFAULT_TOXICITY_RULE_THRESHOLD)
+        if (v < 0) | (v > 1):
+            raise ValueError(f'"threshold" must be between 0 and 1')
+        return v
+
+
+class UserPermission(BaseModel):
+    action: UserPermissionAction
+    resource: UserPermissionResource
+
+    def __hash__(self) -> int:
+        return hash((self.action, self.resource))
+
+    def __eq__(self, other: object) -> bool:
+        return isinstance(other, UserPermission) and self.__hash__() == other.__hash__()
+
+
+class VariableTemplateValue(BaseModel):
+    name: str = Field(..., description="Name of the variable")
+    value: str = Field(..., description="Value of the variable")

arthur_common/models/connectors.py

@@ -3,12 +3,12 @@ from pydantic import BaseModel, Field
 
 class ConnectorPaginationOptions(BaseModel):
     page: int = Field(default=1, ge=1)
-    page_size: int = Field(default=25, gt=0, le=500)
+    page_size: int = Field(default=25, ge=1, le=500)
 
     @property
     def page_params(self) -> tuple[int, int]:
         if self.page is not None:
-            return self.page, self.page_size
+            return (self.page, self.page_size)
         else:
             raise ValueError(
                 "Pagination options must be set to return a page and page size",
@@ -38,6 +38,14 @@ ODBC_CONNECTOR_DRIVER_FIELD = "driver"
 ODBC_CONNECTOR_TABLE_NAME_FIELD = "table_name"
 ODBC_CONNECTOR_DIALECT_FIELD = "dialect"
 
+# Snowflake connector constants
+SNOWFLAKE_CONNECTOR_ACCOUNT_FIELD = "account"
+SNOWFLAKE_CONNECTOR_SCHEMA_FIELD = "schema"
+SNOWFLAKE_CONNECTOR_WAREHOUSE_FIELD = "warehouse"
+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"
 
 # dataset (connector type dependent) constants
 SHIELD_DATASET_TASK_ID_FIELD = "task_id"

arthur_common/models/constants.py

@@ -0,0 +1,24 @@
+# RBAC
+CHAT_USER: str = "CHAT-USER"
+ORG_ADMIN: str = "ORG-ADMIN"
+TASK_ADMIN: str = "TASK-ADMIN"
+DEFAULT_RULE_ADMIN: str = "DEFAULT-RULE-ADMIN"
+VALIDATION_USER: str = "VALIDATION-USER"
+ORG_AUDITOR: str = "ORG-AUDITOR"
+ADMIN_KEY: str = "ADMIN-KEY"
+
+LEGACY_KEYCLOAK_ROLES: dict[str, str] = {
+    "genai_engine_admin_user": TASK_ADMIN,
+}
+
+# Make sure the policy and description match
+GENAI_ENGINE_KEYCLOAK_PASSWORD_LENGTH = 12
+GENAI_ENGINE_KEYCLOAK_PASSWORD_POLICY = f"length({GENAI_ENGINE_KEYCLOAK_PASSWORD_LENGTH}) and specialChars(1) and upperCase(1) and lowerCase(1)"
+ERROR_PASSWORD_POLICY_NOT_MET = f"Password should be at least {GENAI_ENGINE_KEYCLOAK_PASSWORD_LENGTH} characters and contain at least one special character, lowercase character, and uppercase character."
+ERROR_DEFAULT_METRICS_ENGINE = "This metric could not be evaluated"
+
+# Miscellaneous
+DEFAULT_TOXICITY_RULE_THRESHOLD = 0.5
+DEFAULT_PII_RULE_CONFIDENCE_SCORE_THRESHOLD = 0
+NEGATIVE_BLOOD_EXAMPLE = "John has O negative blood group"
+HALLUCINATION_RULE_NAME = "Hallucination Rule"

arthur_common/models/datasets.py

@@ -1,15 +1,6 @@
 from enum import Enum
 
 
-class ModelProblemType(str, Enum):
-    REGRESSION = "regression"
-    BINARY_CLASSIFICATION = "binary_classification"
-    ARTHUR_SHIELD = "arthur_shield"
-    CUSTOM = "custom"
-    MULTICLASS_CLASSIFICATION = "multiclass_classification"
-    AGENTIC_TRACE = "agentic_trace"
-
-
 class DatasetFileType(str, Enum):
     JSON = "json"
     CSV = "csv"

arthur_common/models/enums.py

@@ -0,0 +1,177 @@
+from enum import IntEnum, StrEnum
+
+from arthur_common.models.constants import (
+    DEFAULT_RULE_ADMIN,
+    ORG_ADMIN,
+    ORG_AUDITOR,
+    TASK_ADMIN,
+    VALIDATION_USER,
+)
+
+
+class BaseEnum(StrEnum):
+    @classmethod
+    def values(self) -> list[str]:
+        values: list[str] = [e for e in self]
+        return values
+
+    def __str__(self) -> str:
+        return str(self.value)
+
+
+class APIKeysRolesEnum(BaseEnum):
+    DEFAULT_RULE_ADMIN = DEFAULT_RULE_ADMIN
+    TASK_ADMIN = TASK_ADMIN
+    VALIDATION_USER = VALIDATION_USER
+    ORG_AUDITOR = ORG_AUDITOR
+    ORG_ADMIN = ORG_ADMIN
+
+
+class InferenceFeedbackTarget(BaseEnum):
+    CONTEXT = "context"
+    RESPONSE_RESULTS = "response_results"
+    PROMPT_RESULTS = "prompt_results"
+
+
+class MetricType(BaseEnum):
+    QUERY_RELEVANCE = "QueryRelevance"
+    RESPONSE_RELEVANCE = "ResponseRelevance"
+    TOOL_SELECTION = "ToolSelection"
+
+
+class ModelProblemType(BaseEnum):
+    REGRESSION = "regression"
+    BINARY_CLASSIFICATION = "binary_classification"
+    ARTHUR_SHIELD = "arthur_shield"
+    CUSTOM = "custom"
+    MULTICLASS_CLASSIFICATION = "multiclass_classification"
+    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):
+    CREDIT_CARD = "CREDIT_CARD"
+    CRYPTO = "CRYPTO"
+    DATE_TIME = "DATE_TIME"
+    EMAIL_ADDRESS = "EMAIL_ADDRESS"
+    IBAN_CODE = "IBAN_CODE"
+    IP_ADDRESS = "IP_ADDRESS"
+    NRP = "NRP"
+    LOCATION = "LOCATION"
+    PERSON = "PERSON"
+    PHONE_NUMBER = "PHONE_NUMBER"
+    MEDICAL_LICENSE = "MEDICAL_LICENSE"
+    URL = "URL"
+    US_BANK_NUMBER = "US_BANK_NUMBER"
+    US_DRIVER_LICENSE = "US_DRIVER_LICENSE"
+    US_ITIN = "US_ITIN"
+    US_PASSPORT = "US_PASSPORT"
+    US_SSN = "US_SSN"
+
+    @classmethod
+    def to_string(cls) -> str:
+        return ",".join(member.value for member in cls)
+
+
+class PaginationSortMethod(BaseEnum):
+    ASCENDING = "asc"
+    DESCENDING = "desc"
+
+
+class RuleResultEnum(BaseEnum):
+    PASS = "Pass"
+    FAIL = "Fail"
+    SKIPPED = "Skipped"
+    UNAVAILABLE = "Unavailable"
+    PARTIALLY_UNAVAILABLE = "Partially Unavailable"
+    MODEL_NOT_AVAILABLE = "Model Not Available"
+
+
+class RuleScope(BaseEnum):
+    DEFAULT = "default"
+    TASK = "task"
+
+
+class RuleType(BaseEnum):
+    KEYWORD = "KeywordRule"
+    MODEL_HALLUCINATION_V2 = "ModelHallucinationRuleV2"
+    MODEL_SENSITIVE_DATA = "ModelSensitiveDataRule"
+    PII_DATA = "PIIDataRule"
+    PROMPT_INJECTION = "PromptInjectionRule"
+    REGEX = "RegexRule"
+    TOXICITY = "ToxicityRule"
+
+
+class TaskType(BaseEnum):
+    TRADITIONAL = "traditional"
+    AGENTIC = "agentic"
+
+
+class TokenUsageScope(BaseEnum):
+    RULE_TYPE = "rule_type"
+    TASK = "task"
+
+
+class ToolClassEnum(IntEnum):
+    INCORRECT = 0
+    CORRECT = 1
+    NA = 2
+
+    def __str__(self) -> str:
+        return str(self.value)
+
+
+class ToxicityViolationType(BaseEnum):
+    BENIGN = "benign"
+    HARMFUL_REQUEST = "harmful_request"
+    TOXIC_CONTENT = "toxic_content"
+    PROFANITY = "profanity"
+    UNKNOWN = "unknown"
+
+
+# If you added values here, did you update permission_mappings.py in arthur-engine?
+class UserPermissionAction(BaseEnum):
+    CREATE = "create"
+    READ = "read"
+
+
+# If you added values here, did you update permission_mappings.py in arthur-engine?
+class UserPermissionResource(BaseEnum):
+    PROMPTS = "prompts"
+    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"
+
+
+class StatusCodeEnum(BaseEnum):
+    OK = "Ok"
+    ERROR = "Error"
+    UNSET = "Unset"
+
+
+class ContinuousEvalRunStatus(BaseEnum):
+    PENDING = "pending"
+    PASSED = "passed"
+    RUNNING = "running"
+    FAILED = "failed"
+    SKIPPED = "skipped"
+    ERROR = "error"
+
+
+class AgenticAnnotationType(BaseEnum):
+    HUMAN = "human"
+    CONTINUOUS_EVAL = "continuous_eval"

arthur_common/models/metric_schemas.py

@@ -0,0 +1,63 @@
+from typing import Any, Dict, List, Optional
+
+from pydantic import BaseModel, Field
+
+
+class RelevanceMetricConfig(BaseModel):
+    """Configuration for relevance metrics including QueryRelevance and ResponseRelevance"""
+
+    relevance_threshold: Optional[float] = Field(
+        default=None,
+        description="Threshold for determining relevance when not using LLM judge",
+    )
+    use_llm_judge: bool = Field(
+        default=True,
+        description="Whether to use LLM as a judge for relevance scoring",
+    )
+
+
+class RelevanceMetric(BaseModel):
+    bert_f_score: Optional[float] = None
+    reranker_relevance_score: Optional[float] = None
+    llm_relevance_score: Optional[float] = None
+    reason: Optional[str] = None
+    refinement: Optional[str] = None
+
+
+class QueryRelevanceMetric(RelevanceMetric):
+    """Inherits from RelevanceMetric. This class is left empty so that the openapi response schema remains the same as before, but we have a single source of truth for the relevance metric details."""
+
+
+class ResponseRelevanceMetric(RelevanceMetric):
+    """Inherits from RelevanceMetric. This class is left empty so that the openapi response schema remains the same as before, but we have a single source of truth for the relevance metric details."""
+
+
+class MetricRequest(BaseModel):
+    system_prompt: Optional[str] = Field(
+        description="System prompt to be used by GenAI Engine for computing metrics.",
+        default=None,
+    )
+    user_query: Optional[str] = Field(
+        description="User query to be used by GenAI Engine for computing metrics.",
+        default=None,
+    )
+    context: List[Dict[str, Any]] = Field(
+        description="Conversation history and additional context to be used by GenAI Engine for computing metrics.",
+        default_factory=list,
+        examples=[
+            {"role": "user", "value": "What is the weather in Tokyo?"},
+            {"role": "assistant", "value": "WeatherTool", "args": {"city": "Tokyo"}},
+            {
+                "role": "tool",
+                "value": '[{"name": "WeatherTool", "result": {"temperature": "20°C", "humidity": "50%", "condition": "sunny"}}]',
+            },
+            {
+                "role": "assistant",
+                "value": "The weather in Tokyo is sunny and the temperature is 20°C.",
+            },
+        ],
+    )
+    response: Optional[str] = Field(
+        description="Response to be used by GenAI Engine for computing metrics.",
+        default=None,
+    )

arthur_common/models/metrics.py

@@ -7,7 +7,7 @@ from uuid import UUID
 from pydantic import BaseModel, Field, field_validator, model_validator
 from typing_extensions import Self
 
-from arthur_common.models.datasets import ModelProblemType
+from arthur_common.models.enums import ModelProblemType
 from arthur_common.models.schema_definitions import (
     DType,
     SchemaTypeUnion,
@@ -195,7 +195,7 @@ class MetricsColumnParameterSchema(MetricsParameterSchema, BaseColumnParameterSc
 
 class MetricsColumnListParameterSchema(
     MetricsParameterSchema,
-    BaseColumnParameterSchema,
+    BaseColumnBaseParameterSchema,
 ):
     # list column parameter schema specific to default metrics
     parameter_type: Literal["column_list"] = "column_list"
@@ -295,10 +295,3 @@ class ReportedCustomAggregation(BaseReportedAggregation):
     dimension_columns: list[str] = Field(
         description="Name of any dimension columns returned from the SQL query. Max length is 1.",
     )
-
-    @field_validator("dimension_columns")
-    @classmethod
-    def validate_dimension_columns_length(cls, v: list[str]) -> str:
-        if len(v) > 1:
-            raise ValueError("Only one dimension column can be specified.")
-        return v