Build the arbitration input with a per-ID extras merge

Co-Authored-By: Claude Opus 5 (1M context) <noreply@anthropic.com>

tools/src/pop_fem_audit_tools/commands/compare_codings.py

@@ -14,8 +14,11 @@ agreed.  The quotes supporting a keyword never take part in the
 comparison; they are carried along as the evidence of the disagreed
 keyword they support.  The disagreements are written as a JSON file,
 as :data:`DISAGREEMENTS_JSON`, holding only the songs the two runs
-disagree on; it is what the arbitration step of the coding
-procedure settles.  The agreed half is not written; it is returned
+disagree on, each song's disagreed keywords wrapped in a
+"disagreements" object so that the file merges into the arbitration
+input as the per-song extra parameters; it is what the arbitration
+step of the coding procedure settles.  The agreed half is not
+written; it is returned
 by :func:`compare_codings`, as the keyword names alone, for later
 steps to consume.  The step is fully deterministic; no LLM call is
 made.
@@ -229,10 +232,13 @@ def write_disagreements(path: Path,
     """Write the per-song disagreement JSON file.
 
     Writes a JSON file holding a single object mapping the song ID
-    to the disagreed keywords of that song, each with the quotes of
-    the run that assigned it, in the given order, UTF-8, with a
-    trailing newline.  Only the songs the two runs disagree on are
-    written.
+    to a ``{"disagreements": {...}}`` object holding the disagreed
+    keywords of that song, each with the quotes of the run that
+    assigned it, in the given order, UTF-8, with a trailing newline.
+    The wrapper is what makes the file merge into the arbitration
+    input as the per-song extra parameters, giving a "disagreements"
+    sibling of the lyrics rather than loose keywords.  Only the
+    songs the two runs disagree on are written.
 
     :param path: The path of the disagreement JSON file to write.
     :param disagreements: The disagreed keywords of every song,
@@ -240,9 +246,11 @@ def write_disagreements(path: Path,
     :return: None.
     :raises OSError: When the file cannot be written.
     """
+    wrapped: dict[str, dict[str, Coding]] = {
+        song_id: {"disagreements": coding}
+        for song_id, coding in disagreements.items()}
     path.write_text(
-        json.dumps(disagreements, ensure_ascii=False, indent=2)
-        + "\n",
+        json.dumps(wrapped, ensure_ascii=False, indent=2) + "\n",
         encoding="utf-8")
 
 

tools/src/pop_fem_audit_tools/commands/export_llm_input.py

@@ -16,6 +16,14 @@ object serialized as a string, its ``lyrics`` key holding the
 song's lyrics followed by the keys of the given extras file in
 their file order, so a step that needs parameters alongside the
 lyrics can carry them without this module knowing what they mean.
+
+With ``--extras-per-id``, the same merge happens per song: the
+given file maps a song ID to the extra keys of that one song, and
+the export is restricted to the song IDs the file names, so a step
+that revisits only some of the songs, each with its own parameters,
+gets exactly those records.  The two options may be given together,
+in which case a record's keys are ``lyrics``, the shared extras'
+keys, then that song's own keys, each group in its file order.
 """
 import argparse
 import json
@@ -50,6 +58,14 @@ def parse_args(argv: list[str] | None) -> argparse.Namespace:
              " becomes a JSON object string with a \"lyrics\" key"
              " followed by the extras' keys, instead of the bare"
              " lyrics string")
+    parser.add_argument(
+        "--extras-per-id", type=Path, default=None,
+        help="a JSON file holding a single JSON object that maps a"
+             " song ID, as \"song-<N>\", to a JSON object of extra"
+             " parameters for that one song; the song's object is"
+             " merged into its \"content\" the same way as with"
+             " --extras, and the export is restricted to the song"
+             " IDs the file names")
     return parser.parse_args(argv)
 
 
@@ -72,6 +88,31 @@ def _no_duplicate_keys(
     return result
 
 
+def _load_json_object(path: Path, label: str) -> dict[str, Any]:
+    """Load a single JSON object from a file, in file order.
+
+    :param path: The JSON file.
+    :param label: The kind of file, for the error messages.
+    :return: The object, in file order.
+    :raises OSError: When the file cannot be read.
+    :raises ValueError: When the file is not valid JSON, is not
+        a JSON object, or has duplicate keys.
+    """
+    with open(path, encoding="utf-8") as file:
+        text: str = file.read()
+    try:
+        data: Any = json.loads(
+            text, object_pairs_hook=_no_duplicate_keys)
+    except json.JSONDecodeError as error:
+        raise ValueError(
+            f"invalid JSON in {label} file {path}: {error}") \
+            from error
+    if not isinstance(data, dict):
+        raise ValueError(
+            f"{label} file {path} must contain a JSON object")
+    return data
+
+
 def load_extras(path: Path) -> dict[str, Any]:
     """Load the extras object from a JSON file.
 
@@ -81,64 +122,111 @@ def load_extras(path: Path) -> dict[str, Any]:
     :raises ValueError: When the file is not valid JSON, is not
         a JSON object, has duplicate keys, or has a "lyrics" key.
     """
-    with open(path, encoding="utf-8") as file:
-        text: str = file.read()
-    try:
-        data: Any = json.loads(
-            text, object_pairs_hook=_no_duplicate_keys)
-    except json.JSONDecodeError as error:
-        raise ValueError(
-            f"invalid JSON in extras file {path}: {error}") \
-            from error
-    if not isinstance(data, dict):
-        raise ValueError(
-            f"extras file {path} must contain a JSON object")
+    data: dict[str, Any] = _load_json_object(path, "extras")
     if "lyrics" in data:
         raise ValueError(
             f"extras file {path} must not have a \"lyrics\" key")
     return data
 
 
+def load_extras_per_id(path: Path) -> dict[str, dict[str, Any]]:
+    """Load the per-ID extras object from a JSON file.
+
+    :param path: The per-ID extras JSON file, mapping a song ID,
+        as ``song-<N>``, to the extras of that one song.
+    :return: The extras of each song ID, in file order, every
+        song's own extras in their file order too.
+    :raises OSError: When the file cannot be read.
+    :raises ValueError: When the file is not valid JSON, is not
+        a JSON object, has duplicate keys, has a song whose value
+        is not a JSON object, or has a song with a "lyrics" key.
+    """
+    data: dict[str, Any] = _load_json_object(path, "per-ID extras")
+    song_id: str
+    extras: Any
+    for song_id, extras in data.items():
+        if not isinstance(extras, dict):
+            raise ValueError(
+                f"per-ID extras file {path}: id {song_id} must"
+                " have a JSON object")
+        if "lyrics" in extras:
+            raise ValueError(
+                f"per-ID extras file {path}: id {song_id} must not"
+                " have a \"lyrics\" key")
+    return data
+
+
 def build_lines(
         session: Session,
-        extras: dict[str, Any] | None = None) -> list[str]:
-    """Build the JSONL lines of every song's lyrics.
+        extras: dict[str, Any] | None = None,
+        extras_per_id: dict[str, dict[str, Any]] | None = None) \
+        -> list[str]:
+    """Build the JSONL lines of the exported songs' lyrics.
 
-    Without extras, each record's ``content`` is the bare lyrics
-    string.  With extras, ``content`` is a JSON object serialized
-    as a string, whose first key is ``"lyrics"`` holding the
-    lyrics string, followed by the extras' keys in their given
-    order.
+    Without extras of either kind, each record's ``content`` is
+    the bare lyrics string.  With extras, ``content`` is a JSON
+    object serialized as a string, whose first key is ``"lyrics"``
+    holding the lyrics string, followed by the shared extras' keys
+    and then the song's own per-ID extras' keys, each group in its
+    given order.
+
+    Every song is exported, unless per-ID extras are given, in
+    which case only the songs they name are.
 
     :param session: The database session.
     :param extras: The extra parameters merged into every
         record's content alongside the lyrics, in the order they
-        are to appear, or None for the bare lyrics string.
-    :return: The JSON lines, one per song, ordered by song ID.
-    :raises ValueError: When a song has no lyrics.
+        are to appear, or None for none.
+    :param extras_per_id: The extra parameters merged into the
+        content of one record alone, keyed by that record's song
+        ID and in the order they are to appear, restricting the
+        export to the song IDs they name, or None for no such
+        extras and no such restriction.
+    :return: The JSON lines, one per exported song, ordered by
+        song ID.
+    :raises ValueError: When an exported song has no lyrics, or
+        the per-ID extras name a song the working store does not
+        have.
     """
     lines: list[str] = []
