imacat/pop-fem-audit
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases 1 Wiki Activity

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

Browse Source 
Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
imacat
2026-08-04 15:13:31 +08:00
co-authored by Claude Fable 5
parent 8d223177cd
commit e0eb343ce0
11 changed files with 540 additions and 333 deletions
Download Patch File Download Diff File Expand all files Collapse all files
tools/src/pop_fem_audit_tools/__main__.py
+1 -2
View File
@@ -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:
tools/src/pop_fem_audit_tools/build_db.py
+55 -30
View File
@@ -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."""
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,
credits, and songs with lyrics, under those keys.
:return: The row counts of the working store.
"""
return {
"songs": session.scalar(
sa.select(sa.func.count()).select_from(Song)) or 0,
"chart entries": session.scalar(
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)) or 0,
"artists": session.scalar(
sa.select(sa.func.count()).select_from(Artist)) or 0,
"credits": session.scalar(
.select_from(ChartEntry)),
artists=count(
sa.select(sa.func.count()).select_from(Artist)),
credits=count(
sa.select(sa.func.count())
.select_from(SongArtist)) or 0,
"songs with lyrics": session.scalar(
.select_from(SongArtist)),
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}"
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
tools/src/pop_fem_audit_tools/fetch_artists.py
+123 -59
View File
@@ -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
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: <reason>``.
error yields a snapshot with what was resolved so far and
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}
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}",
tools/src/pop_fem_audit_tools/fetch_lyrics.py
+42 -18
View File
@@ -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
tools/src/pop_fem_audit_tools/models.py
+14 -3
View File
@@ -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
+222 -151
View File
@@ -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:
"""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)
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)
tools/tests/test_build_db.py
+24 -23
View File
@@ -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
tools/tests/test_fetch_artists.py
+1 -1
View File
@@ -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
+2 -1
View File
@@ -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
+5 -4
View File
@@ -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(
tools/tests/test_run_llm.py
+43 -33
View File
@@ -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."""