feat(rubrics_verifier): grade a multi-turn run as of turn k

src/agent_env/task_step/task_steps/verifiers/rubrics_verifier.py

@@ -7,7 +7,7 @@
 import logging
 import re
 import uuid
-from dataclasses import dataclass, field
+from dataclasses import dataclass, field, replace
 from typing import ClassVar, Optional
 
 from agent_env.a2a_agent.object_transfer import (
@@ -254,6 +254,49 @@ def _is_per_criterion(tf: TrajectoryFilter | None) -> bool:
             and tf.frame_strategy == FrameStrategy.PER_CRITERION)
 
 
+
+def _validated_turn(turn: Optional[int]) -> Optional[int]:
+    if turn is None:
+        return None
+    if isinstance(turn, bool) or not isinstance(turn, int) or turn < 1:
+        raise ValueError(f"grade_up_to_turn must be a positive (1-based) turn number, got {turn!r}")
+    return turn
+
+
+def _as_of_turn(prompt_response: PromptResponse, turn: Optional[int]) -> PromptResponse:
+    """A copy of a multi-turn PromptResponse cut off after ``turn``, or the original when there is
+    nothing to cut. The stored response is the conversation's last reply, so a cut copy points the
+    judge at the end of the turn's own trajectory instead of leaking later turns. A failed turn
+    ends the conversation, so every turn before the last completed: the stored error is the last
+    turn's and is dropped from a cut copy."""
+    turn_urls = prompt_response.target_agent_per_turn_trajectory_object_urls or []
+    prompt_parts = prompt_response.source_agent_per_turn_prompt_parts
+    # Per-turn prompt parts are recorded every turn; per-turn trajectories only when the agent
+    # exposes them, so the turn count comes from whichever is longer.
+    n_turns = max(len(turn_urls), len(prompt_parts or []))
+    if turn is None or turn >= n_turns:
+        return prompt_response
+    if len(turn_urls) < n_turns or turn_urls[turn - 1] is None:
+        raise RuntimeError(
+            f"grade_up_to_turn={turn}: the run has {n_turns} turns but no stored trajectory for turn "
+            f"{turn}, so the judge would see either later turns or nothing from turn {turn}."
+        )
+    return replace(
+        prompt_response,
+        target_agent_per_turn_trajectory_object_urls=turn_urls[:turn],
+        agent_trajectory_object_url=turn_urls[turn - 1],
+        source_agent_per_turn_prompt_parts=prompt_parts[:turn] if prompt_parts is not None else None,
+        structured_output=None,
+        error_type=None,
+        error_code=None,
+        error_message=None,
+        response=(
+            f"(Graded as of turn {turn} of {len(turn_urls)}: the agent's reply at this point is the "
+            f"final message of turn {turn}'s trajectory.)"
+        ),
+    )
+
+
 class RubricsVerifierTaskStep(TaskStep):
     type: ClassVar[str] = "rubrics_verifier"
     entity_refs = (EntityRef.agent("judge_a2a_agent_id"),)
@@ -296,10 +339,14 @@ def __init__(
         output_format: JudgeOutputFormat | str = JudgeOutputFormat.RUBRIC_BINARY,
         grading_policy_prompt: Optional[str] = None,
         context_facts_for_judge: Optional[dict[str, str]] = None,
+        grade_up_to_turn: Optional[int] = None,
         depends_on: Optional[list[TaskStepDependency]] = None,
         fail_task_on_error: bool = True,
     ):
         super().__init__(id, version, depends_on=depends_on, fail_task_on_error=fail_task_on_error)
+        # Grade a multi-turn run as it stood after this (1-based) turn: the judge sees turns 1..k only,
+        # so a per-turn env snapshot can be scored on what the agent had done by then. None grades all.
+        self.grade_up_to_turn = _validated_turn(grade_up_to_turn)
         # Task-authored rules the criteria alone do not state, appended as a "## Grading policy" section;
         # with rubric_evidence its findings come back as `checks`.
         self.grading_policy_prompt = grading_policy_prompt
@@ -412,6 +459,8 @@ def to_dict(self) -> dict:
         base["output_format"] = self.output_format.value
         base["grading_policy_prompt"] = self.grading_policy_prompt
         base["context_facts_for_judge"] = self.context_facts_for_judge
+        if self.grade_up_to_turn is not None:
+            base["grade_up_to_turn"] = self.grade_up_to_turn
         return base
 
     @classmethod
@@ -440,6 +489,7 @@ def from_dict(cls, data: dict) -> RubricsVerifierTaskStep:
             output_format=data.get("output_format", "rubric_binary"),
             grading_policy_prompt=data.get("grading_policy_prompt"),
             context_facts_for_judge=data.get("context_facts_for_judge"),
+            grade_up_to_turn=data.get("grade_up_to_turn"),
         )
 
     def _resolved_criteria(self, context: TaskStepContext) -> list[dict]:
