Restructure the remaining commands into the house shape

Every command now runs as one orchestrating class (the ctor
stores, run() executes, helpers and constants private), main a
thin controller; the guards the fixed corpus cannot trigger are
dropped, docstrings say each level's own contract once, and the
build-db summary reports the songs and the artists alone.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
This commit is contained in:
2026-08-19 10:46:21 +08:00
co-authored by Claude Opus 5
parent 676d7788e8
commit 2fa55f29d7
11 changed files with 1919 additions and 1882 deletions
File diff suppressed because it is too large Load Diff
@@ -24,8 +24,7 @@ keyword set for ``export-llm-input --extras`` is written as a JSON
file holding the group name keywords plus every extra a-priori file holding the group name keywords plus every extra a-priori
keyword the caller gives with the repeatable ``--extra-keyword`` keyword the caller gives with the repeatable ``--extra-keyword``
command-line option, as command-line option, as
:attr:`KeywordsToMerge.KEYWORDS_TO_MERGE_JSON`; with no :attr:`KeywordsToMerge.KEYWORDS_TO_MERGE_JSON`. No default
``--extra-keyword``, it holds the group names alone. No default
extra keyword is ever injected; the caller supplies each one extra keyword is ever injected; the caller supplies each one
consciously. Finally, the command-line choices and the consciously. Finally, the command-line choices and the
environment that produced the numbers -- neither recoverable from environment that produced the numbers -- neither recoverable from
@@ -160,15 +159,8 @@ class KeywordPooler:
if line.strip() == "": if line.strip() == "":
continue continue
record: Any = json.loads(line) record: Any = json.loads(line)
if not isinstance(record, dict) or "id" not in record:
raise ValueError(
f"{path}: record without \"id\": {line}")
if "error" in record: if "error" in record:
continue continue
if "text" not in record:
raise ValueError(
f"{path}: id {record['id']}: record without"
" \"text\" or \"error\"")
song_id: int = cls.__parse_song_id(record["id"], path) song_id: int = cls.__parse_song_id(record["id"], path)
try: try:
keywords: Any = json.loads( keywords: Any = json.loads(
@@ -297,7 +289,8 @@ class KeywordGroups:
class KeywordClusterer: class KeywordClusterer:
"""The clusterer of the pooled keywords into coding groups.""" """The clusterer of the pooled keywords into coding groups."""
DEFAULT_MODEL: str = "sentence-transformers/all-mpnet-base-v2" DEFAULT_MODEL: ClassVar[str] \
= "sentence-transformers/all-mpnet-base-v2"
"""The sentence embedding model used when the caller names """The sentence embedding model used when the caller names
none.""" none."""
@@ -668,11 +661,7 @@ def parse_args(argv: list[str] | None) -> argparse.Namespace:
parser.add_argument( parser.add_argument(
"output_dir", type=Path, "output_dir", type=Path,
help="the output directory, created if missing, that" help="the output directory, created if missing, that"
f" receives {PooledKeywords.SOURCE_KEYWORDS_TXT}," " receives the run's output artifacts")
f" {KeywordGroups.RESULT_KEYWORDS_TXT},"
f" {KeywordGroups.RESULT_GROUPS_CSV},"
f" {KeywordsToMerge.KEYWORDS_TO_MERGE_JSON},"
f" and {RunMeta.META_JSON}")
parser.add_argument( parser.add_argument(
"--model", default=model, "--model", default=model,
help=f"the sentence embedding model (default \"{model}\")") help=f"the sentence embedding model (default \"{model}\")")
@@ -698,16 +687,9 @@ def parse_args(argv: list[str] | None) -> argparse.Namespace:
def main(argv: list[str] | None = None) -> int: def main(argv: list[str] | None = None) -> int:
"""Pool the two tagging runs' keywords and cluster them. """Pool the two tagging runs' keywords and cluster them.
Writes the five fixed-named artifacts under the output Creates the output directory (with parents) if it does not
directory, creating it (with parents) if it does not exist: exist. Each output artifact is written as soon as its
the pooled keyword text file; then the group membership CSV content is computed, so when the input is rejected, or an
file, holding the clustering result alone; the group name
keyword text file, holding the same group names as a readable
list; the coding keyword set JSON file, holding the group
names plus every extra keyword given via ``--extra-keyword``;
and the run metadata JSON file, recording the command-line
choices and the environment. Each file is written as soon as
its content is computed, so when the input is rejected, or an
extra keyword duplicates a group name or another extra extra keyword duplicates a group name or another extra
keyword, the output directory holds whatever the steps before keyword, the output directory holds whatever the steps before
the failing one produced, and the error message names what the failing one produced, and the error message names what
@@ -719,8 +701,8 @@ def main(argv: list[str] | None = None) -> int:
""" """
started: float = time.monotonic() started: float = time.monotonic()
args: argparse.Namespace = parse_args(argv) args: argparse.Namespace = parse_args(argv)
args.output_dir.mkdir(parents=True, exist_ok=True)
try: try:
args.output_dir.mkdir(parents=True, exist_ok=True)
source: PooledKeywords = KeywordPooler( source: PooledKeywords = KeywordPooler(
args.run_dir_1, args.run_dir_2, args.output_dir).run() args.run_dir_1, args.run_dir_2, args.output_dir).run()
clusters: KeywordGroups = KeywordClusterer( clusters: KeywordGroups = KeywordClusterer(
@@ -735,7 +717,7 @@ def main(argv: list[str] | None = None) -> int:
f"Done. Clustered {len(source.keywords)} keywords into" f"Done. Clustered {len(source.keywords)} keywords into"
f" {len(clusters.names)}. {elapsed} elapsed.", f" {len(clusters.names)}. {elapsed} elapsed.",
file=sys.stderr) file=sys.stderr)
except ClusterError as error: except (ClusterError, OSError) as error:
print(f"error: {error}", file=sys.stderr) print(f"error: {error}", file=sys.stderr)
return 1 return 1
return 0 return 0
@@ -11,26 +11,17 @@ project's lyrics-only firewall: the output carries only the
lyrics text of each song, identified by an opaque song key; no lyrics text of each song, identified by an opaque song key; no
title, artist, or chart data crosses into the LLM input. title, artist, or chart data crosses into the LLM input.
With ``--extras``, each record's ``content`` becomes a JSON With ``--extras`` and ``--extras-per-id``, a record's content may
object serialized as a string, its ``lyrics`` key holding the carry extra parameters alongside the lyrics, so a step that needs
song's lyrics followed by the keys of the given extras file in them can get them without this module knowing what they mean; see
their file order, so a step that needs parameters alongside the the exporter's content-building step for how the two merge.
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 argparse
import json import json
import sys import sys
import time import time
from pathlib import Path from pathlib import Path
from typing import Any from typing import Any, ClassVar
import sqlalchemy as sa import sqlalchemy as sa
from sqlalchemy.orm import Session from sqlalchemy.orm import Session
@@ -40,6 +31,270 @@ from ..models import Song
from ..utils import format_duration from ..utils import format_duration
class LlmInputExporter:
"""The exporter of the LLM input JSONL file."""
__LYRICS_KEY: ClassVar[str] = "lyrics"
"""The key holding the lyrics in a record's merged content,
and the key forbidden in an extras file."""
def __init__(
self, output_jsonl: Path, extras: Path | None = None,
extras_per_id: Path | None = None) -> None:
"""Set up the exporter of the LLM input JSONL file.
:param output_jsonl: The JSONL output file.
:param extras: The extras JSON file, or None for none.
:param extras_per_id: The per-ID extras JSON file, or
None for none.
"""
self.__output_jsonl: Path = output_jsonl
"""The JSONL output file."""
self.__extras_path: Path | None = extras
"""The extras JSON file, or None for none."""
self.__extras_per_id_path: Path | None = extras_per_id
"""The per-ID extras JSON file, or None for none."""
def run(self) -> int:
"""Export the songs' lyrics to the output JSONL file.
:return: The number of songs exported.
:raises OSError: When a file cannot be read or written.
:raises sqlalchemy.exc.SQLAlchemyError: When the working
store cannot be read.
:raises ValueError: When an extras file is malformed, an
exported song has no lyrics, or the per-ID extras
name a song the working store does not have.
"""
session: Session = ds.get_db()
try:
extras: dict[str, Any] | None = None
if self.__extras_path is not None:
extras = self.__load_extras(self.__extras_path)
extras_per_id: dict[str, dict[str, Any]] | None = None
if self.__extras_per_id_path is not None:
extras_per_id = self.__load_extras_per_id(
self.__extras_per_id_path)
lines: list[str] = self.__build_lines(
session, extras, extras_per_id)
finally:
session.close()
self.__write_output(lines)
return len(lines)
@staticmethod
def __no_duplicate_keys(
pairs: list[tuple[str, Any]]) -> dict[str, Any]:
"""Build a dict from JSON object pairs, rejecting
duplicates.
:param pairs: The key-value pairs of a JSON object, in
file order.
:return: The pairs as a dict, in file order.
:raises ValueError: When a key appears more than once.
"""
result: dict[str, Any] = {}
key: str
value: Any
for key, value in pairs:
if key in result:
raise ValueError(
f"duplicate key \"{key}\" in extras")
result[key] = value
return result
@classmethod
def __load_json_object(
cls, 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=cls.__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
@classmethod
def __load_extras(cls, path: Path) -> dict[str, Any]:
"""Load the extras object from a JSON file.
:param path: The extras JSON file.
:return: The extras, 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, has duplicate keys, or has a
"lyrics" key.
"""
data: dict[str, Any] = cls.__load_json_object(
path, "extras")
if cls.__LYRICS_KEY in data:
raise ValueError(
f"extras file {path} must not have a"
f" \"{cls.__LYRICS_KEY}\" key")
return data
@classmethod
def __load_extras_per_id(
cls, 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] = cls.__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 cls.__LYRICS_KEY in extras:
raise ValueError(
f"per-ID extras file {path}: id {song_id}"
f" must not have a \"{cls.__LYRICS_KEY}\""
" key")
return data
def __build_lines(
self, session: Session,
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.
Every song is exported, unless per-ID extras are given,
in which case only the songs they name are; see
:meth:`__build_content` for how the extras merge into a
record's content.
: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 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")
song_extras: dict[str, Any] | None = None \
if extras_per_id is None \
else extras_per_id[song_id]
content: str = self.__build_content(
song.lyrics, extras, song_extras)
record: dict[str, str] = {
"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"
f" store does not have: {', '.join(missing)}")
return lines
@classmethod
def __build_content(
cls, lyrics: str, extras: dict[str, Any] | None,
song_extras: dict[str, Any] | None) -> str:
"""Build the content of one exported record.
Without extras of either kind, a record's content is the
bare lyrics string. With ``--extras``, the content
becomes a JSON 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. With
``--extras-per-id``, the same merge happens per song: the
song's own extra keys follow the lyrics instead. When
both are given, a record's keys are "lyrics", the shared
extras' keys, then that song's own keys, each group in
its file order.
:param lyrics: The lyrics of the song.
:param extras: The extra parameters shared by every
record, in the order they are to appear, or None for
none.
:param song_extras: The extra parameters of this record
alone, in the order they are to appear, or None for
none.
:return: The bare lyrics when there are no extras of
either kind, or otherwise a JSON object serialized as
a string, whose first key is "lyrics" holding the
lyrics, followed by the shared extras' keys and then
this record's own keys, each group in its given
order.
"""
if extras is None and song_extras is None:
return lyrics
payload: dict[str, Any] = {cls.__LYRICS_KEY: lyrics}
if extras is not None:
payload.update(extras)
if song_extras is not None:
payload.update(song_extras)
return json.dumps(payload, ensure_ascii=False)
def __write_output(self, lines: list[str]) -> None:
"""Write the exported lines to the output JSONL file.
Creates the parent directory when it does not exist.
:param lines: The JSONL lines, in the output order.
:return: None.
:raises OSError: When the file cannot be written.
"""
self.__output_jsonl.parent.mkdir(
parents=True, exist_ok=True)
with open(
self.__output_jsonl, "w",
encoding="utf-8") as file:
line: str
for line in lines:
file.write(line + "\n")
def parse_args(argv: list[str] | None) -> argparse.Namespace: def parse_args(argv: list[str] | None) -> argparse.Namespace:
"""Parse the command-line arguments. """Parse the command-line arguments.
@@ -56,227 +311,33 @@ def parse_args(argv: list[str] | None) -> argparse.Namespace:
parser.add_argument( parser.add_argument(
"--extras", type=Path, default=None, "--extras", type=Path, default=None,
help="a JSON file holding a single JSON object of extra" help="a JSON file holding a single JSON object of extra"
" parameters; when given, each record's \"content\"" " parameters merged into every record's content")
" becomes a JSON object string with a \"lyrics\" key"
" followed by the extras' keys, instead of the bare"
" lyrics string")
parser.add_argument( parser.add_argument(
"--extras-per-id", type=Path, default=None, "--extras-per-id", type=Path, default=None,
help="a JSON file holding a single JSON object that maps a" help="a JSON file holding a single JSON object that maps a"
" song ID, as \"song-<N>\", to a JSON object of extra" " song ID, as \"song-<N>\", to a JSON object of extra"
" parameters for that one song; the song's object is" " parameters for that one song, restricting the"
" merged into its \"content\" the same way as with" " export to the song IDs the file names")
" --extras, and the export is restricted to the song"
" IDs the file names")
return parser.parse_args(argv) return parser.parse_args(argv)
def __no_duplicate_keys(
pairs: list[tuple[str, Any]]) -> dict[str, Any]:
"""Build a dict from JSON object pairs, rejecting duplicates.
:param pairs: The key-value pairs of a JSON object, in file
order.
:return: The pairs as a dict, in file order.
:raises ValueError: When a key appears more than once.
"""
result: dict[str, Any] = {}
key: str
value: Any
for key, value in pairs:
if key in result:
raise ValueError(f"duplicate key \"{key}\" in extras")
result[key] = value
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.
:param path: The extras JSON file.
:return: The extras, 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, has duplicate keys, or has a "lyrics" key.
"""
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_content(
lyrics: str,
extras: dict[str, Any] | None,
song_extras: dict[str, Any] | None) -> str:
"""Build the content of one exported record.
:param lyrics: The lyrics of the song.
:param extras: The extra parameters shared by every record,
in the order they are to appear, or None for none.
:param song_extras: The extra parameters of this record
alone, in the order they are to appear, or None for none.
:return: The bare lyrics when there are no extras of either
kind, or otherwise a JSON object serialized as a string,
whose first key is ``"lyrics"`` holding the lyrics,
followed by the shared extras' keys and then this
record's own keys, each group in its given order.
"""
if extras is None and song_extras is None:
return lyrics
payload: dict[str, Any] = {"lyrics": lyrics}
if extras is not None:
payload.update(extras)
if song_extras is not None:
payload.update(song_extras)
return json.dumps(payload, ensure_ascii=False)
def build_lines(
session: Session,
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 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 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")
song_extras: dict[str, Any] | None = None \
if extras_per_id is None else extras_per_id[song_id]
content: str = __build_content(
song.lyrics, extras, song_extras)
record: dict[str, str] = {
"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: def main(argv: list[str] | None = None) -> int:
"""Export the LLM input JSONL file from the working store. """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 :param argv: The command-line arguments, or None for
``sys.argv``. ``sys.argv``.
:return: The exit status: 0 on success, non-zero on failure. :return: The exit status: 0 on success, non-zero on failure.
""" """
started: float = time.monotonic() started: float = time.monotonic()
args: argparse.Namespace = parse_args(argv) args: argparse.Namespace = parse_args(argv)
session: Session = ds.get_db()
lines: list[str]
try: try:
extras: dict[str, Any] | None = None count: int = LlmInputExporter(
if args.extras is not None: args.output_jsonl, args.extras,
extras = load_extras(args.extras) args.extras_per_id).run()
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: except (OSError, sa.exc.SQLAlchemyError, ValueError) as error:
print(f"error: {error}", file=sys.stderr) print(f"error: {error}", file=sys.stderr)
return 1 return 1
finally:
session.close()
args.output_jsonl.parent.mkdir(parents=True, exist_ok=True)
with open(args.output_jsonl, "w", encoding="utf-8") as file:
line: str
for line in lines:
file.write(line + "\n")
elapsed: str = format_duration(time.monotonic() - started) elapsed: str = format_duration(time.monotonic() - started)
print(f"Done. {len(lines)} songs exported." print(f"Done. {count} songs exported."
f" {elapsed} elapsed.", file=sys.stderr) f" {elapsed} elapsed.", file=sys.stderr)
return 0 return 0
@@ -12,12 +12,10 @@ layer. The working store is only read, never written; the
``build-db`` subcommand assembles the captured files into the ``build-db`` subcommand assembles the captured files into the
store on the next rebuild. store on the next rebuild.
Every fetched row is meant for later human verification: the An unresolved artist or an error on one artist is noted on its
description of the resolved item is recorded in the note column row and does not fail the run. A row whose name is no longer an
so that a bad match can be spotted. An unresolved artist or an artist of the store is dropped from the snapshot and reported on
error on one artist is noted on its row and does not fail the the standard error.
run. A row whose name is no longer an artist of the store is
dropped from the snapshot and reported on the standard error.
""" """
import argparse import argparse
import csv import csv
@@ -33,7 +31,7 @@ import urllib.request
from collections.abc import Container, Sequence from collections.abc import Container, Sequence
from dataclasses import asdict, dataclass, field, fields from dataclasses import asdict, dataclass, field, fields
from pathlib import Path from pathlib import Path
from typing import Any, Literal, TextIO from typing import Any, ClassVar, Literal, TextIO
import sqlalchemy as sa import sqlalchemy as sa
from sqlalchemy.orm import Session from sqlalchemy.orm import Session
@@ -43,71 +41,6 @@ from ..database import ds
from ..models import Artist, Song, SongArtist from ..models import Artist, Song, SongArtist
from ..utils import format_duration from ..utils import format_duration
API_URL: str = "https://www.wikidata.org/w/api.php"
"""The URL of the Wikidata API endpoint."""
SPARQL_URL: str = "https://query.wikidata.org/sparql"
"""The URL of the Wikidata Query Service SPARQL endpoint."""
USER_AGENT: str = (
f"pop-fem-audit-tools/{VERSION}"
" (https://github.com/imacat/pop-fem-audit;"
" mailto:imacat@mail.imacat.idv.tw)")
"""The User-Agent header sent on every HTTP request."""
TIMEOUT: float = 30.0
"""The timeout of an API HTTP request, in seconds."""
SPARQL_TIMEOUT: float = 90.0
"""The timeout of a SPARQL HTTP request, in seconds.
Higher than the API timeout: the WDQS server aborts a slow
query at 60 seconds, and a lower client timeout would race
that server-side abort and misclassify a slow-but-answerable
query as a client-side timeout instead of letting the server's
own HTTP error response arrive and enter the retry path."""
SLEEP_SECONDS: float = 1.0
"""The delay between consecutive HTTP requests, in seconds."""
MAX_ATTEMPTS: int = 5
"""The maximum number of attempts on a transient error."""
RETRY_SECONDS: float = 15.0
"""The back-off unit on a transient error, in seconds;
multiplied by the attempt number already made."""
RETRY_STATUSES: frozenset[int] = frozenset({429, 500, 502, 503})
"""The HTTP statuses that are retried with a back-off."""
MAX_STAGE1_TITLES: int = 3
"""The maximum number of charted titles used for the stage-1 song
corroboration."""
HUMAN_QID: str = "Q5"
"""The Wikidata item ID of "human"."""
ENSEMBLE_QID: str = "Q2088357"
"""The Wikidata item ID of "musical ensemble"."""
ORIGINAL_CAST_QID: str = "Q106497009"
"""The Wikidata item ID of "original cast"."""
GROUP_KEYWORDS: Sequence[str] = ("band", "group", "duo", "trio")
"""The label keywords that suggest a musical ensemble, covering
labels like "boy band" and "girl group"."""
NOTE_NOT_FOUND: str = "not found"
"""The note sentinel of an artist without a resolved Wikidata
item, written to the snapshot and read back for the
classification."""
CORPUS_START_YEAR: int = 2016
"""The first year of the corpus window: a member who left a group
before it never performed a corpus song."""
MIXED_GENDER: str = "mixed"
"""The gender recorded for a group whose members do not share one
gender."""
TIME_YEAR_PATTERN: re.Pattern[str] = re.compile(r"^[+-]?\d+")
"""The leading year of a Wikidata time value."""
PINNED_QIDS: dict[str, str] = {}
"""The last-resort pinned item IDs, keyed by the artist name.
An entry is for an artist the algorithm documented on
``ArtistFetcher`` is structurally unable to resolve, with its
justification recorded here. Currently empty: the only pin ever
needed, "Pinkfong" (typed as a brand, which the type gate
excludes by design), became moot when the store's artist entity
behind that credit was identified as Hope Segoine.
A pinned name skips the candidate retrieval and corroboration
steps; its item ID is used directly."""
class ArtistType(enum.StrEnum): class ArtistType(enum.StrEnum):
"""The decided artist type of a snapshot row.""" """The decided artist type of a snapshot row."""
@@ -147,15 +80,14 @@ class ArtistSnapshot:
return asdict(self) return asdict(self)
SNAPSHOT_FIELDS: Sequence[str] = tuple(
x.name for x in fields(ArtistSnapshot))
"""The header columns of the Wikidata artist snapshot CSV file."""
@dataclass @dataclass
class GroupMember: class GroupMember:
"""One has-part member of a Wikidata group item.""" """One has-part member of a Wikidata group item."""
__CORPUS_START_YEAR: ClassVar[int] = 2016
"""The first year of the corpus window: a member who left a
group before it never performed a corpus song."""
qid: str qid: str
"""The item ID of the member.""" """The item ID of the member."""
start_years: list[int] = field(default_factory=list) start_years: list[int] = field(default_factory=list)
@@ -181,7 +113,7 @@ class GroupMember:
if len(self.end_years) == 0: if len(self.end_years) == 0:
return True return True
last_end: int = max(self.end_years) last_end: int = max(self.end_years)
if last_end >= CORPUS_START_YEAR: if last_end >= self.__CORPUS_START_YEAR:
return True return True
if len(self.start_years) == 0: if len(self.start_years) == 0:
return False return False
@@ -229,22 +161,6 @@ class RetryExhausted(Exception):
""" """
def parse_args(argv: list[str] | None) -> argparse.Namespace:
"""Parse the command-line arguments.
:param argv: The command-line arguments, or None for
``sys.argv``.
:return: The parsed arguments.
"""
parser: argparse.ArgumentParser = argparse.ArgumentParser(
description="Fetch the artist metadata from Wikidata"
" into the capture layer.")
parser.add_argument(
"wikidata_csv", type=Path,
help="the Wikidata artist snapshot CSV file")
return parser.parse_args(argv)
class ArtistFetcher: class ArtistFetcher:
"""A fetcher of artist metadata from Wikidata. """A fetcher of artist metadata from Wikidata.
@@ -291,6 +207,74 @@ class ArtistFetcher:
in the note. in the note.
""" """
__API_URL: ClassVar[str] = "https://www.wikidata.org/w/api.php"
"""The URL of the Wikidata API endpoint."""
__SPARQL_URL: ClassVar[str] \
= "https://query.wikidata.org/sparql"
"""The URL of the Wikidata Query Service SPARQL endpoint."""
__USER_AGENT: ClassVar[str] = (
f"pop-fem-audit-tools/{VERSION}"
" (https://github.com/imacat/pop-fem-audit;"
" mailto:imacat@mail.imacat.idv.tw)")
"""The User-Agent header sent on every HTTP request."""
__TIMEOUT: ClassVar[float] = 30.0
"""The timeout of an API HTTP request, in seconds."""
__SPARQL_TIMEOUT: ClassVar[float] = 90.0
"""The timeout of a SPARQL HTTP request, in seconds.
Higher than the API timeout: the WDQS server aborts a slow
query at 60 seconds, and a lower client timeout would race
that server-side abort and misclassify a slow-but-answerable
query as a client-side timeout instead of letting the
server's own HTTP error response arrive and enter the retry
path."""
__SLEEP_SECONDS: ClassVar[float] = 1.0
"""The delay between consecutive HTTP requests, in
seconds."""
__MAX_ATTEMPTS: ClassVar[int] = 5
"""The maximum number of attempts on a transient error."""
__RETRY_SECONDS: ClassVar[float] = 15.0
"""The back-off unit on a transient error, in seconds;
multiplied by the attempt number already made."""
__RETRY_STATUSES: ClassVar[frozenset[int]] \
= frozenset({429, 500, 502, 503})
"""The HTTP statuses that are retried with a back-off."""
__MAX_STAGE1_TITLES: ClassVar[int] = 3
"""The maximum number of charted titles used for the
stage-1 song corroboration."""
__HUMAN_QID: ClassVar[str] = "Q5"
"""The Wikidata item ID of "human"."""
__ENSEMBLE_QID: ClassVar[str] = "Q2088357"
"""The Wikidata item ID of "musical ensemble"."""
__ORIGINAL_CAST_QID: ClassVar[str] = "Q106497009"
"""The Wikidata item ID of "original cast"."""
__GROUP_KEYWORDS: ClassVar[Sequence[str]] \
= ("band", "group", "duo", "trio")
"""The label keywords that suggest a musical ensemble,
covering labels like "boy band" and "girl group"."""
__NOTE_NOT_FOUND: ClassVar[str] = "not found"
"""The note sentinel of an artist without a resolved
Wikidata item."""
__MIXED_GENDER: ClassVar[str] = "mixed"
"""The gender recorded for a group whose members do not
share one gender."""
__TIME_YEAR_PATTERN: ClassVar[re.Pattern[str]] \
= re.compile(r"^[+-]?\d+")
"""The leading year of a Wikidata time value."""
__PINNED_QIDS: ClassVar[dict[str, str]] = {}
"""The last-resort pinned item IDs, keyed by the artist name.
An entry is for an artist the algorithm documented on
``ArtistFetcher`` is structurally unable to resolve, with its
justification recorded here. Currently empty: the only pin
ever needed, "Pinkfong" (typed as a brand, which the type
gate excludes by design), became moot when the store's
artist entity behind that credit was identified as Hope
Segoine.
A pinned name skips the candidate retrieval and corroboration
steps; its item ID is used directly."""
def __init__(self) -> None: def __init__(self) -> None:
"""Construct the fetcher.""" """Construct the fetcher."""
self.__sent: int = 0 self.__sent: int = 0
@@ -317,7 +301,7 @@ class ArtistFetcher:
try: try:
qid: str | None = self.__resolve_qid(name, titles) qid: str | None = self.__resolve_qid(name, titles)
if qid is None: if qid is None:
snapshot.note = NOTE_NOT_FOUND snapshot.note = self.__NOTE_NOT_FOUND
return snapshot return snapshot
snapshot.qid = qid snapshot.qid = qid
self.__resolve(snapshot) self.__resolve(snapshot)
@@ -342,8 +326,8 @@ class ArtistFetcher:
transient error are exhausted. transient error are exhausted.
:raises ValueError: On a JSON decoding error. :raises ValueError: On a JSON decoding error.
""" """
if name in PINNED_QIDS: if name in self.__PINNED_QIDS:
return PINNED_QIDS[name] return self.__PINNED_QIDS[name]
candidates: list[str] = self.__candidates(name) candidates: list[str] = self.__candidates(name)
if len(candidates) == 0: if len(candidates) == 0:
return None return None
@@ -373,11 +357,11 @@ class ArtistFetcher:
{{ ?item rdfs:label ?name }} {{ ?item rdfs:label ?name }}
UNION {{ ?item skos:altLabel ?name }} UNION {{ ?item skos:altLabel ?name }}
{{ {{
?item wdt:P31 wd:{HUMAN_QID} ?item wdt:P31 wd:{self.__HUMAN_QID}
}} UNION {{ }} UNION {{
?item wdt:P31/wdt:P279* wd:{ENSEMBLE_QID} ?item wdt:P31/wdt:P279* wd:{self.__ENSEMBLE_QID}
}} UNION {{ }} UNION {{
?item wdt:P31 wd:{ORIGINAL_CAST_QID} ?item wdt:P31 wd:{self.__ORIGINAL_CAST_QID}
}} }}
}} }}
""" """
@@ -403,7 +387,7 @@ class ArtistFetcher:
transient error are exhausted. transient error are exhausted.
:raises ValueError: On a JSON decoding error. :raises ValueError: On a JSON decoding error.
""" """
subset: Sequence[str] = titles[:MAX_STAGE1_TITLES] subset: Sequence[str] = titles[:self.__MAX_STAGE1_TITLES]
if len(subset) == 0: if len(subset) == 0:
return None return None
query: str = f""" query: str = f"""
@@ -532,7 +516,7 @@ class ArtistFetcher:
qid: str qid: str
for qid in qids: for qid in qids:
member: MemberClaims = claims.get(qid, MemberClaims()) member: MemberClaims = claims.get(qid, MemberClaims())
if HUMAN_QID not in member.instance_of_ids: if self.__HUMAN_QID not in member.instance_of_ids:
continue continue
if len(member.gender_ids) == 0: if len(member.gender_ids) == 0:
return return
@@ -543,7 +527,7 @@ class ArtistFetcher:
[x[1] for x in genders], any_language=True) [x[1] for x in genders], any_language=True)
unique: set[str] = {x[1] for x in genders} unique: set[str] = {x[1] for x in genders}
snapshot.gender = labels[genders[0][1]] \ snapshot.gender = labels[genders[0][1]] \
if len(unique) == 1 else MIXED_GENDER if len(unique) == 1 else self.__MIXED_GENDER
basis: str = "gender derived from members: " + "; ".join( basis: str = "gender derived from members: " + "; ".join(
f"{x} {labels[y]}" for x, y in genders) f"{x} {labels[y]}" for x, y in genders)
snapshot.note = f"{snapshot.note}; {basis}" \ snapshot.note = f"{snapshot.note}; {basis}" \
@@ -677,7 +661,8 @@ class ArtistFetcher:
or not isinstance(value.get("time"), str): or not isinstance(value.get("time"), str):
continue continue
match: re.Match[str] | None \ match: re.Match[str] | None \
= TIME_YEAR_PATTERN.match(value["time"]) = ArtistFetcher.__TIME_YEAR_PATTERN.match(
value["time"])
if match is not None: if match is not None:
years.append(int(match.group())) years.append(int(match.group()))
return years return years
@@ -806,12 +791,13 @@ class ArtistFetcher:
``ArtistType.GROUP`` for a musical ensemble, or the ``ArtistType.GROUP`` for a musical ensemble, or the
empty string for the human to decide. empty string for the human to decide.
""" """
if HUMAN_QID in type_ids: if ArtistFetcher.__HUMAN_QID in type_ids:
return ArtistType.SOLO return ArtistType.SOLO
qid: str qid: str
for qid in type_ids: for qid in type_ids:
label: str = labels.get(qid, "").lower() label: str = labels.get(qid, "").lower()
if any(x in label for x in GROUP_KEYWORDS): if any(x in label
for x in ArtistFetcher.__GROUP_KEYWORDS):
return ArtistType.GROUP return ArtistType.GROUP
return "" return ""
@@ -827,14 +813,15 @@ class ArtistFetcher:
transient error are exhausted. transient error are exhausted.
:raises ValueError: On a JSON decoding error. :raises ValueError: On a JSON decoding error.
""" """
url: str = (f"{SPARQL_URL}?" url: str = (
f"{self.__SPARQL_URL}?"
f"{urllib.parse.urlencode({'query': query})}") f"{urllib.parse.urlencode({'query': query})}")
request: urllib.request.Request = urllib.request.Request( request: urllib.request.Request = urllib.request.Request(
url, headers={ url, headers={
"User-Agent": USER_AGENT, "User-Agent": self.__USER_AGENT,
"Accept": "application/sparql-results+json"}) "Accept": "application/sparql-results+json"})
body: bytes = self.__send( body: bytes = self.__send(
request, timeout=SPARQL_TIMEOUT) request, timeout=self.__SPARQL_TIMEOUT)
data: Any = json.loads(body) data: Any = json.loads(body)
bindings: Any = None bindings: Any = None
if isinstance(data, dict) \ if isinstance(data, dict) \
@@ -868,13 +855,14 @@ class ArtistFetcher:
transient error are exhausted. transient error are exhausted.
:raises ValueError: On a JSON decoding error. :raises ValueError: On a JSON decoding error.
""" """
url: str = f"{API_URL}?{urllib.parse.urlencode(params)}" url: str \
= f"{self.__API_URL}?{urllib.parse.urlencode(params)}"
request: urllib.request.Request = urllib.request.Request( request: urllib.request.Request = urllib.request.Request(
url, headers={"User-Agent": USER_AGENT}) url, headers={"User-Agent": self.__USER_AGENT})
return json.loads(self.__send(request)) return json.loads(self.__send(request))
def __send(self, request: urllib.request.Request, def __send(self, request: urllib.request.Request,
timeout: float = TIMEOUT) -> bytes: timeout: float = __TIMEOUT) -> bytes:
"""Send an HTTP request, retrying on a transient error. """Send an HTTP request, retrying on a transient error.
Consecutive requests are separated by a fixed delay. A Consecutive requests are separated by a fixed delay. A
@@ -892,7 +880,7 @@ class ArtistFetcher:
transient error are exhausted. transient error are exhausted.
""" """
if self.__sent > 0: if self.__sent > 0:
time.sleep(SLEEP_SECONDS) time.sleep(self.__SLEEP_SECONDS)
self.__sent += 1 self.__sent += 1
attempt: int = 1 attempt: int = 1
reason: str | None reason: str | None
@@ -905,10 +893,10 @@ class ArtistFetcher:
reason = self.__retry_reason(error) reason = self.__retry_reason(error)
if reason is None: if reason is None:
raise raise
if attempt >= MAX_ATTEMPTS: if attempt >= self.__MAX_ATTEMPTS:
raise RetryExhausted( raise RetryExhausted(
f"retries exhausted ({reason})") from error f"retries exhausted ({reason})") from error
time.sleep(RETRY_SECONDS * attempt) time.sleep(self.__RETRY_SECONDS * attempt)
attempt += 1 attempt += 1
@staticmethod @staticmethod
@@ -923,7 +911,7 @@ class ArtistFetcher:
the error is not transient and must not be retried. the error is not transient and must not be retried.
""" """
if isinstance(error, urllib.error.HTTPError): if isinstance(error, urllib.error.HTTPError):
if error.code not in RETRY_STATUSES: if error.code not in ArtistFetcher.__RETRY_STATUSES:
return None return None
return str(error) return str(error)
if isinstance(error, TimeoutError): if isinstance(error, TimeoutError):
@@ -968,7 +956,106 @@ class ArtistFetcher:
return uri.rsplit("/", 1)[-1] return uri.rsplit("/", 1)[-1]
def read_snapshot_rows(file: TextIO) -> list[dict[str, str]]: @dataclass(frozen=True)
class FetchCounts:
"""The outcome counts of one snapshot update run."""
fetched: int
"""The number of artists newly resolved."""
not_found: int
"""The number of artists left unresolved."""
errors: int
"""The number of artists that ended in an error."""
class ArtistSnapshotUpdater:
"""The updater of the Wikidata artist snapshot CSV file.
Fetches the metadata of every artist of the working store
that the snapshot does not resolve yet, appends a row for
each to the snapshot as it is fetched, and rewrites the
snapshot sorted by artist name with its stale rows dropped.
"""
__SNAPSHOT_FIELDS: ClassVar[Sequence[str]] = tuple(
x.name for x in fields(ArtistSnapshot))
"""The header columns of the Wikidata artist snapshot CSV file."""
def __init__(self, wikidata_csv: Path) -> None:
"""Set up the updater.
:param wikidata_csv: The Wikidata artist snapshot CSV
file.
"""
self.__wikidata_csv: Path = wikidata_csv
"""The Wikidata artist snapshot CSV file."""
def run(self) -> FetchCounts:
"""Fetch every unresolved artist and update the snapshot.
:return: The counts of the run.
:raises OSError: When the snapshot file, or its parent
directory, cannot be read or written.
:raises sqlalchemy.exc.SQLAlchemyError: When the working
store cannot be read.
"""
session: Session = ds.get_db()
try:
return self.__run(session)
finally:
session.close()
def __run(self, session: Session) -> FetchCounts:
"""Run the fetch loop with an open database session.
:param session: The database session.
:return: The counts of the run.
:raises OSError: When the snapshot file, or its parent
directory, cannot be read or written.
"""
fetcher: ArtistFetcher = ArtistFetcher()
fetched: int = 0
not_found: int = 0
errors: int = 0
self.__wikidata_csv.parent.mkdir(
parents=True, exist_ok=True)
with open(self.__wikidata_csv, "a+", encoding="utf-8",
newline="") as csv_file:
done: set[str] = {
x["name"] for x in
self.__read_snapshot_rows(csv_file)
if x["gender"] != ""}
self.__ensure_snapshot_header(csv_file)
names: set[str] = set()
artist: Artist
for artist in session.scalars(
sa.select(Artist).order_by(Artist.id)):
names.add(artist.name)
if artist.name in done:
continue
titles: list[str] = self.__read_artist_titles(
session, artist.id)
snapshot: ArtistSnapshot = fetcher.fetch(
artist.name, titles)
self.__append_row(csv_file, snapshot)
status: str = snapshot.qid
if snapshot.note == "not found":
not_found += 1
status = "not found"
elif snapshot.note.startswith("error: "):
errors += 1
status = snapshot.note
else:
fetched += 1
print(f"artist \"{artist.name}\": {status}",
file=sys.stderr)
self.__write_snapshot(csv_file, names)
return FetchCounts(
fetched=fetched, not_found=not_found, errors=errors)
@staticmethod
def __read_snapshot_rows(file: TextIO) \
-> list[dict[str, str]]:
"""Read the current rows of a snapshot CSV file handle. """Read the current rows of a snapshot CSV file handle.
:param file: The open, seekable snapshot CSV file. :param file: The open, seekable snapshot CSV file.
@@ -979,15 +1066,15 @@ def read_snapshot_rows(file: TextIO) -> list[dict[str, str]]:
reader: csv.DictReader[str] = csv.DictReader(file) reader: csv.DictReader[str] = csv.DictReader(file)
return list(reader) return list(reader)
@staticmethod
def read_artist_titles(session: Session, def __read_artist_titles(session: Session,
artist_id: int) -> list[str]: artist_id: int) -> list[str]:
"""Read the charted song titles credited to an artist. """Read the charted song titles credited to an artist.
:param session: The database session. :param session: The database session.
:param artist_id: The artist ID. :param artist_id: The artist ID.
:return: The song titles credited to the artist, ordered by :return: The song titles credited to the artist, ordered
the song ID, with the duplicate titles removed. by the song ID, with the duplicate titles removed.
""" """
titles: Sequence[str] = session.scalars( titles: Sequence[str] = session.scalars(
sa.select(Song.title) sa.select(Song.title)
@@ -996,9 +1083,10 @@ def read_artist_titles(session: Session,
.order_by(Song.id)).all() .order_by(Song.id)).all()
return list(dict.fromkeys(titles)) return list(dict.fromkeys(titles))
@staticmethod
def ensure_snapshot_header(file: TextIO) -> None: def __ensure_snapshot_header(file: TextIO) -> None:
"""Write the snapshot CSV header row if the file is empty. """Write the snapshot CSV header row if the file is
empty.
:param file: The open, seekable snapshot CSV file. :param file: The open, seekable snapshot CSV file.
:return: None. :return: None.
@@ -1006,32 +1094,38 @@ def ensure_snapshot_header(file: TextIO) -> None:
""" """
file.seek(0, os.SEEK_END) file.seek(0, os.SEEK_END)
if file.tell() == 0: if file.tell() == 0:
csv.writer(file).writerow(SNAPSHOT_FIELDS) csv.writer(file).writerow(
ArtistSnapshotUpdater.__SNAPSHOT_FIELDS)
file.flush() file.flush()
@staticmethod
def append_row(file: TextIO, snapshot: ArtistSnapshot) -> None: def __append_row(file: TextIO,
snapshot: ArtistSnapshot) -> None:
"""Append a snapshot row to a snapshot CSV file handle. """Append a snapshot row to a snapshot CSV file handle.
:param file: The open snapshot CSV file, opened for append. :param file: The open snapshot CSV file, opened for
append.
:param snapshot: The snapshot of an artist. :param snapshot: The snapshot of an artist.
:return: None. :return: None.
:raises OSError: When the file cannot be written. :raises OSError: When the file cannot be written.
""" """
csv.DictWriter(file, SNAPSHOT_FIELDS).writerow( csv.DictWriter(
file, ArtistSnapshotUpdater.__SNAPSHOT_FIELDS).writerow(
snapshot.to_row()) snapshot.to_row())
file.flush() file.flush()
@staticmethod
def __write_snapshot(file: TextIO, names: Container[str]) \
-> None:
"""Rewrite a snapshot CSV file handle sorted by artist
name.
def write_snapshot(file: TextIO, names: Container[str]) -> None: The rows are ordered by the case-folded artist name,
"""Rewrite a snapshot CSV file handle sorted by artist name. matching the convention of the derived ``artists.csv``.
An artist keeps one row only, the last one of the file,
The rows are ordered by the case-folded artist name, matching so that a re-fetched artist replaces its earlier row. A
the convention of the derived ``artists.csv``. An artist row whose name is not an artist of the store is dropped
keeps one row only, the last one of the file, so that a and reported on the standard error.
re-fetched artist replaces its earlier row. A row whose name
is not an artist of the store is dropped and reported on the
standard error.
:param file: The open, seekable snapshot CSV file. :param file: The open, seekable snapshot CSV file.
:param names: The artist names of the working store. :param names: The artist names of the working store.
@@ -1040,10 +1134,12 @@ def write_snapshot(file: TextIO, names: Container[str]) -> None:
""" """
kept: dict[str, dict[str, str]] = {} kept: dict[str, dict[str, str]] = {}
row: dict[str, str] row: dict[str, str]
for row in read_snapshot_rows(file): for row in ArtistSnapshotUpdater.__read_snapshot_rows(
file):
if row["name"] not in names: if row["name"] not in names:
print(f"dropped stale row \"{row['name']}\":" print(f"dropped stale row \"{row['name']}\":"
" no such artist in the store", file=sys.stderr) " no such artist in the store",
file=sys.stderr)
continue continue
kept[row["name"]] = row kept[row["name"]] = row
ordered: list[dict[str, str]] = sorted( ordered: list[dict[str, str]] = sorted(
@@ -1051,11 +1147,27 @@ def write_snapshot(file: TextIO, names: Container[str]) -> None:
file.seek(0) file.seek(0)
file.truncate() file.truncate()
writer: csv.DictWriter[str] = csv.DictWriter( writer: csv.DictWriter[str] = csv.DictWriter(
file, SNAPSHOT_FIELDS) file, ArtistSnapshotUpdater.__SNAPSHOT_FIELDS)
writer.writeheader() writer.writeheader()
writer.writerows(ordered) writer.writerows(ordered)
def parse_args(argv: list[str] | None) -> argparse.Namespace:
"""Parse the command-line arguments.
:param argv: The command-line arguments, or None for
``sys.argv``.
:return: The parsed arguments.
"""
parser: argparse.ArgumentParser = argparse.ArgumentParser(
description="Fetch the artist metadata from Wikidata"
" into the capture layer.")
parser.add_argument(
"wikidata_csv", type=Path,
help="the Wikidata artist snapshot CSV file")
return parser.parse_args(argv)
def main(argv: list[str] | None = None) -> int: def main(argv: list[str] | None = None) -> int:
"""Fetch the artist metadata from Wikidata. """Fetch the artist metadata from Wikidata.
@@ -1066,52 +1178,16 @@ def main(argv: list[str] | None = None) -> int:
""" """
started: float = time.monotonic() started: float = time.monotonic()
args: argparse.Namespace = parse_args(argv) args: argparse.Namespace = parse_args(argv)
fetcher: ArtistFetcher = ArtistFetcher()
fetched: int = 0
not_found: int = 0
errors: int = 0
session: Session = ds.get_db()
try: try:
args.wikidata_csv.parent.mkdir( counts: FetchCounts \
parents=True, exist_ok=True) = ArtistSnapshotUpdater(args.wikidata_csv).run()
with open(args.wikidata_csv, "a+", encoding="utf-8",
newline="") as csv_file:
done: set[str] = {x["name"] for x in
read_snapshot_rows(csv_file)
if x["gender"] != ""}
ensure_snapshot_header(csv_file)
names: set[str] = set()
artist: Artist
for artist in session.scalars(
sa.select(Artist).order_by(Artist.id)):
names.add(artist.name)
if artist.name in done:
continue
titles: list[str] = read_artist_titles(
session, artist.id)
snapshot: ArtistSnapshot = fetcher.fetch(
artist.name, titles)
append_row(csv_file, snapshot)
status: str = snapshot.qid
if snapshot.note == NOTE_NOT_FOUND:
not_found += 1
status = "not found"
elif snapshot.note.startswith("error: "):
errors += 1
status = snapshot.note
else:
fetched += 1
print(f"artist \"{artist.name}\": {status}",
file=sys.stderr)
write_snapshot(csv_file, names)
except (OSError, sa.exc.SQLAlchemyError) as error: except (OSError, sa.exc.SQLAlchemyError) as error:
print(f"error: {error}", file=sys.stderr) print(f"error: {error}", file=sys.stderr)
return 1 return 1
finally: attempted: int = counts.fetched + counts.not_found \
session.close() + counts.errors
attempted: int = fetched + not_found + errors
elapsed: str = format_duration(time.monotonic() - started) elapsed: str = format_duration(time.monotonic() - started)
print(f"Done. Resolved {fetched}/{attempted} artists." print(f"Done. Resolved {counts.fetched}/{attempted}"
f" {elapsed} elapsed.", f" artists. {elapsed} elapsed.",
file=sys.stderr) file=sys.stderr)
return 0 return 0
@@ -30,8 +30,9 @@ import time
import urllib.parse import urllib.parse
import urllib.request import urllib.request
from collections.abc import Sequence from collections.abc import Sequence
from dataclasses import dataclass
from pathlib import Path from pathlib import Path
from typing import Any from typing import Any, ClassVar
import sqlalchemy as sa import sqlalchemy as sa
from sqlalchemy.orm import Session from sqlalchemy.orm import Session
@@ -45,60 +46,6 @@ from ..models import (
) )
from ..utils import format_duration from ..utils import format_duration
PROVENANCE_FIELDS: Sequence[str] = (
"song_id", "source", "method", "acquired_at", "note")
"""The header columns of the lyrics provenance CSV file."""
USER_AGENT: str = ("pop-fem-audit-tools"
" (https://github.com/imacat/pop-fem-audit)")
"""The User-Agent header sent on every HTTP request."""
TIMEOUT: float = 30.0
"""The timeout of an HTTP request, in seconds."""
SLEEP_SECONDS: float = 1.0
"""The delay between consecutive HTTP requests, in seconds."""
def __build_normalization() -> dict[int, str | None]:
"""Build the lyrics normalization translation table.
:return: The codepoint-to-replacement mapping, a replacement
of None meaning removal.
"""
table: dict[int, str | None] = {}
codepoint: int
for codepoint in range(0x80, 0xa0):
try:
table[codepoint] = bytes([codepoint]).decode("cp1252")
except UnicodeDecodeError:
table[codepoint] = None
table[0x0435] = "e"
table[0x03cc] = "ó"
for codepoint in (0x2005, 0x205f, 0x200a):
table[codepoint] = " "
for codepoint in (0x200b, 0x200c, 0x200d, 0xfeff):
table[codepoint] = None
return table
NORMALIZATION: dict[int, str | None] = __build_normalization()
"""The codepoint-to-replacement mapping applied to fetched
lyrics: cp1252-mojibake restoration for U+0080-U+009F (with the
five byte values undefined in cp1252 removed), homoglyph
restoration for the Cyrillic "e" and the Greek "o" with tonos,
ASCII-space restoration for exotic space variants, and removal
of zero-width characters. A replacement of None removes the
codepoint."""
def normalize_lyrics(text: str) -> str:
"""Restore or remove watermark and mojibake characters.
:param text: The lyrics text as fetched from an API.
:return: The text with the codepoints in
:data:`NORMALIZATION` replaced or removed; every other
character is unchanged.
"""
return text.translate(NORMALIZATION)
def parse_args(argv: list[str] | None) -> argparse.Namespace: def parse_args(argv: list[str] | None) -> argparse.Namespace:
"""Parse the command-line arguments. """Parse the command-line arguments.
@@ -122,6 +69,16 @@ def parse_args(argv: list[str] | None) -> argparse.Namespace:
class LyricsFetcher: class LyricsFetcher:
"""A fetcher of song lyrics from the public lyrics APIs.""" """A fetcher of song lyrics from the public lyrics APIs."""
__USER_AGENT: ClassVar[str] = (
"pop-fem-audit-tools"
" (https://github.com/imacat/pop-fem-audit)")
"""The User-Agent header sent on every HTTP request."""
__TIMEOUT: ClassVar[float] = 30.0
"""The timeout of an HTTP request, in seconds."""
__SLEEP_SECONDS: ClassVar[float] = 1.0
"""The delay between consecutive HTTP requests, in
seconds."""
def __init__(self) -> None: def __init__(self) -> None:
"""Construct the fetcher.""" """Construct the fetcher."""
self.__sent: int = 0 self.__sent: int = 0
@@ -193,25 +150,154 @@ class LyricsFetcher:
network, or decoding error. network, or decoding error.
""" """
if self.__sent > 0: if self.__sent > 0:
time.sleep(SLEEP_SECONDS) time.sleep(self.__SLEEP_SECONDS)
self.__sent += 1 self.__sent += 1
request: urllib.request.Request = urllib.request.Request( request: urllib.request.Request = urllib.request.Request(
url, headers={"User-Agent": USER_AGENT}) url, headers={"User-Agent": self.__USER_AGENT})
try: try:
with urllib.request.urlopen( with urllib.request.urlopen(
request, timeout=TIMEOUT) as response: request, timeout=self.__TIMEOUT) as response:
return json.load(response) return json.load(response)
except (OSError, ValueError): except (OSError, ValueError):
return None return None
def query_artist(session: Session, song_id: int) -> str: @dataclass(frozen=True)
class LyricsFetchCounts:
"""The outcome of one run of fetching the missing lyrics."""
fetched: int
"""The number of songs newly fetched."""
missed: int
"""The number of songs every API missed."""
class LyricsFetchRunner:
"""The orchestrator of one run of fetching missing lyrics."""
__PROVENANCE_FIELDS: ClassVar[Sequence[str]] = (
"song_id", "source", "method", "acquired_at", "note")
"""The header columns of the lyrics provenance CSV file."""
@staticmethod
def __build_normalization() -> dict[int, str | None]:
"""Build the lyrics normalization translation table.
:return: The codepoint-to-replacement mapping, a
replacement of None meaning removal.
"""
table: dict[int, str | None] = {}
codepoint: int
for codepoint in range(0x80, 0xa0):
try:
table[codepoint] = bytes(
[codepoint]).decode("cp1252")
except UnicodeDecodeError:
table[codepoint] = None
table[0x0435] = "e"
table[0x03cc] = "ó"
for codepoint in (0x2005, 0x205f, 0x200a):
table[codepoint] = " "
for codepoint in (0x200b, 0x200c, 0x200d, 0xfeff):
table[codepoint] = None
return table
__NORMALIZATION: ClassVar[dict[int, str | None]] \
= __build_normalization()
"""The codepoint-to-replacement mapping applied to fetched
lyrics: cp1252-mojibake restoration for U+0080-U+009F (with
the five byte values undefined in cp1252 removed), homoglyph
restoration for the Cyrillic "e" and the Greek "o" with
tonos, ASCII-space restoration for exotic space variants, and
removal of zero-width characters. A replacement of None
removes the codepoint."""
def __init__(self, lyrics_dir: Path,
provenance_csv: Path) -> None:
"""Set up the fetch run.
:param lyrics_dir: The lyrics cache directory.
:param provenance_csv: The lyrics provenance CSV file.
"""
self.__lyrics_dir: Path = lyrics_dir
"""The lyrics cache directory."""
self.__provenance_csv: Path = provenance_csv
"""The lyrics provenance CSV file."""
self.__fetcher: LyricsFetcher = LyricsFetcher()
"""The fetcher of the public lyrics APIs."""
def run(self) -> LyricsFetchCounts:
"""Fetch the missing lyrics of every song in the store.
Every song fetched or missed is reported on the standard
error as an observable side effect.
:return: The number of songs fetched and missed.
:raises OSError: When a cache file or the provenance CSV
cannot be written.
:raises sqlalchemy.exc.SQLAlchemyError: On a database
error.
"""
fetched: int = 0
missed: int = 0
session: Session = ds.get_db()
try:
song: Song
for song in session.scalars(
sa.select(Song).order_by(Song.id)):
if (self.__lyrics_dir
/ f"{song.id}.txt").exists():
continue
if self.__fetch_one(session, song):
fetched += 1
else:
missed += 1
finally:
session.close()
return LyricsFetchCounts(fetched=fetched, missed=missed)
def __fetch_one(self, session: Session, song: Song) -> bool:
"""Fetch and save the lyrics of one song.
The song is queried by its primary-role artist name; when
every API misses and the song's full artist credit
differs from that name, the same APIs are queried again
with the artist credit.
:param session: The database session.
:param song: The song to fetch.
:return: True when a lyrics text was fetched and saved,
False when every API missed on both queries.
:raises OSError: When the cache file or the provenance
CSV cannot be written.
"""
artist: str = self.__query_artist(session, song.id)
result: tuple[str, str] | None = self.__fetcher.fetch(
artist, song.title)
if result is None and song.artist_credit != artist:
result = self.__fetcher.fetch(
song.artist_credit, song.title)
if result is None:
print(f"song {song.id} \"{song.title}\": miss",
file=sys.stderr)
return False
lyrics: str
source: str
lyrics, source = result
self.__save_lyrics(song.id, lyrics)
self.__append_provenance(song.id, source)
print(f"song {song.id} \"{song.title}\": {source}",
file=sys.stderr)
return True
@staticmethod
def __query_artist(session: Session, song_id: int) -> str:
"""Find the artist name to query the APIs with. """Find the artist name to query the APIs with.
:param session: The database session. :param session: The database session.
:param song_id: The song ID. :param song_id: The song ID.
:return: The name of the primary-role artist with the lowest :return: The name of the primary-role artist with the
position. lowest position.
""" """
name: str | None = session.scalar( name: str | None = session.scalar(
sa.select(Artist.name) sa.select(Artist.name)
@@ -223,49 +309,58 @@ def query_artist(session: Session, song_id: int) -> str:
assert name is not None assert name is not None
return name return name
def __save_lyrics(self, song_id: int, lyrics: str) -> None:
def save_lyrics(lyrics_dir: Path, song_id: int,
lyrics: str) -> None:
"""Write the lyrics of a song into the cache directory. """Write the lyrics of a song into the cache directory.
The cache directory is created when missing. The cache directory is created when missing.
The lyrics text is normalized with :func:`normalize_lyrics` The lyrics text is normalized with
before being written. :meth:`normalize_lyrics` before being written.
:param lyrics_dir: The lyrics cache directory.
:param song_id: The song ID. :param song_id: The song ID.
:param lyrics: The lyrics text. :param lyrics: The lyrics text.
:return: None. :return: None.
:raises OSError: When the file cannot be written. :raises OSError: When the file cannot be written.
""" """
lyrics_dir.mkdir(parents=True, exist_ok=True) self.__lyrics_dir.mkdir(parents=True, exist_ok=True)
(lyrics_dir / f"{song_id}.txt").write_text( (self.__lyrics_dir / f"{song_id}.txt").write_text(
normalize_lyrics(lyrics), encoding="utf-8") self.normalize_lyrics(lyrics), encoding="utf-8")
def __append_provenance(self, song_id: int,
def append_provenance(path: Path, song_id: int,
source: str) -> None: source: str) -> None:
"""Append a provenance row for a fetched lyrics file. """Append a provenance row for a fetched lyrics file.
The CSV file is created with the header row when missing. The CSV file is created with the header row when
missing.
:param path: The lyrics provenance CSV file.
:param song_id: The song ID. :param song_id: The song ID.
:param source: The source name of the fetched lyrics. :param source: The source name of the fetched lyrics.
:return: None. :return: None.
:raises OSError: When the file cannot be written. :raises OSError: When the file cannot be written.
""" """
is_new: bool = not path.exists() is_new: bool = not self.__provenance_csv.exists()
path.parent.mkdir(parents=True, exist_ok=True) self.__provenance_csv.parent.mkdir(
with open(path, "a", encoding="utf-8", parents=True, exist_ok=True)
with open(self.__provenance_csv, "a", encoding="utf-8",
newline="") as file: newline="") as file:
writer: Any = csv.writer(file) writer: Any = csv.writer(file)
if is_new: if is_new:
writer.writerow(PROVENANCE_FIELDS) writer.writerow(self.__PROVENANCE_FIELDS)
writer.writerow([song_id, source, "api-fetch", writer.writerow(
[song_id, source, "api-fetch",
datetime.date.today().isoformat(), ""]) datetime.date.today().isoformat(), ""])
@classmethod
def normalize_lyrics(cls, text: str) -> str:
"""Restore or remove watermark and mojibake characters.
:param text: The lyrics text as fetched from an API.
:return: The text with the codepoints of the
normalization table replaced or removed; every other
character is unchanged.
"""
return text.translate(cls.__NORMALIZATION)
def main(argv: list[str] | None = None) -> int: def main(argv: list[str] | None = None) -> int:
"""Fetch the missing song lyrics from the public APIs. """Fetch the missing song lyrics from the public APIs.
@@ -277,44 +372,15 @@ def main(argv: list[str] | None = None) -> int:
""" """
started: float = time.monotonic() started: float = time.monotonic()
args: argparse.Namespace = parse_args(argv) args: argparse.Namespace = parse_args(argv)
fetcher: LyricsFetcher = LyricsFetcher()
fetched: int = 0
missed: int = 0
session: Session = ds.get_db()
try: try:
song: Song counts: LyricsFetchCounts = LyricsFetchRunner(
for song in session.scalars( args.lyrics_dir, args.provenance_csv).run()
sa.select(Song).order_by(Song.id)):
if (args.lyrics_dir / f"{song.id}.txt").exists():
continue
artist: str = query_artist(session, song.id)
result: tuple[str, str] | None = fetcher.fetch(
artist, song.title)
if result is None and song.artist_credit != artist:
result = fetcher.fetch(
song.artist_credit, song.title)
if result is None:
missed += 1
print(f"song {song.id} \"{song.title}\": miss",
file=sys.stderr)
continue
lyrics: str
source: str
lyrics, source = result
save_lyrics(args.lyrics_dir, song.id, lyrics)
append_provenance(args.provenance_csv, song.id,
source)
fetched += 1
print(f"song {song.id} \"{song.title}\": {source}",
file=sys.stderr)
except (OSError, sa.exc.SQLAlchemyError) as error: except (OSError, sa.exc.SQLAlchemyError) as error:
print(f"error: {error}", file=sys.stderr) print(f"error: {error}", file=sys.stderr)
return 1 return 1
finally: attempted: int = counts.fetched + counts.missed
session.close()
attempted: int = fetched + missed
elapsed: str = format_duration(time.monotonic() - started) elapsed: str = format_duration(time.monotonic() - started)
print(f"Done. Fetched lyrics for {fetched}/{attempted}" print(f"Done. Fetched lyrics for {counts.fetched}/"
f" songs. {elapsed} elapsed.", f"{attempted} songs. {elapsed} elapsed.",
file=sys.stderr) file=sys.stderr)
return 0 return 0
+325 -224
View File
@@ -28,27 +28,13 @@ import time
from dataclasses import asdict, dataclass from dataclasses import asdict, dataclass
from datetime import datetime from datetime import datetime
from pathlib import Path from pathlib import Path
from typing import Any, Self from typing import Any, ClassVar, Self
import anthropic import anthropic
from ..config import get_settings from ..config import get_settings
from ..utils import format_duration from ..utils import format_duration
# claude-fable-5 accepts neither "temperature" nor "thinking";
# a model's entry holds exactly the extra request parameters it
# accepts.
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
class InputFormatError(Exception): class InputFormatError(Exception):
"""An error in the JSONL input file.""" """An error in the JSONL input file."""
@@ -136,13 +122,23 @@ class BatchResult:
if x.type == "text") if x.type == "text")
return cls(id=entry.custom_id, text=text, return cls(id=entry.custom_id, text=text,
stop_reason=message.stop_reason, stop_reason=message.stop_reason,
usage=usage_to_dict(message.usage)) usage=cls.__usage_to_dict(message.usage))
case "errored": case "errored":
return cls(id=entry.custom_id, return cls(id=entry.custom_id,
error=result.error.error.type) error=result.error.error.type)
case other: case other:
return cls(id=entry.custom_id, error=str(other)) return cls(id=entry.custom_id, error=str(other))
@staticmethod
def __usage_to_dict(usage: Any) -> dict[str, Any]:
"""Convert a usage object to a plain dictionary.
:param usage: The usage object of a message.
:return: The usage as a dictionary, without null entries.
"""
return {k: v for k, v in usage.model_dump().items()
if v is not None}
def to_record(self) -> dict[str, Any]: def to_record(self) -> dict[str, Any]:
"""Return this result as an archive JSONL record. """Return this result as an archive JSONL record.
@@ -169,88 +165,226 @@ class BatchInfo:
is still processing.""" is still processing."""
def parse_args(argv: list[str] | None) -> argparse.Namespace: @dataclass(frozen=True)
"""Parse the command-line arguments. class ExecutionOutcome:
"""The outcome of submitting and awaiting one batch."""
:param argv: The command-line arguments, or None for ``sys.argv``. batch: BatchInfo
:return: The parsed arguments. """The submitted batch's bookkeeping."""
results: Results
"""The batch's results, keyed by item ID."""
@dataclass(frozen=True)
class RunOutcome:
"""The outcome of one LLM definition file run."""
item_count: int
"""The number of loaded input items."""
dry_run: bool
"""Whether this was a dry run."""
dry_run_request: dict[str, Any] | None
"""The first item's preview request, for a dry run; None for
an actual run."""
failed: list[str]
"""The failed item IDs, in item order; always empty for a dry
run."""
class LLMRunner:
"""The orchestrator of one LLM definition file run."""
# claude-fable-5 accepts neither "temperature" nor "thinking";
# a model's entry holds exactly the extra request parameters
# it accepts.
MODELS: ClassVar[dict[str, dict[str, Any]]] = {
"claude-sonnet-4-6": {
"temperature": 0.0,
"thinking": {"type": "disabled"},
},
"claude-fable-5": {},
}
"""The supported model IDs and their extra request
parameters."""
DEFAULT_MODEL: ClassVar[str] = "claude-sonnet-4-6"
"""The default model ID."""
__SCRIPT_VERSION: ClassVar[str] = "run_llm.py 3.1.0"
"""The script version recorded into the archive metadata."""
__POLL_INTERVAL_SECONDS: ClassVar[float] = 60.0
"""The interval between batch status polls."""
def __init__(self, prompt: Path, input_path: Path,
archive_dir: Path, model: str, max_tokens: int,
dry_run: bool, replace: bool) -> None:
"""Set up the run of one LLM definition file.
:param prompt: The prompt definition file, used as the
system prompt.
:param input_path: The JSONL input file with "id" and
"content".
:param archive_dir: The destination archive directory.
:param model: The model ID, a key of :attr:`MODELS`.
:param max_tokens: The maximum output tokens per request.
:param dry_run: Whether to validate and archive without
calling the API.
:param replace: Whether to replace an already existing
archive directory.
""" """
parser: argparse.ArgumentParser = argparse.ArgumentParser( self.__prompt: Path = prompt
description="Run one LLM definition file against one input" """The prompt definition file."""
" and archive the result.") self.__input: Path = input_path
parser.add_argument( """The JSONL input file."""
"prompt", type=Path, self.__archive_dir: Path = archive_dir
help="the prompt definition file, used as the system prompt") """The destination archive directory."""
parser.add_argument( self.__model: str = model
"input", type=Path, """The model ID."""
help="the JSONL input file with \"id\" and \"content\"") self.__max_tokens: int = max_tokens
parser.add_argument( """The maximum output tokens per request."""
"archive_dir", type=Path, self.__dry_run: bool = dry_run
help="the destination archive directory") """Whether to validate and archive without calling the
parser.add_argument( API."""
"--model", choices=sorted(MODELS), default=DEFAULT_MODEL, self.__replace: bool = replace
help=f"the model ID (default {DEFAULT_MODEL})") """Whether to replace an already existing archive
parser.add_argument( directory."""
"--max-tokens", type=int, default=2048,
help="the maximum output tokens per request (default 2048)")
parser.add_argument(
"--dry-run", action="store_true",
help="validate and archive without calling the API")
parser.add_argument(
"--replace", action="store_true",
help="replace an already existing archive directory")
return parser.parse_args(argv)
def run(self) -> RunOutcome:
"""Load the input, archive the prompt, and run the batch.
def load_items(path: Path) -> list[InputItem]: Always writes ``prompt.md`` and ``meta.json`` into the
archive directory. A dry run stops there, previewing the
first item's request; an actual run also submits the
batch, awaits it, and writes ``output.jsonl``.
:return: The outcome of the run.
:raises InputFormatError: When the input file is
malformed.
:raises OSError: When the input or prompt file cannot be
read, the archive directory already exists without
``replace``, or an output file cannot be written.
"""
items: list[InputItem] = self.__load_items()
prompt_text: str = self.__prompt.read_text(encoding="utf-8")
archive_dir: Path = self.__create_archive_dir()
(archive_dir / "prompt.md").write_bytes(
self.__prompt.read_bytes())
meta: dict[str, Any] = self.__build_meta(items)
meta_path: Path = archive_dir / "meta.json"
if self.__dry_run:
self.__write_json(meta_path, meta)
request: dict[str, Any] = self.__build_request(
items[0], prompt_text)
return RunOutcome(
item_count=len(items), dry_run=True,
dry_run_request=request, failed=[])
client: anthropic.Anthropic = anthropic.Anthropic(
api_key=get_settings().ANTHROPIC_API_KEY)
outcome: ExecutionOutcome = self.__execute_run(
client, items, prompt_text)
item_ids: list[str] = [x.id for x in items]
self.__write_jsonl(
archive_dir / "output.jsonl",
[outcome.results[x].to_record() for x in item_ids
if x in outcome.results])
meta["batch"] = outcome.batch
meta["usage"] = self.__sum_usage(outcome.results)
self.__write_meta(meta_path, meta)
failed: list[str] = self.__find_failures(
item_ids, outcome.results)
return RunOutcome(
item_count=len(items), dry_run=False,
dry_run_request=None, failed=failed)
def __load_items(self) -> list[InputItem]:
"""Load and validate the JSONL input items. """Load and validate the JSONL input items.
:param path: The path of the JSONL input file.
:return: The input items, in file order. :return: The input items, in file order.
:raises InputFormatError: When a line is malformed, an ID is :raises InputFormatError: When a line is malformed, an ID
duplicated, or the file contains no item. is duplicated, or the file contains no item.
:raises OSError: When the file cannot be read. :raises OSError: When the file cannot be read.
""" """
items: list[InputItem] = [] items: list[InputItem] = []
seen: set[str] = set() seen: set[str] = set()
with open(path, encoding="utf-8") as file: with open(self.__input, encoding="utf-8") as file:
for number, line in enumerate(file, start=1): for number, line in enumerate(file, start=1):
if line.strip() == "": if line.strip() == "":
continue continue
data: Any
try: try:
data: Any = json.loads(line) data = json.loads(line)
except json.JSONDecodeError as error: except json.JSONDecodeError as error:
raise InputFormatError( raise InputFormatError(
f"{path}: line {number}: malformed JSON: {error}") f"{self.__input}: line {number}: malformed"
f" JSON: {error}")
item: InputItem = InputItem.get_instance( item: InputItem = InputItem.get_instance(
data, path, number) data, self.__input, number)
if item.id in seen: if item.id in seen:
raise InputFormatError( raise InputFormatError(
f"{path}: line {number}: duplicated ID" f"{self.__input}: line {number}:"
f" \"{item.id}\"") f" duplicated ID \"{item.id}\"")
seen.add(item.id) seen.add(item.id)
items.append(item) items.append(item)
if len(items) == 0: if len(items) == 0:
raise InputFormatError(f"{path}: no input items") raise InputFormatError(f"{self.__input}: no input items")
return items return items
def __create_archive_dir(self) -> Path:
"""Create the archive directory.
def build_request(item: InputItem, system_prompt: str, Only this directory is ever created or removed; no other
max_tokens: int, model: str) -> dict[str, Any]: directory is ever touched.
:return: The created archive directory.
:raises FileExistsError: When the archive directory
already exists and ``replace`` is False.
"""
if self.__archive_dir.exists():
if not self.__replace:
raise FileExistsError(
f"{self.__archive_dir} already exists; pass"
" --replace to replace it")
shutil.rmtree(self.__archive_dir)
self.__archive_dir.mkdir(parents=True)
return self.__archive_dir
def __build_meta(self, items: list[InputItem]) -> dict[str, Any]:
"""Build the initial archive metadata.
:param items: The loaded input items.
:return: The metadata, "batch" and "usage" not yet filled
in for an actual run.
"""
return {
"script_version": self.__SCRIPT_VERSION,
"model": self.__model,
"temperature": self.MODELS[self.__model].get(
"temperature"),
"thinking": self.MODELS[self.__model].get("thinking"),
"max_tokens": self.__max_tokens,
"prompt_path": str(self.__prompt),
"prompt_sha256": self.__sha256_of(self.__prompt),
"input_path": str(self.__input),
"input_sha256": self.__sha256_of(self.__input),
"item_count": len(items),
"dry_run": self.__dry_run,
"started_at": self.__now_iso(),
"batch": None,
"usage": {},
}
def __build_request(self, item: InputItem, system_prompt: str) \
-> dict[str, Any]:
"""Build one Message Batches request for an input item. """Build one Message Batches request for an input item.
:param item: The input item. :param item: The input item.
:param system_prompt: The system prompt text. :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: The batch request with "custom_id" and "params".
""" """
return { return {
"custom_id": item.id, "custom_id": item.id,
"params": { "params": {
"model": model, "model": self.__model,
"max_tokens": max_tokens, "max_tokens": self.__max_tokens,
**MODELS[model], **self.MODELS[self.__model],
"system": system_prompt, "system": system_prompt,
"messages": [ "messages": [
{"role": "user", "content": item.content}, {"role": "user", "content": item.content},
@@ -258,8 +392,33 @@ def build_request(item: InputItem, system_prompt: str,
}, },
} }
def __execute_run(
self, client: anthropic.Anthropic,
items: list[InputItem], system_prompt: str) \
-> ExecutionOutcome:
"""Submit the batch of this run and await its results.
def submit_batch(client: anthropic.Anthropic, :param client: The Anthropic client.
:param items: The input items.
:param system_prompt: The system prompt text.
:return: The submitted batch's bookkeeping and its
results.
"""
requests: list[dict[str, Any]] = [
self.__build_request(x, system_prompt) for x in items]
info: BatchInfo = BatchInfo(
batch_id=self.__submit_batch(client, requests),
submitted_at=self.__now_iso())
print(f"submitted batch {info.batch_id}", file=sys.stderr)
batches: dict[str, Any] = self.__poll_batches(
client, [info.batch_id])
info.ended_at = batches[info.batch_id].ended_at.isoformat()
results: Results = self.__collect_results(
client, info.batch_id)
return ExecutionOutcome(batch=info, results=results)
@staticmethod
def __submit_batch(client: anthropic.Anthropic,
requests: list[dict[str, Any]]) -> str: requests: list[dict[str, Any]]) -> str:
"""Submit one message batch. """Submit one message batch.
@@ -269,8 +428,8 @@ def submit_batch(client: anthropic.Anthropic,
""" """
return client.messages.batches.create(requests=requests).id return client.messages.batches.create(requests=requests).id
@classmethod
def poll_batches(client: anthropic.Anthropic, def __poll_batches(cls, client: anthropic.Anthropic,
batch_ids: list[str]) -> dict[str, Any]: batch_ids: list[str]) -> dict[str, Any]:
"""Poll the batches until every one of them has ended. """Poll the batches until every one of them has ended.
@@ -278,11 +437,13 @@ def poll_batches(client: anthropic.Anthropic,
:param client: The Anthropic client. :param client: The Anthropic client.
:param batch_ids: The batch IDs to poll. :param batch_ids: The batch IDs to poll.
:return: The final batch object of each batch, keyed by batch ID. :return: The final batch object of each batch, keyed by
batch ID.
""" """
while True: while True:
batches: dict[str, Any] = { batches: dict[str, Any] = {
x: client.messages.batches.retrieve(x) for x in batch_ids} x: client.messages.batches.retrieve(x)
for x in batch_ids}
pending: list[str] = [ pending: list[str] = [
x for x in batch_ids x for x in batch_ids
if batches[x].processing_status != "ended"] if batches[x].processing_status != "ended"]
@@ -291,20 +452,40 @@ def poll_batches(client: anthropic.Anthropic,
print(f"batch {batch_id}: {status}", file=sys.stderr) print(f"batch {batch_id}: {status}", file=sys.stderr)
if len(pending) == 0: if len(pending) == 0:
return batches return batches
time.sleep(POLL_INTERVAL_SECONDS) time.sleep(cls.__POLL_INTERVAL_SECONDS)
@staticmethod
def __collect_results(client: anthropic.Anthropic,
batch_id: str) -> Results:
"""Collect the results of an ended batch.
def usage_to_dict(usage: Any) -> dict[str, Any]: :param client: The Anthropic client.
"""Convert a usage object to a plain dictionary. :param batch_id: The batch ID.
:return: The result records, keyed by custom ID.
:param usage: The usage object of a message.
:return: The usage as a dictionary, without null entries.
""" """
return {k: v for k, v in usage.model_dump().items() results: Results = {}
if v is not None} for entry in client.messages.batches.results(batch_id):
results[entry.custom_id] = BatchResult.get_instance(
entry)
return results
@staticmethod
def __find_failures(item_ids: list[str],
results: Results) -> list[str]:
"""Find the item IDs that failed in a result set.
def sum_usage(results: Results) -> dict[str, int]: An item failed when it is missing from the results or when
its record is a failure.
:param item_ids: The item IDs to check, in order.
:param results: The result records, keyed by item ID.
:return: The failed item IDs, in the given order.
"""
return [x for x in item_ids
if x not in results or results[x].is_failure]
@staticmethod
def __sum_usage(results: Results) -> dict[str, int]:
"""Sum the token usage of every succeeded result. """Sum the token usage of every succeeded result.
:param results: The result records, keyed by item ID. :param results: The result records, keyed by item ID.
@@ -319,97 +500,51 @@ def sum_usage(results: Results) -> dict[str, int]:
totals[key] = totals.get(key, 0) + value totals[key] = totals.get(key, 0) + value
return totals return totals
@staticmethod
def collect_results(client: anthropic.Anthropic, def __write_jsonl(path: Path,
batch_id: str) -> Results: records: list[dict[str, Any]]) -> None:
"""Collect the results of an ended batch.
:param client: The Anthropic client.
:param batch_id: The batch ID.
:return: The result records, keyed by custom ID.
"""
results: Results = {}
for entry in client.messages.batches.results(batch_id):
results[entry.custom_id] = BatchResult.get_instance(entry)
return results
def find_failures(item_ids: list[str],
results: Results) -> list[str]:
"""Find the item IDs that failed in a result set.
An item failed when it is missing from the results or when its
record is a failure.
:param item_ids: The item IDs to check, in order.
:param results: The result records, keyed by item ID.
:return: The failed item IDs, in the given order.
"""
return [x for x in item_ids
if x not in results or results[x].is_failure]
def create_archive_dir(directory: Path, replace: bool) -> Path:
"""Create the archive directory.
Only this directory is ever created or removed; no other
directory is ever touched.
:param directory: The destination archive directory.
:param replace: Whether to remove an already existing archive
directory before creating it.
:return: The created archive directory.
:raises FileExistsError: When the archive directory already
exists and ``replace`` is False.
"""
if directory.exists():
if not replace:
raise FileExistsError(
f"{directory} already exists; pass --replace to"
" replace it")
shutil.rmtree(directory)
directory.mkdir(parents=True)
return directory
def write_jsonl(path: Path, records: list[dict[str, Any]]) -> None:
"""Write records to a file as JSON Lines. """Write records to a file as JSON Lines.
:param path: The path of the file to write. :param path: The path of the file to write.
:param records: The records, one per line. :param records: The records, one per line.
:return: None. :return: None.
:raises OSError: When the file cannot be written.
""" """
with open(path, "w", encoding="utf-8") as file: with open(path, "w", encoding="utf-8") as file:
for record in records: for record in records:
file.write(json.dumps(record, ensure_ascii=False) + "\n") file.write(
json.dumps(record, ensure_ascii=False) + "\n")
@staticmethod
def write_json(path: Path, data: dict[str, Any]) -> None: def __write_json(path: Path, data: dict[str, Any]) -> None:
"""Write data to a file as pretty-printed JSON. """Write data to a file as pretty-printed JSON.
:param path: The path of the file to write. :param path: The path of the file to write.
:param data: The data to write. :param data: The data to write.
:return: None. :return: None.
:raises OSError: When the file cannot be written.
""" """
path.write_text( path.write_text(
json.dumps(data, ensure_ascii=False, indent=2) + "\n", json.dumps(data, ensure_ascii=False, indent=2) + "\n",
encoding="utf-8") encoding="utf-8")
@classmethod
def write_meta(path: Path, meta: dict[str, Any]) -> None: def __write_meta(cls, path: Path, meta: dict[str, Any]) -> None:
"""Write the metadata to the ``meta.json`` file. """Write the metadata to the ``meta.json`` file.
The ``BatchInfo`` value under ``batch`` is written as a plain The ``BatchInfo`` value under "batch" is written as a
JSON object. plain JSON object.
:param path: The path of the ``meta.json`` file. :param path: The path of the ``meta.json`` file.
:param meta: The metadata to write. :param meta: The metadata to write.
:return: None. :return: None.
:raises OSError: When the file cannot be written.
""" """
write_json(path, {**meta, "batch": asdict(meta["batch"])}) cls.__write_json(
path, {**meta, "batch": asdict(meta["batch"])})
@staticmethod
def sha256_of(path: Path) -> str: def __sha256_of(path: Path) -> str:
"""Calculate the SHA-256 digest of a file. """Calculate the SHA-256 digest of a file.
:param path: The path of the file. :param path: The path of the file.
@@ -418,111 +553,77 @@ def sha256_of(path: Path) -> str:
with open(path, "rb") as file: with open(path, "rb") as file:
return hashlib.file_digest(file, "sha256").hexdigest() return hashlib.file_digest(file, "sha256").hexdigest()
@staticmethod
def now_iso() -> str: def __now_iso() -> str:
"""Return the current local time in ISO 8601 format. """Return the current local time in ISO 8601 format.
:return: The current local time with the timezone offset. :return: The current local time with the timezone offset.
""" """
return datetime.now().astimezone().isoformat(timespec="seconds") return datetime.now().astimezone().isoformat(
timespec="seconds")
def execute_run( def parse_args(argv: list[str] | None) -> argparse.Namespace:
client: anthropic.Anthropic, items: list[InputItem], """Parse the command-line arguments.
system_prompt: str, max_tokens: int, model: str,
meta: dict[str, Any],
) -> Results:
"""Submit the batch of this run and await its results.
The batch ID and timestamps are recorded into the metadata as an :param argv: The command-line arguments, or None for
observable side effect. ``sys.argv``.
:return: The parsed arguments.
:param client: The Anthropic client.
: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]] = [ parser: argparse.ArgumentParser = argparse.ArgumentParser(
build_request(x, system_prompt, max_tokens, model) description="Run one LLM definition file against one input"
for x in items] " and archive the result.")
info: BatchInfo = BatchInfo( parser.add_argument(
batch_id=submit_batch(client, requests), "prompt", type=Path,
submitted_at=now_iso()) help="the prompt definition file, used as the system"
meta["batch"] = info " prompt")
print(f"submitted batch {info.batch_id}", file=sys.stderr) parser.add_argument(
batches: dict[str, Any] = poll_batches(client, [info.batch_id]) "input", type=Path,
info.ended_at = batches[info.batch_id].ended_at.isoformat() help="the JSONL input file with \"id\" and \"content\"")
return collect_results(client, info.batch_id) parser.add_argument(
"archive_dir", type=Path,
help="the destination archive directory")
parser.add_argument(
"--model", choices=sorted(LLMRunner.MODELS),
default=LLMRunner.DEFAULT_MODEL,
help=f"the model ID (default {LLMRunner.DEFAULT_MODEL})")
parser.add_argument(
"--max-tokens", type=int, default=2048,
help="the maximum output tokens per request (default"
" 2048)")
parser.add_argument(
"--dry-run", action="store_true",
help="validate and archive without calling the API")
parser.add_argument(
"--replace", action="store_true",
help="replace an already existing archive directory")
return parser.parse_args(argv)
def main(argv: list[str] | None = None) -> int: def main(argv: list[str] | None = None) -> int:
"""Run one LLM definition file against one input and archive it. """Run one LLM definition file against one input and archive it.
:param argv: The command-line arguments, or None for ``sys.argv``. :param argv: The command-line arguments, or None for
``sys.argv``.
:return: The exit status: 0 on success, non-zero on failure. :return: The exit status: 0 on success, non-zero on failure.
""" """
started: float = time.monotonic() started: float = time.monotonic()
args: argparse.Namespace = parse_args(argv) args: argparse.Namespace = parse_args(argv)
try: try:
items: list[InputItem] = load_items(args.input) outcome: RunOutcome = LLMRunner(
prompt_text: str = args.prompt.read_text(encoding="utf-8") args.prompt, args.input, args.archive_dir, args.model,
except (OSError, InputFormatError) as error: args.max_tokens, args.dry_run, args.replace).run()
except (InputFormatError, OSError) as error:
print(f"error: {error}", file=sys.stderr) print(f"error: {error}", file=sys.stderr)
return 1 return 1
try: if not outcome.dry_run and len(outcome.failed) > 0:
archive_dir: Path = create_archive_dir( print(f"error: failed items: {', '.join(outcome.failed)}",
args.archive_dir, args.replace)
except FileExistsError as error:
print(f"error: {error}", file=sys.stderr)
return 1
meta_path: Path = archive_dir / "meta.json"
(archive_dir / "prompt.md").write_bytes(args.prompt.read_bytes())
meta: dict[str, Any] = {
"script_version": SCRIPT_VERSION,
"model": args.model,
"temperature": MODELS[args.model].get("temperature"),
"thinking": MODELS[args.model].get("thinking"),
"max_tokens": args.max_tokens,
"prompt_path": str(args.prompt),
"prompt_sha256": sha256_of(args.prompt),
"input_path": str(args.input),
"input_sha256": sha256_of(args.input),
"item_count": len(items),
"dry_run": args.dry_run,
"started_at": now_iso(),
"batch": None,
"usage": {},
}
if args.dry_run:
write_json(meta_path, meta)
print(json.dumps(
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."
f" {elapsed} elapsed.", file=sys.stderr)
return 0
client: anthropic.Anthropic = anthropic.Anthropic(
api_key=get_settings().ANTHROPIC_API_KEY)
results: Results = execute_run(
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",
[results[x].to_record() for x in item_ids if x in results])
meta["usage"] = sum_usage(results)
write_meta(meta_path, meta)
failed: list[str] = find_failures(item_ids, results)
if len(failed) > 0:
print(f"error: failed items: {', '.join(failed)}",
file=sys.stderr) file=sys.stderr)
return 1 return 1
if outcome.dry_run:
print(json.dumps(
outcome.dry_run_request, ensure_ascii=False, indent=2))
elapsed: str = format_duration(time.monotonic() - started) elapsed: str = format_duration(time.monotonic() - started)
print(f"Done. {len(items)} jobs finished." print(f"Done. {outcome.item_count} jobs finished."
f" {elapsed} elapsed.", file=sys.stderr) f" {elapsed} elapsed.", file=sys.stderr)
return 0 return 0
@@ -6,49 +6,16 @@ r"""The majority tally of the three coding runs.
Settles the coding step: the same coding definition file is run Settles the coding step: the same coding definition file is run
three times independently, and this command counts the votes and three times independently, and this command counts the votes and
writes the final coding table the paper cites, as the CSV file writes the final coding table the paper cites. A (song, keyword)
given as the fourth positional command-line argument. Only the pair is written out when at least two of the three runs assign
keyword key sets of the three runs' archived ``output.jsonl`` it, carrying the pooled, deduplicated lyric quotes of the runs
files take part in the tally; the lyric quotes never do. A that assigned it. ``--corrections`` names a CSV file of
(song, keyword) pair is written out when at least two of the researcher-reviewed repairs to a run's records, applied before
three runs assign it, so three votes never tie, and it carries the tally. ``--valid-keywords`` names a plain text file of the
the lyric quotes of every run that assigned it, pooled, allowed keywords that every record's keywords must appear in.
deduplicated, sorted by Unicode code point, and joined with a The songs are named from the working store, so this command runs
single ``|``: the three runs are peers, so the quote order after ``build-db``. When any input is malformed, the tally fails
follows the text alone. A quote carries the lyric line-break and nothing is written; the error message names what failed.
convention ``" / "`` where the lyric has a newline, applied once
when the run records load, so the corrections, the coding table,
and the database all share the one representation and nothing is
ever converted back. No lyric of the 883-song corpus contains
``" / "`` -- a corpus fact checked exhaustively, not a structural
guarantee -- so the convention is unambiguous here. The three
archives must cover exactly the same set of song IDs, every
record must be a successful result, and every record's "text"
must parse to a JSON object; otherwise the tally fails and
nothing is written.
Two optional inputs guard the tally. ``--corrections`` names a
CSV file of researcher-reviewed repairs, applied to each run's
records before anything else happens: a keyword row renames or
drops one keyword assignment of one song in one run, and an
evidence row rewrites or drops one lyric quote string wherever it
appears in that song's record for that run. Its two text fields
carry the two characters ``\n`` where the text has a newline, so
the file holds one row per line. Every row must match, so a
stale row fails the run. ``--valid-keywords`` names a
plain text file of the allowed keywords, one per line; once the
corrections are in, every keyword left in any record must appear
in it. The order is fixed and matters: the corrections come
first, so a repair may reunite the votes of a misspelled keyword
that the check would otherwise reject. With neither option, no
record is touched and no vocabulary is checked.
The archives identify a song as ``song-<ID>``, where ``<ID>`` is
the song's ID in the SQLite working store. The output table does
not carry that ID: every song is looked up in the working store
and written as its title and its stored artist credit instead, so
this command runs after ``build-db``. The step is fully
deterministic; no LLM call is made.
""" """
import argparse import argparse
import csv import csv
@@ -127,11 +94,6 @@ class Correction:
class CorrectionTable: class CorrectionTable:
"""The researcher-reviewed repairs of the runs' records.""" """The researcher-reviewed repairs of the runs' records."""
MANUAL_CORRECTIONS_CSV: ClassVar[str] \
= "coding-corrections.csv"
"""The correction table CSV file's conventional name under
``data/manual/``."""
path: Path path: Path
"""The correction table CSV file the repairs came from.""" """The correction table CSV file the repairs came from."""
corrections: list[Correction] corrections: list[Correction]
@@ -141,7 +103,7 @@ class CorrectionTable:
class CorrectionsLoader: class CorrectionsLoader:
"""The loader of the researcher-reviewed correction table.""" """The loader of the researcher-reviewed correction table."""
__HEADER: tuple[str, str, str, str, str] = ( __HEADER: ClassVar[tuple[str, str, str, str, str]] = (
"Song ID", "Run", "Type", "To Be Replaced", "Correct Term") "Song ID", "Run", "Type", "To Be Replaced", "Correct Term")
"""The header row the correction table CSV file must carry.""" """The header row the correction table CSV file must carry."""
@@ -164,10 +126,8 @@ class CorrectionsLoader:
of the runs the command was given, and a known type. The of the runs the command was given, and a known type. The
file is read with the CSV reader, so a quoted field may file is read with the CSV reader, so a quoted field may
hold a comma or a double quote. No field holds a line hold a comma or a double quote. No field holds a line
break: the two text fields carry the lyric line-break break (see the line-break convention on
convention ``" / "`` where the text has a newline -- the ``CodingTallier``). Nothing is written.
same representation the loaded run records carry -- and
are matched and applied verbatim. Nothing is written.
:return: The repairs, in file order. :return: The repairs, in file order.
:raises TallyError: When the file cannot be read, the :raises TallyError: When the file cannot be read, the
@@ -340,16 +300,16 @@ class TalliedCodings:
class CodingTallier: class CodingTallier:
"""The tallier of the three coding runs' keyword votes.""" """The tallier of the three coding runs' keyword votes."""
__MAJORITY: int = 2 __MAJORITY: ClassVar[int] = 2
"""The number of runs that must assign a keyword to a song for """The number of runs that must assign a keyword to a song for
that code to be settled.""" that code to be settled."""
__MAX_REPORTED_IDS: int = 10 __MAX_REPORTED_IDS: ClassVar[int] = 10
"""The number of song IDs an error message lists before """The number of song IDs an error message lists before
summarizing the rest as a count.""" summarizing the rest as a count."""
__QUOTE_SEPARATOR: str = "|" __QUOTE_SEPARATOR: ClassVar[str] = "|"
"""The separator between the distinct lyric quotes of one """The separator between the distinct lyric quotes of one
settled code.""" settled code."""
__LINE_BREAK: str = " / " __LINE_BREAK: ClassVar[str] = " / "
"""The lyric line-break convention replacing every LF inside a """The lyric line-break convention replacing every LF inside a
quote. Unambiguous for this corpus only: none of the 883 quote. Unambiguous for this corpus only: none of the 883
songs' lyrics contains the three characters, checked songs' lyrics contains the three characters, checked
@@ -616,11 +576,8 @@ class CodingTallier:
if line.strip() == "": if line.strip() == "":
continue continue
record: Any = cls.__parse_json(line, str(path)) record: Any = cls.__parse_json(line, str(path))
if not isinstance(record, dict) or "id" not in record:
raise ValueError(
f"{path}: record without \"id\": {line}")
item_id: Any = record["id"] item_id: Any = record["id"]
if "error" in record or "text" not in record: if "text" not in record:
raise ValueError( raise ValueError(
f"{path}: id {item_id}: not a successful" f"{path}: id {item_id}: not a successful"
" result") " result")
@@ -647,8 +604,8 @@ class CodingTallier:
:param label: The location of the record, for the error :param label: The location of the record, for the error
message. message.
:return: The lyric quotes of every keyword, in the given :return: The lyric quotes of every keyword, in the given
order, every LF inside a quote turned into the lyric order, every LF inside a quote turned into the
line-break convention ``" / "``. line-break convention.
:raises ValueError: When a keyword's value is not a list :raises ValueError: When a keyword's value is not a list
of strings. of strings.
""" """
@@ -739,7 +696,7 @@ class CodingTallier:
first: set[int] = set(runs[0]) first: set[int] = set(runs[0])
index: int index: int
records: dict[int, dict[str, list[str]]] records: dict[int, dict[str, list[str]]]
for index, records in enumerate(runs): for index, records in enumerate(runs[1:], start=1):
song_ids: set[int] = set(records) song_ids: set[int] = set(records)
if song_ids == first: if song_ids == first:
continue continue
@@ -813,9 +770,6 @@ class CodingTallier:
class CodingTable: class CodingTable:
"""The final coding table the paper cites.""" """The final coding table the paper cites."""
RESULT_CODINGS_CSV: ClassVar[str] = "codings.csv"
"""The coding table CSV file's conventional name under
``results/``."""
__HEADER: ClassVar[tuple[str, str, str, str]] \ __HEADER: ClassVar[tuple[str, str, str, str]] \
= ("Song", "Artist Credit", "Keyword", "Quote") = ("Song", "Artist Credit", "Keyword", "Quote")
"""The header row of the coding table CSV file.""" """The header row of the coding table CSV file."""
@@ -833,11 +787,8 @@ class CodingTable:
endings, carrying the header row endings, carrying the header row
``Song,Artist Credit,Keyword,Quote`` and one row per ``Song,Artist Credit,Keyword,Quote`` and one row per
settled keyword, in the row order. Every field is settled keyword, in the row order. Every field is
written verbatim; a quote carries the lyric line-break written verbatim, so the file holds one row per line.
convention ``" / "`` where the lyric has a line break, so The parent directory is created when it does not exist.
no field holds a line break and the file holds one row
per line. The parent directory is created when it does
not exist.
:param output_csv: The output CSV file. :param output_csv: The output CSV file.
:return: None. :return: None.
@@ -970,8 +921,7 @@ def parse_args(argv: list[str] | None) -> argparse.Namespace:
help="the third coding run's archive directory") help="the third coding run's archive directory")
parser.add_argument( parser.add_argument(
"output_csv", type=Path, "output_csv", type=Path,
help="the output CSV file, by convention" help="the output CSV file")
f" results/{CodingTable.RESULT_CODINGS_CSV}")
parser.add_argument( parser.add_argument(
"--valid-keywords", type=Path, default=None, "--valid-keywords", type=Path, default=None,
help="a plain text file of the allowed keywords, one per" help="a plain text file of the allowed keywords, one per"
@@ -980,33 +930,23 @@ def parse_args(argv: list[str] | None) -> argparse.Namespace:
parser.add_argument( parser.add_argument(
"--corrections", type=Path, default=None, "--corrections", type=Path, default=None,
help="the researcher-reviewed correction table CSV file," help="the researcher-reviewed correction table CSV file,"
" by convention"
f" data/manual/{CorrectionTable.MANUAL_CORRECTIONS_CSV},"
" applied to the runs' records before the tally" " applied to the runs' records before the tally"
" (default: no repair)") " (default: no repair)")
return parser.parse_args(argv) return parser.parse_args(argv)
def main(argv: list[str] | None = None) -> int: def main(argv: list[str] | None = None) -> int:
r"""Settle the coding by a majority of the three coding runs. """Settle the coding by a majority of the three coding runs.
Writes the final coding table as the given CSV file, holding Writes the final coding table CSV file described in the
the header row ``Song,Artist Credit,Keyword,Quote`` and one module docstring. The records are repaired from the
row per keyword at least two of the three runs assign, the ``--corrections`` table and then checked against the
song named by its title and its stored artist credit from the ``--valid-keywords`` list, when either is given. Nothing is
SQLite working store, and the keyword carrying the pooled, written when the three archives do not cover the same songs,
deduplicated, and sorted lyric quotes of the runs that a record is not a successful result, a correction is invalid
assigned it, joined with a single ``|`` and carrying the or matches nothing, a keyword is not in the valid keyword
lyric line-break convention ``" / "`` where the lyric has a list, or a song is not in the working store; the error message
newline, so the table holds one row per line. The records are names what failed.
repaired from the ``--corrections`` table and then checked
against the ``--valid-keywords`` list, when either is given.
Nothing is written when the three archives do not cover the
same songs, a record is not a successful result, a record's
"text" does not parse to a JSON object of quote string lists,
a correction is invalid or matches nothing, a keyword is not
in the valid keyword list, or a song is not in the working
store; the error message names what failed.
:param argv: The command-line arguments, or None for :param argv: The command-line arguments, or None for
``sys.argv``. ``sys.argv``.
+20 -44
View File
@@ -285,9 +285,11 @@ class TestBuildDB(unittest.TestCase):
patchers: list[Any] = [ patchers: list[Any] = [
mock.patch.object(build_db, "ds", self.__ds), mock.patch.object(build_db, "ds", self.__ds),
mock.patch.object( mock.patch.object(
build_db.SongImporter, "YEARS", [2016, 2017]), build_db.SongImporter, "_SongImporter__YEARS",
[2016, 2017]),
mock.patch.object( mock.patch.object(
build_db.SongImporter, "RANKS_PER_YEAR", 2)] build_db.SongImporter,
"_SongImporter__RANKS_PER_YEAR", 2)]
for patcher in patchers: for patcher in patchers:
patcher.start() patcher.start()
self.addCleanup(patcher.stop) self.addCleanup(patcher.stop)
@@ -1014,12 +1016,9 @@ class TestBuildDB(unittest.TestCase):
"""Test that the coding CSV imports one row per song and """Test that the coding CSV imports one row per song and
keyword, storing the quote column verbatim.""" keyword, storing the quote column verbatim."""
self.__write_codings(self.CODINGS_CSV) self.__write_codings(self.CODINGS_CSV)
status: int self.assertEqual(
stderr: str self.__run_build("--codings", str(self.__codings))[0],
status, stderr = self.__run_build( 0)
"--codings", str(self.__codings))
self.assertEqual(status, 0)
self.assertIn("3 codings", stderr)
self.assertEqual( self.assertEqual(
self.__stored_codings(), self.__stored_codings(),
{("Hello", "longing"): {("Hello", "longing"):
@@ -1043,11 +1042,7 @@ class TestBuildDB(unittest.TestCase):
def test_omitted_codings_leaves_table_empty(self) -> None: def test_omitted_codings_leaves_table_empty(self) -> None:
"""Test that an omitted coding option leaves no codings.""" """Test that an omitted coding option leaves no codings."""
self.__write_codings(self.CODINGS_CSV) self.__write_codings(self.CODINGS_CSV)
status: int self.assertEqual(self.__run_build()[0], 0)
stderr: str
status, stderr = self.__run_build()
self.assertEqual(status, 0)
self.assertIn("0 codings", stderr)
self.assertEqual(self.__stored_codings(), {}) self.assertEqual(self.__stored_codings(), {})
def test_codings_unknown_song_fails(self) -> None: def test_codings_unknown_song_fails(self) -> None:
@@ -1119,12 +1114,9 @@ class TestBuildDB(unittest.TestCase):
"Song,Artist Credit,Keyword,Quote\n" "Song,Artist Credit,Keyword,Quote\n"
"Shape of You,Ed Sheeran,attraction,I'm in love with" "Shape of You,Ed Sheeran,attraction,I'm in love with"
" your body\n") " your body\n")
status: int self.assertEqual(
stderr: str self.__run_build("--codings", str(self.__codings))[0],
status, stderr = self.__run_build( 0)
"--codings", str(self.__codings))
self.assertEqual(status, 0)
self.assertIn("1 codings", stderr)
self.assertEqual( self.assertEqual(
self.__stored_codings(), self.__stored_codings(),
{("Shape of You", "attraction"): {("Shape of You", "attraction"):
@@ -1159,12 +1151,8 @@ class TestBuildDB(unittest.TestCase):
"""Test that the group CSV imports one row per group and """Test that the group CSV imports one row per group and
keyword, the votes stored as integers.""" keyword, the votes stored as integers."""
self.__write_groups(self.GROUPS_CSV) self.__write_groups(self.GROUPS_CSV)
status: int self.assertEqual(
stderr: str self.__run_build("--groups", str(self.__groups))[0], 0)
status, stderr = self.__run_build(
"--groups", str(self.__groups))
self.assertEqual(status, 0)
self.assertIn("3 group members", stderr)
self.assertEqual( self.assertEqual(
self.__stored_groups(), self.__stored_groups(),
{("masculine", "dominance-and-power"): 3, {("masculine", "dominance-and-power"): 3,
@@ -1180,12 +1168,8 @@ class TestBuildDB(unittest.TestCase):
self.__write_groups( self.__write_groups(
"Group,Keyword,Votes\n" "Group,Keyword,Votes\n"
"vulnerable,longing-and-loss,3\n") "vulnerable,longing-and-loss,3\n")
status: int self.assertEqual(
stderr: str self.__run_build("--groups", str(self.__groups))[0], 0)
status, stderr = self.__run_build(
"--groups", str(self.__groups))
self.assertEqual(status, 0)
self.assertIn("1 group members", stderr)
self.assertEqual( self.assertEqual(
self.__stored_groups(), self.__stored_groups(),
{("vulnerable", "longing-and-loss"): 3}) {("vulnerable", "longing-and-loss"): 3})
@@ -1407,14 +1391,11 @@ class TestBuildDB(unittest.TestCase):
together, reflected in the final counts message.""" together, reflected in the final counts message."""
self.__write_patterns(self.PATTERNS_CSV) self.__write_patterns(self.PATTERNS_CSV)
self.__write_annotations(self.ANNOTATIONS_CSV) self.__write_annotations(self.ANNOTATIONS_CSV)
status: int self.assertEqual(
stderr: str self.__run_build(
status, stderr = self.__run_build(
"--patterns", str(self.__patterns), "--patterns", str(self.__patterns),
"--annotations", str(self.__annotations)) "--annotations", str(self.__annotations))[0],
self.assertEqual(status, 0) 0)
self.assertIn("2 patterns", stderr)
self.assertIn("2 annotations", stderr)
self.assertEqual( self.assertEqual(
self.__stored_patterns(), self.__stored_patterns(),
{"M1": ("male", "Dominance", {"M1": ("male", "Dominance",
@@ -1432,12 +1413,7 @@ class TestBuildDB(unittest.TestCase):
annotation tables empty.""" annotation tables empty."""
self.__write_patterns(self.PATTERNS_CSV) self.__write_patterns(self.PATTERNS_CSV)
self.__write_annotations(self.ANNOTATIONS_CSV) self.__write_annotations(self.ANNOTATIONS_CSV)
status: int self.assertEqual(self.__run_build()[0], 0)
stderr: str
status, stderr = self.__run_build()
self.assertEqual(status, 0)
self.assertIn("0 patterns", stderr)
self.assertIn("0 annotations", stderr)
self.assertEqual(self.__stored_patterns(), {}) self.assertEqual(self.__stored_patterns(), {})
self.assertEqual(self.__stored_annotations(), {}) self.assertEqual(self.__stored_annotations(), {})
+17 -11
View File
@@ -51,10 +51,12 @@ class TestFetchArtists(unittest.TestCase):
self.addCleanup(self.__ds.engine.dispose) self.addCleanup(self.__ds.engine.dispose)
patchers: list[Any] = [ patchers: list[Any] = [
mock.patch.object(fetch_artists, "ds", self.__ds), mock.patch.object(fetch_artists, "ds", self.__ds),
mock.patch.object(fetch_artists, "SLEEP_SECONDS", mock.patch.object(
0.0), fetch_artists.ArtistFetcher,
mock.patch.object(fetch_artists, "RETRY_SECONDS", "_ArtistFetcher__SLEEP_SECONDS", 0.0),
0.0)] mock.patch.object(
fetch_artists.ArtistFetcher,
"_ArtistFetcher__RETRY_SECONDS", 0.0)]
for patcher in patchers: for patcher in patchers:
patcher.start() patcher.start()
self.addCleanup(patcher.stop) self.addCleanup(patcher.stop)
@@ -315,10 +317,12 @@ class TestFetchArtists(unittest.TestCase):
self.assertEqual(status, 0) self.assertEqual(status, 0)
self.assertEqual(urlopen.call_count, 3) self.assertEqual(urlopen.call_count, 3)
first: Any = urlopen.call_args_list[0][0][0] first: Any = urlopen.call_args_list[0][0][0]
self.assertEqual(first.get_header("User-agent"), self.assertEqual(
fetch_artists.USER_AGENT) first.get_header("User-agent"),
self.assertTrue( fetch_artists.ArtistFetcher._ArtistFetcher__USER_AGENT)
first.full_url.startswith(fetch_artists.SPARQL_URL)) self.assertTrue(first.full_url.startswith(
fetch_artists.ArtistFetcher
._ArtistFetcher__SPARQL_URL))
self.assertEqual( self.assertEqual(
first.get_header("Accept"), first.get_header("Accept"),
"application/sparql-results+json") "application/sparql-results+json")
@@ -358,7 +362,8 @@ class TestFetchArtists(unittest.TestCase):
qid, {}, "a brand the type gate excludes") qid, {}, "a brand the type gate excludes")
urlopen: mock.Mock urlopen: mock.Mock
with mock.patch.object( with mock.patch.object(
fetch_artists, "PINNED_QIDS", fetch_artists.ArtistFetcher,
"_ArtistFetcher__PINNED_QIDS",
{"Brandy Brand": qid}), \ {"Brandy Brand": qid}), \
mock.patch( mock.patch(
"urllib.request.urlopen", "urllib.request.urlopen",
@@ -369,7 +374,7 @@ class TestFetchArtists(unittest.TestCase):
self.assertEqual(urlopen.call_count, 1) self.assertEqual(urlopen.call_count, 1)
first: Any = urlopen.call_args_list[0][0][0] first: Any = urlopen.call_args_list[0][0][0]
self.assertTrue(first.full_url.startswith( self.assertTrue(first.full_url.startswith(
fetch_artists.API_URL)) fetch_artists.ArtistFetcher._ArtistFetcher__API_URL))
rows: list[list[str]] = self.__read_rows( rows: list[list[str]] = self.__read_rows(
self.__snapshot) self.__snapshot)
self.assertEqual(rows[1], [ self.assertEqual(rows[1], [
@@ -886,7 +891,8 @@ class TestFetchArtists(unittest.TestCase):
"urllib.request.urlopen", "urllib.request.urlopen",
side_effect=[self.__response(empty)]) as urlopen, side_effect=[self.__response(empty)]) as urlopen,
mock.patch.object( mock.patch.object(
fetch_artists, "read_artist_titles", fetch_artists.ArtistSnapshotUpdater,
"_ArtistSnapshotUpdater__read_artist_titles",
side_effect=[[], OSError("boom")])): side_effect=[[], OSError("boom")])):
status: int = self.__run_fetch()[0] status: int = self.__run_fetch()[0]
self.assertNotEqual(status, 0) self.assertNotEqual(status, 0)
+17 -14
View File
@@ -50,7 +50,9 @@ class TestFetchLyrics(unittest.TestCase):
self.addCleanup(self.__ds.engine.dispose) self.addCleanup(self.__ds.engine.dispose)
patchers: list[Any] = [ patchers: list[Any] = [
mock.patch.object(fetch_lyrics, "ds", self.__ds), mock.patch.object(fetch_lyrics, "ds", self.__ds),
mock.patch.object(fetch_lyrics, "SLEEP_SECONDS", 0.0)] mock.patch.object(
fetch_lyrics.LyricsFetcher,
"_LyricsFetcher__SLEEP_SECONDS", 0.0)]
for patcher in patchers: for patcher in patchers:
patcher.start() patcher.start()
self.addCleanup(patcher.stop) self.addCleanup(patcher.stop)
@@ -158,8 +160,9 @@ class TestFetchLyrics(unittest.TestCase):
self.assertEqual(status, 0) self.assertEqual(status, 0)
self.assertEqual(urlopen.call_count, 1) self.assertEqual(urlopen.call_count, 1)
request: Any = urlopen.call_args[0][0] request: Any = urlopen.call_args[0][0]
self.assertEqual(request.get_header("User-agent"), self.assertEqual(
fetch_lyrics.USER_AGENT) request.get_header("User-agent"),
fetch_lyrics.LyricsFetcher._LyricsFetcher__USER_AGENT)
self.assertEqual( self.assertEqual(
(self.__lyrics / "1.txt") (self.__lyrics / "1.txt")
.read_text(encoding="utf-8"), .read_text(encoding="utf-8"),
@@ -321,39 +324,39 @@ class TestFetchLyrics(unittest.TestCase):
def test_normalize_cp1252_mojibake(self) -> None: def test_normalize_cp1252_mojibake(self) -> None:
"""Test that cp1252 mojibake codepoints are restored.""" """Test that cp1252 mojibake codepoints are restored."""
self.assertEqual( self.assertEqual(
fetch_lyrics.normalize_lyrics("wait…"), fetch_lyrics.LyricsFetchRunner.normalize_lyrics("wait…"),
"wait…") "wait…")
self.assertEqual( self.assertEqual(
fetch_lyrics.normalize_lyrics( fetch_lyrics.LyricsFetchRunner.normalize_lyrics(
"‘quote’"), "‘quote’"),
"quote") "quote")
self.assertEqual( self.assertEqual(
fetch_lyrics.normalize_lyrics( fetch_lyrics.LyricsFetchRunner.normalize_lyrics(
"“quote”"), "“quote”"),
"“quote”") "“quote”")
self.assertEqual( self.assertEqual(
fetch_lyrics.normalize_lyrics("dash—line"), fetch_lyrics.LyricsFetchRunner.normalize_lyrics("dash—line"),
"dash—line") "dash—line")
def test_normalize_undefined_cp1252_removed(self) -> None: def test_normalize_undefined_cp1252_removed(self) -> None:
"""Test that undefined cp1252 byte values are removed.""" """Test that undefined cp1252 byte values are removed."""
text: str = ("abcdef") text: str = ("abcdef")
self.assertEqual( self.assertEqual(
fetch_lyrics.normalize_lyrics(text), "abcdef") fetch_lyrics.LyricsFetchRunner.normalize_lyrics(text), "abcdef")
def test_normalize_homoglyphs(self) -> None: def test_normalize_homoglyphs(self) -> None:
"""Test that watermark homoglyphs are restored.""" """Test that watermark homoglyphs are restored."""
self.assertEqual( self.assertEqual(
fetch_lyrics.normalize_lyrics("likе that"), fetch_lyrics.LyricsFetchRunner.normalize_lyrics("likе that"),
"like that") "like that")
self.assertEqual( self.assertEqual(
fetch_lyrics.normalize_lyrics("lό que soy"), fetch_lyrics.LyricsFetchRunner.normalize_lyrics("lό que soy"),
"ló que soy") "ló que soy")
def test_normalize_space_variants(self) -> None: def test_normalize_space_variants(self) -> None:
"""Test that exotic space variants become ASCII space.""" """Test that exotic space variants become ASCII space."""
self.assertEqual( self.assertEqual(
fetch_lyrics.normalize_lyrics( fetch_lyrics.LyricsFetchRunner.normalize_lyrics(
"abcd"), "abcd"),
"a b c d") "a b c d")
@@ -362,19 +365,19 @@ class TestFetchLyrics(unittest.TestCase):
text: str = ( text: str = (
"abcde") "abcde")
self.assertEqual( self.assertEqual(
fetch_lyrics.normalize_lyrics(text), "abcde") fetch_lyrics.LyricsFetchRunner.normalize_lyrics(text), "abcde")
def test_normalize_ascii_unchanged(self) -> None: def test_normalize_ascii_unchanged(self) -> None:
"""Test that plain ASCII text passes through unchanged.""" """Test that plain ASCII text passes through unchanged."""
text: str = "Hello, it's me\n" text: str = "Hello, it's me\n"
self.assertEqual( self.assertEqual(
fetch_lyrics.normalize_lyrics(text), text) fetch_lyrics.LyricsFetchRunner.normalize_lyrics(text), text)
def test_normalize_legitimate_non_ascii_unchanged(self) -> None: def test_normalize_legitimate_non_ascii_unchanged(self) -> None:
"""Test that legitimate non-ASCII content is unchanged.""" """Test that legitimate non-ASCII content is unchanged."""
text: str = "¿cómo estás? 안녕하세요\n" text: str = "¿cómo estás? 안녕하세요\n"
self.assertEqual( self.assertEqual(
fetch_lyrics.normalize_lyrics(text), text) fetch_lyrics.LyricsFetchRunner.normalize_lyrics(text), text)
def test_fetched_lyrics_saved_normalized(self) -> None: def test_fetched_lyrics_saved_normalized(self) -> None:
"""Test that a fetched lyric is normalized before saving.""" """Test that a fetched lyric is normalized before saving."""
+123 -179
View File
@@ -21,7 +21,7 @@ from pop_fem_audit_tools.commands import run_llm
class RunLLMTestCase(unittest.TestCase): class RunLLMTestCase(unittest.TestCase):
"""The common base test case with the shared helpers.""" """The common base test case with the shared helpers."""
def _make_temp_dir(self) -> Path: def make_temp_dir(self) -> Path:
"""Create a temporary directory removed on test cleanup. """Create a temporary directory removed on test cleanup.
:return: The path of the temporary directory. :return: The path of the temporary directory.
@@ -31,7 +31,7 @@ class RunLLMTestCase(unittest.TestCase):
self.addCleanup(tmp.cleanup) self.addCleanup(tmp.cleanup)
return Path(tmp.name) return Path(tmp.name)
def _make_success_entry(self, custom_id: str, def make_success_entry(self, custom_id: str,
text: str) -> mock.Mock: text: str) -> mock.Mock:
"""Create a mock succeeded batch result entry. """Create a mock succeeded batch result entry.
@@ -47,7 +47,7 @@ class RunLLMTestCase(unittest.TestCase):
return entry return entry
@staticmethod @staticmethod
def _make_error_entry(custom_id: str, def make_error_entry(custom_id: str,
error_type: str) -> mock.Mock: error_type: str) -> mock.Mock:
"""Create a mock errored batch result entry. """Create a mock errored batch result entry.
@@ -90,86 +90,113 @@ class RunLLMTestCase(unittest.TestCase):
class TestLoadItems(RunLLMTestCase): class TestLoadItems(RunLLMTestCase):
"""Test cases for the input JSONL validation.""" """Test cases for the input JSONL validation, driven by a dry
run since the validation happens before any API call."""
def setUp(self) -> None: def setUp(self) -> None:
"""Create a temporary directory for the input files.""" """Create the prompt and input paths for a dry run."""
self.__dir: Path = self._make_temp_dir() directory: Path = self.make_temp_dir()
self.__prompt: Path = directory / "task.md"
self.__prompt.write_text("The task.\n", encoding="utf-8")
self.__input: Path = directory / "items.jsonl"
self.__archive_dir: Path = directory / "runs" / "run1"
def __write_input(self, content: str) -> Path: def __run_dry(self, content: str) -> tuple[int, str]:
"""Write an input file with the given content. """Write the input file and dry-run against it.
:param content: The file content. :param content: The input file content.
:return: The path of the input file. :return: The exit status and the standard error text.
""" """
path: Path = self.__dir / "items.jsonl" self.__input.write_text(content, encoding="utf-8")
path.write_text(content, encoding="utf-8") stderr: io.StringIO = io.StringIO()
return path status: int
with redirect_stderr(stderr):
def test_valid_items(self) -> None: status = run_llm.main([
"""Test that valid items are loaded in file order.""" str(self.__prompt), str(self.__input),
path: Path = self.__write_input( str(self.__archive_dir), "--dry-run"])
'{"id": "a", "content": "one"}\n' return status, stderr.getvalue()
'{"id": "b", "content": "two"}\n')
items: list[run_llm.InputItem] = run_llm.load_items(path)
self.assertEqual(items, [
run_llm.InputItem(id="a", content="one"),
run_llm.InputItem(id="b", content="two")])
def test_malformed_json_names_line(self) -> None: def test_malformed_json_names_line(self) -> None:
"""Test that malformed JSON reports the line number.""" """Test that malformed JSON reports the line number."""
path: Path = self.__write_input( status: int
stderr: str
status, stderr = self.__run_dry(
'{"id": "a", "content": "one"}\n' '{"id": "a", "content": "one"}\n'
'not json\n') 'not json\n')
with self.assertRaises(run_llm.InputFormatError) as context: self.assertEqual(status, 1)
run_llm.load_items(path) self.assertIn("line 2", stderr)
self.assertIn("line 2", str(context.exception))
def test_missing_key_names_line(self) -> None: def test_missing_key_names_line(self) -> None:
"""Test that a missing key reports the line number.""" """Test that a missing key reports the line number."""
path: Path = self.__write_input('{"id": "a"}\n') status: int
with self.assertRaises(run_llm.InputFormatError) as context: stderr: str
run_llm.load_items(path) status, stderr = self.__run_dry('{"id": "a"}\n')
self.assertIn("line 1", str(context.exception)) self.assertEqual(status, 1)
self.assertIn("line 1", stderr)
def test_extra_key_rejected(self) -> None: def test_extra_key_rejected(self) -> None:
"""Test that an extra key is rejected.""" """Test that an extra key is rejected."""
path: Path = self.__write_input( status: int
status, _ = self.__run_dry(
'{"id": "a", "content": "one", "extra": 1}\n') '{"id": "a", "content": "one", "extra": 1}\n')
with self.assertRaises(run_llm.InputFormatError): self.assertEqual(status, 1)
run_llm.load_items(path)
def test_non_string_content_rejected(self) -> None: def test_non_string_content_rejected(self) -> None:
"""Test that a non-string content is rejected.""" """Test that a non-string content is rejected."""
path: Path = self.__write_input('{"id": "a", "content": 3}\n') status: int
with self.assertRaises(run_llm.InputFormatError): status, _ = self.__run_dry('{"id": "a", "content": 3}\n')
run_llm.load_items(path) self.assertEqual(status, 1)
def test_duplicated_id_names_line(self) -> None: def test_duplicated_id_names_line(self) -> None:
"""Test that a duplicated ID reports the line number.""" """Test that a duplicated ID reports the line number."""
path: Path = self.__write_input( status: int
stderr: str
status, stderr = self.__run_dry(
'{"id": "a", "content": "one"}\n' '{"id": "a", "content": "one"}\n'
'{"id": "a", "content": "two"}\n') '{"id": "a", "content": "two"}\n')
with self.assertRaises(run_llm.InputFormatError) as context: self.assertEqual(status, 1)
run_llm.load_items(path) self.assertIn("line 2", stderr)
self.assertIn("line 2", str(context.exception)) self.assertIn("a", stderr)
self.assertIn("a", str(context.exception))
def test_empty_file_rejected(self) -> None: def test_empty_file_rejected(self) -> None:
"""Test that an empty input file is rejected.""" """Test that an empty input file is rejected."""
path: Path = self.__write_input("") status: int
with self.assertRaises(run_llm.InputFormatError): status, _ = self.__run_dry("")
run_llm.load_items(path) self.assertEqual(status, 1)
self.assertFalse(self.__archive_dir.exists())
class TestRequestBuilding(RunLLMTestCase): class TestRequestBuilding(RunLLMTestCase):
"""Test cases for the request construction.""" """Test cases for the request preview, driven by a dry run."""
def test_build_request(self) -> None: def setUp(self) -> None:
"""Test the shape of a batch request.""" """Create the prompt and input files for a dry run."""
request: dict[str, Any] = run_llm.build_request( directory: Path = self.make_temp_dir()
run_llm.InputItem(id="song-1", content="the lyrics"), self.__prompt: Path = directory / "task.md"
"the system prompt", 2048, "claude-sonnet-4-6") self.__prompt.write_text(
"the system prompt", encoding="utf-8")
self.__input: Path = directory / "items.jsonl"
self.__input.write_text(
'{"id": "song-1", "content": "the lyrics"}\n',
encoding="utf-8")
self.__archive_dir: Path = directory / "runs" / "run1"
def __preview(self, extra_argv: list[str]) -> dict[str, Any]:
"""Dry-run and parse the previewed request.
:param extra_argv: The extra command-line arguments.
:return: The parsed request.
"""
stdout: io.StringIO = io.StringIO()
with redirect_stdout(stdout):
run_llm.main([
str(self.__prompt), str(self.__input),
str(self.__archive_dir), "--dry-run"] + extra_argv)
return json.loads(stdout.getvalue())
def test_default_model_request(self) -> None:
"""Test the request shape for the default model."""
request: dict[str, Any] = self.__preview([])
self.assertEqual(request["custom_id"], "song-1") self.assertEqual(request["custom_id"], "song-1")
params: dict[str, Any] = request["params"] params: dict[str, Any] = request["params"]
self.assertEqual(params["model"], "claude-sonnet-4-6") self.assertEqual(params["model"], "claude-sonnet-4-6")
@@ -177,125 +204,19 @@ class TestRequestBuilding(RunLLMTestCase):
self.assertEqual(params["thinking"], {"type": "disabled"}) self.assertEqual(params["thinking"], {"type": "disabled"})
self.assertEqual(params["max_tokens"], 2048) self.assertEqual(params["max_tokens"], 2048)
self.assertEqual(params["system"], "the system prompt") self.assertEqual(params["system"], "the system prompt")
self.assertEqual(params["messages"], self.assertEqual(
params["messages"],
[{"role": "user", "content": "the lyrics"}]) [{"role": "user", "content": "the lyrics"}])
def test_build_request_fable_5(self) -> None: def test_fable_5_request(self) -> None:
"""Test the request shape for the claude-fable-5 model.""" """Test the request shape for the claude-fable-5 model."""
request: dict[str, Any] = run_llm.build_request( request: dict[str, Any] = self.__preview(
run_llm.InputItem(id="group-1", content="the groups"), ["--model", "claude-fable-5", "--max-tokens", "8192"])
"the system prompt", 8192, "claude-fable-5")
params: dict[str, Any] = request["params"] params: dict[str, Any] = request["params"]
self.assertEqual(params["model"], "claude-fable-5") self.assertEqual(params["model"], "claude-fable-5")
self.assertNotIn("temperature", params) self.assertNotIn("temperature", params)
self.assertNotIn("thinking", params) self.assertNotIn("thinking", params)
self.assertEqual(params["max_tokens"], 8192) 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."""
def test_collect_success_and_error(self) -> None:
"""Test collecting succeeded and errored results."""
client: mock.Mock = mock.Mock()
client.messages.batches.results.return_value = iter([
self._make_success_entry("a", "output a"),
self._make_error_entry("b", "invalid_request_error")])
results: run_llm.Results = run_llm.collect_results(
client, "batch_x")
self.assertEqual(results["a"].text, "output a")
self.assertEqual(results["a"].stop_reason, "end_turn")
self.assertEqual(results["a"].usage,
{"input_tokens": 10, "output_tokens": 5})
self.assertEqual(results["b"], run_llm.BatchResult(
id="b", error="invalid_request_error"))
client.messages.batches.results.assert_called_once_with(
"batch_x")
def test_find_failures(self) -> None:
"""Test finding failed and missing items."""
results: run_llm.Results = {
"a": run_llm.BatchResult(id="a", text="fine"),
"b": run_llm.BatchResult(id="b", error="errored")}
self.assertEqual(
run_llm.find_failures(["a", "b", "c"], results),
["b", "c"])
def test_sum_usage(self) -> None:
"""Test summing the token usage across results."""
results: run_llm.Results = {
"a": run_llm.BatchResult(
id="a", text="fine",
usage={"input_tokens": 10, "output_tokens": 5}),
"b": run_llm.BatchResult(
id="b", text="fine",
usage={"input_tokens": 3, "output_tokens": 2}),
"c": run_llm.BatchResult(id="c", error="errored")}
self.assertEqual(
run_llm.sum_usage(results),
{"input_tokens": 13, "output_tokens": 7})
class TestArchive(RunLLMTestCase):
"""Test cases for the archive directory handling."""
def setUp(self) -> None:
"""Create a temporary directory as the runs root."""
self.__dir: Path = self._make_temp_dir()
def test_create_archive_dir(self) -> None:
"""Test the archive directory creation."""
target: Path = self.__dir / "01-01-tag" / "run1"
directory: Path = run_llm.create_archive_dir(target, False)
self.assertTrue(directory.is_dir())
self.assertEqual(directory, target)
def test_existing_archive_dir_rejected_without_replace(
self) -> None:
"""Test that an existing archive is rejected by default."""
target: Path = self.__dir / "01-01-tag" / "run1"
run_llm.create_archive_dir(target, False)
with self.assertRaises(FileExistsError):
run_llm.create_archive_dir(target, False)
def test_existing_archive_dir_replaced(self) -> None:
"""Test that --replace replaces an existing archive."""
target: Path = self.__dir / "01-01-tag" / "run1"
first: Path = run_llm.create_archive_dir(target, False)
(first / "stale.txt").write_text("stale", encoding="utf-8")
second: Path = run_llm.create_archive_dir(target, True)
self.assertEqual(first, second)
self.assertFalse((second / "stale.txt").exists())
def test_replace_leaves_sibling_dir_untouched(self) -> None:
"""Test that replacing run2 does not touch run1."""
run1: Path = run_llm.create_archive_dir(
self.__dir / "01-01-tag" / "run1", False)
(run1 / "output.jsonl").write_text(
"run1 data", encoding="utf-8")
run2: Path = run_llm.create_archive_dir(
self.__dir / "01-01-tag" / "run2", False)
(run2 / "stale.jsonl").write_text("stale", encoding="utf-8")
run_llm.create_archive_dir(
self.__dir / "01-01-tag" / "run2", True)
self.assertEqual(
(run1 / "output.jsonl").read_text(encoding="utf-8"),
"run1 data")
def test_write_jsonl(self) -> None:
"""Test writing records as JSON Lines."""
path: Path = self.__dir / "out.jsonl"
run_llm.write_jsonl(path, [{"id": "a", "text": "中文"},
{"id": "b", "text": "two"}])
lines: list[str] = path.read_text(
encoding="utf-8").splitlines()
self.assertEqual(len(lines), 2)
self.assertEqual(json.loads(lines[0]),
{"id": "a", "text": "中文"})
self.assertIn("中文", lines[0])
class TestMainFlow(RunLLMTestCase): class TestMainFlow(RunLLMTestCase):
@@ -303,7 +224,7 @@ class TestMainFlow(RunLLMTestCase):
def setUp(self) -> None: def setUp(self) -> None:
"""Create a temporary directory with the input files.""" """Create a temporary directory with the input files."""
directory: Path = self._make_temp_dir() directory: Path = self.make_temp_dir()
self.__runs: Path = directory / "runs" self.__runs: Path = directory / "runs"
self.__archive_dir: Path = self.__runs / "task_v1" / "run1" self.__archive_dir: Path = self.__runs / "task_v1" / "run1"
self.__prompt: Path = directory / "task_v1.md" self.__prompt: Path = directory / "task_v1.md"
@@ -322,8 +243,7 @@ class TestMainFlow(RunLLMTestCase):
ANTHROPIC_API_KEY="test-key") ANTHROPIC_API_KEY="test-key")
config.set_settings(self.__settings) config.set_settings(self.__settings)
@staticmethod def __make_client(self, entries: list[Any]) -> mock.Mock:
def __make_client(entries: list[Any]) -> mock.Mock:
"""Create a mock Anthropic client serving canned results. """Create a mock Anthropic client serving canned results.
:param entries: The result entries of the single run batch. :param entries: The result entries of the single run batch.
@@ -386,10 +306,12 @@ class TestMainFlow(RunLLMTestCase):
"Done. 2 jobs finished. 02:05 elapsed.")) "Done. 2 jobs finished. 02:05 elapsed."))
def test_run_produces_output_file(self) -> None: def test_run_produces_output_file(self) -> None:
"""Test that a run submits one batch and writes output.""" """Test that a run submits one batch and writes output,
the item order preserved and the token usage summed, the
output text written unescaped."""
client: mock.Mock = self.__make_client( client: mock.Mock = self.__make_client(
[self._make_success_entry("a", "answer a"), [self.make_success_entry("a", "answer 中文 a"),
self._make_success_entry("b", "answer b")]) self.make_success_entry("b", "answer b")])
status: int status: int
stderr: str stderr: str
with mock.patch( with mock.patch(
@@ -400,10 +322,13 @@ class TestMainFlow(RunLLMTestCase):
client.messages.batches.create.call_count, 1) client.messages.batches.create.call_count, 1)
run_dir: Path = self.__archive_dir run_dir: Path = self.__archive_dir
self.assertTrue((run_dir / "output.jsonl").exists()) self.assertTrue((run_dir / "output.jsonl").exists())
output_text: str = (run_dir / "output.jsonl").read_text(
encoding="utf-8")
self.assertIn("中文", output_text)
output: list[dict[str, Any]] = [ output: list[dict[str, Any]] = [
json.loads(x) for x in (run_dir / "output.jsonl") json.loads(x) for x in output_text.splitlines()]
.read_text(encoding="utf-8").splitlines()] self.assertEqual(output[0]["text"], "answer 中文 a")
self.assertEqual(output[0]["text"], "answer a") self.assertEqual(output[1]["text"], "answer b")
meta: dict[str, Any] = json.loads( meta: dict[str, Any] = json.loads(
(run_dir / "meta.json").read_text(encoding="utf-8")) (run_dir / "meta.json").read_text(encoding="utf-8"))
self.assertNotIn("run", meta) self.assertNotIn("run", meta)
@@ -431,8 +356,8 @@ class TestMainFlow(RunLLMTestCase):
(run_dir / "stale.jsonl").write_text( (run_dir / "stale.jsonl").write_text(
"stale", encoding="utf-8") "stale", encoding="utf-8")
client: mock.Mock = self.__make_client( client: mock.Mock = self.__make_client(
[self._make_success_entry("a", "answer a"), [self.make_success_entry("a", "answer a"),
self._make_success_entry("b", "answer b")]) self.make_success_entry("b", "answer b")])
status: int = self.__run_main( status: int = self.__run_main(
self.__argv + ["--replace"], client)[0] self.__argv + ["--replace"], client)[0]
self.assertEqual(status, 0) self.assertEqual(status, 0)
@@ -450,8 +375,8 @@ class TestMainFlow(RunLLMTestCase):
(run2_dir / "output.jsonl").write_text( (run2_dir / "output.jsonl").write_text(
"stale run2 data", encoding="utf-8") "stale run2 data", encoding="utf-8")
client: mock.Mock = self.__make_client( client: mock.Mock = self.__make_client(
[self._make_success_entry("a", "answer a"), [self.make_success_entry("a", "answer a"),
self._make_success_entry("b", "answer b")]) self.make_success_entry("b", "answer b")])
argv: list[str] = [ argv: list[str] = [
str(self.__prompt), str(self.__input), str(self.__prompt), str(self.__input),
str(run2_dir), "--replace"] str(run2_dir), "--replace"]
@@ -465,11 +390,14 @@ class TestMainFlow(RunLLMTestCase):
"stale run2 data") "stale run2 data")
def test_run_failure_exits_non_zero(self) -> None: def test_run_failure_exits_non_zero(self) -> None:
"""Test that a failed item aborts with a non-zero status.""" """Test that a failed item aborts with a non-zero status,
the summed usage counting only the succeeded item."""
client: mock.Mock = self.__make_client( client: mock.Mock = self.__make_client(
[self._make_success_entry("a", "answer a"), [self.make_success_entry("a", "answer a"),
self._make_error_entry("b", "invalid_request_error")]) self.make_error_entry("b", "invalid_request_error")])
status: int = self.__run_main(self.__argv, client)[0] status: int
stderr: str
status, _, stderr = self.__run_main(self.__argv, client)
self.assertEqual(status, 1) self.assertEqual(status, 1)
run_dir: Path = self.__archive_dir run_dir: Path = self.__archive_dir
self.assertTrue((run_dir / "output.jsonl").exists()) self.assertTrue((run_dir / "output.jsonl").exists())
@@ -479,6 +407,22 @@ class TestMainFlow(RunLLMTestCase):
self.assertEqual(json.loads(output_lines[1]), self.assertEqual(json.loads(output_lines[1]),
{"id": "b", {"id": "b",
"error": "invalid_request_error"}) "error": "invalid_request_error"})
meta: dict[str, Any] = json.loads(
(run_dir / "meta.json").read_text(encoding="utf-8"))
self.assertEqual(meta["usage"],
{"input_tokens": 10, "output_tokens": 5})
self.assertNotIn("Done.", stderr)
def test_missing_result_item_is_a_failure(self) -> None:
"""Test that an item missing from the batch results is
reported as a failed item."""
client: mock.Mock = self.__make_client(
[self.make_success_entry("a", "answer a")])
status: int
stderr: str
status, _, stderr = self.__run_main(self.__argv, client)
self.assertEqual(status, 1)
self.assertIn("b", stderr)
def test_invalid_input_exits_non_zero(self) -> None: def test_invalid_input_exits_non_zero(self) -> None:
"""Test that an invalid input file aborts before archiving.""" """Test that an invalid input file aborts before archiving."""