arthur-common 2.1.58__tar.gz → 2.1.59__tar.gz

{arthur_common-2.1.58 → arthur_common-2.1.59}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.3
 Name: arthur-common
-Version: 2.1.58
+Version: 2.1.59
 Summary: Utility code common to Arthur platform components.
 License: MIT
 Author: Arthur
@@ -62,6 +62,14 @@ This project uses [pytest](https://pytest.org/) for testing. To run the tests, e
 poetry run pytest
 ```
 
+## Release process
+1. Merge changes into **main** branch
+2. Go to **Actions** -> **Arthur Common Version Bump**
+3. Manually trigger workflow there, it will create a PR with version bumping
+4. Go to **Pull requests** and check PR for version bump, accept it if everything is okay
+5. Version bump commit will be merged to **main** branch and it will start release process
+6. Update package version in your project (arthur-engine)
+
 ## License
 
 This project is licensed under the MIT License.

{arthur_common-2.1.58 → arthur_common-2.1.59}/README.md

@@ -37,6 +37,14 @@ This project uses [pytest](https://pytest.org/) for testing. To run the tests, e
 poetry run pytest
 ```
 
+## Release process
+1. Merge changes into **main** branch
+2. Go to **Actions** -> **Arthur Common Version Bump**
+3. Manually trigger workflow there, it will create a PR with version bumping
+4. Go to **Pull requests** and check PR for version bump, accept it if everything is okay
+5. Version bump commit will be merged to **main** branch and it will start release process
+6. Update package version in your project (arthur-engine)
+
 ## License
 
 This project is licensed under the MIT License.

{arthur_common-2.1.58 → arthur_common-2.1.59}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "arthur-common"
-version = "2.1.58"
+version = "2.1.59"
 description = "Utility code common to Arthur platform components."
 authors = ["Arthur <engineering@arthur.ai>"]
 license = "MIT"
@@ -27,7 +27,53 @@ responses = "0.25.7"
 pytest-xdist = "3.6.1"
 pytest-cov = "^6.1.1"
 pre-commit = "^4.2.0"
-mypy = "^1.16.1"
+
+
+[tool.poetry.group.linters.dependencies]
+autoflake = "^2.3.1"
+isort = "^6.0.1"
+black = "^25.1.0"
+mypy = "^1.17.0"
+
+
+[tool.autoflake]
+remove-all-unused-imports = true
+in-place = true
+recursive = true
+
+
+[tool.isort]
+profile = "black"
+src_paths = ["src"]
+
+
+[tool.black]
+target-version = ['py312', 'py313']
+include = '\.pyi?$'
+extend-exclude = '''
+/(
+  # directories
+  \.eggs
+  | \.git
+  06
+  | \.hg
+  | \.mypy_cache
+  | \.tox
+  | \.venv
+  | build
+  | dist
+)/
+'''
+
+[tool.mypy]
+ignore_missing_imports = true
+implicit_reexport = true
+explicit_package_bases = true
+strict = true
+exclude = ["clients/python", "alembic_app_db", "alembic_ts_db", "tests"]
+namespace_packages = true
+mypy_path = "src"
+
 
 [tool.pytest.ini_options]
 pythonpath = ["src"]

{arthur_common-2.1.58 → arthur_common-2.1.59}/src/arthur_common/aggregations/functions/agentic_aggregations.py

@@ -1,6 +1,6 @@
 import json
 import logging
-from typing import Annotated
+from typing import Annotated, Any
 from uuid import UUID
 
 import pandas as pd
@@ -27,7 +27,10 @@ TOOL_SCORE_NO_TOOL_VALUE = 2
 logger = logging.getLogger(__name__)
 
 
-def extract_spans_with_metrics_and_agents(root_spans):
+# TODO: create TypedDict for span
+def extract_spans_with_metrics_and_agents(
+    root_spans: list[str | dict[str, Any]],
+) -> list[tuple[dict[str, Any], str]]:
     """Recursively extract all spans with metrics and their associated agent names from the span tree.
 
     Returns:
@@ -35,14 +38,21 @@ def extract_spans_with_metrics_and_agents(root_spans):
     """
     spans_with_metrics_and_agents = []
 
-    def traverse_spans(spans, current_agent_name="unknown"):
-        for span_str in spans:
-            span = json.loads(span_str) if type(span_str) == str else span_str
+    # 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",
+    ) -> 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
 
             # Update current agent name if this span is an AGENT
