refactor: introduce JsonPathStr type for JSONPath validation

libs/giskard-checks/src/giskard/checks/builtin/comparison.py

@@ -7,7 +7,7 @@
 from pydantic import Field, model_validator
 
 from ..core.check import Check
-from ..core.extraction import NoMatch, provided_or_resolve, resolve
+from ..core.extraction import JsonPathStr, NoMatch, provided_or_resolve, resolve
 from ..core.result import CheckResult
 from ..core.trace import Trace
 from ..utils.normalization import NormalizationForm, normalize_data
@@ -34,14 +34,14 @@ class ComparisonCheck[InputType, OutputType, TraceType: Trace, ExpectedType](  #
     - Result formatting
     """
 
-    key: str = Field(
+    key: JsonPathStr = Field(
         ..., description="The key to extract the actual value from the trace"
     )
     expected_value: ExpectedType | None = Field(
         default=None,
         description="The expected value to compare against. If None, the expected value is extracted from the trace using the expected_value_key.",
     )
-    expected_value_key: str | NotProvided = Field(
+    expected_value_key: JsonPathStr | NotProvided = Field(
         default=NOT_PROVIDED,
         description="The key to extract the expected value from the trace. If None, the expected value is used as is. If provided, the expected value is extracted from the trace using the expected_value_key.",
     )

libs/giskard-checks/src/giskard/checks/builtin/semantic_similarity.py

@@ -5,7 +5,7 @@
 from pydantic import Field
 
 from ..core.check import Check
-from ..core.extraction import NoMatch, provided_or_resolve, resolve
+from ..core.extraction import JsonPathStr, NoMatch, provided_or_resolve, resolve
 from ..core.mixin import WithEmbeddingMixin
 from ..core.result import CheckResult
 from ..core.trace import Trace
@@ -90,11 +90,11 @@ class SemanticSimilarity[InputType, OutputType, TraceType: Trace](  # pyright: i
     reference_text: str | None = Field(
         default=None, description="The reference text to compare the output with"
     )
-    reference_text_key: str = Field(
+    reference_text_key: JsonPathStr = Field(
         default="trace.last.metadata.reference_text",
         description="The key to extract the reference text from the trace",
     )
-    actual_answer_key: str = Field(
+    actual_answer_key: JsonPathStr = Field(
         default="trace.last.outputs",
         description="The key to extract the actual answer from the trace",
     )

libs/giskard-checks/src/giskard/checks/builtin/string_matching.py

@@ -13,7 +13,7 @@
 from pydantic import Field, model_validator
 
 from ..core.check import Check
-from ..core.extraction import NoMatch, provided_or_resolve
+from ..core.extraction import JsonPathStr, NoMatch, provided_or_resolve
 from ..core.result import CheckResult
 from ..core.trace import Trace
 from ..utils.normalization import NormalizationForm, normalize_string
@@ -91,15 +91,15 @@ class StringMatching[InputType, OutputType, TraceType: Trace](  # pyright: ignor
         default=None,
         description="The text string to search within. If None, extracted from trace using text_key.",
     )
-    text_key: str = Field(
+    text_key: JsonPathStr = Field(
         default="trace.last.outputs",
         description="JSONPath expression to extract the text from the trace (e.g., 'trace.last.outputs.response').",
     )
     keyword: str | None = Field(
         default=None,
         description="The keyword to search for within the text. Either this or keyword_key must be provided.",
     )
-    keyword_key: str | None = Field(
+    keyword_key: JsonPathStr | None = Field(
         default=None,
         description="JSONPath expression to extract the keyword from the trace (e.g., 'trace.last.inputs.expected'). Either this or keyword must be provided.",
     )

libs/giskard-checks/src/giskard/checks/core/extraction.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import Any, override
+from typing import Annotated, Any, override
 
 from giskard.core import NOT_PROVIDED, NotProvided
 from jsonpath_ng import (
@@ -15,11 +15,36 @@
     WhereNot,
     parse,
 )
-from pydantic import BaseModel, Field
+from jsonpath_ng.exceptions import JsonPathLexerError, JsonPathParserError
+from pydantic import AfterValidator, BaseModel, Field
 
 from .trace import Trace
 
 
+class _JsonPathStrMarker:
+    """Marker placed in JsonPathStr metadata. Used by the enforcement test."""
+
+
+_REQUIRED_JSONPATH_PREFIX = "trace."
+
+
+def _validate_jsonpath_syntax(v: str) -> str:
+    if not v.startswith(_REQUIRED_JSONPATH_PREFIX):
+        raise ValueError(
+            f"Invalid JSONPath expression {v!r}: path must start with 'trace.'"
+        )
+    try:
+        parse(v)
+        return v
+    except (JsonPathLexerError, JsonPathParserError) as e:
+        raise ValueError(f"Invalid JSONPath expression {v!r}: {e}") from e
+
+
+JsonPathStr = Annotated[
+    str, AfterValidator(_validate_jsonpath_syntax), _JsonPathStrMarker()
+]
+
+
 class NoMatch(BaseModel):
     """Indicates that a key was provided but no match was found during extraction.
 

libs/giskard-checks/src/giskard/checks/judges/groundedness.py

@@ -5,7 +5,7 @@
 from pydantic import Field
 
 from ..core.check import Check
-from ..core.extraction import provided_or_resolve
+from ..core.extraction import JsonPathStr, provided_or_resolve
 from ..core.trace import Trace
 from .base import BaseLLMCheck
 
@@ -51,14 +51,14 @@ class Groundedness[InputType, OutputType, TraceType: Trace](  # pyright: ignore[
     answer: str | None = Field(
         default=None, description="Input source for the answer to evaluate"
     )
-    answer_key: str = Field(
+    answer_key: JsonPathStr = Field(
         default="trace.last.outputs",
         description="Key to extract the answer from the trace",
     )
     context: str | list[str] | None = Field(
         default=None, description="Input source for the reference context"
     )
-    context_key: str = Field(
+    context_key: JsonPathStr = Field(
         default="trace.last.metadata.context",
         description="Key to extract the context from the trace",
     )

libs/giskard-checks/tests/core/test_extraction.py

@@ -0,0 +1,116 @@
+"""Unit tests for JSONPath validation in extraction.py."""
+
+import pytest
+from giskard.checks.core.extraction import JsonPathStr, _validate_jsonpath_syntax
+from pydantic import BaseModel, ValidationError
+
+
+class TestValidateJsonpathSyntax:
+    """Tests for the _validate_jsonpath_syntax validator function."""
+
+    def test_valid_simple_path(self):
+        assert _validate_jsonpath_syntax("trace.last.outputs") == "trace.last.outputs"
+
+    def test_valid_index_path(self):
+        assert (
+            _validate_jsonpath_syntax("trace.interactions[-1].outputs")
+            == "trace.interactions[-1].outputs"
+        )
+
+    def test_valid_wildcard_path(self):
+        assert (
+            _validate_jsonpath_syntax("trace.interactions[*].outputs")
+            == "trace.interactions[*].outputs"
+        )
+
+    def test_valid_nested_path(self):
+        assert (
+            _validate_jsonpath_syntax("trace.last.metadata.context")
+            == "trace.last.metadata.context"
+        )
+
+    def test_valid_metadata_reference(self):
+        assert (
+            _validate_jsonpath_syntax("trace.last.metadata.reference_text")
+            == "trace.last.metadata.reference_text"
+        )
+
+    def test_syntax_error_unclosed_bracket(self):
+        with pytest.raises(ValueError, match="Invalid JSONPath expression"):
+            _validate_jsonpath_syntax("trace.last.outputs[")
+
+    def test_syntax_error_double_bracket(self):
+        with pytest.raises(ValueError, match="Invalid JSONPath expression"):
+            _validate_jsonpath_syntax("trace.last.outputs[[")
+
+    def test_missing_trace_prefix(self):
+        with pytest.raises(ValueError, match="path must start with 'trace\\.'"):
+            _validate_jsonpath_syntax("last.outputs")
+
+    def test_root_typo(self):
+        with pytest.raises(ValueError, match="path must start with 'trace\\.'"):
+            _validate_jsonpath_syntax("tras.last.outputs")
+
+    def test_empty_string_missing_prefix(self):
+        with pytest.raises(ValueError, match="path must start with 'trace\\.'"):
+            _validate_jsonpath_syntax("")
+
+    def test_just_prefix_is_invalid(self):
+        """JSONPath 'trace.' alone should fail parsing as it's incomplete."""
+        with pytest.raises(ValueError, match="Invalid JSONPath expression"):
+            _validate_jsonpath_syntax("trace.")
+
+    def test_prefix_is_case_sensitive(self):
+        """Prefix must be lowercase 'trace.', not 'Trace.' or 'TRACE.'"""
+        with pytest.raises(ValueError, match="path must start with 'trace\\.'"):
+            _validate_jsonpath_syntax("Trace.last.outputs")
+
+    def test_whitespace_prefix_is_invalid(self):
+        """Leading whitespace should be rejected."""
+        with pytest.raises(ValueError, match="path must start with 'trace\\.'"):
+            _validate_jsonpath_syntax(" trace.last.outputs")
+
+    def test_valid_descendants_operator(self):
+        """JSONPath descendants operator (..) should be valid."""
+        assert _validate_jsonpath_syntax("trace..outputs") == "trace..outputs"
+
+
+class TestJsonPathStrAnnotatedType:
+    """Tests for JsonPathStr as a Pydantic Annotated field type."""
+
+    def test_valid_path_accepted(self):
+        class Model(BaseModel):
+            key: JsonPathStr
+
+        m = Model(key="trace.last.outputs")
+        assert m.key == "trace.last.outputs"
+
+    def test_invalid_syntax_raises_validation_error(self):
+        class Model(BaseModel):
+            key: JsonPathStr
+
+        with pytest.raises(ValidationError, match="Invalid JSONPath expression"):
+            Model(key="trace.last.outputs[")
+
+    def test_missing_trace_prefix_raises_validation_error(self):
+        class Model(BaseModel):
+            key: JsonPathStr
+
+        with pytest.raises(ValidationError, match="path must start with 'trace\\.'"):
+            Model(key="last.outputs")
+
+    def test_optional_jsonpath_str_accepts_none(self):
+        class Model(BaseModel):
+            key: JsonPathStr | None = None
+
+        m = Model(key=None)
+        assert m.key is None
+
+    def test_optional_jsonpath_str_validates_when_provided(self):
+        """Optional fields should still validate when a string value is provided."""
+
+        class Model(BaseModel):
+            key: JsonPathStr | None
+
+        with pytest.raises(ValidationError, match="Invalid JSONPath expression"):
+            Model(key="trace.last.outputs[")

libs/giskard-checks/tests/core/test_jsonpath_enforcement.py

@@ -0,0 +1,79 @@
+"""Enforcement test: all JSONPath fields in Check subclasses must use JsonPathStr."""
+
+import re
+import types
+from typing import Annotated, Union, get_args, get_origin
+
+import giskard.checks.builtin  # noqa: F401 - triggers all @Check.register imports
+from giskard.checks.core.check import Check
+from giskard.checks.core.extraction import _JsonPathStrMarker
+
+JSONPATH_FIELD = re.compile(r"^(key|.+_key)$")
+
+
+def _all_check_subclasses(cls):
+    """Recursively yield all concrete and abstract subclasses of cls.
+
+    This ensures all Check subclasses follow the JsonPathStr convention,
+    even if they're abstract base classes.
+    """
+    for sub in cls.__subclasses__():
+        yield sub
+        yield from _all_check_subclasses(sub)
+
+
+def _annotation_has_marker(annotation) -> bool:
+    """Recursively check if an annotation contains _JsonPathStrMarker.
+
+    This handles complex type annotations like:
+    - Annotated[str, AfterValidator(...), _JsonPathStrMarker()]
+    - JsonPathStr | None
+    - JsonPathStr | NotProvided
+    """
+    if get_origin(annotation) is Annotated:
+        return any(isinstance(m, _JsonPathStrMarker) for m in get_args(annotation)[1:])
+    origin = get_origin(annotation)
+    if origin is Union or isinstance(annotation, types.UnionType):
+        return any(_annotation_has_marker(arg) for arg in get_args(annotation))
+    return False
+
+
+def _has_jsonpath_marker(field_info) -> bool:
+    """Return True if the field uses JsonPathStr.
+
+    Pydantic v2 stores Annotated metadata in two places depending on the type:
+    - Simple `JsonPathStr`: marker is in field_info.metadata
+    - Union `JsonPathStr | None` / `JsonPathStr | NotProvided`: marker is
+      inside field_info.annotation (the Annotated[str, ...] is preserved within
+      the Union at the annotation level)
+    """
+    if any(isinstance(m, _JsonPathStrMarker) for m in field_info.metadata):
+        return True
+    return _annotation_has_marker(field_info.annotation)
+
+
+def test_all_jsonpath_fields_use_jsonpath_str():
+    """Enforce that all JSONPath fields use the JsonPathStr type.
+
+    This architectural constraint ensures:
+    1. JSONPath syntax is validated at model creation time
+    2. All JSONPath expressions start with 'trace.' prefix
+    3. API consistency across all Check subclasses
+    4. Better error messages for users when they provide invalid paths
+    """
+    violations = []
+    for cls in _all_check_subclasses(Check):
+        if not hasattr(cls, "model_fields"):
+            continue
+        for field_name, field_info in cls.model_fields.items():
+            if JSONPATH_FIELD.match(field_name):
+                if not _has_jsonpath_marker(field_info):
+                    violations.append(
+                        f"{cls.__name__}.{field_name}: {field_info.annotation}"
+                    )
+    assert not violations, (
+        "The following JSONPath fields do not use JsonPathStr.\n"
+        "All fields named 'key' or ending in '_key' must be annotated as JsonPathStr "
+        "(or JsonPathStr | None / JsonPathStr | NotProvided):\n"
+        + "\n".join(f"  - {v}" for v in violations)
+    )