Update mainapp.py

mainapp.py

@@ -1,43 +1,67 @@
-from fastapi import FastAPI, UploadFile, File, HTTPException 
+from fastapi import FastAPI, UploadFile, File, HTTPException
 from fastapi.responses import JSONResponse
 from PIL import Image
 import io
 import asyncio
-from llm import  call_llm
+import fitz  # PyMuPDF
+from llm import call_llm
+
 app = FastAPI()
- # Initialize once
 
+@app.get("/")
+def home():
+    return {"message": "home page"}
 
 
+# 🔥 Process single image
+async def process_image(image: Image.Image, page_num: int):
+    prompt = """
+    Extract all text and return ONLY valid JSON structured.
+    
+    """
+
+    result = await call_llm(image, prompt)
+
+    return {
+        "page": page_num,
+        "llm": result
+    }
 
 
-@app.get("/")
-def home ():
-    return JSONResponse(content={'message':'home page'})
 @app.post("/ocr-llm")
 async def ocr_llm_endpoint(file: UploadFile = File(...)):
+    
     if not file.filename.lower().endswith((".pdf", ".png", ".jpg", ".jpeg")):
         raise HTTPException(status_code=400, detail="File must be PDF or image")
-    
+
     try:
         results = []
 
+        # ✅ HANDLE PDF
         if file.filename.lower().endswith(".pdf"):
             file_bytes = await file.read()
-            # doc = fitz.open(stream=file_bytes, filetype="pdf")
+            doc = fitz.open(stream=file_bytes, filetype="pdf")
+
+            tasks = []
 
-            # # Process each page sequentially or concurrently
-            # tasks = [process_pdf_page(page) for page in doc]
-            # results = await asyncio.gather(*tasks)
+            for i, page in enumerate(doc):
+                pix = page.get_pixmap()
+                img_bytes = pix.tobytes("png")
+                image = Image.open(io.BytesIO(img_bytes)).convert("RGB")
 
+                tasks.append(process_image(image, i + 1))
+
+            results = await asyncio.gather(*tasks)
+
+        # ✅ HANDLE IMAGE
         else:
-            # Single image file
             image_bytes = await file.read()
             image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
-            llm_results=await call_llm(image,"")
-            results.append({"page": 1,  "llm": llm_results})
 
-        return JSONResponse(content={"results": results})
+            result = await process_image(image, 1)
+            results.append(result)
+
+        return {"results": results}
 
     except Exception as e:
         raise HTTPException(status_code=500, detail=str(e))