feat(checks): add Trace.steps() helper for step-based grouping

libs/giskard-checks/src/giskard/checks/core/interaction/trace.py

@@ -127,7 +127,35 @@ async def with_interaction(
             if generator is not None:
                 await generator.aclose()
 
-    # TODO def steps() -> list[list[Interaction[InputType, OutputType]]]: # Index based
+    def steps(self) -> list[list[Interaction[InputType, OutputType]]]:
+        """Return interactions grouped by logical scenario step.
+
+        Step grouping is derived from the ``step_index`` value stored in each
+        interaction's ``metadata``. Interactions recorded without a
+        ``step_index`` are treated as belonging to step ``0``. The helper
+        preserves interaction order within each step and emits groups in the
+        order the step indices first appear in the trace.
+
+        Returns
+        -------
+        list[list[Interaction[InputType, OutputType]]]
+            Interactions grouped by logical step. Empty traces return ``[]``.
+        """
+        grouped: list[list[Interaction[InputType, OutputType]]] = []
+        current_step_index: int | None = None
+
+        for interaction in self.interactions:
+            step_index = interaction.metadata.get("step_index", 0)
+            if not isinstance(step_index, int) or isinstance(step_index, bool):
+                step_index = 0
+
+            if current_step_index != step_index:
+                grouped.append([])
+                current_step_index = step_index
+
+            grouped[-1].append(interaction)
+
+        return grouped
 
     def __rich_console__(
         self, console: Console, options: ConsoleOptions

libs/giskard-checks/src/giskard/checks/scenarios/runner.py

@@ -87,6 +87,30 @@ def _resolve_trace_type[InputType, OutputType, TraceType: Trace[Any, Any]](
     return cast(type[TraceType], inferred if inferred is not None else Trace)
 
 
+def _tag_interactions_with_step_index[InputType, OutputType, TraceType: Trace[Any, Any]](
+    trace: TraceType,
+    *,
+    step_index: int,
+    previous_count: int,
+) -> TraceType:
+    if len(trace.interactions) <= previous_count:
+        return trace
+
+    interactions = list(trace.interactions[:previous_count])
+    interactions.extend(
+        interaction.model_copy(
+            update={
+                "metadata": {
+                    **interaction.metadata,
+                    "step_index": step_index,
+                }
+            }
+        )
+        for interaction in trace.interactions[previous_count:]
+    )
+    return cast(TraceType, trace.model_copy(update={"interactions": interactions}))
+
+
 class ScenarioRunner:
     """Execute scenarios by running their steps sequentially.
 
@@ -149,8 +173,14 @@ async def _run_once[InputType, OutputType, TraceType: Trace[Any, Any]](
             properties=shape_props,
         )
 
-        for step in steps:
+        for step_index, step in enumerate(steps):
+            previous_count = len(trace.interactions)
             trace = await trace.with_interactions(*step.interacts)
+            trace = _tag_interactions_with_step_index(
+                trace,
+                step_index=step_index,
+                previous_count=previous_count,
+            )
 
             test_case = TestCase(
                 trace=trace,

libs/giskard-checks/tests/trace/test_trace.py

@@ -0,0 +1,54 @@
+import pytest
+from giskard.checks import Equals, Interaction, Scenario, Trace
+from giskard.checks.scenarios.runner import ScenarioRunner
+
+
+def test_trace_steps_returns_empty_list_for_empty_trace():
+    trace = Trace[str, str]()
+
+    assert trace.steps() == []
+
+
+def test_trace_steps_groups_single_step_without_metadata():
+    trace = Trace[str, str](
+        interactions=[
+            Interaction(inputs="a", outputs="A"),
+            Interaction(inputs="b", outputs="B"),
+        ]
+    )
+
+    assert trace.steps() == [trace.interactions]
+
+
+def test_trace_steps_groups_multiple_steps_by_step_index():
+    first = Interaction(inputs="a", outputs="A", metadata={"step_index": 0})
+    second = Interaction(inputs="b", outputs="B", metadata={"step_index": 0})
+    third = Interaction(inputs="c", outputs="C", metadata={"step_index": 1})
+    fourth = Interaction(inputs="d", outputs="D", metadata={"step_index": 2})
+
+    trace = Trace[str, str](interactions=[first, second, third, fourth])
+
+    assert trace.steps() == [
+        [first, second],
+        [third],
+        [fourth],
+    ]
+
+
+@pytest.mark.asyncio
+async def test_scenario_runner_tags_trace_interactions_with_step_index():
+    scenario = (
+        Scenario("multi_step")
+        .interact("hello", "HELLO")
+        .check(Equals(expected_value="HELLO", key="trace.last.outputs"))
+        .interact("world", "WORLD")
+    )
+
+    result = await ScenarioRunner().run(scenario)
+
+    assert result.final_trace.steps() == [
+        [result.final_trace.interactions[0]],
+        [result.final_trace.interactions[1]],
+    ]
+    assert result.final_trace.interactions[0].metadata["step_index"] == 0
+    assert result.final_trace.interactions[1].metadata["step_index"] == 1