Update main.py

main.py

@@ -1,12 +1,13 @@
 import os
 import io
+import time
 import tempfile
 import PIL.Image
 from fastapi import FastAPI, File, UploadFile, HTTPException
-from fastapi.responses import FileResponse
 from pdf2image import convert_from_bytes
 from google import genai
 from google.genai import types
+from google.genai.errors import ClientError
 
 app = FastAPI(title="PDF/Image Text Extraction API")
 
@@ -21,14 +22,38 @@ client = genai.Client(api_key=API_KEY)
 def extract_text_from_image(img):
     """
     Extracts text from a PIL image using the Google GenAI API.
+    Includes error handling for RESOURCE_EXHAUSTED errors.
     """
-    response = client.models.generate_content(
-        model="gemini-2.0-flash",
-        contents=["Extract the text from the image. Do not write anything except the extracted content", img]
-    )
-    return response.text
+    max_retries = 3
+    for attempt in range(max_retries):
+        try:
+            response = client.models.generate_content(
+                model="gemini-2.0-flash",
+                contents=[
+                    "Extract the text from the image. Do not write anything except the extracted content",
+                    img,
+                ]
+            )
+            return response.text
+        except ClientError as e:
+            # Check if error code is 429 (RESOURCE_EXHAUSTED)
+            if e.status_code == 429:
+                if attempt < max_retries - 1:
+                    # Wait for an exponentially increasing delay before retrying.
+                    time.sleep(2 ** attempt)
+                    continue
+                else:
+                    raise HTTPException(
+                        status_code=503,
+                        detail="API resource exhausted. Please try again later."
+                    )
+            else:
+                raise HTTPException(
+                    status_code=500,
+                    detail=f"Error processing image: {str(e)}"
+                )
 
-@app.post("/upload", summary="Upload a PDF or image file", response_description="Returns a Markdown file with the extracted text")
+@app.post("/upload", summary="Upload a PDF or image file", response_description="Returns a JSON response with the extracted text")
 async def upload_file(file: UploadFile = File(...)):
     if not file.filename:
         raise HTTPException(status_code=400, detail="No file provided")
@@ -57,16 +82,9 @@ async def upload_file(file: UploadFile = File(...)):
         
         output_text += extract_text_from_image(img) + "\n\n"
     
-    # Save the extracted text to a temporary Markdown file.
-    temp_md = tempfile.NamedTemporaryFile(delete=False, suffix=".md")
-    with open(temp_md.name, "w", encoding="utf-8") as md_file:
-        md_file.write(output_text)
-    
-    # Return the file as a downloadable response.
-    return FileResponse(temp_md.name, filename="output.md", media_type="text/markdown")
+    # Return the extracted text as JSON.
+    return {"extracted_text": output_text}
 
-# --- API Endpoints ---
 @app.get("/", summary="Health Check")
 async def root():
     return {"message": "API is up and running."}
-