Script.org file and batching utilities

fetch_batched_results.py

@@ -0,0 +1,63 @@
+import os
+import sys
+import argparse
+from pathlib import Path
+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()
+
+    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.")
+
+    client = genai.Client()
+
+    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)
+
+    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.")
+
+    # 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()