diff --git a/tools/src/pop_fem_audit_tools/__main__.py b/tools/src/pop_fem_audit_tools/__main__.py index 61242ee..d728b06 100644 --- a/tools/src/pop_fem_audit_tools/__main__.py +++ b/tools/src/pop_fem_audit_tools/__main__.py @@ -81,8 +81,7 @@ def main(argv: list[str] | None = None) -> int: return 2 prog_backup: str = sys.argv[0] main_module: ModuleType = sys.modules["__main__"] - spec_backup: ModuleSpec | None = getattr( - main_module, "__spec__", None) + spec_backup: ModuleSpec | None = main_module.__spec__ sys.argv[0] = f"{prog()} {args[0]}" main_module.__spec__ = None try: diff --git a/tools/src/pop_fem_audit_tools/build_db.py b/tools/src/pop_fem_audit_tools/build_db.py index 1f41dd0..854d8e3 100644 --- a/tools/src/pop_fem_audit_tools/build_db.py +++ b/tools/src/pop_fem_audit_tools/build_db.py @@ -27,8 +27,9 @@ import csv import re import sys from collections.abc import Iterable, Sequence +from dataclasses import dataclass from pathlib import Path -from typing import Any +from typing import Self import sqlalchemy as sa from sqlalchemy.orm import Session @@ -37,6 +38,7 @@ from .database import Base, ds from .models import ( Artist, ChartEntry, + Role, Song, SongArtist, ) @@ -56,7 +58,7 @@ RANKS_PER_YEAR: int = 100 ARTIST_FIELDS: dict[str, str] = { "qid": "wikidata_qid", "gender": "gender", - "artist_type": "artist_type", + "type": "type", "genre": "genre", "country": "country", } @@ -86,7 +88,7 @@ def parse_args(argv: list[str] | None) -> argparse.Namespace: return parser.parse_args(argv) -def parse_artist_credit(credit: str) -> list[tuple[str, str]]: +def parse_artist_credit(credit: str) -> list[tuple[str, Role]]: """Parse a combined artist credit into artists and roles. The credit splits into a primary side and a featured side on @@ -101,13 +103,14 @@ def parse_artist_credit(credit: str) -> list[tuple[str, str]]: :param credit: The combined artist credit string. :return: The (name, role) pairs in credit order, primary side - first, with the role "primary" or "featured". + first, with the role ``Role.PRIMARY`` or + ``Role.FEATURED``. """ sides: list[str] = FEATURING_PATTERN.split(credit, maxsplit=1) - pairs: list[tuple[str, str]] = [] - role: str + pairs: list[tuple[str, Role]] = [] + role: Role side: str - for side, role in zip(sides, ("primary", "featured")): + for side, role in zip(sides, (Role.PRIMARY, Role.FEATURED)): token: str for token in DELIMITER_PATTERN.split(side): name: str = token.strip() @@ -139,7 +142,7 @@ def create_song(session: Session, song_id: int, title: str, seen: set[str] = set() position: int = 0 name: str - role: str + role: Role for name, role in parse_artist_credit(credit): if name in seen: print(f"warning: {credit}: duplicated artist" @@ -220,7 +223,7 @@ def apply_artist_csv(session: Session, path: Path) -> None: :param session: The database session, with the artists flushed. :param path: The CSV file with the columns name, qid, gender, - artist_type, genre, country, and note. + type, genre, country, and note. :return: None. :raises BuildError: When a name matches no artist. :raises OSError: When the file cannot be read. @@ -275,7 +278,7 @@ def find_violations(session: Session, years: Iterable[int], f"unexpected chart entry: year {year} rank {rank}") primary_ids: set[int] = set(session.scalars( sa.select(SongArtist.song_id) - .where(SongArtist.role == "primary"))) + .where(SongArtist.role == Role.PRIMARY))) song: Song for song in session.scalars(sa.select(Song).order_by(Song.id)): if song.id not in primary_ids: @@ -290,29 +293,48 @@ def find_violations(session: Session, years: Iterable[int], return violations -def count_rows(session: Session) -> dict[str, int]: - """Count the loaded rows for the build summary. +@dataclass +class StoreCounts: + """The row counts of the working store, for the build summary.""" - :param session: The database session with the loaded data - flushed. - :return: The counts of the songs, chart entries, artists, - credits, and songs with lyrics, under those keys. - """ - return { - "songs": session.scalar( - sa.select(sa.func.count()).select_from(Song)) or 0, - "chart entries": session.scalar( - sa.select(sa.func.count()) - .select_from(ChartEntry)) or 0, - "artists": session.scalar( - sa.select(sa.func.count()).select_from(Artist)) or 0, - "credits": session.scalar( - sa.select(sa.func.count()) - .select_from(SongArtist)) or 0, - "songs with lyrics": session.scalar( - sa.select(sa.func.count()).select_from(Song) - .where(Song.lyrics.is_not(None))) or 0, - } + songs: int + """The number of the songs.""" + chart_entries: int + """The number of the chart entries.""" + artists: int + """The number of the artists.""" + credits: int + """The number of the song-artist credits.""" + songs_with_lyrics: int + """The number of the songs with lyrics.""" + + @classmethod + def get_instance(cls, session: Session) -> Self: + """Counts the loaded rows and returns the counts. + + :param session: The database session with the loaded data + flushed. + :return: The row counts of the working store. + """ + def count(selectable: sa.Select[tuple[int]]) -> int: + value: int | None = session.scalar(selectable) + assert value is not None + return value + + return cls( + songs=count( + sa.select(sa.func.count()).select_from(Song)), + chart_entries=count( + sa.select(sa.func.count()) + .select_from(ChartEntry)), + artists=count( + sa.select(sa.func.count()).select_from(Artist)), + credits=count( + sa.select(sa.func.count()) + .select_from(SongArtist)), + songs_with_lyrics=count( + sa.select(sa.func.count()).select_from(Song) + .where(Song.lyrics.is_not(None)))) def prepare_engine(engine: sa.Engine) -> None: @@ -355,7 +377,7 @@ def main(argv: list[str] | None = None) -> int: prepare_engine(engine) Base.metadata.create_all(engine) session: Session = ds.get_db() - counts: dict[str, int] + counts: StoreCounts try: reset_store(session) load_chart(session, CHART_CSV) @@ -371,7 +393,7 @@ def main(argv: list[str] | None = None) -> int: for violation in violations: print(f"error: {violation}", file=sys.stderr) return 1 - counts = count_rows(session) + counts = StoreCounts.get_instance(session) session.commit() except (OSError, BuildError) as error: session.rollback() @@ -379,7 +401,10 @@ def main(argv: list[str] | None = None) -> int: return 1 finally: session.close() - print("done: " + ", ".join(f"{count} {name}" - for name, count in counts.items()), + print(f"done: {counts.songs} songs," + f" {counts.chart_entries} chart entries," + f" {counts.artists} artists," + f" {counts.credits} credits," + f" {counts.songs_with_lyrics} songs with lyrics", file=sys.stderr) return 0 diff --git a/tools/src/pop_fem_audit_tools/fetch_artists.py b/tools/src/pop_fem_audit_tools/fetch_artists.py index 61a9e2e..dd072b6 100644 --- a/tools/src/pop_fem_audit_tools/fetch_artists.py +++ b/tools/src/pop_fem_audit_tools/fetch_artists.py @@ -20,14 +20,16 @@ current working directory. """ import argparse import csv +import enum import json import sys import time import urllib.parse import urllib.request from collections.abc import Sequence +from dataclasses import asdict, dataclass, field, fields from pathlib import Path -from typing import Any +from typing import Any, Literal import sqlalchemy as sa from sqlalchemy.orm import Session @@ -37,10 +39,6 @@ from .models import Artist WIKIDATA_CSV: Path = Path("data/artists_wikidata.csv") """The Wikidata artist snapshot CSV file.""" -SNAPSHOT_FIELDS: Sequence[str] = ( - "name", "qid", "gender", "artist_type", "genre", "country", - "note") -"""The header columns of the Wikidata artist snapshot CSV file.""" API_URL: str = "https://www.wikidata.org/w/api.php" """The URL of the Wikidata API endpoint.""" USER_AGENT: str = ("pop-fem-audit-tools" @@ -55,6 +53,71 @@ HUMAN_QID: str = "Q5" 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 Wikidata search hit, +written to the snapshot and read back for the classification.""" + + +class ArtistType(enum.StrEnum): + """The decided artist type of a snapshot row.""" + + SOLO = "solo" + """A solo artist: a human.""" + GROUP = "group" + """A musical ensemble.""" + MIXED = "mixed" + """A mixed act, assigned manually via the overrides; never + derived by the fetcher.""" + + +@dataclass +class ArtistSnapshot: + """One row of the Wikidata artist snapshot CSV file.""" + + name: str + """The artist name.""" + qid: str = "" + """The Wikidata item ID, or empty when unresolved.""" + gender: str = "" + """The gender label, or empty when unresolved.""" + type: str = "" + """The artist type, an ``ArtistType`` value, or empty for the + human to decide.""" + genre: str = "" + """The genre labels, joined with ``; ``.""" + country: str = "" + """The country label, or empty when unresolved.""" + note: str = "" + """The note for human verification: the description of the + search hit, ``not found``, or ``error: ``.""" + + def to_row(self) -> dict[str, str]: + """Return this snapshot as a CSV row. + + :return: The row values, keyed by the column name. + """ + 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 +class ArtistClaims: + """The item-ID claim targets of a Wikidata artist item.""" + + gender_ids: list[str] = field(default_factory=list) + """The item IDs of the gender targets.""" + instance_of_ids: list[str] = field(default_factory=list) + """The item IDs of the instance-of targets.""" + genre_ids: list[str] = field(default_factory=list) + """The item IDs of the genre targets.""" + country_ids: list[str] = field(default_factory=list) + """The item IDs of the country-of-citizenship targets.""" + origin_country_ids: list[str] = field(default_factory=list) + """The item IDs of the country-of-origin targets.""" def parse_args(argv: list[str] | None) -> argparse.Namespace: @@ -78,30 +141,29 @@ class ArtistFetcher: self.__sent: int = 0 """The number of the HTTP requests already sent.""" - def fetch(self, name: str) -> dict[str, str]: + def fetch(self, name: str) -> ArtistSnapshot: """Fetch the metadata of an artist. - The note column carries the description of the search hit - for human verification. A search miss yields a row with + The note carries the description of the search hit for + human verification. A search miss yields a snapshot with the note ``not found``. An HTTP, network, or decoding - error yields a row with what was resolved so far and the - note ``error: ``. + error yields a snapshot with what was resolved so far and + the note ``error: ``. :param name: The artist name to query with. - :return: The snapshot CSV row of the artist. + :return: The snapshot of the artist. """ - row: dict[str, str] = {x: "" for x in SNAPSHOT_FIELDS} - row["name"] = name + snapshot: ArtistSnapshot = ArtistSnapshot(name=name) try: hit: tuple[str, str] | None = self.__search(name) if hit is None: - row["note"] = "not found" - return row - row["qid"], row["note"] = hit - self.__resolve(row) + snapshot.note = NOTE_NOT_FOUND + return snapshot + snapshot.qid, snapshot.note = hit + self.__resolve(snapshot) except (OSError, ValueError) as error: - row["note"] = f"error: {error}" - return row + snapshot.note = f"error: {error}" + return snapshot def __search(self, name: str) -> tuple[str, str] | None: """Search Wikidata for an artist. @@ -127,39 +189,36 @@ class ArtistFetcher: return hit["id"], \ description if isinstance(description, str) else "" - def __resolve(self, row: dict[str, str]) -> None: - """Resolve the claims of an artist into the row fields. + def __resolve(self, snapshot: ArtistSnapshot) -> None: + """Resolve the claims of an artist into the snapshot. - :param row: The snapshot CSV row, with the QID set. + :param snapshot: The snapshot, with the QID set. :return: None. :raises OSError: On an HTTP or network error. :raises ValueError: On a JSON decoding error. """ - claims: dict[str, list[str]] \ - = self.__get_claims(row["qid"]) - gender_ids: list[str] = claims.get("P21", []) - type_ids: list[str] = claims.get("P31", []) - genre_ids: list[str] = claims.get("P136", []) - country_ids: list[str] = claims.get("P27", []) + claims: ArtistClaims = self.__get_claims(snapshot.qid) + country_ids: list[str] = claims.country_ids if len(country_ids) == 0: - country_ids = claims.get("P495", []) + country_ids = claims.origin_country_ids labels: dict[str, str] = self.__get_labels( - gender_ids + type_ids + genre_ids + country_ids) - if len(gender_ids) > 0: - row["gender"] = labels.get(gender_ids[0], "") - row["artist_type"] = self.__artist_type(type_ids, labels) - row["genre"] = "; ".join( - labels[x] for x in genre_ids if x in labels) + claims.gender_ids + claims.instance_of_ids + + claims.genre_ids + country_ids) + if len(claims.gender_ids) > 0: + snapshot.gender = labels.get(claims.gender_ids[0], "") + snapshot.type = self.__artist_type( + claims.instance_of_ids, labels) + snapshot.genre = "; ".join( + labels[x] for x in claims.genre_ids if x in labels) if len(country_ids) > 0: - row["country"] = labels.get(country_ids[0], "") + snapshot.country = labels.get(country_ids[0], "") - def __get_claims(self, qid: str) -> dict[str, list[str]]: + def __get_claims(self, qid: str) -> ArtistClaims: """Fetch the item-ID claim targets of a Wikidata item. :param qid: The item ID. :return: The item-ID targets of the gender, instance-of, - genre, and country properties, keyed by the property - ID. + genre, and country properties. :raises OSError: On an HTTP or network error. :raises ValueError: On a JSON decoding error. """ @@ -172,9 +231,13 @@ class ArtistFetcher: and isinstance(data["entities"].get(qid), dict): claims = data["entities"][qid].get("claims") if not isinstance(claims, dict): - return {} - return {x: self.__targets(claims.get(x)) - for x in ("P21", "P31", "P136", "P27", "P495")} + return ArtistClaims() + return ArtistClaims( + gender_ids=self.__targets(claims.get("P21")), + instance_of_ids=self.__targets(claims.get("P31")), + genre_ids=self.__targets(claims.get("P136")), + country_ids=self.__targets(claims.get("P27")), + origin_country_ids=self.__targets(claims.get("P495"))) @staticmethod def __targets(statements: Any) -> list[str]: @@ -238,21 +301,23 @@ class ArtistFetcher: @staticmethod def __artist_type(type_ids: Sequence[str], - labels: dict[str, str]) -> str: + labels: dict[str, str]) \ + -> ArtistType | Literal[""]: """Derive the artist type from the instance-of targets. :param type_ids: The item IDs of the instance-of targets. :param labels: The English labels, keyed by the item ID. - :return: ``solo`` for a human, ``group`` for a musical - ensemble, or empty for the human to decide. + :return: ``ArtistType.SOLO`` for a human, + ``ArtistType.GROUP`` for a musical ensemble, or the + empty string for the human to decide. """ if HUMAN_QID in type_ids: - return "solo" + return ArtistType.SOLO qid: str for qid in type_ids: label: str = labels.get(qid, "").lower() if any(x in label for x in GROUP_KEYWORDS): - return "group" + return ArtistType.GROUP return "" def __get_json(self, params: dict[str, str]) -> Any: @@ -288,17 +353,16 @@ def read_snapshot_names() -> set[str]: with open(WIKIDATA_CSV, encoding="utf-8", newline="") as file: reader: csv.DictReader[str] = csv.DictReader(file) - return {x["name"] for x in reader - if x.get("name") is not None} + return {x["name"] for x in reader} -def append_row(row: dict[str, str]) -> None: - """Append a row to the snapshot CSV file. +def append_row(snapshot: ArtistSnapshot) -> None: + """Append a snapshot row to the snapshot CSV file. The CSV file is created with the header row when missing; the existing rows are preserved. - :param row: The snapshot CSV row. + :param snapshot: The snapshot of an artist. :return: None. :raises OSError: When the file cannot be written. """ @@ -310,7 +374,7 @@ def append_row(row: dict[str, str]) -> None: file, SNAPSHOT_FIELDS) if is_new: writer.writeheader() - writer.writerow(row) + writer.writerow(snapshot.to_row()) def main(argv: list[str] | None = None) -> int: @@ -336,15 +400,15 @@ def main(argv: list[str] | None = None) -> int: if name in done: skipped += 1 continue - row: dict[str, str] = fetcher.fetch(name) - append_row(row) - status: str = row["qid"] - if row["note"] == "not found": + snapshot: ArtistSnapshot = fetcher.fetch(name) + append_row(snapshot) + status: str = snapshot.qid + if snapshot.note == NOTE_NOT_FOUND: not_found += 1 status = "not found" - elif row["note"].startswith("error: "): + elif snapshot.note.startswith("error: "): errors += 1 - status = row["note"] + status = snapshot.note else: fetched += 1 print(f"artist \"{name}\": {status}", diff --git a/tools/src/pop_fem_audit_tools/fetch_lyrics.py b/tools/src/pop_fem_audit_tools/fetch_lyrics.py index 1614975..0b9e302 100644 --- a/tools/src/pop_fem_audit_tools/fetch_lyrics.py +++ b/tools/src/pop_fem_audit_tools/fetch_lyrics.py @@ -27,6 +27,7 @@ import time import urllib.parse import urllib.request from collections.abc import Sequence +from dataclasses import dataclass from pathlib import Path from typing import Any @@ -36,6 +37,7 @@ from sqlalchemy.orm import Session from .database import ds from .models import ( Artist, + Role, Song, SongArtist, ) @@ -61,6 +63,28 @@ SLEEP_SECONDS: float = 1.0 """The delay between consecutive HTTP requests, in seconds.""" +@dataclass +class MissingLyrics: + """One row of the missing lyrics report CSV file.""" + + song_id: int + """The song ID.""" + title: str + """The song title.""" + artist_credit: str + """The artist credit of the song.""" + reason: str + """The reason the lyrics are missing.""" + + def to_row(self) -> list[str]: + """Return this report entry as a CSV row. + + :return: The row values, in the column order. + """ + return [str(self.song_id), self.title, + self.artist_credit, self.reason] + + def parse_args(argv: list[str] | None) -> argparse.Namespace: """Parse the command-line arguments. @@ -160,21 +184,23 @@ class LyricsFetcher: return None -def query_artist(session: Session, song_id: int) -> str | None: +def query_artist(session: Session, song_id: int) -> str: """Find the artist name to query the APIs with. :param session: The database session. :param song_id: The song ID. :return: The name of the primary-role artist with the lowest - position, or None when the song has no primary artist. + position. """ - return session.scalar( + name: str | None = session.scalar( sa.select(Artist.name) .join(SongArtist, SongArtist.artist_id == Artist.id) .where(SongArtist.song_id == song_id, - SongArtist.role == "primary") + SongArtist.role == Role.PRIMARY) .order_by(SongArtist.position) .limit(1)) + assert name is not None + return name def save_lyrics(song_id: int, lyrics: str) -> None: @@ -213,15 +239,14 @@ def append_provenance(song_id: int, source: str) -> None: datetime.date.today().isoformat(), ""]) -def write_missing(misses: Sequence[Sequence[Any]]) -> None: +def write_missing(misses: Sequence[MissingLyrics]) -> None: """Rewrite the missing lyrics report CSV file. The previous content is replaced, so the file reflects the current misses only. - :param misses: The rows of the songs still without lyrics, - each with the song ID, the title, the artist credit, and - the reason. + :param misses: The report entries of the songs still without + lyrics. :return: None. :raises OSError: When the file cannot be written. """ @@ -230,7 +255,7 @@ def write_missing(misses: Sequence[Sequence[Any]]) -> None: newline="") as file: writer: Any = csv.writer(file) writer.writerow(MISSING_FIELDS) - writer.writerows(misses) + writer.writerows(x.to_row() for x in misses) def main(argv: list[str] | None = None) -> int: @@ -244,7 +269,7 @@ def main(argv: list[str] | None = None) -> int: parse_args(argv) fetcher: LyricsFetcher = LyricsFetcher() fetched: int = 0 - misses: list[tuple[int, str, str, str]] = [] + misses: list[MissingLyrics] = [] session: Session = ds.get_db() try: song: Song @@ -252,15 +277,14 @@ def main(argv: list[str] | None = None) -> int: sa.select(Song).order_by(Song.id)): if (LYRICS_DIR / f"{song.id}.txt").exists(): continue - artist: str | None = query_artist(session, song.id) - result: tuple[str, str] | None = None - reason: str = "no primary artist" - if artist is not None: - result = fetcher.fetch(artist, song.title) - reason = "not found" + artist: str = query_artist(session, song.id) + result: tuple[str, str] | None = fetcher.fetch( + artist, song.title) if result is None: - misses.append((song.id, song.title, - song.artist_credit, reason)) + misses.append(MissingLyrics( + song_id=song.id, title=song.title, + artist_credit=song.artist_credit, + reason="not found")) print(f"song {song.id} \"{song.title}\": miss", file=sys.stderr) continue diff --git a/tools/src/pop_fem_audit_tools/models.py b/tools/src/pop_fem_audit_tools/models.py index 66d98e3..71a03c0 100644 --- a/tools/src/pop_fem_audit_tools/models.py +++ b/tools/src/pop_fem_audit_tools/models.py @@ -9,12 +9,23 @@ lyrics, their yearly chart entries, the individual artists, and the song-artist credits with the role and order. """ +import enum + import sqlalchemy as sa from sqlalchemy.orm import Mapped, mapped_column, relationship from .database import Base +class Role(enum.StrEnum): + """The role of an artist on a song credit.""" + + PRIMARY = "primary" + """The primary artist role.""" + FEATURED = "featured" + """The featured artist role.""" + + class Song(Base): """A song, identified by its title and combined artist credit.""" __tablename__ = "songs" @@ -66,7 +77,7 @@ class Artist(Base): """The Wikidata QID of the artist.""" gender: Mapped[str | None] """The gender of the artist.""" - artist_type: Mapped[str | None] + type: Mapped[str | None] """The artist type: solo, group, or mixed.""" genre: Mapped[str | None] """The music genre of the artist.""" @@ -89,7 +100,7 @@ class SongArtist(Base): sa.ForeignKey(Artist.id), primary_key=True) """The ID of the credited artist.""" role: Mapped[str] = mapped_column() - """The role of the artist: primary or featured.""" + """The role of the artist, a ``Role`` value.""" position: Mapped[int] """The 0-based position of the artist in the credit order.""" song: Mapped[Song] = relationship(back_populates="song_artists") @@ -98,6 +109,6 @@ class SongArtist(Base): = relationship(back_populates="song_artists") """The credited artist.""" __table_args__ = ( - sa.CheckConstraint(role.in_(["primary", "featured"]), + sa.CheckConstraint(role.in_([x.value for x in Role]), name="ck_song_artists_role"),) """The table-level constraints.""" diff --git a/tools/src/pop_fem_audit_tools/run_llm.py b/tools/src/pop_fem_audit_tools/run_llm.py index 87c616a..9e325a8 100644 --- a/tools/src/pop_fem_audit_tools/run_llm.py +++ b/tools/src/pop_fem_audit_tools/run_llm.py @@ -12,27 +12,20 @@ arbitration batch, and archives every artifact self-contained under ``runs//-/``. """ import argparse +import enum import hashlib import json import sys import time +from dataclasses import asdict, dataclass from datetime import datetime from pathlib import Path -from typing import Any +from typing import Any, Self import anthropic from .config import get_settings -type Item = dict[str, str] -"""An input item with "id" and "content".""" - -type Result = dict[str, Any] -"""One batch result record.""" - -type Results = dict[str, Result] -"""The batch result records, keyed by item ID.""" - MODEL: str = "claude-sonnet-4-6" TEMPERATURE: float = 0.0 THINKING: dict[str, str] = {"type": "disabled"} @@ -44,10 +37,134 @@ ARBITRATION_TEMPLATE: str = ( "\n{run2}\n") +class Source(enum.StrEnum): + """The source of a final record.""" + + AGREED = "agreed" + """The record takes the run text the two runs agreed on.""" + ARBITRATION = "arbitration" + """The record takes the arbitration output.""" + + class InputFormatError(Exception): """An error in the JSONL input file.""" +@dataclass(frozen=True) +class InputItem: + """An input item of the LLM step.""" + + id: str + """The item ID.""" + content: str + """The item content.""" + + @classmethod + def get_instance(cls, data: Any, path: Path, + number: int) -> Self: + """Validate one parsed JSONL record as an input item. + + :param data: The parsed JSON value of the line. + :param path: The path of the JSONL input file, for the + messages. + :param number: The line number, for the messages. + :return: The validated input item. + :raises InputFormatError: When the record is malformed. + """ + if not isinstance(data, dict): + raise InputFormatError( + f"{path}: line {number}: not a JSON object") + if set(data.keys()) != {"id", "content"}: + raise InputFormatError( + f"{path}: line {number}: keys must be exactly" + " \"id\" and \"content\"") + if not isinstance(data["id"], str) or data["id"] == "": + raise InputFormatError( + f"{path}: line {number}: \"id\" must be a" + " non-empty string") + if not isinstance(data["content"], str): + raise InputFormatError( + f"{path}: line {number}: \"content\" must be a" + " string") + return cls(id=data["id"], content=data["content"]) + + +@dataclass +class BatchResult: + """One batch result record.""" + + id: str + """The item ID.""" + text: str | None = None + """The output text of a succeeded result.""" + stop_reason: str | None = None + """The stop reason of a succeeded result.""" + usage: dict[str, Any] | None = None + """The token usage of a succeeded result.""" + error: str | None = None + """The error code of a failed result.""" + + @property + def is_failure(self) -> bool: + """Whether this result is a failure. + + :return: True when the result carries an error, or False + when it succeeded. + """ + return self.error is not None + + @classmethod + def get_instance(cls, entry: Any) -> Self: + """Create the result record of a batch result entry. + + A succeeded entry yields the text, the stop reason, and + the usage; any other entry yields the error code. + + :param entry: The batch result entry. + :return: The result record. + """ + result: Any = entry.result + match result.type: + case "succeeded": + message: Any = result.message + text: str = "".join( + x.text for x in message.content + if x.type == "text") + return cls(id=entry.custom_id, text=text, + stop_reason=message.stop_reason, + usage=usage_to_dict(message.usage)) + case "errored": + return cls(id=entry.custom_id, + error=result.error.error.type) + case other: + return cls(id=entry.custom_id, error=str(other)) + + def to_record(self) -> dict[str, Any]: + """Return this result as an archive JSONL record. + + :return: The record, with the None fields omitted. + """ + return {k: v for k, v in asdict(self).items() + if v is not None} + + +type Results = dict[str, BatchResult] +"""The batch result records, keyed by item ID.""" + + +@dataclass +class BatchInfo: + """The bookkeeping of one submitted message batch.""" + + batch_id: str + """The batch ID.""" + submitted_at: str + """The submission time, in ISO 8601 format.""" + ended_at: str | None = None + """The end time, in ISO 8601 format, or None while the batch + is still processing.""" + + def parse_args(argv: list[str] | None) -> argparse.Namespace: """Parse the command-line arguments. @@ -77,43 +194,16 @@ def parse_args(argv: list[str] | None) -> argparse.Namespace: return parser.parse_args(argv) -def validate_item(data: Any, path: Path, number: int) -> Item: - """Validate one parsed JSONL record as an input item. - - :param data: The parsed JSON value of the line. - :param path: The path of the JSONL input file, for the messages. - :param number: The line number, for the messages. - :return: The validated item with "id" and "content". - :raises InputFormatError: When the record is malformed. - """ - if not isinstance(data, dict): - raise InputFormatError( - f"{path}: line {number}: not a JSON object") - if set(data.keys()) != {"id", "content"}: - raise InputFormatError( - f"{path}: line {number}: keys must be exactly" - " \"id\" and \"content\"") - if not isinstance(data["id"], str) or data["id"] == "": - raise InputFormatError( - f"{path}: line {number}: \"id\" must be a" - " non-empty string") - if not isinstance(data["content"], str): - raise InputFormatError( - f"{path}: line {number}: \"content\" must be a" - " string") - return {"id": data["id"], "content": data["content"]} - - -def load_items(path: Path) -> list[Item]: +def load_items(path: Path) -> list[InputItem]: """Load and validate the JSONL input items. :param path: The path of the JSONL input file. - :return: The items, each with "id" and "content", in file order. + :return: The input items, in file order. :raises InputFormatError: When a line is malformed, an ID is duplicated, or the file contains no item. :raises OSError: When the file cannot be read. """ - items: list[Item] = [] + items: list[InputItem] = [] seen: set[str] = set() with open(path, encoding="utf-8") as file: for number, line in enumerate(file, start=1): @@ -124,29 +214,30 @@ def load_items(path: Path) -> list[Item]: except json.JSONDecodeError as error: raise InputFormatError( f"{path}: line {number}: malformed JSON: {error}") - item: Item = validate_item(data, path, number) - if item["id"] in seen: + item: InputItem = InputItem.get_instance( + data, path, number) + if item.id in seen: raise InputFormatError( f"{path}: line {number}: duplicated ID" - f" \"{item['id']}\"") - seen.add(item["id"]) + f" \"{item.id}\"") + seen.add(item.id) items.append(item) if len(items) == 0: raise InputFormatError(f"{path}: no input items") return items -def build_request(item: Item, system_prompt: str, +def build_request(item: InputItem, system_prompt: str, max_tokens: int) -> dict[str, Any]: """Build one Message Batches request for an input item. - :param item: The input item with "id" and "content". + :param item: The input item. :param system_prompt: The system prompt text. :param max_tokens: The maximum output tokens. :return: The batch request with "custom_id" and "params". """ return { - "custom_id": item["id"], + "custom_id": item.id, "params": { "model": MODEL, "max_tokens": max_tokens, @@ -154,7 +245,7 @@ def build_request(item: Item, system_prompt: str, "thinking": THINKING, "system": system_prompt, "messages": [ - {"role": "user", "content": item["content"]}, + {"role": "user", "content": item.content}, ], }, } @@ -214,45 +305,21 @@ def usage_to_dict(usage: Any) -> dict[str, Any]: :param usage: The usage object of a message. :return: The usage as a dictionary, without null entries. """ - if hasattr(usage, "model_dump"): - return {k: v for k, v in usage.model_dump().items() - if v is not None} - return dict(usage) + return {k: v for k, v in usage.model_dump().items() + if v is not None} def collect_results(client: anthropic.Anthropic, batch_id: str) -> Results: """Collect the results of an ended batch. - A succeeded result carries "text", "stop_reason", and "usage"; - any other result carries "error" instead. - :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): - result: Any = entry.result - record: Result - match result.type: - case "succeeded": - message: Any = result.message - text: str = "".join( - x.text for x in message.content - if x.type == "text") - record = {"id": entry.custom_id, "text": text, - "stop_reason": message.stop_reason, - "usage": usage_to_dict(message.usage)} - case "errored": - error_type: Any = getattr( - result.error, "type", "unknown") - record = {"id": entry.custom_id, - "error": str(error_type)} - case other: - record = {"id": entry.custom_id, - "error": str(other)} - results[entry.custom_id] = record + results[entry.custom_id] = BatchResult.get_instance(entry) return results @@ -261,18 +328,18 @@ def find_failures(item_ids: 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 carries an "error" field. + 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 "error" in results[x]] + if x not in results or results[x].is_failure] def split_by_agreement( - items: list[Item], run1: Results, run2: Results, + items: list[InputItem], run1: Results, run2: Results, ) -> tuple[list[str], list[str]]: """Split the item IDs into agreed and disagreeing ones. @@ -287,18 +354,18 @@ def split_by_agreement( agreed: list[str] = [] disagreed: list[str] = [] for item in items: - item_id: str = item["id"] - text1: str = run1[item_id]["text"].strip() - text2: str = run2[item_id]["text"].strip() - if text1 == text2: - agreed.append(item_id) + text1: str | None = run1[item.id].text + text2: str | None = run2[item.id].text + assert text1 is not None and text2 is not None + if text1.strip() == text2.strip(): + agreed.append(item.id) else: - disagreed.append(item_id) + disagreed.append(item.id) return agreed, disagreed def build_final_records( - items: list[Item], run1: Results, arbitration: Results, + items: list[InputItem], run1: Results, arbitration: Results, ) -> list[dict[str, str]]: """Assemble the final records, one per item, in input order. @@ -312,15 +379,17 @@ def build_final_records( """ records: list[dict[str, str]] = [] for item in items: - item_id: str = item["id"] - if item_id in arbitration: - records.append({"id": item_id, - "text": arbitration[item_id]["text"], - "source": "arbitration"}) + text: str | None + if item.id in arbitration: + text = arbitration[item.id].text + assert text is not None + records.append({"id": item.id, "text": text, + "source": Source.ARBITRATION}) else: - records.append({"id": item_id, - "text": run1[item_id]["text"].strip(), - "source": "agreed"}) + text = run1[item.id].text + assert text is not None + records.append({"id": item.id, "text": text.strip(), + "source": Source.AGREED}) return records @@ -335,9 +404,7 @@ def create_archive_dir(runs_root: Path, phase: str, prompt_path: Path, :return: The created archive directory. :raises FileExistsError: When the directory already exists. """ - stem: str = prompt_path.name - if stem.endswith(".md"): - stem = stem[:-len(".md")] + stem: str = prompt_path.stem name: str = f"{now.strftime('%Y%m%d-%H%M')}-{stem}" directory: Path = runs_root / phase / name if directory.exists(): @@ -352,6 +419,7 @@ def write_jsonl(path: Path, records: list[dict[str, Any]]) -> None: :param path: The path of the file to write. :param records: The records, one per line. + :return: None. """ with open(path, "w", encoding="utf-8") as file: for record in records: @@ -363,12 +431,29 @@ def write_json(path: Path, data: dict[str, Any]) -> None: :param path: The path of the file to write. :param data: The data to write. + :return: None. """ path.write_text( json.dumps(data, ensure_ascii=False, indent=2) + "\n", encoding="utf-8") +def write_meta(path: Path, meta: dict[str, Any]) -> None: + """Write the metadata to the ``meta.json`` file. + + The ``BatchInfo`` values under ``batches`` are written as + plain JSON objects. + + :param path: The path of the ``meta.json`` file. + :param meta: The metadata to write. + :return: None. + """ + write_json(path, { + **meta, + "batches": {k: asdict(v) + for k, v in meta["batches"].items()}}) + + def sha256_of(path: Path) -> str: """Calculate the SHA-256 digest of a file. @@ -387,24 +472,8 @@ def now_iso() -> str: return datetime.now().astimezone().isoformat(timespec="seconds") -def format_timestamp(value: Any) -> str: - """Format a timestamp value from a batch object as a string. - - :param value: The timestamp: a datetime, a string, or None. - :return: The timestamp as a string, or the current local time - when the value is missing. - """ - match value: - case datetime(): - return value.isoformat() - case str(): - return value - case _: - return now_iso() - - def execute_runs( - client: anthropic.Anthropic, items: list[Item], + client: anthropic.Anthropic, items: list[InputItem], system_prompt: str, max_tokens: int, meta: dict[str, Any], ) -> tuple[Results, Results]: """Submit the two identical runs and await their results. @@ -421,26 +490,25 @@ def execute_runs( """ requests: list[dict[str, Any]] = [ build_request(x, system_prompt, max_tokens) for x in items] - batch_ids: dict[str, str] = {} + infos: dict[str, BatchInfo] = {} for run_name in ("run1", "run2"): - batch_id: str = submit_batch(client, requests) - batch_ids[run_name] = batch_id - meta["batches"][run_name] = { - "batch_id": batch_id, "submitted_at": now_iso(), - "ended_at": None} - print(f"{run_name}: submitted batch {batch_id}", + info: BatchInfo = BatchInfo( + batch_id=submit_batch(client, requests), + submitted_at=now_iso()) + infos[run_name] = info + meta["batches"][run_name] = info + print(f"{run_name}: submitted batch {info.batch_id}", file=sys.stderr) batches: dict[str, Any] = poll_batches( - client, list(batch_ids.values())) - for run_name, batch_id in batch_ids.items(): - meta["batches"][run_name]["ended_at"] = format_timestamp( - getattr(batches[batch_id], "ended_at", None)) - return (collect_results(client, batch_ids["run1"]), - collect_results(client, batch_ids["run2"])) + client, [x.batch_id for x in infos.values()]) + for info in infos.values(): + info.ended_at = batches[info.batch_id].ended_at.isoformat() + return (collect_results(client, infos["run1"].batch_id), + collect_results(client, infos["run2"].batch_id)) def execute_arbitration( - client: anthropic.Anthropic, items: list[Item], + client: anthropic.Anthropic, items: list[InputItem], disagreed: list[str], run1: Results, run2: Results, system_prompt: str, max_tokens: int, meta: dict[str, Any], ) -> Results: @@ -460,23 +528,26 @@ def execute_arbitration( :return: The arbitration results, keyed by item ID. """ content_by_id: dict[str, str] = { - x["id"]: x["content"] for x in items} - requests: list[dict[str, Any]] = [ - build_request( - {"id": x, - "content": build_arbitration_content( - content_by_id[x], run1[x]["text"], run2[x]["text"])}, - system_prompt, max_tokens) - for x in disagreed] - batch_id: str = submit_batch(client, requests) - meta["batches"]["arbitration"] = { - "batch_id": batch_id, "submitted_at": now_iso(), - "ended_at": None} - print(f"arbitration: submitted batch {batch_id}", file=sys.stderr) - batches: dict[str, Any] = poll_batches(client, [batch_id]) - meta["batches"]["arbitration"]["ended_at"] = format_timestamp( - getattr(batches[batch_id], "ended_at", None)) - return collect_results(client, batch_id) + x.id: x.content for x in items} + requests: list[dict[str, Any]] = [] + for item_id in disagreed: + text1: str | None = run1[item_id].text + text2: str | None = run2[item_id].text + assert text1 is not None and text2 is not None + requests.append(build_request( + InputItem(id=item_id, + content=build_arbitration_content( + content_by_id[item_id], text1, text2)), + system_prompt, max_tokens)) + info: BatchInfo = BatchInfo( + batch_id=submit_batch(client, requests), + submitted_at=now_iso()) + meta["batches"]["arbitration"] = info + print(f"arbitration: submitted batch {info.batch_id}", + file=sys.stderr) + batches: dict[str, Any] = poll_batches(client, [info.batch_id]) + info.ended_at = batches[info.batch_id].ended_at.isoformat() + return collect_results(client, info.batch_id) def main(argv: list[str] | None = None) -> int: @@ -487,7 +558,7 @@ def main(argv: list[str] | None = None) -> int: """ args: argparse.Namespace = parse_args(argv) try: - items: list[Item] = load_items(args.input) + items: list[InputItem] = load_items(args.input) prompt_text: str = args.prompt.read_text(encoding="utf-8") arbitration_text: str = args.arbitration_prompt.read_text( encoding="utf-8") @@ -522,7 +593,7 @@ def main(argv: list[str] | None = None) -> int: "script_version": SCRIPT_VERSION, } if args.dry_run: - write_json(meta_path, meta) + write_meta(meta_path, meta) print(json.dumps( build_request(items[0], prompt_text, args.max_tokens), ensure_ascii=False, indent=2)) @@ -535,15 +606,15 @@ def main(argv: list[str] | None = None) -> int: run2: Results run1, run2 = execute_runs( client, items, prompt_text, args.max_tokens, meta) - item_ids: list[str] = [x["id"] for x in items] + item_ids: list[str] = [x.id for x in items] write_jsonl(run_dir / "run1.jsonl", - [run1[x] for x in item_ids if x in run1]) + [run1[x].to_record() for x in item_ids if x in run1]) write_jsonl(run_dir / "run2.jsonl", - [run2[x] for x in item_ids if x in run2]) + [run2[x].to_record() for x in item_ids if x in run2]) failed: set[str] = (set(find_failures(item_ids, run1)) | set(find_failures(item_ids, run2))) if len(failed) > 0: - write_json(meta_path, meta) + write_meta(meta_path, meta) names: str = ", ".join(x for x in item_ids if x in failed) print(f"error: failed items: {names}", file=sys.stderr) return 1 @@ -558,18 +629,18 @@ def main(argv: list[str] | None = None) -> int: client, items, disagreed, run1, run2, arbitration_text, args.max_tokens, meta) write_jsonl(run_dir / "arbitration.jsonl", - [arbitration[x] for x in disagreed + [arbitration[x].to_record() for x in disagreed if x in arbitration]) arb_failed: list[str] = find_failures(disagreed, arbitration) if len(arb_failed) > 0: - write_json(meta_path, meta) + write_meta(meta_path, meta) names = ", ".join(arb_failed) print(f"error: failed arbitration items: {names}", file=sys.stderr) return 1 write_jsonl(run_dir / "final.jsonl", build_final_records(items, run1, arbitration)) - write_json(meta_path, meta) + write_meta(meta_path, meta) print(f"done: {len(items)} items, {len(agreed)} agreed," f" {len(disagreed)} arbitrated; archived to {run_dir}", file=sys.stderr) diff --git a/tools/tests/test_build_db.py b/tools/tests/test_build_db.py index ff47893..7357715 100644 --- a/tools/tests/test_build_db.py +++ b/tools/tests/test_build_db.py @@ -20,6 +20,7 @@ from pop_fem_audit_tools.database import DataSource from pop_fem_audit_tools.models import ( Artist, ChartEntry, + Role, Song, ) @@ -30,47 +31,47 @@ class TestParseArtistCredit(unittest.TestCase): def test_plain_solo(self) -> None: """Test a plain solo artist credit.""" self.assertEqual(build_db.parse_artist_credit("Adele"), - [("Adele", "primary")]) + [("Adele", Role.PRIMARY)]) def test_featuring_with_and(self) -> None: """Test a featuring credit with an "and" delimiter.""" self.assertEqual( build_db.parse_artist_credit( "Drake featuring Wizkid and Kyla"), - [("Drake", "primary"), - ("Wizkid", "featured"), - ("Kyla", "featured")]) + [("Drake", Role.PRIMARY), + ("Wizkid", Role.FEATURED), + ("Kyla", Role.FEATURED)]) def test_comma_and_ampersand(self) -> None: """Test a credit with comma and ampersand delimiters.""" self.assertEqual( build_db.parse_artist_credit( "Lady Gaga, Bradley Cooper & BloodPop"), - [("Lady Gaga", "primary"), - ("Bradley Cooper", "primary"), - ("BloodPop", "primary")]) + [("Lady Gaga", Role.PRIMARY), + ("Bradley Cooper", Role.PRIMARY), + ("BloodPop", Role.PRIMARY)]) def test_x_delimiter(self) -> None: """Test the "x" delimiter.""" self.assertEqual( build_db.parse_artist_credit("KAROL G x Nicki Minaj"), - [("KAROL G", "primary"), - ("Nicki Minaj", "primary")]) + [("KAROL G", Role.PRIMARY), + ("Nicki Minaj", Role.PRIMARY)]) def test_plus_delimiter(self) -> None: """Test the "+" delimiter.""" self.assertEqual( build_db.parse_artist_credit("Marshmello + Halsey"), - [("Marshmello", "primary"), - ("Halsey", "primary")]) + [("Marshmello", Role.PRIMARY), + ("Halsey", Role.PRIMARY)]) def test_with_delimiter(self) -> None: """Test the "with" delimiter.""" self.assertEqual( build_db.parse_artist_credit( "Kane Brown with Lauren Alaina"), - [("Kane Brown", "primary"), - ("Lauren Alaina", "primary")]) + [("Kane Brown", Role.PRIMARY), + ("Lauren Alaina", Role.PRIMARY)]) def test_feat_abbreviation(self) -> None: """Test that "Feat." splits the featured side.""" @@ -78,17 +79,17 @@ class TestParseArtistCredit(unittest.TestCase): build_db.parse_artist_credit( "Ariana Grande Feat. Doja Cat" " & Megan Thee Stallion"), - [("Ariana Grande", "primary"), - ("Doja Cat", "featured"), - ("Megan Thee Stallion", "featured")]) + [("Ariana Grande", Role.PRIMARY), + ("Doja Cat", Role.FEATURED), + ("Megan Thee Stallion", Role.FEATURED)]) def test_case_insensitive_featuring(self) -> None: """Test that "Featuring" splits case-insensitively.""" self.assertEqual( build_db.parse_artist_credit( "24kGoldn Featuring iann dior"), - [("24kGoldn", "primary"), - ("iann dior", "featured")]) + [("24kGoldn", Role.PRIMARY), + ("iann dior", Role.FEATURED)]) class TestBuildDB(unittest.TestCase): @@ -183,8 +184,8 @@ class TestBuildDB(unittest.TestCase): [(2016, 2), (2017, 1)]) self.assertEqual([(x.artist.name, x.role, x.position) for x in song.song_artists], - [("Drake", "primary", 0), - ("Wizkid", "featured", 1)]) + [("Drake", Role.PRIMARY, 0), + ("Wizkid", Role.FEATURED, 1)]) self.assertIn("3 songs", stderr) self.assertIn("4 chart entries", stderr) self.assertIn("4 artists", stderr) @@ -244,11 +245,11 @@ class TestBuildDB(unittest.TestCase): def test_overrides_apply_over_wikidata(self) -> None: """Test that the overrides win over the Wikidata snapshot.""" Path("data/artists_wikidata.csv").write_text( - "name,qid,gender,artist_type,genre,country,note\n" + "name,qid,gender,type,genre,country,note\n" "Adele,Q2831,female,solo,pop,GB,\n", encoding="utf-8") Path("data/artists_overrides.csv").write_text( - "name,qid,gender,artist_type,genre,country,note\n" + "name,qid,gender,type,genre,country,note\n" "Adele,,,,soul,,manually checked\n", encoding="utf-8") self.assertEqual(self.__run_build()[0], 0) @@ -264,7 +265,7 @@ class TestBuildDB(unittest.TestCase): def test_unknown_override_name_fails(self) -> None: """Test that an unknown override name fails the build.""" Path("data/artists_overrides.csv").write_text( - "name,qid,gender,artist_type,genre,country,note\n" + "name,qid,gender,type,genre,country,note\n" "Adel,,female,,,,typo\n", encoding="utf-8") status: int stderr: str diff --git a/tools/tests/test_fetch_artists.py b/tools/tests/test_fetch_artists.py index 9c449cb..504fc37 100644 --- a/tools/tests/test_fetch_artists.py +++ b/tools/tests/test_fetch_artists.py @@ -26,7 +26,7 @@ class TestFetchArtists(unittest.TestCase): """Test cases for the artist metadata fetcher.""" HEADER: list[str] = [ - "name", "qid", "gender", "artist_type", "genre", + "name", "qid", "gender", "type", "genre", "country", "note"] """The expected header row of the snapshot CSV file.""" diff --git a/tools/tests/test_fetch_lyrics.py b/tools/tests/test_fetch_lyrics.py index 7aa9cea..aac105b 100644 --- a/tools/tests/test_fetch_lyrics.py +++ b/tools/tests/test_fetch_lyrics.py @@ -21,6 +21,7 @@ from pop_fem_audit_tools import config, fetch_lyrics from pop_fem_audit_tools.database import Base, DataSource from pop_fem_audit_tools.models import ( Artist, + Role, Song, SongArtist, ) @@ -81,7 +82,7 @@ class TestFetchLyrics(unittest.TestCase): session.add(song) session.add(SongArtist(song=song, artist=artists[artist], - role="primary", + role=Role.PRIMARY, position=0)) session.commit() finally: diff --git a/tools/tests/test_models.py b/tools/tests/test_models.py index 5d7692c..51a4cd6 100644 --- a/tools/tests/test_models.py +++ b/tools/tests/test_models.py @@ -14,6 +14,7 @@ from pop_fem_audit_tools.database import Base, DataSource from pop_fem_audit_tools.models import ( Artist, ChartEntry, + Role, Song, SongArtist, ) @@ -42,9 +43,9 @@ class TestModels(unittest.TestCase): song.chart_entries = [ChartEntry(year=2016, rank=4)] song.song_artists = [ SongArtist(artist=Artist(name="Drake"), - role="primary", position=0), + role=Role.PRIMARY, position=0), SongArtist(artist=Artist(name="Wizkid"), - role="featured", position=1)] + role=Role.FEATURED, position=1)] song.lyrics = "Baby, I like your style" self.__session.add(song) self.__session.commit() @@ -63,8 +64,8 @@ class TestModels(unittest.TestCase): [(2016, 4)]) self.assertEqual([(x.artist.name, x.role, x.position) for x in song.song_artists], - [("Drake", "primary", 0), - ("Wizkid", "featured", 1)]) + [("Drake", Role.PRIMARY, 0), + ("Wizkid", Role.FEATURED, 1)]) self.assertEqual(song.lyrics, "Baby, I like your style") artist: Artist | None = self.__session.scalar( diff --git a/tools/tests/test_run_llm.py b/tools/tests/test_run_llm.py index f5b0a15..c762773 100644 --- a/tools/tests/test_run_llm.py +++ b/tools/tests/test_run_llm.py @@ -51,8 +51,12 @@ class RunLLMTestCase(unittest.TestCase): error_type: str) -> mock.Mock: """Create a mock errored batch result entry. + The error object is shaped as the SDK envelope: the outer + error carries the constant type ``error``, and the specific + error code lives at ``error.error.type``. + :param custom_id: The custom ID of the entry. - :param error_type: The error type. + :param error_type: The specific error code. :return: The mock result entry. """ entry: mock.Mock = mock.Mock() @@ -60,7 +64,9 @@ class RunLLMTestCase(unittest.TestCase): entry.result = mock.Mock() entry.result.type = "errored" entry.result.error = mock.Mock() - entry.result.error.type = error_type + entry.result.error.type = "error" + entry.result.error.error = mock.Mock() + entry.result.error.error.type = error_type return entry @staticmethod @@ -105,9 +111,10 @@ class TestLoadItems(RunLLMTestCase): path: Path = self.__write_input( '{"id": "a", "content": "one"}\n' '{"id": "b", "content": "two"}\n') - items: list[run_llm.Item] = run_llm.load_items(path) - self.assertEqual(items, [{"id": "a", "content": "one"}, - {"id": "b", "content": "two"}]) + 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: """Test that malformed JSON reports the line number.""" @@ -161,7 +168,7 @@ class TestRequestBuilding(RunLLMTestCase): def test_build_request(self) -> None: """Test the shape of a batch request.""" request: dict[str, Any] = run_llm.build_request( - {"id": "song-1", "content": "the lyrics"}, + run_llm.InputItem(id="song-1", content="the lyrics"), "the system prompt", 2048) self.assertEqual(request["custom_id"], "song-1") params: dict[str, Any] = request["params"] @@ -188,18 +195,18 @@ class TestAgreement(RunLLMTestCase): def test_split_by_agreement(self) -> None: """Test splitting items into agreed and disagreeing ones.""" - items: list[run_llm.Item] = [ - {"id": "a", "content": "one"}, - {"id": "b", "content": "two"}, - {"id": "c", "content": "three"}] + items: list[run_llm.InputItem] = [ + run_llm.InputItem(id="a", content="one"), + run_llm.InputItem(id="b", content="two"), + run_llm.InputItem(id="c", content="three")] run1: run_llm.Results = { - "a": {"id": "a", "text": "same\n"}, - "b": {"id": "b", "text": "left"}, - "c": {"id": "c", "text": " padded "}} + "a": run_llm.BatchResult(id="a", text="same\n"), + "b": run_llm.BatchResult(id="b", text="left"), + "c": run_llm.BatchResult(id="c", text=" padded ")} run2: run_llm.Results = { - "a": {"id": "a", "text": "same"}, - "b": {"id": "b", "text": "right"}, - "c": {"id": "c", "text": "padded"}} + "a": run_llm.BatchResult(id="a", text="same"), + "b": run_llm.BatchResult(id="b", text="right"), + "c": run_llm.BatchResult(id="c", text="padded")} agreed: list[str] disagreed: list[str] agreed, disagreed = run_llm.split_by_agreement( @@ -213,14 +220,15 @@ class TestFinalRecords(RunLLMTestCase): def test_build_final_records(self) -> None: """Test assembling final records from runs and arbitration.""" - items: list[run_llm.Item] = [ - {"id": "a", "content": "one"}, - {"id": "b", "content": "two"}] + items: list[run_llm.InputItem] = [ + run_llm.InputItem(id="a", content="one"), + run_llm.InputItem(id="b", content="two")] run1: run_llm.Results = { - "a": {"id": "a", "text": "agreed text\n"}, - "b": {"id": "b", "text": "left"}} + "a": run_llm.BatchResult(id="a", text="agreed text\n"), + "b": run_llm.BatchResult(id="b", text="left")} arbitration: run_llm.Results = { - "b": {"id": "b", "text": "arbitrated text"}} + "b": run_llm.BatchResult(id="b", + text="arbitrated text")} records: list[dict[str, str]] = run_llm.build_final_records( items, run1, arbitration) self.assertEqual(records, [ @@ -237,23 +245,23 @@ class TestCollectResults(RunLLMTestCase): 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")]) + 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"], + 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"], - {"id": "b", "error": "invalid_request"}) + 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": {"id": "a", "text": "fine"}, - "b": {"id": "b", "error": "errored"}} + "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"]) @@ -345,7 +353,7 @@ class TestMainFlow(RunLLMTestCase): mock.Mock(id=x) for x in batch_ids] ended: mock.Mock = mock.Mock() ended.processing_status = "ended" - ended.ended_at = "2026-07-30T20:00:00+08:00" + ended.ended_at = datetime(2026, 7, 30, 20, 0) client.messages.batches.retrieve.return_value = ended results: dict[str, list[Any]] = { "batch_run1": run1, "batch_run2": run2, @@ -482,7 +490,8 @@ class TestMainFlow(RunLLMTestCase): """Test that a failed item aborts with a non-zero status.""" client: mock.Mock = self.__make_client( run1=[self._make_success_entry("a", "answer a"), - self._make_error_entry("b", "invalid_request")], + self._make_error_entry( + "b", "invalid_request_error")], run2=[self._make_success_entry("a", "answer a"), self._make_success_entry("b", "answer b")]) status: int = self.__run_main(self.__argv, client)[0] @@ -495,7 +504,8 @@ class TestMainFlow(RunLLMTestCase): run1_lines: list[str] = (run_dir / "run1.jsonl") \ .read_text(encoding="utf-8").splitlines() self.assertEqual(json.loads(run1_lines[1]), - {"id": "b", "error": "invalid_request"}) + {"id": "b", + "error": "invalid_request_error"}) def test_invalid_input_exits_non_zero(self) -> None: """Test that an invalid input file aborts before archiving."""