Record ordinary judge call provenance

tests/v1/test_e2e.py

@@ -322,7 +322,7 @@ async def test_rubric_judge(run_v1, tmp_path):
     assert trace.ok
     assert trace.rewards["rubric"].score > 0  # the judge's verdict landed in the reward
     assert trace.metrics["rubric/always_yes"] == 1.0
-    assert trace.info["judge"]  # the call was recorded onto the trace
+    assert trace.judge_calls  # the call was recorded onto the trace
 
 
 @pytest.mark.e2e

tests/v1/test_judges.py

@@ -1,16 +1,21 @@
-"""Pluggable judges: plugin resolution, base-`TaskConfig.judges` narrowing, the built-in
-`reference` / `rubric` judges, and `Task.score` running plugged judges after the decorated
-rewards. Judge model calls are faked at `Judge.complete` — no network."""
+"""Pluggable judges, scoring integration, and the typed provider-call lifecycle."""
 
+import asyncio
 import json
 import re
 
+import httpx
 import pytest
+from openai import APITimeoutError, ContentFilterFinishReasonError
+from openai.resources.chat.completions import AsyncCompletions
+from openai.types.chat import ChatCompletion
+from pydantic import BaseModel
 
 import verifiers.v1 as vf
 from verifiers.v1.graph import MessageNode
-from verifiers.v1.judge import Judge, JudgeResponse
+from verifiers.v1.judge import JudgeResponse
 from verifiers.v1.loaders import judge_class, judge_config_type, load_judge
