Update main.py

main.py

@@ -4,22 +4,22 @@ except ImportError: # pip < 10.0
 
 pkgs = freeze.freeze()
 for pkg in pkgs: print(pkg)
-import os 
-from fastapi import FastAPI, HTTPException, File, UploadFile
-from fastapi.middleware.cors import CORSMiddleware
-from PyPDF2 import PdfReader
-
-
-import google.generativeai as genai
-import json
 
-secret = os.environ["key"]
-genai.configure(api_key=secret)
-model_vision = genai.GenerativeModel('gemini-pro-vision')
-model_text = genai.GenerativeModel('gemini-pro')
+import io
+import asyncio
+import time
+import aiohttp
+from PyPDF2 import PdfReader, PdfWriter
+import os
+from pathlib import Path
+from aiohttp import FormData
+from fastapi import FastAPI, File, UploadFile
+from fastapi.responses import JSONResponse
+from fastapi.middleware.cors import CORSMiddleware
 
 app = FastAPI()
 
+# Configure CORS
 app.add_middleware(
     CORSMiddleware,
     allow_origins=["*"],
@@ -28,30 +28,108 @@ app.add_middleware(
     allow_headers=["*"],
 )
 
-@app.post("/get_ocr_data/")
-async def get_data(pdf: UploadFile = File(...)):
-    try:
-        # Read PDF file using PyPDF2
-        pdf_reader = PdfReader(pdf.file)
-        text = ""
-        
-        # Extract text from each page
-        for page in pdf_reader.pages:
-            text += page.extract_text()
+async def execute_pdfscraper_async(file_contents: bytes, file_name: str, pages_per_chunk: int):
+    split_pdf(file_contents, file_name, pages_per_chunk)
+    response_list = []
+    async with aiohttp.ClientSession() as session:
+        tasks = [
+            call_pdfscraper(session, chunk_data, chunk_name)
+            for chunk_data, chunk_name in load_chunks(file_name)
+        ]
+        responses = await asyncio.gather(*tasks)
+        for response in responses:
+            response_list.append(response[0])
+
+    return response_list
+
+async def call_pdfscraper(session, file_contents, pdf_name):
+    headers = {"Origin": "http://localhost:8080"}
+    url = "https://us-central1-neuralgap-1.cloudfunctions.net/scraperPDFDocxTables_v3"
+    # Create a FormData object
+    data = FormData()
+    data.add_field(
+        "pdf",
+        file_contents,
+        filename=os.path.basename(pdf_name),
+        content_type="application/pdf",
+    )
+    data.add_field("processTables", "True")
+
+    async with session.post(url, data=data, headers=headers) as resp:
+        if resp.status == 200:
+            response = await resp.json()
+        else:
+            print(f"Failed to get response: {resp.status}")
+            return {}
+
+    return response, pdf_name
+
+def collect_pdfscraper_response(scrape_response_list):
+    content_list = []
+    tables_dict = {}
+    table_count = 1
+    for response in scrape_response_list:
+        content = response["corpus"]
+        table_content = response["tables_raw"]
+
+        content_list.append(content)
+        try:
+            for table_key in table_content.keys():
+                tables_dict[str(table_count)] = table_content[table_key]
+                table_count += 1
+
+        except AttributeError:
+            pass
+
+    content_str = "\n".join(content_list)
+
+    return content_str, tables_dict
+
+def split_pdf(file_contents, file_name, pages_per_chunk):
+    file_bytes = io.BytesIO(file_contents)
+    reader = PdfReader(file_bytes)
+    total_pages = len(reader.pages)
+
+    output_dir = Path(file_name).parent / "chunks"
+    os.makedirs(output_dir, exist_ok=True)
+
+    num_chunks = (total_pages + pages_per_chunk - 1) // pages_per_chunk
+
+    for i in range(num_chunks):
+        writer = PdfWriter()
+        start_page = i * pages_per_chunk
+        end_page = min(start_page + pages_per_chunk, total_pages)
+
+        for page_number in range(start_page, end_page):
+            writer.add_page(reader.pages[page_number])
 
-        # call gemini to get required data extracted text
-        prompt = f"""this is cv data : {text.strip()} 
-         i want only 
+        chunk_file_name = f"{Path(file_name).stem}_{i + 1}.pdf"
+        output_path = output_dir / chunk_file_name
+        with open(output_path, "wb") as output_pdf:
+            writer.write(output_pdf)
 
-         fisrtname,lastname,contact number,total years of experince,linkdn link,experinece,skils 
+def load_chunks(file_name):
+    output_dir = Path(file_name).parent / "chunks"
+    chunk_list = os.listdir(output_dir)
+    chunk_byte_list = [
+        (open(f"{output_dir}/{file}", "rb").read(), file) for file in chunk_list
+    ]
+    return chunk_byte_list
 
-         in json format only"""
-        
-        response = model_text.generate_content(prompt)
-        data = json.loads(response.text.replace("```json","").replace("```",""))
-        return {"data":data}
+@app.post("/process-pdf/")
+async def process_pdf(file: UploadFile = File(...), pages_per_chunk: int = 2):
+    file_contents = await file.read()
+    file_name = file.filename
 
-    except Exception as e:
-        raise HTTPException(status_code=500, detail=f"Error processing PDF: {str(e)}")
+    start_time = time.time()
+    scrape_response_list = await execute_pdfscraper_async(file_contents, file_name, pages_per_chunk)
+    content, table_string = collect_pdfscraper_response(scrape_response_list)
+    end_time = time.time()
 
+    time_taken = end_time - start_time
+    return JSONResponse(content={"content": content, "tables": table_string, "time_taken": time_taken})
 
+# Start the FastAPI app
+# if __name__ == "__main__":
+#     import uvicorn
+#     uvicorn.run(app,port=7000,workers=2)