1 files changed, 60 insertions, 0 deletions
diff --git a/scripts/check_batch_status.py b/scripts/check_batch_status.py
new file mode 100644
index 0000000..8c3ea5d
--- /dev/null
+++ b/scripts/check_batch_status.py
@@ -0,0 +1,60 @@
+import json
+import os
+import time
+from openai import OpenAI
+
+BATCH_IDS_FILE = "data/putnam_eval/submitted_batch_ids.json"
+
+def check_status():
+    api_key = os.getenv("OPENAI_API_KEY")
+    if not api_key:
+        print("Error: OPENAI_API_KEY not set.")
+        return
+    client = OpenAI(api_key=api_key)
+
+    if not os.path.exists(BATCH_IDS_FILE):
+        print(f"Error: {BATCH_IDS_FILE} not found. Did you submit batches?")
+        return
+
+    with open(BATCH_IDS_FILE, "r") as f:
+        batch_ids = json.load(f)
+    
+    print(f"Checking status for {len(batch_ids)} batches...\n")
+    
+    all_completed = True
+    completed_count = 0
+    
+    print(f"{'BATCH ID':<35} | {'STATUS':<15} | {'REQ COUNT':<10} | {'COMPLETED':<10} | {'FAILED':<10}")
+    print("-" * 95)
+    
+    for b_id in batch_ids:
+        try:
+            batch = client.batches.retrieve(b_id)
+            status = batch.status
+            counts = batch.request_counts
+            total = counts.total if counts else 0
+            comp = counts.completed if counts else 0
+            fail = counts.failed if counts else 0
+            
+            print(f"{b_id:<35} | {status:<15} | {total:<10} | {comp:<10} | {fail:<10}")
+            
+            if status != "completed":
+                all_completed = False
+            else:
+                completed_count += 1
+                
+        except Exception as e:
+            print(f"{b_id:<35} | {'ERROR':<15} | {str(e)}")
+            all_completed = False
+
+    print("-" * 95)
+    print(f"\nProgress: {completed_count}/{len(batch_ids)} batches completed.")
+    
+    if all_completed:
+        print("\nSUCCESS: All batches finished! You can now run scripts/retrieve_results.py")
+    else:
+        print("\nSome batches are still processing. Check again later.")
+
+if __name__ == "__main__":
+    check_status()
+