+from verifiers.v1.push import trace_to_sample
 from verifiers.v1.types import AssistantMessage, UserMessage
 
 RUBRIC_TOML = """
@@ -52,20 +57,185 @@ def make_trace(
     )
 
 
+def model_response(
+    text: str, *, finish_reason: str = "stop", refusal: str | None = None
+) -> ChatCompletion:
+    return ChatCompletion.model_validate(
+        {
+            "id": "fake",
+            "object": "chat.completion",
+            "created": 0,
+            "model": "fake",
+            "choices": [
+                {
+                    "index": 0,
+                    "finish_reason": finish_reason,
+                    "message": {
+                        "role": "assistant",
+                        "content": text,
+                        "refusal": refusal,
+                    },
+                }
+            ],
+            "usage": {
+                "prompt_tokens": 2,
+                "completion_tokens": 1,
+                "total_tokens": 3,
+            },
+        }
+    )
+
+
+async def test_complete_uses_sdk_strict_schema(monkeypatch):
+    request = {}
+
+    class Verdict(BaseModel):
+        score: int | None = None
+
+    async def fake_response(self, **kwargs):
+        request.update(kwargs)
+        return model_response('{"score": 1}')
+
+    monkeypatch.setattr(AsyncCompletions, "create", fake_response)
+    trace = make_trace()
+    response = await vf.Judge().complete("grade this", trace=trace, schema=Verdict)
+    schema = request["response_format"]["json_schema"]["schema"]
+    assert schema["additionalProperties"] is False
+    assert schema["required"] == ["score"]
+    assert response.parsed == Verdict(score=1)
+    assert trace.judge_calls[0].call.sampling is not None
+    assert "response_format" not in trace.judge_calls[0].call.sampling.model_dump()
+    assert trace.judge_calls[0].call.usage is not None
+    assert trace.judge_calls[0].call.usage.total_tokens == 3
+    restored = vf.Trace[QAData].model_validate_json(trace.model_dump_json())
+    assert restored.judge_calls == trace.judge_calls
+    assert trace_to_sample(trace)["info"]["judge_calls"] == [
+        trace.judge_calls[0].model_dump(mode="json", exclude_none=True)
+    ]
+
+
+async def test_complete_rejects_truncated_structured_output(monkeypatch):
+    class Verdict(BaseModel):
+        score: int
+
+    async def fake_response(self, **kwargs):
+        return model_response('{"score": 1}', finish_reason="length")
+
+    monkeypatch.setattr(AsyncCompletions, "create", fake_response)
+    trace = make_trace()
+    with pytest.raises(ValueError, match="structured output was truncated"):
+        await vf.Judge().complete("grade this", trace=trace, schema=Verdict)
+    assert trace.judge_calls[-1].call.finish_reason == "length"
+    assert trace.judge_calls[-1].error is not None
+
+
+async def test_complete_rejects_content_filtered_output(monkeypatch):
+    class Verdict(BaseModel):
+        score: int
+
+    async def fake_response(self, **kwargs):
+        return model_response('{"score": 1}', finish_reason="content_filter")
+
+    monkeypatch.setattr(AsyncCompletions, "create", fake_response)
+    trace = make_trace()
+    with pytest.raises(ContentFilterFinishReasonError):
+        await vf.Judge().complete("grade this", trace=trace, schema=Verdict)
+    judge_call = trace.judge_calls[-1]
+    assert judge_call.call.error is None
+    assert judge_call.error is not None
+    assert judge_call.error.type == "ContentFilterFinishReasonError"
+
+
+async def test_complete_records_refusal_and_cancellation(monkeypatch):
+    async def refuse(self, **kwargs):
+        return model_response("", refusal="cannot grade")
+
+    monkeypatch.setattr(AsyncCompletions, "create", refuse)
+    trace = make_trace()
+    with pytest.raises(ValueError, match="refused"):
+        await vf.Judge().complete("grade this", trace=trace)
+    assert trace.judge_calls[-1].refusal == "cannot grade"
+    assert trace.judge_calls[-1].call.error is None
+    assert trace.judge_calls[-1].error is not None
+
+    async def cancel(self, **kwargs):
+        raise asyncio.CancelledError
+
+    monkeypatch.setattr(AsyncCompletions, "create", cancel)
+    trace = make_trace()
+    with pytest.raises(asyncio.CancelledError):
+        await vf.Judge().complete("grade this", trace=trace)
+    provider_error = trace.judge_calls[-1].call.error
+    assert provider_error is not None and provider_error.type == "CancelledError"
+    assert trace.judge_calls[-1].error is None
+
+
+@pytest.mark.parametrize(
+    ("configured", "sampling", "expected"),
+    [
+        ({"max_completion_tokens": 30}, {}, {"max_completion_tokens": 30}),
+        (
+            {"max_tokens": 100},
+            {"max_completion_tokens": 20},
+            {"max_completion_tokens": 20},
+        ),
+    ],
+)
+async def test_complete_preserves_provider_sampling_aliases(
+    monkeypatch, configured, sampling, expected
+):
+    request = {}
+
+    async def fake_response(self, **kwargs):
+        request.update(kwargs)
+        return model_response("yes")
+
+    monkeypatch.setattr(AsyncCompletions, "create", fake_response)
+    await vf.Judge(vf.JudgeConfig(sampling=configured)).complete(
+        "grade this", model="ignored", **sampling
+    )
+    assert request["model"] == "openai/gpt-5.4-nano"
+    limits = {
+        key: request[key]
+        for key in ("max_tokens", "max_completion_tokens")
+        if key in request
+    }
+    assert limits == expected
+
+
+async def test_judge_provider_error_uses_shared_retry_type(monkeypatch):
+    class JudgeTask(vf.Task[QAData]):
+        @vf.reward
+        async def judged(self, trace: vf.Trace) -> float:
+            await vf.Judge().complete("grade this", trace=trace)
+            return 1.0
+
+    async def fail(self, **kwargs):
+        raise APITimeoutError(request=httpx.Request("POST", "https://judge.test"))
+
+    monkeypatch.setattr(AsyncCompletions, "create", fail)
+    trace = make_trace()
+    with pytest.raises(vf.ProviderError):
+        await JudgeTask(trace.task.data).score(trace)
+    assert trace.judge_calls[-1].error is None
+    provider_error = trace.judge_calls[-1].call.error
+    assert provider_error is not None and provider_error.type == "ProviderError"
+    assert "judged" not in trace.rewards
+
+
 @pytest.fixture
 def fake_judge_model(monkeypatch):
     """Fake the judge's model call, recording each prompt for assertions. Rubric calls (a JSON
     `verdicts` instruction or a `schema`) reply one reasoned verdict per `- name: text` criterion,
     "yes" iff it mentions Paris; other judges reply plain "yes"/"no" by the response block."""
     prompts: list[str] = []
 
-    async def fake_complete(
-        self, messages, *, trace=None, schema=None, parse=None, **sampling
-    ):
+    async def fake_response(self, **kwargs):
+        messages = kwargs["messages"][0]["content"]
         prompts.append(messages)
         # Rubric calls carry criteria lines + a JSON `verdicts` instruction; reply one verdict per
         # criterion (yes iff its text mentions Paris) with a reason. Other judges get plain yes/no.
-        if schema is not None or '"verdicts"' in messages:
+        if "response_format" in kwargs or '"verdicts"' in messages:
             verdicts = [
                 {
                     "name": name,
@@ -74,23 +244,12 @@ async def fake_complete(
                 }
                 for name, text in re.findall(r"^- ([^:]+): (.+)$", messages, re.M)
             ]
-            response = JudgeResponse(
-                text=json.dumps({"verdicts": verdicts}),
-                parsed=schema.model_validate({"verdicts": verdicts})
-                if schema
-                else None,
-            )
+            text = json.dumps({"verdicts": verdicts})
         else:
-            response = JudgeResponse(
-                text="yes" if "Paris" in messages.split("Response:")[-1] else "no"
-            )
-        if parse is not None:
-            response.parsed = parse(response)
-        if trace is not None:
-            trace.record_judge(response)
-        return response
-
-    monkeypatch.setattr(Judge, "complete", fake_complete)
+            text = "yes" if "Paris" in messages.split("Response:")[-1] else "no"
+        return model_response(text)
+
+    monkeypatch.setattr(AsyncCompletions, "create", fake_response)
     return prompts
 
 
@@ -215,10 +374,12 @@ async def test_reference_score(fake_judge_model):
     assert (
         "Capital of France?" in fake_judge_model[0]
     )  # the task prompt is in the judge prompt
-    assert len(trace.info["judge"]) == 1  # the call is recorded onto the trace
+    assert len(trace.judge_calls) == 1  # the call is recorded onto the trace
+    assert trace.judge_calls[0].error is None
 
     trace = make_trace(reply="It is Rome.")
     assert await vf.ReferenceJudge().score(trace.task.data, trace) == 0.0
+    assert trace.judge_calls[0].error is None
 
     judge = vf.ReferenceJudge(vf.ReferenceJudgeConfig(answer_field="gold"))
     with pytest.raises(ValueError, match="no 'gold' field"):  # misconfig raises, not 0
@@ -432,21 +593,12 @@ async def test_reference_choices(fake_judge_model):
 
 async def test_error_attribution(monkeypatch, tmp_path):
     # The policy: a MODEL failure scores 0.0; a JUDGE failure errors the rollout (raises
-    # out of Task.score as a TaskError) so training skips the sample instead of
+    # out of Task.score) so training skips the sample instead of
     # punishing the model for a broken judge.
-    async def gibberish_judge(
-        self, messages, *, trace=None, schema=None, parse=None, **s
-    ):
-        response = JudgeResponse(text="as an AI language model I cannot grade this")
-        try:
-            if parse is not None:
-                response.parsed = parse(response)
-            return response
-        finally:
-            if trace is not None:
-                trace.record_judge(response)
-
-    monkeypatch.setattr(Judge, "complete", gibberish_judge)
+    async def gibberish_judge(self, **kwargs):
+        return model_response("as an AI language model I cannot grade this")
+
+    monkeypatch.setattr(AsyncCompletions, "create", gibberish_judge)
     taskset = JudgedTaskset(
         JudgedConfig.model_validate({"task": {"judges": [{"id": "reference"}]}})
     )
@@ -461,7 +613,7 @@ async def gibberish_judge(
             trace, runtime=None
         )
     assert "reference" not in trace.rewards
-    assert len(trace.info["judge"]) == 1  # the billed call is still recorded
+    assert trace.judge_calls[0].error is not None
 
 
 # --- rubric --------------------------------------------------------------------------------
@@ -521,33 +673,34 @@ async def test_rubric_score(tmp_path, fake_judge_model):
     trace = make_trace()
     assert await judge.score(trace.task.data, trace) == 0.75
     assert trace.metrics == {"rubric/mentions_paris": 1.0, "rubric/is_polite": 0.0}
-    assert len(trace.info["judge"]) == 1  # one call for the whole rubric
+    assert len(trace.judge_calls) == 1  # one call for the whole rubric
 
 
 async def test_rubric_verdict_mismatch_raises(tmp_path, monkeypatch):
     # A reply that doesn't verdict exactly the rubric's criteria is a judge failure: raise
     # (-> rollout error), don't guess or silently score 0.
-    async def wrong_names(self, messages, *, trace=None, schema=None, parse=None, **s):
+    async def wrong_names(self, **kwargs):
         verdicts = {"verdicts": [{"name": "typo", "reason": "x", "verdict": "yes"}]}
-        return JudgeResponse(text=json.dumps(verdicts), parsed=None)
+        return model_response(json.dumps(verdicts))
 
-    monkeypatch.setattr(Judge, "complete", wrong_names)
+    monkeypatch.setattr(AsyncCompletions, "create", wrong_names)
     judge = rubric_judge(tmp_path)
     trace = make_trace()
     with pytest.raises(ValueError, match="expected the batch"):
         await judge.score(trace.task.data, trace)
+    assert trace.judge_calls[0].error is not None
 
 
 CHOICES_TOML = '[[criteria]]\nname = "depth"\ntext = "How thorough?"\nchoices = ["none", "partial", "good"]\n'
 
 
 async def test_rubric_choices_normalize(tmp_path, monkeypatch):
     # Ordered choices (worst→best) score by rank: "partial" of ["none","partial","good"] -> 0.5.
-    async def graded(self, messages, *, trace=None, schema=None, parse=None, **s):
+    async def graded(self, **kwargs):
         v = {"verdicts": [{"name": "depth", "reason": "r", "verdict": "partial"}]}
-        return JudgeResponse(text=json.dumps(v), parsed=None)
+        return model_response(json.dumps(v))
 
-    monkeypatch.setattr(Judge, "complete", graded)
+    monkeypatch.setattr(AsyncCompletions, "create", graded)
     judge = rubric_judge(tmp_path, body=CHOICES_TOML, name="q")
     trace = make_trace()
     assert await judge.score(trace.task.data, trace) == 0.5
@@ -556,11 +709,11 @@ async def graded(self, messages, *, trace=None, schema=None, parse=None, **s):
 
 async def test_rubric_off_menu_answer_raises(tmp_path, monkeypatch):
     # A verdict that isn't one of the criterion's choices is a judge failure, not a 0.
-    async def off_menu(self, messages, *, trace=None, schema=None, parse=None, **s):
+    async def off_menu(self, **kwargs):
         v = {"verdicts": [{"name": "depth", "reason": "r", "verdict": "maybe"}]}
-        return JudgeResponse(text=json.dumps(v), parsed=None)
+        return model_response(json.dumps(v))
 
-    monkeypatch.setattr(Judge, "complete", off_menu)
+    monkeypatch.setattr(AsyncCompletions, "create", off_menu)
     judge = rubric_judge(tmp_path, body=CHOICES_TOML)
     trace = make_trace()
     with pytest.raises(ValueError, match="expected one of"):
@@ -630,9 +783,7 @@ async def test_task_score_runs_plugged_judges(tmp_path, fake_judge_model):
     assert (
         trace.rewards["quality"].score == 0.75
     )  # the rubric's aggregate, under its `name`
-    assert (
-        len(trace.info["judge"]) == 2
-    )  # every judge call recorded (rubric = one call)
+    assert len(trace.judge_calls) == 2  # every judge call recorded (rubric = one call)
 
 
 async def test_task_without_judges_scores_as_before():

verifiers/v1/__init__.py

@@ -106,6 +106,7 @@
     Error,
     EvalRunInfo,
     GenerationSpan,
+    JudgeCall,
     ModelCall,
     Reward,
     RunInfo,
@@ -179,6 +180,7 @@
     "AgentInfo",
     "RunInfo",
     "EvalRunInfo",
+    "JudgeCall",
     "ModelCall",
     "TrainRunInfo",
     "VersionInfo",