-            if span.get("span_kind") == "AGENT":
+            if parsed_span.get("span_kind") == "AGENT":
                 try:
-                    raw_data = span.get("raw_data", {})
+                    raw_data = parsed_span.get("raw_data", {})
                     if isinstance(raw_data, str):
                         raw_data = json.loads(raw_data)
 
@@ -52,29 +62,31 @@ def extract_spans_with_metrics_and_agents(root_spans):
                         current_agent_name = 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')}",
+                        f"Error parsing attributes from span (span_id: {parsed_span.get('span_id')}) in trace {parsed_span.get('trace_id')}",
                     )
 
             # Check if this span has metrics
-            if span.get("metric_results") and len(span.get("metric_results", [])) > 0:
-                spans_with_metrics_and_agents.append((span, current_agent_name))
+            if parsed_span.get("metric_results", []):
+                spans_with_metrics_and_agents.append(
+                    (parsed_span, current_agent_name),
+                )
 
             # Recursively traverse children with the current agent name
-            if span.get("children", []):
-                traverse_spans(span["children"], current_agent_name)
+            if children_span := parsed_span.get("children", []):
+                traverse_spans(children_span, current_agent_name)
 
     traverse_spans(root_spans)
     return spans_with_metrics_and_agents
 
 
-def determine_relevance_pass_fail(score):
+def determine_relevance_pass_fail(score: float | None) -> str | None:
     """Determine pass/fail for relevance scores using global threshold"""
     if score is None:
         return None
     return "pass" if score >= RELEVANCE_SCORE_THRESHOLD else "fail"
 
 
-def determine_tool_pass_fail(score):
+def determine_tool_pass_fail(score: int | None) -> str | None:
     """Determine pass/fail for tool scores using global threshold"""
     if score is None:
         return None
@@ -177,7 +189,7 @@ class AgenticMetricsOverTimeAggregation(SketchAggregationFunction):
 
                 for metric_result in metric_results:
                     metric_type = metric_result.get("metric_type")
-                    details = json.loads(metric_result.get("details", '{}'))
+                    details = json.loads(metric_result.get("details", "{}"))
 
                     if metric_type == "ToolSelection":
                         tool_selection = details.get("tool_selection", {})
@@ -430,7 +442,7 @@ class AgenticRelevancePassFailCountAggregation(NumericAggregationFunction):
 
                 for metric_result in metric_results:
                     metric_type = metric_result.get("metric_type")
