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

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
This commit is contained in:
2026-08-04 15:13:31 +08:00
co-authored by Claude Fable 5
parent 8d223177cd
commit e0eb343ce0
11 changed files with 540 additions and 333 deletions
+43 -33
View File
@@ -51,8 +51,12 @@ class RunLLMTestCase(unittest.TestCase):
error_type: str) -> mock.Mock:
"""Create a mock errored batch result entry.
The error object is shaped as the SDK envelope: the outer
error carries the constant type ``error``, and the specific
error code lives at ``error.error.type``.
:param custom_id: The custom ID of the entry.
:param error_type: The error type.
:param error_type: The specific error code.
:return: The mock result entry.
"""
entry: mock.Mock = mock.Mock()
@@ -60,7 +64,9 @@ class RunLLMTestCase(unittest.TestCase):
entry.result = mock.Mock()
entry.result.type = "errored"
entry.result.error = mock.Mock()
entry.result.error.type = error_type
entry.result.error.type = "error"
entry.result.error.error = mock.Mock()
entry.result.error.error.type = error_type
return entry
@staticmethod
@@ -105,9 +111,10 @@ class TestLoadItems(RunLLMTestCase):
path: Path = self.__write_input(
'{"id": "a", "content": "one"}\n'
'{"id": "b", "content": "two"}\n')
items: list[run_llm.Item] = run_llm.load_items(path)
self.assertEqual(items, [{"id": "a", "content": "one"},
{"id": "b", "content": "two"}])
items: list[run_llm.InputItem] = run_llm.load_items(path)
self.assertEqual(items, [
run_llm.InputItem(id="a", content="one"),
run_llm.InputItem(id="b", content="two")])
def test_malformed_json_names_line(self) -> None:
"""Test that malformed JSON reports the line number."""
@@ -161,7 +168,7 @@ class TestRequestBuilding(RunLLMTestCase):
def test_build_request(self) -> None:
"""Test the shape of a batch request."""
request: dict[str, Any] = run_llm.build_request(
{"id": "song-1", "content": "the lyrics"},
run_llm.InputItem(id="song-1", content="the lyrics"),
"the system prompt", 2048)
self.assertEqual(request["custom_id"], "song-1")
params: dict[str, Any] = request["params"]
@@ -188,18 +195,18 @@ class TestAgreement(RunLLMTestCase):
def test_split_by_agreement(self) -> None:
"""Test splitting items into agreed and disagreeing ones."""
items: list[run_llm.Item] = [
{"id": "a", "content": "one"},
{"id": "b", "content": "two"},
{"id": "c", "content": "three"}]
items: list[run_llm.InputItem] = [
run_llm.InputItem(id="a", content="one"),
run_llm.InputItem(id="b", content="two"),
run_llm.InputItem(id="c", content="three")]
run1: run_llm.Results = {
"a": {"id": "a", "text": "same\n"},
"b": {"id": "b", "text": "left"},
"c": {"id": "c", "text": " padded "}}
"a": run_llm.BatchResult(id="a", text="same\n"),
"b": run_llm.BatchResult(id="b", text="left"),
"c": run_llm.BatchResult(id="c", text=" padded ")}
run2: run_llm.Results = {
"a": {"id": "a", "text": "same"},
"b": {"id": "b", "text": "right"},
"c": {"id": "c", "text": "padded"}}
"a": run_llm.BatchResult(id="a", text="same"),
"b": run_llm.BatchResult(id="b", text="right"),
"c": run_llm.BatchResult(id="c", text="padded")}
agreed: list[str]
disagreed: list[str]
agreed, disagreed = run_llm.split_by_agreement(
@@ -213,14 +220,15 @@ class TestFinalRecords(RunLLMTestCase):
def test_build_final_records(self) -> None:
"""Test assembling final records from runs and arbitration."""
items: list[run_llm.Item] = [
{"id": "a", "content": "one"},
{"id": "b", "content": "two"}]
items: list[run_llm.InputItem] = [
run_llm.InputItem(id="a", content="one"),
run_llm.InputItem(id="b", content="two")]
run1: run_llm.Results = {
"a": {"id": "a", "text": "agreed text\n"},
"b": {"id": "b", "text": "left"}}
"a": run_llm.BatchResult(id="a", text="agreed text\n"),
"b": run_llm.BatchResult(id="b", text="left")}
arbitration: run_llm.Results = {
"b": {"id": "b", "text": "arbitrated text"}}
"b": run_llm.BatchResult(id="b",
text="arbitrated text")}
records: list[dict[str, str]] = run_llm.build_final_records(
items, run1, arbitration)
self.assertEqual(records, [
@@ -237,23 +245,23 @@ class TestCollectResults(RunLLMTestCase):
client: mock.Mock = mock.Mock()
client.messages.batches.results.return_value = iter([
self._make_success_entry("a", "output a"),
self._make_error_entry("b", "invalid_request")])
self._make_error_entry("b", "invalid_request_error")])
results: run_llm.Results = run_llm.collect_results(
client, "batch_x")
self.assertEqual(results["a"]["text"], "output a")
self.assertEqual(results["a"]["stop_reason"], "end_turn")
self.assertEqual(results["a"]["usage"],
self.assertEqual(results["a"].text, "output a")
self.assertEqual(results["a"].stop_reason, "end_turn")
self.assertEqual(results["a"].usage,
{"input_tokens": 10, "output_tokens": 5})
self.assertEqual(results["b"],
{"id": "b", "error": "invalid_request"})
self.assertEqual(results["b"], run_llm.BatchResult(
id="b", error="invalid_request_error"))
client.messages.batches.results.assert_called_once_with(
"batch_x")
def test_find_failures(self) -> None:
"""Test finding failed and missing items."""
results: run_llm.Results = {
"a": {"id": "a", "text": "fine"},
"b": {"id": "b", "error": "errored"}}
"a": run_llm.BatchResult(id="a", text="fine"),
"b": run_llm.BatchResult(id="b", error="errored")}
self.assertEqual(
run_llm.find_failures(["a", "b", "c"], results),
["b", "c"])
@@ -345,7 +353,7 @@ class TestMainFlow(RunLLMTestCase):
mock.Mock(id=x) for x in batch_ids]
ended: mock.Mock = mock.Mock()
ended.processing_status = "ended"
ended.ended_at = "2026-07-30T20:00:00+08:00"
ended.ended_at = datetime(2026, 7, 30, 20, 0)
client.messages.batches.retrieve.return_value = ended
results: dict[str, list[Any]] = {
"batch_run1": run1, "batch_run2": run2,
@@ -482,7 +490,8 @@ class TestMainFlow(RunLLMTestCase):
"""Test that a failed item aborts with a non-zero status."""
client: mock.Mock = self.__make_client(
run1=[self._make_success_entry("a", "answer a"),
self._make_error_entry("b", "invalid_request")],
self._make_error_entry(
"b", "invalid_request_error")],
run2=[self._make_success_entry("a", "answer a"),
self._make_success_entry("b", "answer b")])
status: int = self.__run_main(self.__argv, client)[0]
@@ -495,7 +504,8 @@ class TestMainFlow(RunLLMTestCase):
run1_lines: list[str] = (run_dir / "run1.jsonl") \
.read_text(encoding="utf-8").splitlines()
self.assertEqual(json.loads(run1_lines[1]),
{"id": "b", "error": "invalid_request"})
{"id": "b",
"error": "invalid_request_error"})
def test_invalid_input_exits_non_zero(self) -> None:
"""Test that an invalid input file aborts before archiving."""