@@ -456,6 +506,19 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext:
         prompt_response = next((pr for pr in context.prompt_responses if pr.prompt_id == self.prompt_id), None)
         if prompt_response is None:
             raise RuntimeError(f"PromptResponse with prompt_id='{self.prompt_id}' not found in context")
+        overrides_for_step = self.step_param_overrides(context)
+        grade_up_to_turn = (
+            _validated_turn(overrides_for_step["grade_up_to_turn"])
+            if "grade_up_to_turn" in overrides_for_step else self.grade_up_to_turn
+        )
+        as_of = _as_of_turn(prompt_response, grade_up_to_turn)
+        if as_of is not prompt_response and not self.use_trajectory:
+            # A cut copy's reply only points at turn k's trajectory, which use_trajectory=False never loads.
+            raise RuntimeError(
+                f"Verifier '{self.verifier_id}': grade_up_to_turn={grade_up_to_turn} needs "
+                "use_trajectory=True — turn k's reply is only in its trajectory."
+            )
+        prompt_response = as_of
 
         if prompt_response.error_type:
             logger.warning(f"Skipping verification '{self.verifier_id}': prompt had error_type={prompt_response.error_type}")
@@ -697,6 +760,8 @@ async def execute(self, context: TaskStepContext) -> TaskStepContext:
             judge_call_fields["judge_output_retries"] = judge_output_retries
         if judge_output_discrepancies:
             judge_call_fields["judge_output_discrepancies"] = judge_output_discrepancies
+        if grade_up_to_turn is not None:
+            judge_call_fields["graded_up_to_turn"] = grade_up_to_turn
         self._write_entries(context, verification_results, judge_call_fields)
         return context
 

tst/unit/task_step/verifiers/test_rubrics_verifier_grade_up_to_turn.py

