arthur-common 2.1.59__py3-none-any.whl → 2.1.61__py3-none-any.whl

arthur_common/aggregations/aggregator.py

@@ -172,7 +172,9 @@ class SketchAggregationFunction(AggregationFunction, ABC):
         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),
+                SketchAggregationFunction._group_to_series(
+                    group, timestamp_col, dim_columns, value_col
+                ),
             )
 
         return calculated_metrics

arthur_common/aggregations/functions/agentic_aggregations.py

@@ -10,7 +10,7 @@ from arthur_common.aggregations.aggregator import (
     NumericAggregationFunction,
     SketchAggregationFunction,
 )
-from arthur_common.models.datasets import ModelProblemType
+from arthur_common.models.enums import ModelProblemType
 from arthur_common.models.metrics import (
     BaseReportedAggregation,
     DatasetReference,

arthur_common/aggregations/functions/confusion_matrix.py

@@ -4,7 +4,7 @@ from uuid import UUID
 from duckdb import DuckDBPyConnection
 
 from arthur_common.aggregations.aggregator import NumericAggregationFunction
-from arthur_common.models.datasets import ModelProblemType
+from arthur_common.models.enums import ModelProblemType
 from arthur_common.models.metrics import (
     BaseReportedAggregation,
     DatasetReference,

arthur_common/aggregations/functions/inference_count_by_class.py

@@ -4,7 +4,7 @@ from uuid import UUID
 from duckdb import DuckDBPyConnection
 
 from arthur_common.aggregations.aggregator import NumericAggregationFunction
-from arthur_common.models.datasets import ModelProblemType
+from arthur_common.models.enums import ModelProblemType
 from arthur_common.models.metrics import (
     BaseReportedAggregation,
     DatasetReference,

arthur_common/aggregations/functions/mean_absolute_error.py

@@ -4,7 +4,7 @@ from uuid import UUID
 from duckdb import DuckDBPyConnection
 
 from arthur_common.aggregations.aggregator import NumericAggregationFunction
-from arthur_common.models.datasets import ModelProblemType
+from arthur_common.models.enums import ModelProblemType
 from arthur_common.models.metrics import (
     BaseReportedAggregation,
     DatasetReference,

arthur_common/aggregations/functions/mean_squared_error.py

@@ -4,7 +4,7 @@ from uuid import UUID
 from duckdb import DuckDBPyConnection
 
 from arthur_common.aggregations.aggregator import NumericAggregationFunction
-from arthur_common.models.datasets import ModelProblemType
+from arthur_common.models.enums import ModelProblemType
 from arthur_common.models.metrics import (
     BaseReportedAggregation,
     DatasetReference,

arthur_common/aggregations/functions/multiclass_confusion_matrix.py

@@ -4,7 +4,7 @@ from uuid import UUID
 from duckdb import DuckDBPyConnection
 
 from arthur_common.aggregations.aggregator import NumericAggregationFunction
-from arthur_common.models.datasets import ModelProblemType
+from arthur_common.models.enums import ModelProblemType
 from arthur_common.models.metrics import (
     BaseReportedAggregation,
     DatasetReference,

arthur_common/aggregations/functions/multiclass_inference_count_by_class.py

@@ -6,7 +6,7 @@ from duckdb import DuckDBPyConnection
 from arthur_common.aggregations.functions.inference_count_by_class import (
     BinaryClassifierCountByClassAggregationFunction,
 )
-from arthur_common.models.datasets import ModelProblemType
+from arthur_common.models.enums import ModelProblemType
 from arthur_common.models.metrics import (
     BaseReportedAggregation,
     DatasetReference,

arthur_common/aggregations/functions/shield_aggregations.py

@@ -9,7 +9,7 @@ from arthur_common.aggregations.aggregator import (
     NumericAggregationFunction,
     SketchAggregationFunction,
 )
-from arthur_common.models.datasets import ModelProblemType
+from arthur_common.models.enums import ModelProblemType
 from arthur_common.models.metrics import (
     BaseReportedAggregation,
     DatasetReference,

arthur_common/models/common_schemas.py

@@ -0,0 +1,208 @@
+from __future__ import annotations
+
+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 total_items_count // self.page_size + 1
+
+
+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__()

arthur_common/models/connectors.py

@@ -1,14 +1,14 @@
-from pydantic import BaseModel, Field
+from pydantic import BaseModel, ConfigDict, Field, computed_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",

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,144 @@
+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"
+
+
+# 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):
+    WRONG_TOOL_SELECTED = 0
+    CORRECT_TOOL_SELECTED = 1
+    NO_TOOL_SELECTED = 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"

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,