Refactor internal data passing to dataclasses and enums, and remove unneeded defensive code

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

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(
+    spec_backup: ModuleSpec | None = main_module.__spec__
-        main_module, "__spec__", None)
     sys.argv[0] = f"{prog()} {args[0]}"
     main_module.__spec__ = None
     try:

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]] = []
+    pairs: list[tuple[str, Role]] = []
-    role: str
+    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]:
+@dataclass
-    """Count the loaded rows for the build summary.
+class StoreCounts:
+    """The row counts of the working store, for the build summary."""
+
+    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 counts of the songs, chart entries, artists,
+        :return: The row counts of the working store.
-        credits, and songs with lyrics, under those keys.
         """
-    return {
+        def count(selectable: sa.Select[tuple[int]]) -> int:
-        "songs": session.scalar(
+            value: int | None = session.scalar(selectable)
-            sa.select(sa.func.count()).select_from(Song)) or 0,
+            assert value is not None
-        "chart entries": session.scalar(
+            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)) or 0,
+                .select_from(ChartEntry)),
-        "artists": session.scalar(
+            artists=count(
-            sa.select(sa.func.count()).select_from(Artist)) or 0,
+                sa.select(sa.func.count()).select_from(Artist)),
-        "credits": session.scalar(
+            credits=count(
                 sa.select(sa.func.count())
-            .select_from(SongArtist)) or 0,
+                .select_from(SongArtist)),
-        "songs with lyrics": session.scalar(
+            songs_with_lyrics=count(
                 sa.select(sa.func.count()).select_from(Song)
-            .where(Song.lyrics.is_not(None))) or 0,
+                .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}"
+    print(f"done: {counts.songs} songs,"
-                               for name, count in counts.items()),
+          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

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: <reason>``."""
+
+    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
+        The note carries the description of the search hit for
-        for human verification.  A search miss yields a row with
+        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
+        error yields a snapshot with what was resolved so far and
-        note ``error: <reason>``.
+        the note ``error: <reason>``.
 
         :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}
+        snapshot: ArtistSnapshot = ArtistSnapshot(name=name)
-        row["name"] = name
         try:
             hit: tuple[str, str] | None = self.__search(name)
             if hit is None:
-                row["note"] = "not found"
+                snapshot.note = NOTE_NOT_FOUND
-                return row
+                return snapshot
-            row["qid"], row["note"] = hit
+            snapshot.qid, snapshot.note = hit
-            self.__resolve(row)
+            self.__resolve(snapshot)
         except (OSError, ValueError) as error:
-            row["note"] = f"error: {error}"
+            snapshot.note = f"error: {error}"
-        return row
+        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:
+    def __resolve(self, snapshot: ArtistSnapshot) -> None:
-        """Resolve the claims of an artist into the row fields.
+        """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]] \
+        claims: ArtistClaims = self.__get_claims(snapshot.qid)
-            = self.__get_claims(row["qid"])
+        country_ids: list[str] = claims.country_ids
-        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", [])
         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)
+            claims.gender_ids + claims.instance_of_ids
-        if len(gender_ids) > 0:
+            + claims.genre_ids + country_ids)
-            row["gender"] = labels.get(gender_ids[0], "")
+        if len(claims.gender_ids) > 0:
-        row["artist_type"] = self.__artist_type(type_ids, labels)
+            snapshot.gender = labels.get(claims.gender_ids[0], "")
-        row["genre"] = "; ".join(
+        snapshot.type = self.__artist_type(
-            labels[x] for x in genre_ids if x in labels)
+            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
+            genre, and country properties.
-            ID.
         :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 ArtistClaims()
-        return {x: self.__targets(claims.get(x))
+        return ArtistClaims(
-                for x in ("P21", "P31", "P136", "P27", "P495")}
+            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
+        :return: ``ArtistType.SOLO`` for a human,
-            ensemble, or empty for the human to decide.
+            ``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
+        return {x["name"] for x in reader}
-                if x.get("name") is not None}
 
 
-def append_row(row: dict[str, str]) -> None:
+def append_row(snapshot: ArtistSnapshot) -> None:
-    """Append a row to the snapshot CSV file.
+    """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)
+            snapshot: ArtistSnapshot = fetcher.fetch(name)
-            append_row(row)
+            append_row(snapshot)
-            status: str = row["qid"]
+            status: str = snapshot.qid
-            if row["note"] == "not found":
+            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}",

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,
+    :param misses: The report entries of the songs still without
-        each with the song ID, the title, the artist credit, and
+        lyrics.
-        the reason.
     :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)
+            artist: str = query_artist(session, song.id)
-            result: tuple[str, str] | None = None
+            result: tuple[str, str] | None = fetcher.fetch(
-            reason: str = "no primary artist"
+                artist, song.title)
-            if artist is not None:
-                result = fetcher.fetch(artist, song.title)
-                reason = "not found"
             if result is None:
-                misses.append((song.id, song.title,
+                misses.append(MissingLyrics(
-                               song.artist_credit, reason))
+                    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

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."""

