Standardisation 8

fetch_batched_results.py

@@ -1,63 +1,83 @@
-import os
-import sys
+from __future__ import annotations
+
 import argparse
-from pathlib import Path
+from collections.abc import Sequence
+
 from google import genai
 
-def main():
-    parser = argparse.ArgumentParser(description="Download and combine completed batch jobs for a directory.")
-    parser.add_argument("root_dir", type=str, help="Directory containing the original batches")
-    args = parser.parse_args()
+import config
+from copienator import (
+    CliError,
+    EvaluationWorkspace,
+    ExitCode,
+    atomic_write_bytes,
+    evaluation_parser,
+    execute,
+    read_json,
+    workspace_from_args,
+)
 
-    target_dir = Path(args.root_dir)
-    dir_name = target_dir.name
-    output_path = target_dir / "batched_correction_result.jsonl"
 
-    if "GEMINI_API_KEY" not in os.environ:
-        sys.exit("Error: GEMINI_API_KEY environment variable not set.")
+def run(workspace: EvaluationWorkspace, *, client=None) -> ExitCode:
+    if client is None:
+        if not config.API_KEY:
+            raise CliError("GEMINI_API_KEY is not configured")
+        client = genai.Client(api_key=config.API_KEY)
+    matching = []
+    if workspace.batch_jobs_file.is_file():
+        manifest = read_json(workspace.batch_jobs_file)
+        jobs = manifest.get("jobs") if isinstance(manifest, dict) else None
+        if not isinstance(jobs, dict):
+            raise CliError(f"Invalid batch manifest: {workspace.batch_jobs_file}")
+        matching = [
+            client.batches.get(name=entry["name"])
+            for entry in jobs.values()
+            if isinstance(entry, dict) and isinstance(entry.get("name"), str)
+        ]
+    else:
+        matching = [
+            job
+            for job in client.batches.list()
+            if workspace.name in str(getattr(job, "display_name", ""))
+        ]
+    if not matching:
+        raise CliError(
+            f"No batch jobs found for evaluation {workspace.name!r}"
+        )
+    for job in matching:
+        state = job.state.name if hasattr(job.state, "name") else job.state
+        print(f"{job.display_name}: {state}")
+        if state != "JOB_STATE_SUCCEEDED":
+            print("Not all matching jobs have succeeded yet.")
+            return ExitCode.PARTIAL
 
-    client = genai.Client()
+    chunks = []
+    incomplete = False
+    for job in matching:
+        destination = getattr(job, "dest", None)
+        file_name = getattr(destination, "file_name", None)
+        if not file_name:
+            print(f"Warning: {job.display_name} has no output file.")
+            incomplete = True
+            continue
+        payload = client.files.download(file=file_name)
+        chunks.append(payload.rstrip(b"\n"))
+    if not chunks:
+        return ExitCode.PARTIAL
+    output_path = workspace.batched_correction_result_file
+    atomic_write_bytes(output_path, b"\n".join(chunks) + b"\n")
+    print(f"Saved combined batch results to {output_path}")
+    return ExitCode.PARTIAL if incomplete else ExitCode.SUCCESS
 
-    print(f"Fetching jobs matching '{dir_name}'...")
-    all_jobs = client.batches.list()
-    matching_jobs = []
 
-    # 1. Find jobs associated with this directory
-    for job in all_jobs:
-        if hasattr(job, 'display_name') and job.display_name and dir_name in job.display_name:
-            matching_jobs.append(job)
+def build_parser() -> argparse.ArgumentParser:
+    return evaluation_parser("Download and combine correction batch results")
 
-    if not matching_jobs:
-        sys.exit(f"No batch jobs found containing '{dir_name}' in their display name.")
 
-    # 2. Check that all matching jobs are complete
-    for job in matching_jobs:
-        state = job.state.name if hasattr(job.state, 'name') else job.state
-        print(f"Found Job: {job.display_name} | State: {state}")
-        if state != 'JOB_STATE_SUCCEEDED':
-            sys.exit(f"Error: Job '{job.display_name}' has not succeeded yet. Try again later.")
+def main(argv: Sequence[str] | None = None) -> int:
+    parser = build_parser()
+    return execute(parser, argv, lambda args: run(workspace_from_args(args)))
 
-    # 3. Download and concatenate
-    print("\nAll jobs succeeded. Downloading results...")
-    combined_data = b""
-
-    for job in matching_jobs:
-        if hasattr(job, 'dest') and job.dest and hasattr(job.dest, 'file_name') and job.dest.file_name:
-            print(f"Downloading output for {job.display_name}...")
-            file_content_bytes = client.files.download(file=job.dest.file_name)
-
-            combined_data += file_content_bytes
-            # Ensure proper line separation between files in JSONL
-            if combined_data and not combined_data.endswith(b'\n'):
-                combined_data += b'\n'
-        else:
-            print(f"Warning: Job {job.display_name} succeeded but has no output file.")
-
-    # 4. Save to destination
-    with open(output_path, "wb") as f:
-        f.write(combined_data)
-
-    print(f"\nSuccess! All results concatenated and saved to:\n{output_path}")
 
 if __name__ == "__main__":
-    main()
+    raise SystemExit(main())