@@ -0,0 +1,173 @@
+"""grade_up_to_turn: grade a multi-turn run as it stood after turn k.
+
+A per-turn env snapshot is only gradable if the judge also sees just turns 1..k: otherwise a
+criterion about what the agent said is credited from a later turn. Same seams as
+test_rubrics_verifier_multiturn.py: patch the trajectory read, capture the judge's prompt.
+"""
+from __future__ import annotations
+
+import pytest
+
+from agent_env.task_step.context import PromptResponse, TaskStepContext
+from agent_env.task_step.task_steps.verifiers.rubrics_verifier import RubricsVerifierTaskStep
+
+TURNS = {"s3://t/1.json": "TURN_ONE_EVENTS", "s3://t/2.json": "TURN_TWO_EVENTS", "s3://t/3.json": "TURN_THREE_EVENTS"}
+
+
+def _verifier(**kw) -> RubricsVerifierTaskStep:
+    params = dict(
+        id="verify", version=1, criteria=[{"id": "c1", "description": "outcome"}],
+        prompt_id="p1", use_agent_judge=False, use_trajectory=True, verifier_id="vt",
+    )
+    params.update(kw)
+    return RubricsVerifierTaskStep(**params)
+
+
+def _three_turn_context() -> TaskStepContext:
+    return TaskStepContext(prompt_responses=[PromptResponse(
+        prompt_id="p1", response="FINAL_REPLY", prompt_text="do it",
+        target_agent_per_turn_trajectory_object_urls=list(TURNS),
+        agent_trajectory_object_url="s3://t/3.json",
+        source_agent_per_turn_prompt_parts=[None, [{"kind": "text", "text": "a"}], [{"kind": "text", "text": "b"}]],
+        structured_output={"from": "the last turn"},
+    )])
+
+
+def _judge_prompt(monkeypatch, verifier) -> dict:
+    captured: dict = {}
+
+    async def fake_run(*, eval_prompt, **kwargs):
+        captured["eval_prompt"] = eval_prompt
+        return ([{**verifier.criteria[0], "score": 1.0, "result": True}], 0, [], None)
+
+    monkeypatch.setattr(verifier, "_read_trajectory_text", lambda uri: TURNS[uri])
+    monkeypatch.setattr(verifier, "_run_judge_with_output_retries", fake_run)
+    return captured
+
+
+@pytest.mark.asyncio
+async def test_the_judge_sees_only_the_turns_up_to_k(monkeypatch):
+    v = _verifier(grade_up_to_turn=2)
+    captured = _judge_prompt(monkeypatch, v)
+    ctx = _three_turn_context()
+
+    await v.execute(ctx)
+
+    prompt = captured["eval_prompt"]
+    assert "TURN_ONE_EVENTS" in prompt and "TURN_TWO_EVENTS" in prompt
+    assert "TURN_THREE_EVENTS" not in prompt
+    assert "FINAL_REPLY" not in prompt
+    assert "as of turn 2 of 3" in prompt
+    assert ctx.metadata["verifications"]["vt"]["graded_up_to_turn"] == 2
+    stored = ctx.prompt_responses[0]
+    assert stored.response == "FINAL_REPLY"
+    assert len(stored.target_agent_per_turn_trajectory_object_urls) == 3
+    assert stored.structured_output == {"from": "the last turn"}
+
+
+@pytest.mark.asyncio
+async def test_turn_one_is_graded_from_its_own_trajectory(monkeypatch):
+    v = _verifier(grade_up_to_turn=1)
+    captured = _judge_prompt(monkeypatch, v)
+
+    await v.execute(_three_turn_context())
+
+    prompt = captured["eval_prompt"]
+    assert "TURN_ONE_EVENTS" in prompt
+    assert "TURN_TWO_EVENTS" not in prompt and "TURN_THREE_EVENTS" not in prompt
+    assert "FINAL_REPLY" not in prompt
+
+
+@pytest.mark.asyncio
+async def test_a_per_run_override_sets_the_turn(monkeypatch):
+    v = _verifier()
+    captured = _judge_prompt(monkeypatch, v)
+    ctx = _three_turn_context()
+    ctx.metadata["user_overrides"] = {"step_params": {"verify": {"grade_up_to_turn": 1}}}
+
+    await v.execute(ctx)
+
+    assert "TURN_TWO_EVENTS" not in captured["eval_prompt"]
+    assert ctx.metadata["verifications"]["vt"]["graded_up_to_turn"] == 1
+    assert v.grade_up_to_turn is None
+
+
+@pytest.mark.asyncio
+async def test_the_last_turn_or_later_grades_the_whole_run(monkeypatch):
+    v = _verifier(grade_up_to_turn=3)
+    captured = _judge_prompt(monkeypatch, v)
+
+    await v.execute(_three_turn_context())
+
+    prompt = captured["eval_prompt"]
+    assert "TURN_THREE_EVENTS" in prompt and "FINAL_REPLY" in prompt
+
+
+@pytest.mark.asyncio
+async def test_without_it_every_turn_is_graded(monkeypatch):
+    v = _verifier()
+    captured = _judge_prompt(monkeypatch, v)
+    ctx = _three_turn_context()
+
+    await v.execute(ctx)
+
+    assert "TURN_THREE_EVENTS" in captured["eval_prompt"]
+    assert "graded_up_to_turn" not in ctx.metadata["verifications"]["vt"]
+
+
+@pytest.mark.parametrize("bad", [0, -1, True, "2", 1.5])
+def test_rejects_a_turn_that_is_not_a_positive_integer(bad):
+    with pytest.raises(ValueError, match="grade_up_to_turn"):
+        _verifier(grade_up_to_turn=bad)
+
+
+def test_round_trips_and_is_omitted_when_unset():
+    v = _verifier(grade_up_to_turn=2)
+    assert v.to_dict()["grade_up_to_turn"] == 2
+    assert RubricsVerifierTaskStep.from_dict(v.to_dict()).grade_up_to_turn == 2
+    assert "grade_up_to_turn" not in _verifier().to_dict()
+
+
+@pytest.mark.asyncio
+async def test_a_later_turns_failure_does_not_skip_an_earlier_turn(monkeypatch):
+    v = _verifier(grade_up_to_turn=2)
+    captured = _judge_prompt(monkeypatch, v)
+    ctx = _three_turn_context()
+    ctx.prompt_responses[0].error_type = "agent_failed"
+
+    await v.execute(ctx)
+
+    assert "TURN_TWO_EVENTS" in captured["eval_prompt"]
+    assert ctx.metadata["verifications"]["vt"]["graded_up_to_turn"] == 2
+    assert ctx.prompt_responses[0].error_type == "agent_failed"
+
+
+@pytest.mark.asyncio
+async def test_a_cut_without_the_trajectory_is_refused(monkeypatch):
+    v = _verifier(grade_up_to_turn=2, use_trajectory=False)
+    _judge_prompt(monkeypatch, v)
+
+    with pytest.raises(RuntimeError, match="use_trajectory=True"):
+        await v.execute(_three_turn_context())
+
+
+@pytest.mark.asyncio
+async def test_a_cut_without_per_turn_trajectories_is_refused(monkeypatch):
+    v = _verifier(grade_up_to_turn=2)
+    _judge_prompt(monkeypatch, v)
+    ctx = _three_turn_context()
+    ctx.prompt_responses[0].target_agent_per_turn_trajectory_object_urls = None
+
+    with pytest.raises(RuntimeError, match="no stored trajectory for turn 2"):
+        await v.execute(ctx)
+
+
+@pytest.mark.asyncio
+async def test_a_cut_at_a_turn_whose_trajectory_was_lost_is_refused(monkeypatch):
+    v = _verifier(grade_up_to_turn=2)
+    _judge_prompt(monkeypatch, v)
+    ctx = _three_turn_context()
+    ctx.prompt_responses[0].target_agent_per_turn_trajectory_object_urls = ["s3://t/1.json", None, "s3://t/3.json"]
+
+    with pytest.raises(RuntimeError, match="no stored trajectory for turn 2"):
+        await v.execute(ctx)