tools/src/pop_fem_audit_tools/run_llm.py

@@ -12,27 +12,20 @@ arbitration batch, and archives every artifact self-contained under
 ``runs/<phase>/<YYYYMMDD-HHMM>-<prompt-stem>/``.
 """
 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 = (
     "<run2>\n{run2}\n</run2>")
 
 
+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:
+def load_items(path: Path) -> list[InputItem]:
-    """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]:
     """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)
+            item: InputItem = InputItem.get_instance(
-            if item["id"] in seen:
+                data, path, number)
+            if item.id in seen:
                 raise InputFormatError(
                     f"{path}: line {number}: duplicated ID"
-                    f" \"{item['id']}\"")
+                    f" \"{item.id}\"")
-            seen.add(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)
 
 
 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
+        results[entry.custom_id] = BatchResult.get_instance(entry)
-        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
     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 | None = run1[item.id].text
-        text1: str = run1[item_id]["text"].strip()
+        text2: str | None = run2[item.id].text
-        text2: str = run2[item_id]["text"].strip()
+        assert text1 is not None and text2 is not None
-        if text1 == text2:
+        if text1.strip() == text2.strip():
-            agreed.append(item_id)
+            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"]
+        text: str | None
-        if item_id in arbitration:
+        if item.id in arbitration:
-            records.append({"id": item_id,
+            text = arbitration[item.id].text
-                            "text": arbitration[item_id]["text"],
+            assert text is not None
-                            "source": "arbitration"})
+            records.append({"id": item.id, "text": text,
+                            "source": Source.ARBITRATION})
         else:
-            records.append({"id": item_id,
+            text = run1[item.id].text
-                            "text": run1[item_id]["text"].strip(),
+            assert text is not None
-                            "source": "agreed"})
+            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
+    stem: str = prompt_path.stem
-    if stem.endswith(".md"):
-        stem = stem[:-len(".md")]
     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)
+        info: BatchInfo = BatchInfo(
-        batch_ids[run_name] = batch_id
+            batch_id=submit_batch(client, requests),
-        meta["batches"][run_name] = {
+            submitted_at=now_iso())
-            "batch_id": batch_id, "submitted_at": now_iso(),
+        infos[run_name] = info
-            "ended_at": None}
+        meta["batches"][run_name] = info
-        print(f"{run_name}: submitted batch {batch_id}",
+        print(f"{run_name}: submitted batch {info.batch_id}",
               file=sys.stderr)
     batches: dict[str, Any] = poll_batches(
-        client, list(batch_ids.values()))
+        client, [x.batch_id for x in infos.values()])
-    for run_name, batch_id in batch_ids.items():
+    for info in infos.values():
-        meta["batches"][run_name]["ended_at"] = format_timestamp(
+        info.ended_at = batches[info.batch_id].ended_at.isoformat()
-            getattr(batches[batch_id], "ended_at", None))
+    return (collect_results(client, infos["run1"].batch_id),
-    return (collect_results(client, batch_ids["run1"]),
+            collect_results(client, infos["run2"].batch_id))
-            collect_results(client, batch_ids["run2"]))
 
 
 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}
+        x.id: x.content for x in items}
-    requests: list[dict[str, Any]] = [
+    requests: list[dict[str, Any]] = []
-        build_request(
+    for item_id in disagreed:
-            {"id": x,
+        text1: str | None = run1[item_id].text
-             "content": build_arbitration_content(
+        text2: str | None = run2[item_id].text
-                 content_by_id[x], run1[x]["text"], run2[x]["text"])},
+        assert text1 is not None and text2 is not None
-            system_prompt, max_tokens)
+        requests.append(build_request(
-        for x in disagreed]
+            InputItem(id=item_id,
-    batch_id: str = submit_batch(client, requests)
+                      content=build_arbitration_content(
-    meta["batches"]["arbitration"] = {
+                          content_by_id[item_id], text1, text2)),
-        "batch_id": batch_id, "submitted_at": now_iso(),
+            system_prompt, max_tokens))
-        "ended_at": None}
+    info: BatchInfo = BatchInfo(
-    print(f"arbitration: submitted batch {batch_id}", file=sys.stderr)
+        batch_id=submit_batch(client, requests),
-    batches: dict[str, Any] = poll_batches(client, [batch_id])
+        submitted_at=now_iso())
-    meta["batches"]["arbitration"]["ended_at"] = format_timestamp(
+    meta["batches"]["arbitration"] = info
-        getattr(batches[batch_id], "ended_at", None))
+    print(f"arbitration: submitted batch {info.batch_id}",
-    return collect_results(client, 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)

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"),
+            [("Drake", Role.PRIMARY),
-             ("Wizkid", "featured"),
+             ("Wizkid", Role.FEATURED),
-             ("Kyla", "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"),
+            [("Lady Gaga", Role.PRIMARY),
-             ("Bradley Cooper", "primary"),
+             ("Bradley Cooper", Role.PRIMARY),
-             ("BloodPop", "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"),
+            [("KAROL G", Role.PRIMARY),
-             ("Nicki Minaj", "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"),
+            [("Marshmello", Role.PRIMARY),
-             ("Halsey", "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"),
+            [("Kane Brown", Role.PRIMARY),
-             ("Lauren Alaina", "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"),
+            [("Ariana Grande", Role.PRIMARY),
-             ("Doja Cat", "featured"),
+             ("Doja Cat", Role.FEATURED),
-             ("Megan Thee Stallion", "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"),
+            [("24kGoldn", Role.PRIMARY),
-             ("iann dior", "featured")])
+             ("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),
+                         [("Drake", Role.PRIMARY, 0),
-                          ("Wizkid", "featured", 1)])
+                          ("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

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."""
 

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:

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),
+                         [("Drake", Role.PRIMARY, 0),
-                          ("Wizkid", "featured", 1)])
+                          ("Wizkid", Role.FEATURED, 1)])
         self.assertEqual(song.lyrics,
                          "Baby, I like your style")
         artist: Artist | None = self.__session.scalar(

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)
+        items: list[run_llm.InputItem] = run_llm.load_items(path)
-        self.assertEqual(items, [{"id": "a", "content": "one"},
+        self.assertEqual(items, [
-                                 {"id": "b", "content": "two"}])
+            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] = [
+        items: list[run_llm.InputItem] = [
-            {"id": "a", "content": "one"},
+            run_llm.InputItem(id="a", content="one"),
-            {"id": "b", "content": "two"},
+            run_llm.InputItem(id="b", content="two"),
-            {"id": "c", "content": "three"}]
+            run_llm.InputItem(id="c", content="three")]
         run1: run_llm.Results = {
-            "a": {"id": "a", "text": "same\n"},
+            "a": run_llm.BatchResult(id="a", text="same\n"),
-            "b": {"id": "b", "text": "left"},
+            "b": run_llm.BatchResult(id="b", text="left"),
-            "c": {"id": "c", "text": " padded "}}
+            "c": run_llm.BatchResult(id="c", text=" padded ")}
         run2: run_llm.Results = {
-            "a": {"id": "a", "text": "same"},
+            "a": run_llm.BatchResult(id="a", text="same"),
-            "b": {"id": "b", "text": "right"},
+            "b": run_llm.BatchResult(id="b", text="right"),
-            "c": {"id": "c", "text": "padded"}}
+            "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] = [
+        items: list[run_llm.InputItem] = [
-            {"id": "a", "content": "one"},
+            run_llm.InputItem(id="a", content="one"),
-            {"id": "b", "content": "two"}]
+            run_llm.InputItem(id="b", content="two")]
         run1: run_llm.Results = {
-            "a": {"id": "a", "text": "agreed text\n"},
+            "a": run_llm.BatchResult(id="a", text="agreed text\n"),
-            "b": {"id": "b", "text": "left"}}
+            "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"].text, "output a")
-        self.assertEqual(results["a"]["stop_reason"], "end_turn")
+        self.assertEqual(results["a"].stop_reason, "end_turn")
-        self.assertEqual(results["a"]["usage"],
+        self.assertEqual(results["a"].usage,
                          {"input_tokens": 10, "output_tokens": 5})
-        self.assertEqual(results["b"],
+        self.assertEqual(results["b"], run_llm.BatchResult(
-                         {"id": "b", "error": "invalid_request"})
+            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"},
+            "a": run_llm.BatchResult(id="a", text="fine"),
-            "b": {"id": "b", "error": "errored"}}
+            "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."""