Update main.py

main.py

@@ -1,12 +1,11 @@
 import os
 import io
 import time
-import tempfile
 import PIL.Image
 from fastapi import FastAPI, File, UploadFile, HTTPException
+from fastapi.responses import JSONResponse
 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")
@@ -36,11 +35,11 @@ def extract_text_from_image(img):
             )
             return response.text
         except ClientError as e:
-            # Check if error code is 429 (RESOURCE_EXHAUSTED)
-            if e.status_code == 429:
+            # Extract error code from the exception arguments
+            error_code = e.args[0] if e.args and isinstance(e.args[0], int) else None
+            if error_code == 429:
                 if attempt < max_retries - 1:
-                    # Wait for an exponentially increasing delay before retrying.
-                    time.sleep(2 ** attempt)
+                    time.sleep(2 ** attempt)  # Exponential backoff before retrying
                     continue
                 else:
                     raise HTTPException(
@@ -53,7 +52,7 @@ def extract_text_from_image(img):
                     detail=f"Error processing image: {str(e)}"
                 )
 
-@app.post("/upload", summary="Upload a PDF or image file", response_description="Returns a JSON response with the extracted text")
+@app.post("/upload", summary="Upload a PDF or image file", response_description="Returns extracted text as JSON")
 async def upload_file(file: UploadFile = File(...)):
     if not file.filename:
         raise HTTPException(status_code=400, detail="No file provided")
@@ -82,9 +81,9 @@ async def upload_file(file: UploadFile = File(...)):
         
         output_text += extract_text_from_image(img) + "\n\n"
     
-    # Return the extracted text as JSON.
-    return {"extracted_text": output_text}
+    # Return the extracted text in a JSON response.
+    return JSONResponse(content={"extracted_text": output_text})
 
 @app.get("/", summary="Health Check")
 async def root():
-    return {"message": "API is up and running."}
+    return JSONResponse(content={"message": "API is up and running."})