Add a model option to run-llm for the claude-fable-5 model

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>

tools/src/pop_fem_audit_tools/commands/run_llm.py

@@ -35,10 +35,18 @@ import anthropic
 from ..config import get_settings
 from ..utils import format_duration
 
-MODEL: str = "claude-sonnet-4-6"
+# claude-fable-5 accepts neither "temperature" nor "thinking";
-TEMPERATURE: float = 0.0
+# a model's entry holds exactly the extra request parameters it
-THINKING: dict[str, str] = {"type": "disabled"}
+# accepts.
-SCRIPT_VERSION: str = "run_llm.py 3.0.0"
+MODELS: dict[str, dict[str, Any]] = {
+    "claude-sonnet-4-6": {
+        "temperature": 0.0,
+        "thinking": {"type": "disabled"},
+    },
+    "claude-fable-5": {},
+}
+DEFAULT_MODEL: str = "claude-sonnet-4-6"
+SCRIPT_VERSION: str = "run_llm.py 3.1.0"
 POLL_INTERVAL_SECONDS: float = 60.0
 
 
@@ -179,6 +187,9 @@ def parse_args(argv: list[str] | None) -> argparse.Namespace:
     parser.add_argument(
         "archive_dir", type=Path,
         help="the destination archive directory")
+    parser.add_argument(
+        "--model", choices=sorted(MODELS), default=DEFAULT_MODEL,
+        help=f"the model ID (default {DEFAULT_MODEL})")
     parser.add_argument(
         "--max-tokens", type=int, default=2048,
         help="the maximum output tokens per request (default 2048)")
@@ -225,21 +236,21 @@ def load_items(path: Path) -> list[InputItem]:
 
 
 def build_request(item: InputItem, system_prompt: str,
-                  max_tokens: int) -> dict[str, Any]:
+                  max_tokens: int, model: str) -> dict[str, Any]:
     """Build one Message Batches request for an input item.
 
     :param item: The input item.
     :param system_prompt: The system prompt text.
     :param max_tokens: The maximum output tokens.
+    :param model: The model ID, a key of ``MODELS``.
     :return: The batch request with "custom_id" and "params".
     """
     return {
         "custom_id": item.id,
         "params": {
-            "model": MODEL,
+            "model": model,
             "max_tokens": max_tokens,
-            "temperature": TEMPERATURE,
+            **MODELS[model],
-            "thinking": THINKING,
             "system": system_prompt,
             "messages": [
                 {"role": "user", "content": item.content},
@@ -418,7 +429,7 @@ def now_iso() -> str:
 
 def execute_run(
         client: anthropic.Anthropic, items: list[InputItem],
-        system_prompt: str, max_tokens: int,
+        system_prompt: str, max_tokens: int, model: str,
         meta: dict[str, Any],
 ) -> Results:
     """Submit the batch of this run and await its results.
@@ -430,11 +441,13 @@ def execute_run(
     :param items: The input items.
     :param system_prompt: The system prompt text.
     :param max_tokens: The maximum output tokens per request.
+    :param model: The model ID, a key of ``MODELS``.
     :param meta: The metadata to record the batch bookkeeping into.
     :return: The results of this run, keyed by item ID.
     """
     requests: list[dict[str, Any]] = [
-        build_request(x, system_prompt, max_tokens) for x in items]
+        build_request(x, system_prompt, max_tokens, model)
+        for x in items]
     info: BatchInfo = BatchInfo(
         batch_id=submit_batch(client, requests),
         submitted_at=now_iso())
@@ -469,9 +482,9 @@ def main(argv: list[str] | None = None) -> int:
     (archive_dir / "prompt.md").write_bytes(args.prompt.read_bytes())
     meta: dict[str, Any] = {
         "script_version": SCRIPT_VERSION,
-        "model": MODEL,
+        "model": args.model,
-        "temperature": TEMPERATURE,
+        "temperature": MODELS[args.model].get("temperature"),
-        "thinking": THINKING,
+        "thinking": MODELS[args.model].get("thinking"),
         "max_tokens": args.max_tokens,
         "prompt_path": str(args.prompt),
         "prompt_sha256": sha256_of(args.prompt),
@@ -486,7 +499,8 @@ def main(argv: list[str] | None = None) -> int:
     if args.dry_run:
         write_json(meta_path, meta)
         print(json.dumps(
-            build_request(items[0], prompt_text, args.max_tokens),
+            build_request(items[0], prompt_text, args.max_tokens,
+                          args.model),
             ensure_ascii=False, indent=2))
         elapsed: str = format_duration(time.monotonic() - started)
         print(f"Done.  {len(items)} jobs finished."
@@ -495,7 +509,8 @@ def main(argv: list[str] | None = None) -> int:
     client: anthropic.Anthropic = anthropic.Anthropic(
         api_key=get_settings().ANTHROPIC_API_KEY)
     results: Results = execute_run(
-        client, items, prompt_text, args.max_tokens, meta)
+        client, items, prompt_text, args.max_tokens, args.model,
+        meta)
     item_ids: list[str] = [x.id for x in items]
     write_jsonl(
         archive_dir / "output.jsonl",

tools/tests/test_run_llm.py

@@ -169,7 +169,7 @@ class TestRequestBuilding(RunLLMTestCase):
         """Test the shape of a batch request."""
         request: dict[str, Any] = run_llm.build_request(
             run_llm.InputItem(id="song-1", content="the lyrics"),
-            "the system prompt", 2048)
+            "the system prompt", 2048, "claude-sonnet-4-6")
         self.assertEqual(request["custom_id"], "song-1")
         params: dict[str, Any] = request["params"]
         self.assertEqual(params["model"], "claude-sonnet-4-6")
@@ -180,6 +180,20 @@ class TestRequestBuilding(RunLLMTestCase):
         self.assertEqual(params["messages"],
                          [{"role": "user", "content": "the lyrics"}])
 
+    def test_build_request_fable_5(self) -> None:
+        """Test the request shape for the claude-fable-5 model."""
+        request: dict[str, Any] = run_llm.build_request(
+            run_llm.InputItem(id="group-1", content="the groups"),
+            "the system prompt", 8192, "claude-fable-5")
+        params: dict[str, Any] = request["params"]
+        self.assertEqual(params["model"], "claude-fable-5")
+        self.assertNotIn("temperature", params)
+        self.assertNotIn("thinking", params)
+        self.assertEqual(params["max_tokens"], 8192)
+        self.assertEqual(params["system"], "the system prompt")
+        self.assertEqual(params["messages"],
+                         [{"role": "user", "content": "the groups"}])
+
 
 class TestCollectResults(RunLLMTestCase):
     """Test cases for the batch result collection."""