-                    details = json.loads(metric_result.get("details", '{}'))
+                    details = json.loads(metric_result.get("details", "{}"))
 
                     if metric_type in ["QueryRelevance", "ResponseRelevance"]:
                         relevance_data = details.get(
@@ -555,7 +567,7 @@ class AgenticToolPassFailCountAggregation(NumericAggregationFunction):
 
                 for metric_result in metric_results:
                     if metric_result.get("metric_type") == "ToolSelection":
-                        details = json.loads(metric_result.get("details", '{}'))
+                        details = json.loads(metric_result.get("details", "{}"))
                         tool_selection = details.get("tool_selection", {})
 
                         tool_selection_score = tool_selection.get("tool_selection")
@@ -723,10 +735,13 @@ class AgenticLLMCallCountAggregation(NumericAggregationFunction):
                 root_spans = json.loads(root_spans)
 
             # Count LLM spans in the tree
-            def count_llm_spans(spans):
+            def count_llm_spans(spans: list[str | dict[str, Any]]) -> int:
                 count = 0
-                for span_str in spans:
-                    span = json.loads(span_str) if type(span_str) == str else span_str
+                for span_to_parse in spans:
+                    if isinstance(span_to_parse, str):
+                        span = json.loads(span_to_parse)
+                    else:
+                        span = span_to_parse
 
                     # Check if this span is an LLM span
                     if span.get("span_kind") == "LLM":
@@ -830,7 +845,7 @@ class AgenticToolSelectionAndUsageByAgentAggregation(NumericAggregationFunction)
 
                 for metric_result in metric_results:
                     if metric_result.get("metric_type") == "ToolSelection":
-                        details = json.loads(metric_result.get("details", '{}'))
+                        details = json.loads(metric_result.get("details", "{}"))
                         tool_selection = details.get("tool_selection", {})
 
                         tool_selection_score = tool_selection.get("tool_selection")

{arthur_common-2.1.58 → arthur_common-2.1.59}/src/arthur_common/models/connectors.py

@@ -38,6 +38,15 @@ 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-2.1.58 → arthur_common-2.1.59}/src/arthur_common/models/metrics.py

@@ -122,6 +122,20 @@ class BaseAggregationParameterSchema(BaseModel):
         description="Description of the parameter.",
     )
 
+    @field_validator("parameter_key")
+    @classmethod
+    def validate_parameter_key_allowed_characters(cls, v: str) -> str:
+        if not v.replace("_", "").isalpha():
+            raise ValueError("Parameter key can only contain letters and underscores.")
+        return v
+
+    @field_validator("friendly_name")
+    @classmethod
+    def validate_friendly_name_allowed_characters(cls, v: str) -> str:
+        if not v.replace("_", "").replace(" ", "").isalpha():
+            raise ValueError("Friendly name can only contain letters and underscores.")
+        return v
+
 
 class MetricsParameterSchema(BaseAggregationParameterSchema):
     # specific to default metrics/Python metrics—not available to custom aggregations
@@ -195,7 +209,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"
@@ -298,7 +312,7 @@ class ReportedCustomAggregation(BaseReportedAggregation):
 
     @field_validator("dimension_columns")
     @classmethod
-    def validate_dimension_columns_length(cls, v: list[str]) -> str:
+    def validate_dimension_columns_length(cls, v: list[str]) -> list[str]:
         if len(v) > 1:
             raise ValueError("Only one dimension column can be specified.")
         return v

{arthur_common-2.1.58 → arthur_common-2.1.59}/src/arthur_common/models/shield.py

@@ -32,7 +32,7 @@ class MetricType(str, Enum):
     RESPONSE_RELEVANCE = "ResponseRelevance"
     TOOL_SELECTION = "ToolSelection"
 
-    def __str__(self):
+    def __str__(self) -> str:
         return self.value
 
 
@@ -575,20 +575,20 @@ class NewMetricRequest(BaseModel):
         },
     )
 
-    @field_validator("type")
-    def validate_metric_type(cls, value):
-        if value not in MetricType:
-            raise ValueError(
-                f"Invalid metric type: {value}. Valid types are: {', '.join([t.value for t in MetricType])}",
-            )
-        return value
-
     @model_validator(mode="before")
-    def set_config_type(cls, values):
+    def set_config_type(cls, values: dict[str, Any] | None) -> dict[str, Any] | None:
         if not isinstance(values, dict):
             return values
 
-        metric_type = values.get("type")
+        try:
+            metric_type = MetricType(values.get("type", "empty_value"))
+        except ValueError:
+            raise HTTPException(
+                status_code=400,
+                detail=f"Invalid metric type: {values.get('type', 'empty_value')}. Must be one of {[t.value for t in MetricType]}",
+                headers={"full_stacktrace": "false"},
+            )
+
         config_values = values.get("config")
 
         # Map metric types to their corresponding config classes