Update main.py

main.py

@@ -1,51 +1,19 @@
-try: from pip._internal.operations import freeze
-except ImportError: # pip < 10.0
-    from pip.operations import freeze
-
-pkgs = freeze.freeze()
-for pkg in pkgs: print(pkg)
-
 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
+import aiohttp
+from PyPDF2 import PdfReader, PdfWriter
+from fastapi import FastAPI, UploadFile, Form
 from fastapi.responses import JSONResponse
-from fastapi.middleware.cors import CORSMiddleware
+from aiohttp import FormData
+import shutil
 
 app = FastAPI()
 
-# Configure CORS
-app.add_middleware(
-    CORSMiddleware,
-    allow_origins=["*"],
-    allow_credentials=True,
-    allow_methods=["*"],
-    allow_headers=["*"],
-)
-
-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):
+async def call_pdfscraper(session, file_contents, pdf_name, processTables):
     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",
@@ -53,24 +21,42 @@ async def call_pdfscraper(session, file_contents, pdf_name):
         filename=os.path.basename(pdf_name),
         content_type="application/pdf",
     )
-    data.add_field("processTables", "True")
+    data.add_field("processTables", processTables)
 
     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 {}, pdf_name
 
     return response, pdf_name
 
+
+async def execute_pdfscraper_async(file_path: str, processTables: str):
+    chunk_list = os.listdir(file_path)
+    chunk_byte_list = [
+        (open(f"{file_path}/{file}", "rb").read(), file) for file in chunk_list
+    ]
+    response_list = []
+    async with aiohttp.ClientSession() as session:
+        tasks = [
+            call_pdfscraper(session, file_all[0], file_all[1], processTables)
+            for file_all in chunk_byte_list
+        ]
+        responses = await asyncio.gather(*tasks)
+        for i, response in enumerate(responses):
+            response_list.append(response[0])
+
+    return response_list
+
+
 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 = response.get("corpus", "")
+        table_content = response.get("tables_raw", {})
 
         content_list.append(content)
         try:
@@ -85,11 +71,11 @@ def collect_pdfscraper_response(scrape_response_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)
 
@@ -108,27 +94,23 @@ def split_pdf(file_contents, file_name, pages_per_chunk):
         with open(output_path, "wb") as output_pdf:
             writer.write(output_pdf)
 
-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
+    return str(output_dir)
+
 
 @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
+async def process_pdf(pdf_file: UploadFile, pages_per_chunk: int = Form(2), processTables: str = Form("True")):
+    file_contents = await pdf_file.read()
 
-    scrape_response_list = await execute_pdfscraper_async(file_contents, file_name, pages_per_chunk)
+    chunks_dir = split_pdf(file_contents, pdf_file.filename, pages_per_chunk)
+    scrape_response_list = await execute_pdfscraper_async(chunks_dir, processTables)
     content, table_string = collect_pdfscraper_response(scrape_response_list)
-
-    shutil.rmtree("chunks")
-    print("old chunks removed") 
+    
+    shutil.rmtree(chunks_dir)  # Clean up chunks after processing
+    
     return JSONResponse(content={"content": content, "tables": table_string})
 
-# Start the FastAPI app
+
+# Starting point for running the FastAPI app
 # if __name__ == "__main__":
 #     import uvicorn
-#     uvicorn.run(app,port=7000,workers=2)
+#     uvicorn.run(app, host="0.0.0.0", port=8000)