+    exported: set[str] = set()
     song: Song
     for song in session.scalars(sa.select(Song).order_by(Song.id)):
+        song_id: str = f"song-{song.id}"
+        if extras_per_id is not None and song_id not in extras_per_id:
+            continue
         if song.lyrics is None:
             raise ValueError(
                 f"song {song.id} \"{song.title}\": no lyrics")
         content: str
-        if extras is None:
+        if extras is None and extras_per_id is None:
             content = song.lyrics
         else:
             payload: dict[str, Any] = {"lyrics": song.lyrics}
+            if extras is not None:
                 payload.update(extras)
+            if extras_per_id is not None:
+                payload.update(extras_per_id[song_id])
             content = json.dumps(payload, ensure_ascii=False)
         record: dict[str, str] = {
-            "id": f"song-{song.id}", "content": content}
+            "id": song_id, "content": content}
         lines.append(json.dumps(record, ensure_ascii=False))
+        exported.add(song_id)
+    if extras_per_id is not None:
+        missing: list[str] = sorted(set(extras_per_id) - exported)
+        if len(missing) > 0:
+            raise ValueError(
+                "the per-ID extras name songs the working store"
+                f" does not have: {', '.join(missing)}")
     return lines
 
 
 def main(argv: list[str] | None = None) -> int:
     """Export the LLM input JSONL file from the working store.
 
+    Every song is exported, unless ``--extras-per-id`` is given,
+    in which case only the songs its file names are.
+
     :param argv: The command-line arguments, or None for
         ``sys.argv``.
     :return: The exit status: 0 on success, non-zero on failure.
@@ -150,7 +238,10 @@ def main(argv: list[str] | None = None) -> int:
         extras: dict[str, Any] | None = None
         if args.extras is not None:
             extras = load_extras(args.extras)
-        lines = build_lines(session, extras)
+        extras_per_id: dict[str, dict[str, Any]] | None = None
+        if args.extras_per_id is not None:
+            extras_per_id = load_extras_per_id(args.extras_per_id)
+        lines = build_lines(session, extras, extras_per_id)
     except (OSError, sa.exc.SQLAlchemyError, ValueError) as error:
         print(f"error: {error}", file=sys.stderr)
         return 1

tools/tests/test_compare_codings.py

@@ -79,11 +79,20 @@ class TestCompareCodings(unittest.TestCase):
     def __read_disagreements(self) -> dict[str, Any]:
         """Read the disagreement JSON file.
 
