From fb5f2c5441ad8e5cb16ced507734b56c20795af7 Mon Sep 17 00:00:00 2001
From: Fadil Rahman <87557055+fadil4u@users.noreply.github.com>
Date: Fri, 6 Jun 2025 23:19:23 +0530
Subject: [PATCH] Add batch polling to python code in batches docs (#11286)

---
 docs/my-website/docs/batches.md | 34 ++++++++++++++++++++++++++++++---
 1 file changed, 31 insertions(+), 3 deletions(-)

diff --git a/docs/my-website/docs/batches.md b/docs/my-website/docs/batches.md
index 7029697699c..d5fbc53c080 100644
--- a/docs/my-website/docs/batches.md
+++ b/docs/my-website/docs/batches.md
@@ -116,10 +116,38 @@ print("response from litellm.create_batch=", create_batch_response)
 **Retrieve the Specific Batch and File Content**
 
 ```python
+    # Maximum wait time before we give up
+    MAX_WAIT_TIME = 300  
+
+    # Time to wait between each status check
+    POLL_INTERVAL = 5
+    
+    #Time waited till now 
+    waited = 0
+
+    # Wait for the batch to finish processing before trying to retrieve output
+    # This loop checks the batch status every few seconds (polling)
+
+    while True:
+        retrieved_batch = await litellm.aretrieve_batch(
+            batch_id=create_batch_response.id,
+            custom_llm_provider="openai"
+        )
+        
+        status = retrieved_batch.status
+        print(f"⏳ Batch status: {status}")
+        
+        if status == "completed" and retrieved_batch.output_file_id:
+            print("✅ Batch complete. Output file ID:", retrieved_batch.output_file_id)
+            break
+        elif status in ["failed", "cancelled", "expired"]:
+            raise RuntimeError(f"❌ Batch failed with status: {status}")
+        
+        await asyncio.sleep(POLL_INTERVAL)
+        waited += POLL_INTERVAL
+        if waited > MAX_WAIT_TIME:
+            raise TimeoutError("❌ Timed out waiting for batch to complete.")
 
-retrieved_batch = await litellm.aretrieve_batch(
-    batch_id=create_batch_response.id, custom_llm_provider="openai"
-)
 print("retrieved batch=", retrieved_batch)
 # just assert that we retrieved a non None batch