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

Restructure the remaining commands into the house shape

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

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