diff --git a/tools/src/pop_fem_audit_tools/commands/compare_codings.py b/tools/src/pop_fem_audit_tools/commands/compare_codings.py index 926af11..02c5b54 100644 --- a/tools/src/pop_fem_audit_tools/commands/compare_codings.py +++ b/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") diff --git a/tools/src/pop_fem_audit_tools/commands/export_llm_input.py b/tools/src/pop_fem_audit_tools/commands/export_llm_input.py index 8891085..458d280 100644 --- a/tools/src/pop_fem_audit_tools/commands/export_llm_input.py +++ b/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-\", 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-``, 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} - payload.update(extras) + 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 diff --git a/tools/tests/test_compare_codings.py b/tools/tests/test_compare_codings.py index e08afce..e26afde 100644 --- a/tools/tests/test_compare_codings.py +++ b/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 diff --git a/tools/tests/test_export_llm_input.py b/tools/tests/test_export_llm_input.py index df62377..37dd094 100644 --- a/tools/tests/test_export_llm_input.py +++ b/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())