-        :return: The parsed disagreements.
+        :return: The parsed disagreements, as written, with every
+            song's keywords in its "disagreements" wrapper.
         """
         return json.loads(
             self.__disagreements_json.read_text(encoding="utf-8"))
 
+    def __read_keywords(self, song_id: str) -> dict[str, Any]:
+        """Read one song's disagreed keywords from the file.
+
+        :param song_id: The song ID.
+        :return: The keywords of that song, unwrapped.
+        """
+        return self.__read_disagreements()[song_id]["disagreements"]
+
     def __read_disagreement_text(self) -> str:
         """Read the disagreement JSON file verbatim.
 
@@ -103,7 +112,23 @@ class TestCompareCodings(unittest.TestCase):
         status, _ = self.__run_compare()
         self.assertEqual(status, 0)
         self.assertEqual(self.__read_disagreements(), {
-            "song-1": {"only-1": ["q2"], "only-2": ["q4"]}})
+            "song-1": {"disagreements": {
+                "only-1": ["q2"], "only-2": ["q4"]}}})
+
+    def test_keywords_wrapped_for_merging(self) -> None:
+        """Test that each song's value is the object merged into
+        the arbitration input, holding "disagreements" alone."""
+        self.__write_codings(self.__run1, {
+            "song-1": {"only-1": ["q"]}})
+        self.__write_codings(self.__run2, {"song-1": {}})
+        status: int
+        status, _ = self.__run_compare()
+        self.assertEqual(status, 0)
+        self.assertEqual(
+            list(self.__read_disagreements()["song-1"].keys()),
+            ["disagreements"])
+        self.assertEqual(
+            self.__read_keywords("song-1"), {"only-1": ["q"]})
 
     def test_quotes_do_not_take_part_in_comparison(self) -> None:
         """Test that identical key sets with different quotes
@@ -162,7 +187,7 @@ class TestCompareCodings(unittest.TestCase):
         status, _ = self.__run_compare()
         self.assertEqual(status, 0)
         self.assertEqual(
-            list(self.__read_disagreements()["song-1"].keys()),
+            list(self.__read_keywords("song-1").keys()),
             ["alpha", "beta", "mu", "zeta"])
 
     def test_summary_line_counts(self) -> None:
@@ -195,8 +220,7 @@ class TestCompareCodings(unittest.TestCase):
         status, _ = self.__run_compare()
         self.assertEqual(status, 0)
         self.assertEqual(
-            self.__read_disagreements(),
-            {"song-1": {"only-1": [quote]}})
+            self.__read_keywords("song-1"), {"only-1": [quote]})
 
     def test_non_ascii_written_verbatim(self) -> None:
         """Test that the output is UTF-8 with the non-ASCII text
