rewardkit: add REWARDKIT_MODEL override for agent judges

docs/content/docs/rewardkit/index.mdx

@@ -276,16 +276,22 @@ All flags are optional with sensible defaults. Passing multiple test directories
 
 ## Provider routing
 
-Judges call LiteLLM, which reads credentials from host environment variables. Two flags help you avoid pinning a rubric to one provider:
+Judges call LiteLLM, which reads credentials from host environment variables. A few flags help you avoid pinning a rubric to one provider:
 
 - `--je KEY=VALUE` sets env vars for the run (repeatable). Same shape as Harbor's `--ve`.
 - `--judge MODEL_OR_AGENT` overwrites the rubric's `[judge].judge` field. This can also be done by setting `REWARDKIT_JUDGE` in the environment, so Harbor users can pass `--ve REWARDKIT_JUDGE=...` to do the same thing.
+- `--model MODEL` overwrites the rubric's `[judge].model` field when the judge is an agent (e.g. `claude-code`, `codex`). Equivalent to setting `REWARDKIT_MODEL`; Harbor users can pass `--ve REWARDKIT_MODEL=...`.
 
 ```bash
 rewardkit /tests \
   --judge bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0 \
   --je AWS_ACCESS_KEY_ID=$AWS_ACCESS_KEY_ID \
   --je AWS_REGION_NAME=us-east-1
+
+# Agent judge with an overridden model
+rewardkit /tests \
+  --judge claude-code \
+  --model anthropic/claude-sonnet-4-6
 ```
 
 The [LiteLLM provider docs](https://docs.litellm.ai/docs/providers) list the env var configurations for each provider.

packages/rewardkit/src/rewardkit/__main__.py

@@ -60,12 +60,20 @@ def main() -> None:
     )
     parser.add_argument(
         "--judge",
-        "--j",
+        "-j",
         default=None,
         metavar="MODEL_OR_AGENT",
         help="Override the rubric's [judge].judge field. "
         "Equivalent to setting REWARDKIT_JUDGE.",
     )
+    parser.add_argument(
+        "--model",
+        "-m",
+        default=None,
+        metavar="MODEL",
+        help="Override the rubric's [judge].model field (used when the judge is an agent). "
+        "Equivalent to setting REWARDKIT_MODEL.",
+    )
 
     args = parser.parse_args()
 
@@ -76,6 +84,8 @@ def main() -> None:
         os.environ[key] = value
     if args.judge:
         os.environ["REWARDKIT_JUDGE"] = args.judge
+    if args.model:
+        os.environ["REWARDKIT_MODEL"] = args.model
     concurrency_kwargs = dict(
         max_concurrent_programmatic=args.max_concurrent_programmatic,
         max_concurrent_llm=args.max_concurrent_llm,

packages/rewardkit/src/rewardkit/runner.py

@@ -87,7 +87,7 @@ def _build_judge_from_toml(judge_config: dict[str, Any]) -> LLMJudge | AgentJudg
     if judge_name in known_agents():
         return AgentJudge(
             agent=judge_name,
-            model=judge_config.get("model"),
+            model=os.environ.get("REWARDKIT_MODEL") or judge_config.get("model"),
             timeout=timeout,
             cwd=judge_config.get("cwd"),
             isolated=isolated,

packages/rewardkit/tests/conftest.py

@@ -11,7 +11,12 @@
 
 from rewardkit.session import Session, _factory_registry, set_current
 
-_TEST_ENV_KEYS = ("REWARDKIT_JUDGE", "TEST_REWARDKIT_VAR", "TEST_REWARDKIT_URL")
+_TEST_ENV_KEYS = (
+    "REWARDKIT_JUDGE",
+    "REWARDKIT_MODEL",
+    "TEST_REWARDKIT_VAR",
+    "TEST_REWARDKIT_URL",
+)
 
 
 @pytest.fixture(autouse=True)

packages/rewardkit/tests/unit/test_cli.py

@@ -179,3 +179,23 @@ def test_judge_flag_sets_rewardkit_judge(self, tmp_path, monkeypatch):
                     os.environ["REWARDKIT_JUDGE"]
                     == "bedrock/anthropic.claude-3-5-sonnet-20240620-v1:0"
                 )
+
+    @pytest.mark.unit
+    def test_model_flag_sets_rewardkit_model(self, tmp_path, monkeypatch):
+        """--model sets REWARDKIT_MODEL on os.environ."""
+        monkeypatch.delenv("REWARDKIT_MODEL", raising=False)
+        tests_dir = str(tmp_path / "tests")
+        with patch("rewardkit.__main__.run", return_value={}):
+            with patch(
+                "sys.argv",
+                [
+                    "rewardkit",
+                    tests_dir,
+                    "--model",
+                    "anthropic/claude-sonnet-4-6",
+                ],
+            ):
+                from rewardkit.__main__ import main
+
+                main()
+                assert os.environ["REWARDKIT_MODEL"] == "anthropic/claude-sonnet-4-6"

packages/rewardkit/tests/unit/test_runner.py

@@ -646,6 +646,35 @@ def test_no_override_uses_rubric_judge(self, monkeypatch):
         assert isinstance(judge, LLMJudge)
         assert judge.model == "openai/gpt-4o"
 
+    @pytest.mark.unit
+    def test_rewardkit_model_env_overrides_agent_model(self, monkeypatch):
+        """REWARDKIT_MODEL replaces the rubric's [judge].model for agent judges."""
+        monkeypatch.setenv("REWARDKIT_MODEL", "anthropic/claude-sonnet-4-6")
+        judge = _build_judge_from_toml(
+            {"judge": "claude-code", "model": "anthropic/claude-haiku-4-5"}
+        )
+        assert isinstance(judge, AgentJudge)
+        assert judge.agent == "claude-code"
+        assert judge.model == "anthropic/claude-sonnet-4-6"
+
+    @pytest.mark.unit
+    def test_rewardkit_model_env_sets_model_when_toml_unset(self, monkeypatch):
+        """REWARDKIT_MODEL applies even when the TOML has no model field."""
+        monkeypatch.setenv("REWARDKIT_MODEL", "anthropic/claude-sonnet-4-6")
+        judge = _build_judge_from_toml({"judge": "claude-code"})
+        assert isinstance(judge, AgentJudge)
+        assert judge.model == "anthropic/claude-sonnet-4-6"
+
+    @pytest.mark.unit
+    def test_no_rewardkit_model_preserves_toml_model(self, monkeypatch):
+        """With REWARDKIT_MODEL unset, the rubric's [judge].model wins."""
+        monkeypatch.delenv("REWARDKIT_MODEL", raising=False)
+        judge = _build_judge_from_toml(
+            {"judge": "claude-code", "model": "anthropic/claude-haiku-4-5"}
+        )
+        assert isinstance(judge, AgentJudge)
+        assert judge.model == "anthropic/claude-haiku-4-5"
+
 
 class TestRewardScore:
     """Tests for the Reward.score property (weighted mean of criterion scores)."""