@@ -362,8 +386,7 @@ class TestCompareCodings(unittest.TestCase):
         status, _ = self.__run_compare()
         self.assertEqual(status, 0)
         self.assertEqual(
-            self.__read_disagreements(),
-            {"song-1": {"only-1": ["q"]}})
+            self.__read_keywords("song-1"), {"only-1": ["q"]})
 
     def test_compare_returns_agreed_keyword_names(self) -> None:
         """Test that the comparison function returns the agreed

tools/tests/test_export_llm_input.py

@@ -76,16 +76,21 @@ class TestExportLlmInput(unittest.TestCase):
             session.close()
 
     def __run_export(
-            self, extras: Path | None = None) -> tuple[int, str]:
+            self, extras: Path | None = None,
+            extras_per_id: Path | None = None) -> tuple[int, str]:
         """Run the exporter with the standard error captured.
 
         :param extras: The extras JSON file, or None for none.
+        :param extras_per_id: The per-ID extras JSON file, or None
+            for none.
         :return: A tuple of the exit status and the standard
             error.
         """
         argv: list[str] = [str(self.__output)]
         if extras is not None:
             argv += ["--extras", str(extras)]
+        if extras_per_id is not None:
+            argv += ["--extras-per-id", str(extras_per_id)]
         stderr: io.StringIO = io.StringIO()
         with redirect_stderr(stderr):
             status: int = export_llm_input.main(argv)
@@ -101,6 +106,16 @@ class TestExportLlmInput(unittest.TestCase):
         path.write_text(text, encoding="utf-8")
         return path
 
+    def __write_extras_per_id(self, text: str) -> Path:
+        """Write a per-ID extras file with the given raw text.
+
+        :param text: The raw file content.
+        :return: The path of the written per-ID extras file.
+        """
+        path: Path = self.__dir / "extras-per-id.json"
+        path.write_text(text, encoding="utf-8")
+        return path
+
     @staticmethod
     def __read_records(path: Path) -> list[dict[str, str]]:
         """Read the JSONL records of a file.
@@ -228,3 +243,115 @@ class TestExportLlmInput(unittest.TestCase):
         self.assertEqual(status, 1)
         self.assertIn("error:", stderr)
         self.assertFalse(self.__output.exists())
+
+    def test_extras_per_id_merges_each_song_its_own(self) -> None:
+        """Test that with per-ID extras, each record's content is
+        a JSON object of the lyrics followed by that song's own
+        keys in their file order."""
+        self.__seed([
+            ("Hello", "Adele", "hello lyrics\n"),
+            ("Umbrella", "Rihanna", "umbrella lyrics\n")])
+        per_id: Path = self.__write_extras_per_id(
+            '{"song-1": {"b": 2, "a": 1},'
+            ' "song-2": {"disagreements": {"k": ["q"]}}}')
+        status: int
+        stderr: str
+        status, stderr = self.__run_export(None, per_id)
+        self.assertEqual(status, 0)
+        records: list[dict[str, str]] = self.__read_records(
+            self.__output)
+        self.assertEqual(len(records), 2)
+        first: dict[str, Any] = json.loads(records[0]["content"])
+        self.assertEqual(
+            list(first.keys()), ["lyrics", "b", "a"])
+        self.assertEqual(first["lyrics"], "hello lyrics\n")
+        self.assertEqual(first["b"], 2)
+        self.assertEqual(first["a"], 1)
+        second: dict[str, Any] = json.loads(records[1]["content"])
+        self.assertEqual(
+            list(second.keys()), ["lyrics", "disagreements"])
+        self.assertEqual(second["lyrics"], "umbrella lyrics\n")
+        self.assertEqual(
+            second["disagreements"], {"k": ["q"]})
+
+    def test_extras_per_id_restricts_the_export(self) -> None:
+        """Test that the songs the per-ID extras do not name are
+        not exported."""
+        self.__seed([
+            ("Hello", "Adele", "hello lyrics\n"),
+            ("Umbrella", "Rihanna", "umbrella lyrics\n"),
+            ("Halo", "Beyonce", "halo lyrics\n")])
+        per_id: Path = self.__write_extras_per_id(
+            '{"song-3": {"a": 1}, "song-1": {"a": 2}}')
+        status: int
+        stderr: str
+        status, stderr = self.__run_export(None, per_id)
+        self.assertEqual(status, 0)
+        records: list[dict[str, str]] = self.__read_records(
+            self.__output)
+        self.assertEqual(
+            [x["id"] for x in records], ["song-1", "song-3"])
+        self.assertIn("done: 2 songs exported", stderr)
+
+    def test_extras_per_id_unknown_id_fails(self) -> None:
+        """Test that a per-ID extras file naming a song the
+        working store does not have is rejected."""
+        self.__seed([("Hello", "Adele", "hello lyrics\n")])
+        per_id: Path = self.__write_extras_per_id(
+            '{"song-1": {"a": 1}, "song-9": {"a": 2}}')
+        status: int
+        stderr: str
+        status, stderr = self.__run_export(None, per_id)
+        self.assertEqual(status, 1)
+        self.assertIn("song-9", stderr)
+        self.assertFalse(self.__output.exists())
+
+    def test_extras_per_id_after_the_shared_extras(self) -> None:
+        """Test that with both extras options, a record's keys are
+        the lyrics, the shared extras, then that song's own."""
+        self.__seed([
+            ("Hello", "Adele", "hello lyrics\n"),
+            ("Umbrella", "Rihanna", "umbrella lyrics\n")])
+        extras: Path = self.__write_extras('{"shared": "s"}')
+        per_id: Path = self.__write_extras_per_id(
+            '{"song-2": {"own": "o"}}')
+        status: int
+        stderr: str
+        status, stderr = self.__run_export(extras, per_id)
+        self.assertEqual(status, 0)
+        records: list[dict[str, str]] = self.__read_records(
+            self.__output)
+        self.assertEqual([x["id"] for x in records], ["song-2"])
+        content: dict[str, Any] = json.loads(
+            records[0]["content"])
+        self.assertEqual(
+            list(content.keys()), ["lyrics", "shared", "own"])
+        self.assertEqual(content["lyrics"], "umbrella lyrics\n")
+        self.assertEqual(content["shared"], "s")
+        self.assertEqual(content["own"], "o")
+
+    def test_extras_per_id_non_object_value_fails(self) -> None:
+        """Test that a per-ID extras file whose song does not have
+        a JSON object is rejected."""
+        self.__seed([("Hello", "Adele", "hello lyrics\n")])
+        per_id: Path = self.__write_extras_per_id(
+            '{"song-1": [1, 2]}')
+        status: int
+        stderr: str
+        status, stderr = self.__run_export(None, per_id)
+        self.assertEqual(status, 1)
+        self.assertIn("error:", stderr)
+        self.assertFalse(self.__output.exists())
+
+    def test_extras_per_id_with_lyrics_key_fails(self) -> None:
+        """Test that a per-ID extras file whose song carries a
+        "lyrics" key is rejected."""
+        self.__seed([("Hello", "Adele", "hello lyrics\n")])
+        per_id: Path = self.__write_extras_per_id(
+            '{"song-1": {"lyrics": "not allowed"}}')
+        status: int
+        stderr: str
+        status, stderr = self.__run_export(None, per_id)
+        self.assertEqual(status, 1)
+        self.assertIn("error:", stderr)
+        self.assertFalse(self.__output.exists())